# small island functions


import xarray as xr
import numpy as np
import pandas as pd 
import seaborn as sns
from numpy.polynomial import Polynomial
import xesmf as xe 

import healpy as hp
import cmocean
import regionmask
from scipy.interpolate import griddata

import easygems.healpix as egh
import matplotlib.tri as tri

import tqdm
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import cartopy.mpl.ticker as cticker



def convert_mtpr_tp(input_data):
    # convert from kg / m2 s to mm
    factor = 0.2777777777777778
    input_data = input_data/factor
    input_data =  1000*input_data
    
    return input_data

def convert_prlr_tp(input_data):
    # convert from m / s to mm
    input_data =  (input_data) * 1000  * 60 * 60
    return input_data


# defining a function to give the closest lat lon indicies based on a target lat lon
def closest_value(input_list, input_value):
    arr = np.asarray(input_list)
    i = (np.abs(arr - input_value)).argmin()
    return arr[i]


def interpolate_timeseries(ref_lat, ref_lon, data, method = "cubic"):
    """Takes a reference lat/lon from an observational station and uses
    cubic interpolation to create a time series from the gridded model
    data"""
    # Extract lats and lons from the data array
    
    try:
        data['lon_adjusted'] = xr.where(data.lon > 180, data.lon - 360, data.lon)
        data['lon'] = data['lon_adjusted']
        data = data.reset_coords(names='lon_adjusted', drop=True)

        lats = data['lat'].values
        lons = data['lon'].values
        
    except AttributeError:
        data['lon_adjusted'] = xr.where(data.rlon > 180, data.rlon - 360, data.rlon)
        data['rlon'] = data['lon_adjusted']
        data = data.reset_coords(names='lon_adjusted', drop=True)

        lats = data['rlat'].values
        lons = data['rlon'].values

    target_lon_nn = closest_value(lons, ref_lon)
    target_index_lon = np.where(lons == target_lon_nn)[0][0]
    target_lat_nn = closest_value(lats, ref_lat)
    target_index_lat = np.where(lats == target_lat_nn)[0][0]

    # time series based on nearest neighbour
    ts =  data[:,target_index_lat, target_index_lon]

    station_ts = xr.zeros_like(ts)

    points = np.array([(lat, lon) for lat in lats for lon in lons])

    # doing the interpolation on each time step
    for t in tqdm.tqdm(range(len(data.time))):

        values = data.values[t,:]  # Your 2D climate data
        # Flatten lat, lon, and value arrays for interpolation
        values_flat = values.flatten()
        
        station_ts[t] = griddata(points, values_flat, (ref_lat, ref_lon), method=method)
        
    return station_ts


def select_box_area(data, latitude, longitude):
    
    #coordinates of Canary islands whole domain 
    latitude = [27.2, 29.5]
    longitude = [-19, -12.5]

    # converting co-ords into degrees from radian
    # data['lon'] = data['lon']*180/(np.pi)
    # data['lat'] = data['lat']*180/(np.pi)

    data['lon_adjusted'] = xr.where(data.lon > 180, data.lon - 360, data.lon)
    data['lon'] = data['lon_adjusted']
    data = data.reset_coords(names='lon_adjusted', drop=True)
    
    indexer = ((data.lat>=latitude[0]) & (data.lat<=latitude[1]) & (data.lon>=longitude[0]) & (data.lon<=longitude[1])).compute()
    data = data.where(indexer, drop=True)

    return data


def create_land_mask(input_data):
    
    country_codes = [37, 47, 54, 78, 103, 121, 122, 137, 156, 160, 224, 234, 235]

    try:
        lat_land = input_data.lat.values 
        lon_land = input_data.lon.values
    except AttributeError:
        lat_land = input_data.rlat.values 
        lon_land = input_data.rlon.values
        
    #The mask has different flag values for each country and nans for the ocean
    land = regionmask.defined_regions.natural_earth_v5_0_0.land_50
    countries = regionmask.defined_regions.natural_earth_v5_0_0.countries_50

    mask_land = land.mask(lon_land, lat_land)
    mask_land
    mask_countries = countries.mask(lon_land, lat_land)
    mask_countries = mask_countries.astype(int)

    output_data = input_data.where(np.isin(mask_countries, country_codes), np.nan)
    
    return output_data

