import xarray as xr
import numpy as np
import pandas as pd
import geopandas as gpd
import calendar
import matplotlib.pyplot as plt
from itertools import groupby

temperatureERA5Path='/esarchive/recon/ecmwf/era5/monthly_mean/tas_f1h-r1440x721cds/'
dewPointERA5Path='/esarchive/recon/ecmwf/era5/monthly_mean/tdps_f1h-r1440x721cds/'
precipitationERA5Path='/esarchive/recon/ecmwf/era5/monthly_mean/prlr_f1h-r1440x721cds/'

anophelesCategories=['Crops','SeasonalWater','PermanentWater','Tree']
ixodesCategories=['Crops','Shrub','Tree','MossLichen']
hyalommaCategories=['Crops','Shrub','Tree','MossLichen','Bare']

def get_t_hr_data_ready_single_month(date):
    # Load datasets using Xarray
    temperature = (xr.load_dataset(temperatureERA5Path + f'tas_{date}.nc')
                   .drop_vars(['bnds', 'time_bnds', 'height'], errors='ignore')
                   .isel(time=0).drop_vars('time'))

    dewPoint = (xr.load_dataset(dewPointERA5Path + f'tdps_{date}.nc')
                .drop_vars(['bnds', 'time_bnds', 'height'], errors='ignore')
                .isel(time=0).drop_vars('time'))

    temperature['tas'] -= 273.15
    dewPoint['tdps'] -= 273.15
    merged = xr.merge([temperature, dewPoint])
    merged['rh'] = 100 * (np.exp((17.625 * merged['tdps']) / (243.04 + merged['tdps'])) /
                           np.exp((17.625 * merged['tas']) / (243.04 + merged['tas'])))
    df = merged[['tas', 'rh']].to_dataframe().reset_index()
    df['lon'] = np.where(df['lon'] > 180, df['lon'] - 360, df['lon'])
    dateExample=pd.to_datetime(f"{date}15")
    df['date'] = dateExample
    df['geometry'] = gpd.points_from_xy(df['lon'] + 0.125, df['lat'] - 0.125)
    gdf = gpd.GeoDataFrame(df, geometry='geometry', crs='EPSG:4326')
    
    world=gpd.read_file('NUTS_RG_20M_2024_4326.shp').loc[lambda dd: dd.LEVL_CODE==3]
    gdf = gdf.sjoin(world[['geometry']], predicate="within").drop(columns=['index_right'])
    gdf['dayLength'] = daylength1(gdf['lat'], dateExample.timetuple().tm_yday)  

    return gdf

def get_t_hr_tp_data_ready_single_month(date):
# Load datasets using Xarray
    temperature = (xr.load_dataset(temperatureERA5Path + f'tas_{date}.nc')
                   .drop_vars(['bnds', 'time_bnds', 'height'], errors='ignore')
                   .isel(time=0).drop_vars('time'))

    dewPoint = (xr.load_dataset(dewPointERA5Path + f'tdps_{date}.nc')
                .drop_vars(['bnds', 'time_bnds', 'height'], errors='ignore')
                .isel(time=0).drop_vars('time'))
    
    precipitation = (xr.load_dataset(precipitationERA5Path + f'prlr_{date}.nc')
                .drop_vars(['bnds', 'time_bnds', 'height'], errors='ignore')
                .isel(time=0).drop_vars('time'))
    
    temperature['tas'] -= 273.15
    dewPoint['tdps'] -= 273.15
    precipitation['prlr'] = precipitation['prlr']*3600 * 24 * 30.44 * 1000
    merged = xr.merge([temperature, dewPoint, precipitation])
    merged['rh'] = 100 * (np.exp((17.625 * merged['tdps']) / (243.04 + merged['tdps'])) /
                           np.exp((17.625 * merged['tas']) / (243.04 + merged['tas'])))
    df = merged[['tas', 'rh', 'prlr']].to_dataframe().reset_index()
    df['lon'] = np.where(df['lon'] > 180, df['lon'] - 360, df['lon'])
    dateExample=pd.to_datetime(f"{date}15")
    df['date'] = dateExample
    df['geometry'] = gpd.points_from_xy(df['lon'] + 0.125, df['lat'] - 0.125)
    gdf = gpd.GeoDataFrame(df, geometry='geometry', crs='EPSG:4326')
    
    world=gpd.read_file('NUTS_RG_20M_2024_4326.shp').loc[lambda dd: dd.LEVL_CODE==3]
    gdf = gdf.sjoin(world[['geometry']], predicate="within").drop(columns=['index_right'])
    gdf['dayLength'] = daylength1(gdf['lat'], dateExample.timetuple().tm_yday)  

    return gdf


def daylength1(lat, dayOfYear):
    """Computes daylength for arrays of latitudes and days of the year using NumPy."""
    latInRad = np.deg2rad(lat)
    declinationOfEarth = 23.45 * np.sin(np.deg2rad(360.0 * (283.0 + dayOfYear) / 365.0))
    
    # Compute condition arrays
    tan_product = -np.tan(latInRad) * np.tan(np.deg2rad(declinationOfEarth))
    
    # Clamp values to valid range [-1, 1] to prevent NaNs in arccos
    tan_product = np.clip(tan_product, -1.0, 1.0)
    
    # Vectorized conditionals
    daylength = np.where(
        tan_product <= -1.0, 24.0, 
        np.where(tan_product >= 1.0, 0.0, 
                 2.0 * np.rad2deg(np.arccos(tan_product)) / 15.0)
    )
    
    return daylength


