import numpy as np
import os 
import json
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
from matplotlib.patches import Polygon
from matplotlib import path
import numpy as np 
import netCDF4

# Path to the JSON file
file_path = 'Basisdata_0000_Norge_4258_Fylker_GeoJSON.geojson' # from https://kartkatalog.geonorge.no/nedlasting

# Function to read JSON file 
def read_json_file(file_path):
    with open(file_path, 'r') as file:
        data = json.load(file)
    return data

# Function to check if a point is inside polygon
def is_point_in_polygon(point, polygon):
    """
    Check if a point is inside a polygon using the ray-casting algorithm.
    
    Parameters:
    - point: A tuple representing the (longitude, latitude) of the point.
    - polygon: A list of tuples representing the vertices of the polygon.
    
    Returns:
    - True if the point is inside the polygon, False otherwise.
    """
    x, y = point
    poly_path = path.Path(polygon)
    return poly_path.contains_point((x, y))

# Reading the data from the JSON file
data = read_json_file(file_path)
coords_counties = []
name_counties = []
for county in data['Fylke']['features']:
    county_name = county['properties']['fylkesnavn']
    if county_name.find('-') >= 0: 
        county_name = county_name[0:county_name.find('-')-1]
    print(county_name)
    name_counties.append(county_name)
    coords_counties.append(county['geometry']['coordinates'][0][0][:][:])

# Define names of regions that include multiple counties 
name_counties.append('Østlandet') # 1, 2, 3, 5, 8, 9, 10  Innlandet, Buskerud, Telemark, Akershus, Østfold, Vestfold
name_counties.append('Sørlandet') # 7 Adger 
name_counties.append('Vestlandet') # 4, 6 Rogaland, Vestland, Møre og Romsdal
name_counties.append('Midt-Norge') # 13 Trøndelag
name_counties.append('Nord-Norge') # 11, 12, 14 Nordland, Troms, Finnmark 
name_counties.append('Norge') # 11, 12, 14


# Define county/region masks and plot 

colors = [
    (255, 0, 0),    # Red
    (0, 255, 0),    # Green
    (0, 0, 255),    # Blue
    (255, 255, 0),  # Yellow
    (255, 192, 203),# Pink
    (255, 0, 255),  # Magenta
    (192, 192, 192),# Silver
    (128, 128, 128),# Gray
    (128, 0, 0),    # Maroon
    (128, 128, 0),  # Olive
    (0, 128, 0),    # Dark Green
    (128, 0, 128),  # Purple
    (0, 128, 128),  # Teal
    (0, 0, 128),    # Navy
    (255, 165, 0)   # Orange
]

# Create a new figure
plt.figure(figsize=(8, 8))

# Create a Basemap instance
m = Basemap(projection='merc', llcrnrlat=55, urcrnrlat=72, llcrnrlon=0, urcrnrlon=35, resolution='i')

# Draw coastlines and countries for context
m.drawcoastlines()
m.drawcountries()

res = [1,1]
lon = np.arange(0,360,res[0])
lat = np.arange(90,-91,-res[1])
lon2, lat2 = np.meshgrid(lon,lat)
area = np.cos(lat2*np.pi/180.)
mask = np.zeros((21,lon2.shape[0],lon2.shape[1]))
weights = np.zeros((21,lon2.shape[0],lon2.shape[1]))
for count,coords in enumerate(coords_counties): 
    # Convert coordinates to the map projection
    poly_coords = [m(x, y) for x, y in coords]

    # Create a polygon patch
    polygon = Polygon(poly_coords, closed=True, edgecolor='k', facecolor='cyan', alpha=0.5)
    
    # Add the polygon to the map
    plt.gca().add_patch(polygon)


    for j, latval in enumerate(lat):
        if latval < 55 or latval > 75: continue
        for i, lonval in enumerate(lon):
            if lonval < 5 or lonval >35: continue
            x1,y1 = m(lonval,latval)
            if is_point_in_polygon((x1, y1),poly_coords):
                if count < 15:
                    mask[count,j,i] = 1
                    plt.gca().plot(x1,y1,color=np.array(colors[count])/255,marker='.',markersize=15)
                if count in {1, 2, 3, 5, 8, 9, 10}:
                    mask[15,j,i] = 1
                    #plt.gca().plot(x1,y1,color='r',marker='+',markersize=10)
                if count in {7}:
                    mask[16,j,i] = 1
                    #plt.gca().plot(x1,y1,color='b',marker='o',markersize=10)
                if count in {0, 4, 6}:
                    mask[17,j,i] = 1
                    #plt.gca().plot(x1,y1,color='g',marker='<',markersize=10)
                if count in {13}:
                    mask[18,j,i] = 1
                    #plt.gca().plot(x1,y1,color='y',marker='>',markersize=10)
                if count in {11, 12, 14}:
                    mask[19,j,i] = 1
                    #plt.gca().plot(x1,y1,color='k',marker='d',markersize=10)
                mask[-1,j,i] = 1
for k in range(weights.shape[0]):
    weights[k,:,:] = area * mask[k,:,:]
    if np.sum(mask[k,:,:]) > 0:
        weights[k,:,:] = weights[k,:,:] / np.sum(weights[k,:,:])

# Save plot
plt.savefig('mask_regions_1x1.png',format='png',dpi=150)

# save masks 
filePath = 'mask_regions_1x1.nc'   
if os.path.exists(filePath): os.remove(filePath)
nc = netCDF4.Dataset(filePath, 'w', format='NETCDF4_CLASSIC')
nc.createDimension('lon',len(lon))
nc.createDimension('lat',len(lat))
nc.createDimension('county_index',len(name_counties))
str_len = 128
nc.createDimension('str_len',str_len)
nclon = nc.createVariable('lon','f4',['lon']) 
nclat = nc.createVariable('lat','f4',['lat']) 
ncmask = nc.createVariable('mask','i4',['county_index','lat','lon']) 
ncweights = nc.createVariable('weights','f4',['county_index','lat','lon']) 
ncname = nc.createVariable('county_name','S1',['county_index','str_len']) 
nclon[:] = lon
nclat[:] = lat
ncmask[:] = mask
ncweights[:] = weights
for county_index in range(len(name_counties)): 
    my_string = name_counties[county_index].replace("ø","oe").replace("Ø","Oe") 
    my_array = np.array(list(my_string.ljust(str_len)))
    ncname[county_index,:] = my_array
nc = nc.close()