## gloria functions

import numpy as np 
import os 
import xarray as xr
import pandas as pd
from numpy import ceil
from pandas import Timedelta, Series
import scipy.stats as scs
from scipy.stats import norm
from pylab import cm
import regionmask 
import geopandas as gp
import seaborn as sns

import matplotlib
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import matplotlib.colors as colors 
import cartopy.mpl.ticker as cticker
from mpl_toolkits.axes_grid1 import make_axes_locatable
from matplotlib.projections import PolarAxes
import mpl_toolkits.axisartist.floating_axes as fa
import mpl_toolkits.axisartist.grid_finder as gf
from matplotlib.animation import FuncAnimation

from dask.distributed import Client, LocalCluster
import matplotlib.gridspec as gs
from matplotlib.colors import LogNorm

import time
#import easygems.healpix as egh
import cartopy.feature as cf
import healpy as hp

def convert_tp(input_data):
    # convert from unites of m/hr to mm / day
    input_data =  (input_data * 1000) * 24
    return input_data

def convert_mtpr(input_data):
    # convert from kg / m2 s to mm/day
    factor = 0.2777777777777778
    input_data = input_data/factor
    input_data = convert_tp(input_data)
    return input_data

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

def convert_prlr(input_data):
    # convert from m / s to mm / day
    input_data =  (input_data) * 1000  * 60 * 60 * 24
    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

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

    Args:
        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 fix_coords_era5(input_data,lons):
    if input_data.lon[-1] > 180:
        res = 0.25
        # Roll the longitude dimension to move the lon_lim-360 values to the beginning
        roll_amount = int(abs(lons[0])/res)
        data_array_rolled = input_data.roll(lon=(roll_amount))
        # Update the longitude values accordingly
        data_array_rolled['lon'] = np.linspace(lons[0], lons[1], len(input_data.lon))

    return data_array_rolled

def fix_aqua_coords(input_data, lat_lim, lon_lim):

    # Update the longitude values accordingly
    input_data['lon'] = np.linspace(lon_lim[0], lon_lim[1], len(input_data.lon))
    input_data['lat'] = np.linspace(lat_lim[0], lat_lim[1], len(input_data.lat))

    return input_data


def create_spain_mask(input_data):
    lat_land = input_data.lat.values 
    lon_land = input_data.lon.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

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

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

    lat_land = input_data.lat.values 
    lon_land = input_data.lon.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_catalan_mask(input_data, hpc):
    
    if hpc == "lumi":
        shpfile = "/users/kgrayson/shape_files/recintos_provinciales_inspire_peninbal_etrs89.shp"
    elif hpc == "levante":
        shpfile = "/home/b/b382291/shape_files/spain/recintos_provinciales_inspire_peninbal_etrs89.shp"        
    region_name = (['Barcelona','Girona','Lleida','Tarragona'])
    maskcoords = xr.Dataset({'lat' : (['lat'],input_data['lat'].values)}, {'lon' : (['lon'],input_data['lon'].values)})
    mask = np.zeros([maskcoords.lat.size, maskcoords.lon.size])

    # read in shapefile
    shpcontents = gp.read_file(shpfile)

    # loop over states to mask
    for i in range(0,len(region_name),1):
    #        print("masking "+regionname[i])
        region = shpcontents[shpcontents.NAMEUNIT == region_name[i]]
        maskt = regionmask.mask_geopandas(region, maskcoords["lon"], maskcoords["lat"])
        maskt = np.where(np.isnan(maskt), 0, 1)
        mask[:,:] = mask[:,:] + maskt[:,:]

    # ensure unmasked region is set to 1, rest set to nan's
    mask = np.where(mask == 0, np.nan, 1)
    mask = xr.DataArray(mask, coords=maskcoords.coords)

    output_data = input_data.where(mask == 1, np.nan)

    return output_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]