def transmission_lenght(df):
    groups=[]
    for k, g in groupby(df):
        if k==True:
            groups.append(len(list(g)))
    if groups:
        return(max(groups))
    else:
        return 0
    
class suitableMonthsThresholds:
    def __init__(self, species: str, year: int):
        
        self.year=year
        self.species=species
        
        if self.species=='Hyalomma':
            conditions=[16,39,10,79,12.01]
            landUseCategories=hyalommaCategories
        elif self.species=='Ixodes':
            conditions=[10,26,45,100,12.01]
            landUseCategories=ixodesCategories
        elif self.species=='Vivax':
            conditions=[14,32,60,100,80,11]
            landUseCategories=anophelesCategories
        elif self.species=='Falciparum':
            conditions=[16,33,60,100,80,11]
            landUseCategories=anophelesCategories    
        else:
            raise Exception("I only know about Hyalomma, Ixodes, Vivax and Falciparum :(") 
        
        self.categories=landUseCategories
        
        yearDF = pd.DataFrame()
        if species in ['Hyalomma', 'Ixodes']:
            for month in ['01','02','03','04','05','06','07','08','09','10','11','12']:
                df = get_t_hr_data_ready_single_month(str(self.year) + month)
                yearDF = pd.concat([yearDF, df], ignore_index=True)

            yearDF['suitable'] = False
            yearDF.loc[(yearDF.tas >= conditions[0]) & (yearDF.tas <= conditions[1]) & (yearDF.rh >= conditions[2]) & (yearDF.rh <= conditions[3]) & (yearDF.dayLength >= conditions[4]), 'suitable'] = True
        
        if species in ['Vivax','Falciparum']:
            for month in ['01','02','03','04','05','06','07','08','09','10','11','12']:
                df = get_t_hr_tp_data_ready_single_month(str(self.year) + month)
                yearDF = pd.concat([yearDF, df], ignore_index=True)

            yearDF['suitable'] = False
            yearDF.loc[(yearDF.tas >= conditions[0]) & (yearDF.tas <= conditions[1]) & (yearDF.rh >= conditions[2]) & (yearDF.rh <= conditions[3]) & (yearDF.prlr>=conditions[4]) & (yearDF.dayLength >= conditions[5]), 'suitable'] = True

        self.df = yearDF

    def plot_num_months_suitable_climate(self):
        fig, ax=plt.subplots(figsize=(16,8))
        
        (self.df
        .groupby(['lat','lon']).agg({'suitable': transmission_lenght}).reset_index()
        .reset_index()
        .assign(geometry= lambda dd: gpd.points_from_xy(dd.lon+0.25/2,dd.lat-0.25/2))
        .pipe(gpd.GeoDataFrame, crs='EPSG:4326')
        .assign(geometry= lambda dd: dd.geometry.buffer(0.25/2, cap_style='square'))
        ).plot('suitable', ax=ax, alpha=0.8, legend=True, vmax=6, vmin=0, legend_kwds={'shrink':0.6})
        world=gpd.read_file('NUTS_RG_20M_2024_4326.shp').loc[lambda dd: dd.LEVL_CODE==3]
        world.boundary.plot(ax=ax, color='black')
        ax.axis('off')
        ax.set_title(self.species + ' Lenght Tranmission season (months) ' + str(self.year))
        
    def num_months_suitable_climate_landuse(self):
        return (self.df.groupby(['lat','lon']).agg({'suitable': 'sum'}).reset_index()
        .reset_index()
        .merge(pd.read_csv('landuse_025.csv').drop(columns=['Unnamed: 0']).rename(columns={'x': 'lon', 'y':'lat'}), how='left', on=['lat','lon'])
         .assign(correctedSuitable= lambda dd: dd[self.categories].sum(axis=1)*dd.suitable*0.01)
        )[['lat','lon','correctedSuitable']]
        
    def plot_num_months_suitable_climate_landuse(self):
        fig, ax=plt.subplots(figsize=(16,8))
        
        (self.num_months_suitable_climate_landuse()
         .assign(geometry= lambda dd: gpd.points_from_xy(dd.lon+0.25/2,dd.lat-0.25/2))
        .pipe(gpd.GeoDataFrame, crs='EPSG:4326')
        .assign(geometry= lambda dd: dd.geometry.buffer(0.25/2, cap_style='square'))
        ).plot('correctedSuitable', ax=ax, alpha=0.8, legend=True, vmax=5, vmin=0, legend_kwds={'shrink':0.6})
        
        world=gpd.read_file('NUTS_RG_20M_2024_4326.shp').loc[lambda dd: dd.LEVL_CODE==3]
        world.boundary.plot(ax=ax, color='black')
        ax.axis('off')
        ax.set_title(self.species + ' Lenght Tranmission season (months) ' + str(self.year))
        

    def plot_single_month(self, month):
        fig, ax=plt.subplots(figsize=(16,8))
        
        (self.df
         .loc[lambda dd: (dd.date==pd.to_datetime(str(self.year)+'-'+str(month)+'-15'))&(dd.suitable==True)]
         .assign(geometry=lambda dd: gpd.points_from_xy(dd.lon+0.25/2,dd.lat-0.25/2))
         .pipe(gpd.GeoDataFrame, crs='EPSG:4326')
         .assign(geometry= lambda dd: dd.geometry.buffer(0.25/2, cap_style='square'))
         .plot(ax=ax, alpha=0.8)
        )
        world=gpd.read_file('NUTS_RG_20M_2024_4326.shp').loc[lambda dd: dd.LEVL_CODE==3]
        world.boundary.plot(ax=ax, color='black')
        ax.axis('off')
        ax.set_title(self.species+' climatic suitability on '+str(self.year)+'-'+str(month))