import iris
from esmvalcore.preprocessor import climate_statistics, extract_time, extract_region
import glob
import matplotlib.pyplot as plt
import iris.analysis.stats
import iris.analysis.cartography
from matplotlib.gridspec import GridSpec 
import iris.plot as iplt
import cartopy.crs as ccrs
import cartopy
import json

file_paths = glob.glob(f"/esarchive/scratch/jcos/output/recipe_light_observational_constrains_20220621_095313/preproc/obs_constrains/tos/*.nc")
file_paths.sort()
ref_path = file_paths[-1]
ref_cube = extract_time(iris.load_cube(ref_path), 1961, 1, 1, 1991, 12, 31)
ref_std_clim = climate_statistics(ref_cube, "std_dev")

weights = iris.analysis.cartography.area_weights(ref_std_clim)
N = len(file_paths)-1

# corr_dict = {}
# cube_dict = {}
# for file in file_paths:
#     if file == ref_path:
#         continue
#     code_alias=file.split("/")[-1].split("_") 
#     dataset = f"{code_alias[1]}_{code_alias[-3]}"
#     cube = extract_time(iris.load_cube(file), 1961, 1, 1, 1991, 12, 31)
#     std_clim = climate_statistics(cube, "std_dev")
#     cube_dict[dataset] = std_clim
#     corr = iris.analysis.stats.pearsonr(std_clim, ref_std_clim, weights=weights, common_mask=True)
#     corr_dict[dataset] = corr.data
# corr_dict_ann = {k:v for k,v in sorted(corr_dict.items(), key=lambda item: item[1], reverse=True)}


# fig = plt.figure(figsize=(5, 100/3*6), dpi=300)
# gs = GridSpec(N+1, 1, figure=fig)
# ax = fig.add_subplot(gs[0,0], projection=ccrs.PlateCarree())
# iplt.pcolormesh(ref_std_clim, cmap=plt.cm.Reds, vmin=0, vmax=1)
# ax.coastlines()
# plt.title(f"ERA5 STD 1961 1991\n1y rolling mean")
# for i, (label, value) in enumerate(corr_dict_ann.items(),1):
#     ax = fig.add_subplot(gs[i,0], projection=ccrs.PlateCarree())
#     iplt.pcolormesh(cube_dict[label], cmap=plt.cm.Reds, vmin=0, vmax=1)
#     ax.coastlines()
#     plt.title(f"{label}\ncorr: {value:.03f}")
# plt.tight_layout()
# plt.savefig(f"/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/plots/1yrolling_mean_SST_STD.png")

regions = {"NAtl": [0,70,-80,0], "Glob": [-180, 180, -90, 90]}
region = "NAtl"

ref_cube = extract_region(ref_cube, regions[region][0], regions[region][1], regions[region][2], regions[region][3]).rolling_window("time", iris.analysis.MEAN, 9)
ref_std_clim = climate_statistics(ref_cube, "std_dev")
weights = iris.analysis.cartography.area_weights(ref_std_clim)
corr_dict = {}
cube_dict = {}
for i, file in enumerate(file_paths):
    if file == ref_path:
        continue
    code_alias=file.split("/")[-1].split("_") 
    dataset = f"{code_alias[1]}_{code_alias[-3]}"
    cube = extract_time(extract_region(iris.load_cube(file), regions[region][0], regions[region][1], regions[region][2], regions[region][3]), 1961, 1, 1, 1991, 12, 31).rolling_window("time", iris.analysis.MEAN, 9)
    std_clim = climate_statistics(cube, "std_dev")
    cube_dict[dataset] = std_clim
    corr = iris.analysis.stats.pearsonr(std_clim, ref_std_clim, weights=weights, common_mask=True)
    corr_dict[dataset] = corr.data
corr_dict_rolled = {k:v for k,v in sorted(corr_dict.items(), key=lambda item: item[1], reverse=True)}
with open(f"/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/plots/9yrolling_mean_SST_STD_{region}.png", "w") as f:
    json.dump(corr_dict_rolled, f)


# fig = plt.figure(figsize=(5, 100/3*6), dpi=300)
# gs = GridSpec(N+1, 1, figure=fig)
# ax = fig.add_subplot(gs[0,0], projection=ccrs.PlateCarree())
# iplt.pcolormesh(ref_std_clim, cmap=plt.cm.Reds, vmin=0, vmax=1)
# ax.coastlines()
# plt.title(f"ERA5 STD 1961 1991\n9y rolling mean")
# for i, (label, value) in enumerate(corr_dict_rolled.items(),1):
#     ax = fig.add_subplot(gs[i,0], projection=ccrs.PlateCarree())
#     iplt.pcolormesh(cube_dict[label], cmap=plt.cm.Reds, vmin=0, vmax=1)
#     ax.coastlines()
#     plt.title(f"{label}\ncorr: {value:.03f}")
# plt.tight_layout()
# plt.savefig(f"/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/plots/9yrolling_mean_SST_STD_{region}.png")