def create_island_mask(input_data):
    
    country_codes = [54, 160] # SPain 54, france 160
    try:
        lat_land = input_data.lat.values 
        lon_land = input_data.lon.values
    except AttributeError:
        lat_land = input_data.rlat.values 
        lon_land = input_data.rlon.values
        

    #The mask has different flag values for each country and nans for the ocean
    land = regionmask.defined_regions.natural_earth_v5_0_0.land_50
    countries = regionmask.defined_regions.natural_earth_v5_0_0.countries_50

    mask_land = land.mask(lon_land, lat_land)
    mask_land
    mask_countries = countries.mask(lon_land, lat_land)
    mask_countries = mask_countries.astype(int)

    output_data = input_data.where(np.isin(mask_countries, country_codes), np.nan)
    
    return output_data

def ticks_round(ticks: list, round_to: int = None):
    """
    Round a tick to the nearest round_to value.

    tick (list):          Tick value.
    round_to (int, opt):  Round to value.

    Returns:
        tick (list):  Rounded tick value.
    """
    if round_to is None:
        # define round_to
        tick_span = ticks[1] - ticks[0]
        if tick_span <= 1:
            round_to = 2
        elif tick_span > 1 and tick_span <= 10:
            round_to = 1
        else:
            round_to = 0

    return np.round(ticks, round_to)