# set the colormap and centre the colorbar
def shiftedColorMap(cmap, start=0, midpoint=0.5, shift = 0, stop=1.0, name='shiftedcmap'):
    '''
    Function to offset the "center" of a colormap. Useful for
    data with a negative min and positive max and you want the
    middle of the colormap's dynamic range to be at zero.

    Input
    -----
      cmap : The matplotlib colormap to be altered
      start : Offset from lowest point in the colormap's range.
          Defaults to 0.0 (no lower offset). Should be between
          0.0 and `midpoint`.
      midpoint : The new center of the colormap. Defaults to 
          0.5 (no shift). Should be between 0.0 and 1.0. In
          general, this should be  1 - vmax / (vmax + abs(vmin))
          For example if your data range from -15.0 to +5.0 and
          you want the center of the colormap at 0.0, `midpoint`
          should be set to  1 - 5/(5 + 15)) or 0.75
      stop : Offset from highest point in the colormap's range.
          Defaults to 1.0 (no upper offset). Should be between
          `midpoint` and 1.0.
    '''
    cdict = {
        'red': [],
        'green': [],
        'blue': [],
        'alpha': []
    }

    # regular index to compute the colors
    reg_index = np.linspace(start, stop, 257)

    # shifted index to match the data
    shift_index = np.hstack([
        np.linspace(0.0, midpoint - shift, 128, endpoint=False), 
        np.linspace(midpoint - shift, 1.0, 129, endpoint=True)
    ])

    for ri, si in zip(reg_index, shift_index):
        r, g, b, a = cmap(ri)

        cdict['red'].append((si, r, r))
        cdict['green'].append((si, g, g))
        cdict['blue'].append((si, b, b))
        cdict['alpha'].append((si, a, a))

    newcmap = colors.LinearSegmentedColormap(name, cdict)
    #plt.register_cmap(cmap=newcmap)
    return newcmap

def compute_integrated_winds(u, v, p, w):
    """
    Compute integrated winds.
    
    Parameters:
    u : numpy.ndarray
        Zonal wind component at the same pressure levels (m/s).
    v : numpy.ndarray
        Meridional wind component at the same pressure levels (m/s).
    p : numpy.ndarray
        Pressure levels (Pa).

    Returns:
    integrated_winds : float
        Integrated winds (kg m^-1 s^-1).
    """
    # Constants
    g = 9.81  # Acceleration due to gravity (m/s^2)

    # Convert pressure from hPa to Pa
    p = p * 100

    # Calculate integrated u and v components using np.trapz
    integrated_u = np.trapz(u, x=p, axis = 1) / g
    integrated_v = np.trapz(v, x=p, axis = 1) / g
    integrated_v = np.trapz(w, x=p, axis = 1) / g
    
    # print("Integrated u-component (m^2/s):", integrated_u)
    # print("Integrated v-component (m^2/s):", integrated_v)

    return integrated_u, integrated_v

def compute_ivt(q, u, v, p, w):
    """
    Compute Integrated Water Vapor Transport (IVT).
    
    Parameters:
    q : numpy.ndarray
        Specific humidity at pressure levels (kg/kg).
    u : numpy.ndarray
        Zonal wind component at the same pressure levels (m/s).
    v : numpy.ndarray
        Meridional wind component at the same pressure levels (m/s).
    p : numpy.ndarray
        Pressure levels (hPa).

    Returns:
    ivt : float
        Integrated Water Vapor Transport (kg m^-1 s^-1).
    """
    g = 9.81 

    # Ensure pressure is in ascending order
    if p[0] > p[-1]:
        p = p[::-1]
        q = q[:,::-1,:,:] # want to reverse the order of pressure along the second dimension
        u = u[:,::-1,:,:] 
        v = v[:,::-1,:,:]
        w = w[:,::-1,:,:]

    # Compute the integrals using trapezoidal integration
    qu =  q*u 
    qv =  q*v 

    # Integrated values
    int_qu_dp = qu.integrate("level") / g # integrating over pressure levels
    int_qv_dp = qv.integrate("level") / g # integrating over pressure levels
    int_w_dp = w.integrate("level") / g # integrating over pressure levels
    
    # # Convert pressure from hPa to Pa
    int_qu_dp = int_qu_dp*100 
    int_qv_dp = int_qv_dp*100 
    int_w_dp = int_w_dp*100 
    
    # Calculate IVT magnitude
    ivt = np.sqrt(int_qu_dp**2 + int_qv_dp**2)

    return ivt, int_qu_dp, int_qv_dp, int_w_dp

#!/usr/bin/env python
# Copyright: This document has been placed in the public domain.

