import pandas as pd
import geopandas as gpd
import matplotlib.pyplot as plt
import numpy as np
import rasterstats
import rioxarray as rxr
import rasterio
import adlfs
import contextily as ctx
import sys
sys.path.append('../sar_snowmelt_timing')
import s1_rtc_bs_utils
import fsspec
from pathlib import Path

import dask_gateway
cluster = dask_gateway.GatewayCluster()
client = cluster.get_client()
cluster.adapt(minimum=90, maximum=100)
print(client.dashboard_link)

f,ax=plt.subplots()
snotel_gf.plot(ax=ax)
ctx.add_basemap(ax=ax, crs=snotel_gf.crs, source=ctx.providers.Stamen.Terrain, attribution=False)

%%time
my_file = Path("snotel_gf.geojson")
if my_file.is_file():
    snotel_gf = gpd.read_file('snotel_gf.geojson')
else:
    for i,row in snotel_gf.iterrows():
        try:
            snotel_data = pd.read_parquet(f'https://github.com/scottyhq/snotel/raw/main/parquet/{row.code}.parquet').loc[slice(f'{years[0]}-02-01',f'{years[-1]}-09-30')]
            dates = snotel_data.groupby(snotel_data.index.year).idxmax()['WTEQ']
        #print(dates)
            if len(dates)==len(years):
                snotel_gf.loc[i,max_swe_years] = dates.values
        except:
            pass
        
        for year in years:
            snotel_gf[f'max_swe_{year}'] = pd.to_datetime(snotel_gf[f'max_swe_{year}']).dt.dayofyear
            
        snotel_gf.to_file("snotel_gf.geojson", driver='GeoJSON')

snotel_row_proj_buff

f,ax = plt.subplots(len(years),1,figsize=(20,20))

for i,year in enumerate(years):
    aoi = rxr.open_rasterio(f"https://snowmelt.blob.core.windows.net/snowmelt/eric/MGRS/{snotel_row.mgrs.values[0]}/runoff_onset_{snotel_row.mgrs.values[0]}_{year}_median_40m.tif")

    snotel_row_proj_buff.boundary.plot(ax=ax[i],color='red')
    aoi.where(aoi>0).plot(ax=ax[i],vmin=80,vmax=200)
    ctx.add_basemap(ax=ax[i], crs=snotel_row_proj_buff.crs, source=ctx.providers.Stamen.Terrain, attribution=False)

    ax[i].set_xlim([snotel_row_proj_buff.total_bounds[0]-buffer,snotel_row_proj_buff.total_bounds[2]+buffer])
    ax[i].set_ylim([snotel_row_proj_buff.total_bounds[1]-buffer,snotel_row_proj_buff.total_bounds[3]+buffer])
    ax[i].set_title(f'{year}')
    ax[i].axis('off')

%%time 
ts_ds_clip = ts_ds_extent.rio.clip(snotel_row_proj_buff.geometry).compute() # maybe.compute() here

for year in years:
    f,ax=plt.subplots(3,1,figsize=(11,8),sharex=True)
    
    f.suptitle(f'{code} SNOTEL Peak SWE with Backscatter Minima Comparison for {year}')
    colors1 = iter(['red','maroon','tomato','lightcoral','rosybrown'])
    colors2 = iter(['green','lime','yellowgreen','mediumspringgreen','aquamarine'])


    snotel_data_thisyear = snotel_data['WTEQ'].loc[slice(f'{year}-01-01',f'{year}-12-31')]
    snotel_data_thisyear.plot(ax=ax[0],color='blue',label=f'SWE [m]')
    
    swe_max = snotel_data_thisyear.loc[slice(f'{year}-02-01',f'{year}-07-31')].idxmax()
    
    ax[0].axvline(x=swe_max,linestyle='dashed',color='blue',label=f'SNOTEL max SWE')


    ts_ds_clip_year = ts_ds_clip.loc[slice(f'{year}-01-01',f'{year}-12-31')].dropna('time',how='all')

    orbits = s1_rtc_bs_utils.get_orbits_with_melt_season_coverage(ts_ds_clip_year,num_acquisitions_during_melt_season=1)
    print(orbits)
    ts_ds_clip_year = ts_ds_clip_year[ts_ds_clip_year['sat:relative_orbit'].isin(orbits)]
    
    estimates = []
    
    for orbit in np.unique(ts_ds_clip_year['sat:relative_orbit']):
        single_orbit = ts_ds_clip_year[ts_ds_clip_year['sat:relative_orbit']==orbit].median(dim=['x','y']) # change to mean or median?##########################################################################################################
        single_orbit_db = 10*np.log10(single_orbit)
        #single_orbit_db = single_orbit

        #color1 = next(colors1)
        #color2 = next(colors2) 

        asc_or_desc = single_orbit['sat:orbit_state'][0]

        if asc_or_desc == 'ascending':
            color = next(colors1)
        else:
            color = next(colors2)

        single_orbit_db.sel(band='vv').plot.line(ax=ax[1],label=f'Orbit {orbit} VV ({asc_or_desc})',linestyle='-',marker='o',color=color)
        single_orbit_db.sel(band='vh').plot.line(ax=ax[2],label=f'Orbit {orbit} VH ({asc_or_desc})',linestyle='-',marker='o',color=color)

        vv_estimate = single_orbit_db.sel(band='vv').loc[slice(f'{year}-02-01',f'{year}-07-31')].idxmin().values
        vh_estimate = single_orbit_db.sel(band='vh').loc[slice(f'{year}-02-01',f'{year}-07-31')].idxmin().values
        
        estimates.append(vv_estimate)
        estimates.append(vh_estimate)
        
        ax[1].axvline(x=vv_estimate,linestyle='dashed',color=color) #,ymin=-12,ymax=-4,
        ax[2].axvline(x=vh_estimate,linestyle='dashed',color=color) #,ymin=-17,ymax=-12 #using axvline gives seperate ticks
        
    median_estimate = pd.to_datetime(pd.Series(estimates)).median()
    ax[0].axvline(x=median_estimate,linestyle='dashed',color='black',label=f'SAR-derived median estimate of backscatter minima')
    
    ax[0].legend()
    ax[1].legend()
    ax[2].legend()

    ax[1].set_title('VV Backscatter Time series')
    ax[2].set_title('VH Backscatter Time series')
    
    print(f'For {year}, SNOTEL max SWE observed on {swe_max} and SAR derived backscatter minima of {median_estimate}')

        #peak_swe = pd.to_datetime(f'{year*1000}{snotel_row_proj_buff[f'max_swe_{year}']}', format = "%Y%j")

    plt.tight_layout()