def plot_islands(to_plot, grid_type, min_value, max_value, latitude, longitude, title, cbar_label, diff_or_cont = "diff"):

    projection = ccrs.PlateCarree()

    fsize = 14
    #Add tick marks with values for latitude and longitude axes
    nxticks = 5
    nyticks = 7
    
    # coordinates of La Palma southern tip 
    # this part was to see where the same 'real' co-ordinate was plotted on different re-grids
    # lat_palma = 28.455732
    # lon_palma = -17.838980

    # lat_cruz = 28.100639
    # lon_cruz = -14.499162

    # if grid_type == 'lat_lon':
    #     closest_point = to_plot.sel(lat=lat_palma, lon=lon_palma, method='nearest')
    #     loni = np.where(to_plot.lon == closest_point.lon.values)[0][0]
    #     lati = np.where(to_plot.lat == closest_point.lat.values)[0][0]
    #     print('closest point palma lat:', closest_point.lat.values, 'lon:', closest_point.lon.values,
    #         'real: ', lat_palma, lon_palma)

    #     closest_point = to_plot.sel(lat=lat_cruz, lon=lon_cruz, method='nearest')
    #     loni2 = np.where(to_plot.lon == closest_point.lon.values)[0][0]
    #     lati2 = np.where(to_plot.lat == closest_point.lat.values)[0][0]
    #     print('closest point cruz lat:', closest_point.lat.values, 'lon:', closest_point.lon.values,
    #         'real: ', lat_cruz, lon_cruz)

    if diff_or_cont == "cont":
        cmap = cmocean.cm.thermal
        cmap = cmap(np.linspace(0, 1, 20))
        cmap = ListedColormap(cmap)
    elif diff_or_cont == "diff":
        cmap = plt.get_cmap('RdBu_r', 20)
        
    fig, ax = plt.subplots(1, 1, constrained_layout=False, figsize=(10, 10), 
                            subplot_kw={'projection': projection}) 


    if grid_type == "lat_lon":

        try:
            lons, lats = np.meshgrid(to_plot.rlon.values, to_plot.rlat.values)
        except AttributeError:    
            lons, lats = np.meshgrid(to_plot.lon.values, to_plot.lat.values)
        
        pcm = ax.pcolormesh(lons, lats, to_plot, transform= projection, 
            cmap=cmap, vmax = max_value, vmin = min_value) 
        
        # ax.plot(to_plot.lon[loni], to_plot.lat[lati], marker='x', color='red', markersize=6, transform=projection)
        # ax.plot(to_plot.lon[loni2], to_plot.lat[lati2], marker='x', color='red', markersize=6, transform=projection)
    
    elif grid_type == "healpix":
        
        ax.set_global()
        ax.set_extent([longitude[0], longitude[1], latitude[0], latitude[1]], crs=projection)
        pcm = egh.healpix_show(to_plot, ax=ax, vmax = max_value, cmap = cmap)

    elif grid_type == "native":
        triang = tri.Triangulation(to_plot.lon, to_plot.lat)
        pcm = plt.tricontourf(triang, to_plot.data, alpha=1, levels=np.linspace(0, max_value, 20),
            cmap=cmap, transform= projection) 

    ax.coastlines(resolution='10m', color='black', linewidth=1)
    ax.add_feature(cfeature.BORDERS, linewidth=0.5)
        
    plt.title(f'{title}')

    # Latitude labels
    if grid_type == "healpix":
        lat_min = latitude[0]
        lat_max = latitude[1]
    else:
        try:
            lat_min = to_plot['lat'].values.min()
            lat_max = to_plot['lat'].values.max()
        except KeyError:
            lat_min = to_plot['rlat'].values.min()
            lat_max = to_plot['rlat'].values.max()
            
    step = (lat_max - lat_min) / (nyticks - 1)
    yticks = np.arange(lat_min, lat_max, step)
    yticks = ticks_round(ticks=yticks, round_to=None)
    ax.set_yticks(yticks, crs=projection)
    lat_formatter = cticker.LatitudeFormatter()
    ax.yaxis.set_major_formatter(lat_formatter)

    # Latitude labels
    if grid_type == "healpix":
        lon_min = longitude[0]
        lon_max = longitude[1]
    else:
        try:
            lon_min = to_plot['lon'].values.min()
            lon_max = to_plot['lon'].values.max()
        except KeyError:
            lon_min = to_plot['rlon'].values.min()
            lon_max = to_plot['rlon'].values.max()
    step = (lon_max - lon_min) / (nxticks - 1)
    xticks = np.arange(lon_min, lon_max + 1, step)
    xticks = ticks_round(ticks=xticks, round_to=None)
    ax.set_xticks(xticks, crs=projection)
    lon_formatter = cticker.LongitudeFormatter()
    ax.xaxis.set_major_formatter(lon_formatter)

    ax.tick_params(axis='y', labelsize=fsize) 
    ax.tick_params(axis='x', labelsize=fsize) 

    try:
        to_plot['lon'].values.max()
        cbar_ax = fig.add_axes([0.36, 0.3, 0.28, 0.01]) 
    except KeyError:
        cbar_ax = fig.add_axes([0.36, 0.27, 0.28, 0.01]) 
        
    cbar = fig.colorbar(pcm, cax = cbar_ax, orientation='horizontal', shrink=0.4)
    
    cbar.set_ticks([np.round(min_value,2),
                    np.round(min_value + (max_value - min_value)/3,2), 
                    np.round(min_value + 2*(max_value - min_value)/3,2), 
                    np.round(max_value,2)])

    cbar.update_ticks()  # Updates the colorbar with the new ticks

    cbar.ax.tick_params(labelsize=fsize)
    cbar.set_label(f"{cbar_label}", fontsize = fsize)

    plt.show()
    
    

# # loop to fix cfsdatetimes 
# # Function to find last valid day of the month
# def last_valid_day(year, month):
#     """Returns the last valid day of a given year and month in the Gregorian calendar."""
#     next_month = datetime(year, month, 28) + timedelta(days=4)  # Move to the next month
#     return (next_month - timedelta(days=next_month.day)).day  # Go back to last valid day

# # Conversion function
# def convert_cftime_to_npdate(cftime_dates):
#     np_dates = []
#     for date in cftime_dates:
#         # Get the last valid day of the month in Gregorian calendar
#         max_day = last_valid_day(date.year, date.month)
        
#         # Adjust day if it exceeds the max valid day
#         safe_day = min(date.day, max_day)
        
#         # Convert to numpy datetime64
#         np_date = np.datetime64(f"{date.year:04d}-{date.month:02d}-{safe_day:02d}T{date.hour:02d}:{date.minute:02d}:{date.second:02d}")
#         np_dates.append(np_date)
    
#     return np.array(np_dates)


def adjust_coords(data):

    data['lon_adjusted'] = xr.where(data.lon > 180, data.lon - 360, data.lon)
    data['lon'] = data['lon_adjusted']
    data = data.reset_coords(names='lon_adjusted', drop=True)
    
    return data