def crop_to_box(data_incoming, data_type):
    """
    Crops the incoming DestinE data into the region defined by the box
    """
    
    # lon_minr = 100 #-3.086207
    # lat_minr = 70 #37.039106
    # lon_maxr = 180 #4.74569
    # lat_maxr = 140 #43.977654

    lat_box = 43.977654
    lon_box = -3.086207
    
    if data_type == "ifsfesom":
        # coordinates for the box
        loni_minr = 100
        lati_minr = 70
        loni_maxr = 180
        lati_maxr = 140

    elif data_type == "era5":
        addition_lon = 32
        addition_lat = 28
        
    elif data_type == "mswep":
        addition_lon = 80
        addition_lat = 70

    if data_type != "ifsfesom":
        # Find the nearest lat/lon in the dataset to the target coordinates
        closest_point = data_incoming.sel(lat=lat_box, lon=lon_box, method='nearest')
        loni_city = np.where(data_incoming.lon == closest_point.lon.values)[0][0]
        lati_city = np.where(data_incoming.lat == closest_point.lat.values)[0][0]

        # print(pr_mswep[0, val_lati, val_loni])
        # # Output the result
        # print(closest_point.lon.values)

        # coordinates for the box
        loni_minr = loni_city
        loni_maxr = loni_city + addition_lon

        if data_type in ["mswep", "ifsfesom"]:
            lati_maxr = lati_city + 1
            lati_minr = lati_maxr - addition_lat
        elif data_type == "era5":
            lati_minr = lati_city + 1
            lati_maxr = lati_minr + addition_lat


    try:
        data_incoming.time
        data_cropped = data_incoming[:,lati_minr:lati_maxr,loni_minr:loni_maxr]
    except AttributeError:
        data_cropped = data_incoming[lati_minr:lati_maxr,loni_minr:loni_maxr]
        
    return data_cropped

def crop_to_city(data_incoming, data_type):
    
    "Crops incoming data to a 50 km square covering 4 cities"
    
    cities = ("Valencia", "Girona", "Ibiza", "Alicante")
    
    if data_type == "era5":
        addition = 2
    elif data_type in ["mswep", "ifsfesom"]:
        addition = 5

    # gir_lat = 42.167597765363126
    # gir_lon = 2.5646551724137936
    # ali_lat = 38.64804469273743
    # ali_lon = -0.9051724137931032
    # val_lat = 39.452513966480446
    # val_lon = -0.706896551724137
    # ibi_lat = 39.25139664804469
    # ibi_lon = 1.1767241379310356

    for count, city in enumerate(cities):
        if city == "Valencia":
            lat_city = 39.452513966480446 #
            lon_city = -0.706896551724137

        elif city == "Girona":
            lat_city = 42.167597765363126 # 122
            lon_city = 2.5646551724137936 # 154
 
        elif city == "Alicante":
            lat_city = 38.64804469273743#top
            lon_city = -0.9051724137931032#left
       
        elif city == "Ibiza":
            lat_city = 39.25139664804469
            lon_city = 1.1767241379310356 # 142

        # Find the nearest lat/lon in the dataset to the target coordinates
        closest_point = data_incoming.sel(lat=lat_city, lon=lon_city, method='nearest')
        loni_city = np.where(data_incoming.lon == closest_point.lon.values)[0][0]
        lati_city = np.where(data_incoming.lat == closest_point.lat.values)[0][0]

        # print(pr_mswep[0, val_lati, val_loni])
        # # Output the result
        # print(closest_point.lon.values)

        # coordinates for the box
        loni_minr = loni_city
        loni_maxr = loni_city + addition

        if data_type in ["mswep", "ifsfesom"]:
            lati_maxr = lati_city + 1
            lati_minr = lati_maxr - addition
        elif data_type == "era5":
            lati_minr = lati_city + 1
            lati_maxr = lati_minr + addition
    
        try:
            data_incoming.time
            data_crop = data_incoming[:,lati_minr:lati_maxr,loni_minr:loni_maxr]
        except AttributeError:
            data_crop = data_incoming[lati_minr:lati_maxr,loni_minr:loni_maxr]
        
        if city == "Valencia":
            data_val = data_crop

        elif city == "Girona":
            data_gir = data_crop
 
        elif city == "Alicante":
            data_ali = data_crop

        elif city == "Ibiza":
            data_ibi = data_crop
            
    return data_val, data_ali, data_ibi, data_gir

#def crop_to_city(data):
    # """
    # Crops the incoming DestinE data into the four cities we defined
    # """

    # cities = ("Valencia", "Girona", "Ibiza", "Alicante")
    # addition = 5

    # ## grid cells with maximum precip
    # # 2k lon -0.4094827586206886 lat 38.44692737430168
    # # 2k lon index 127 lat index 84
    # # 2K 2020-01-19T13:00:00.000000000
    # # 2k lon -0.6077586206896548 lat 38.44692737430168
    # # 2k lon index 125 lat index 84
    # # 2K 2020-01-19T13:00:00.000000000
    # # 2k lon -0.6077586206896548 lat 38.547486033519554
    # # 2k lon index 125 lat index 85

    # for count, city in enumerate(cities):
    #     if city == "Valencia":
    #         lat_city = 94 #
    #         lon_city = 124

    #     elif city == "Girona":
    #         lat_city = 121 # 122
    #         lon_city = 157 # 154
 
    #     elif city == "Alicante":
    #         lat_city = 86 #top
    #         lon_city = 122 #left
       
    #     elif city == "Ibiza":
    #         lat_city = 92
    #         lon_city = 143 # 142

    #     # coordinates for the box
    #     lon_minr = lon_city
    #     lat_maxr = lat_city + 1
    #     lon_maxr = lon_city + addition
    #     lat_minr = lat_maxr - addition

    #     try:
    #         data.time
    #         data_crop = data[:,lat_minr:lat_maxr,lon_minr:lon_maxr]
    #     except AttributeError:
    #         data_crop = data[lat_minr:lat_maxr,lon_minr:lon_maxr]
        
    #     if city == "Valencia":
    #         data_val = data_crop

    #     elif city == "Girona":
    #         data_gir = data_crop
 
    #     elif city == "Alicante":
    #         data_ali = data_crop

    #     elif city == "Ibiza":
    #         data_ibi = data_crop
            
    # return data_val, data_ali, data_ibi, data_gir

def create_time_series(data):
    """
    Creates time series by spatially averaging over the whole domain
    
    Arguments
    -----------
    data : xr.DataArray 
    """

    data_ts = data.mean(axis = 1)
    data_ts = data_ts.mean(axis = 1)

    return data_ts

def create_violin_plot(cont_data, hist_data, twoK_data, var, units, region):
    """
    Creates violin plot of data to compare the 3 scenarios
    
    Arguments
    -----------
    cont_data : xr.DataArray 
    var : string 
    units : string
    region : string
    """

    cont_ts = cont_data.mean(axis = 1)
    cont_ts = cont_ts.mean(axis = 1)

    hist_ts = hist_data.mean(axis = 1)
    hist_ts = hist_ts.mean(axis = 1)

    twoK_ts = twoK_data.mean(axis = 1)
    twoK_ts = twoK_ts.mean(axis = 1)

    data = pd.DataFrame({
        'Date': twoK_ts.time,
        'Control': cont_ts,
        'Historical': hist_ts,
        '+2K': twoK_ts
    })

    # Melt the DataFrame for easier plotting with seaborn
    melted_data = pd.melt(data, id_vars=['Date'], value_vars=['Control', 'Historical', '+2K'],
                        var_name='Time Series', value_name='Value')

    #Plot the violin plots
    plt.figure(figsize=(10, 6))
    sns.violinplot(x='Time Series', y='Value', data=melted_data)
    if var in ["pr", "precipitation", "wind speed"]:
        plt.ylim(0, None)  # Set the lower limit of the y-axis to 0

    #plt.ylim(0, None)  # Set the lower limit of the y-axis to 0
    plt.title(f'{var} - {region}')
    plt.xlabel('Time Series of different scenarios')
    plt.ylabel(fr'{var} [{units}]')
    plt.grid(True)
    #plt.show()
    
    return

"""
Taylor diagram (Taylor, 2001) implementation.
Note: If you have found these software useful for your research, I would
appreciate an acknowledgment.
"""

__version__ = "Time-stamp: <2018-12-06 11:43:41 ycopin>"
__author__ = "Yannick Copin <yannick.copin@laposte.net>"


class TaylorDiagram(object):
    """
    Taylor diagram.
    Plot model standard deviation and correlation to reference (data)
    sample in a single-quadrant polar plot, with r=stddev and
    theta=arccos(correlation).
    """

    def __init__(self, refstd,
                 fig=None, rect=111, label='_', srange=(0, 1.5), extend=False):
        """
        Set up Taylor diagram axes, i.e. single quadrant polar
        plot, using `mpl_toolkits.axisartist.floating_axes`.
        Parameters:
        * refstd: reference standard deviation to be compared to
        * fig: input Figure or None
        * rect: subplot definition
        * label: reference label
        * srange: stddev axis extension, in units of *refstd*
        * extend: extend diagram to negative correlations
        """

        self.refstd = refstd            # Reference standard deviation

        tr = PolarAxes.PolarTransform()

        # Correlation labels
        rlocs = np.array([0, 0.2, 0.4, 0.6, 0.7, 0.8, 0.9, 0.95, 0.99, 1])
        if extend:
            # Diagram extended to negative correlations
            self.tmax = np.pi
            rlocs = np.concatenate((-rlocs[:0:-1], rlocs))
        else:
            # Diagram limited to positive correlations
            self.tmax = np.pi/2
        tlocs = np.arccos(rlocs)        # Conversion to polar angles
        gl1 = gf.FixedLocator(tlocs)    # Positions
        tf1 = gf.DictFormatter(dict(zip(tlocs, map(str, rlocs))))

        # Standard deviation axis extent (in units of reference stddev)
        self.smin = srange[0] * self.refstd
        self.smax = srange[1] * self.refstd

        ghelper = fa.GridHelperCurveLinear(
            tr,
            extremes=(0, self.tmax, self.smin, self.smax),
            grid_locator1=gl1, tick_formatter1=tf1)

        if fig is None:
            fig = plt.figure()

        ax = fa.FloatingSubplot(fig, rect, grid_helper=ghelper)
        fig.add_subplot(ax)

        # Adjust axes
        ax.axis["top"].set_axis_direction("bottom")   # "Angle axis"
        ax.axis["top"].toggle(ticklabels=True, label=True)
        ax.axis["top"].major_ticklabels.set_axis_direction("top")
        ax.axis["top"].label.set_axis_direction("top")
        ax.axis["top"].label.set_text("Correlation")

        ax.axis["left"].set_axis_direction("bottom")  # "X axis"
        ax.axis["left"].label.set_text("Standard deviation")

        ax.axis["right"].set_axis_direction("top")    # "Y-axis"
        ax.axis["right"].toggle(ticklabels=True)
        ax.axis["right"].major_ticklabels.set_axis_direction(
            "bottom" if extend else "left")

        if self.smin:
            ax.axis["bottom"].toggle(ticklabels=False, label=False)
        else:
            ax.axis["bottom"].set_visible(False)          # Unused

        self._ax = ax                   # Graphical axes
        self.ax = ax.get_aux_axes(tr)   # Polar coordinates

        # Add reference point and stddev contour
        l, = self.ax.plot([0], self.refstd, 'k*',
                          ls='', ms=10, label=label)
        t = np.linspace(0, self.tmax)
        r = np.zeros_like(t) + self.refstd
        self.ax.plot(t, r, 'k--', label='_')

        # Collect sample points for latter use (e.g. legend)
        self.samplePoints = [l]

    def add_sample(self, stddev, corrcoef, *args, **kwargs):
        """
        Add sample (*stddev*, *corrcoeff*) to the Taylor
        diagram. *args* and *kwargs* are directly propagated to the
        `Figure.plot` command.
        """

        l, = self.ax.plot(np.arccos(corrcoef), stddev,
                          *args, **kwargs)  # (theta, radius)
        self.samplePoints.append(l)

        return l

    def add_grid(self, *args, **kwargs):
        """Add a grid."""

        self._ax.grid(*args, **kwargs)


    def add_contours(self, levels=5, **kwargs):
        """
        Add constant centered RMS difference contours, defined by *levels*.
        """

        rs, ts = np.meshgrid(np.linspace(self.smin, self.smax),
                             np.linspace(0, self.tmax))
        # Compute centered RMS difference
        rms = np.sqrt(self.refstd**2 + rs**2 - 2*self.refstd*rs*np.cos(ts))

        contours = self.ax.contour(ts, rs, rms, levels, **kwargs)

        return contours
