import iris
import matplotlib.pyplot as plt
import iris.plot as iplt
import json
import glob
import os
import numpy as np
import cartopy.crs as ccrs
from matplotlib.gridspec import GridSpec
import iris.plot as iplt
from PIL import Image
import dask.array as da
from iris.coords import DimCoord, AuxCoord
from iris.cube import Cube
import xarray as xr
from numpy import genfromtxt

from esmvalcore.preprocessor import area_statistics, extract_point, extract_region

"""
scripts:
    -  /esarchive/scratch/jcos/es-esmvaltool/common-tools/projections/diagnostics/scripts/plot_pattern_correlation_maps.py
        - plot_2methods_regional_timeseries -> S6 & 7, remove RPSS boes and so on
    - /esarchive/scratch/jcos/es-esmvaltool/common-tools/projections/diagnostics/shared_constraints_diagnostics/member_selection_heatmap.py
        - heatmap S2 (inplace) add redbox, subtract first decade
    - /esarchive/scratch/jcos/es-esmvaltool/common-tools/projections/diagnostics/scripts/plot_relres.py
        - S4 & 5 (inplace?)
    S3 (done)
    S1 probably done in inkscape

"""


member_alias = iris.load_cube("/esarchive/scratch/jcos/output/recipe_decadal_constrains_20220718_053727/work/constr_s1968/s1968/PC/1-9y/members_field_Glob_1968.nc").coord("member_alias").points
canesm_ind = [i for i, alias in enumerate(member_alias) if "CanESM" in alias]
series = np.arange(1961,2001)
period = [1961,2000]
indices = [i for i in np.where(series>=period[0])[0] if i in np.where(series<=period[1])[0]]

rpss = {"Eastern":
        {"1961-2000": {"OBS": {"Glob": [-0.08,0,-0.50,0.07], "NAtl": [-0.15,1,-0.46,0.06]},
                       "DP": {"Glob": [0.14,1,-0.59,0.06], "NAtl": [0.04,0,-0.49,0.05]},
                       "DPts": {"Glob": [0.92,1,-0.56,0.05], "NAtl": [0.93,1,-0.45,0.06], "SPG": [0.01,0,-0.44,0.04]},
                       "ERA5": [-0.56,0.05], "ALL": [-0.44,0.05]}, 
        "1970-2000": {"OBS": {"Glob": [0.20,1,-0.3,0.22], "NAtl": [0.08,1,-0.3,0.22]},
                      "DP": {"Glob": [0.21,1,-0.36,0.21], "NAtl": [0.05,0,-0.32,0.20]},
                      "DPts": {"Glob": [-0.12,0,-0.25,0.18], "NAtl": [-0.03,0,-0.30,0.19], "SPG": [0.06,1,-0.27,0.18]},
                      "ERA5": [-0.36,0.21], "ALL": [-0.28,0.18]}},
        "Western":
        {"1961-2000": {"OBS": {"Glob": [0.04,1,-0.46,0.05], "NAtl": [-0.14,0,-0.41,0.06]},
                       "DP": {"Glob": [0.4,1,-0.59,0.04], "NAtl": [0.06,1,-0.49,0.04]},
                       "DPts": {"Glob": [0.15,1,-0.39,0.05], "NAtl": [0.04,0,-0.42,0.05], "SPG": [0.01,0,-0.41,0.04]},
                       "ERA5": [-0.73,0.09], "ALL": [-0.41,0.05]}, 
        "1970-2000": {"OBS": {"Glob": [-0.04,0,-0.28,0.2], "NAtl": [-0.08,0,-0.27,0.2]},
                      "DP": {"Glob": [0.26,1,-0.34,0.17], "NAtl": [0.15,1,-0.30,0.17]},
                      "DPts": {"Glob": [-0.18,0,-0.25,0.16], "NAtl": [0.31,1,-0.27,0.15], "SPG": [0.07,0,-0.26,0.17]},
                      "ERA5": [-0.34,0.19], "ALL": [-0.25,0.16]}}}



def plot_regional_timeseries(path, region, title, cwip=False):
    """compares the OBS, BEST, ALL forecasts for the 1-20 timeseries in an input area avgd region"""
    "tas_fcst_cube_BEST_Glob_1-20.nc"
    "tas_members_fcst_cube_BEST_Glob_1-20.nc"
    selection_region = title.split(" ")[-2]
    method = title.split(" ")[-3]
    subregion = title.split(" ")[0]
    if selection_region == "SPG":
        fcst_files = glob.glob(os.path.join(path, f"tas_fcst_cube_*_Gyre_1-20.nc"))+glob.glob(os.path.join(path, f"tas_fcst_cube_ALL_1-20.nc")) # fcst, members º1-20 BEST&ALL&OBS
        members_files = glob.glob(os.path.join(path, f"tas_members_fcst_cube_*_Gyre_1-20.nc"))+glob.glob(os.path.join(path, f"tas_members_fcst_cube_ALL_1-20.nc"))
    else:
        fcst_files = glob.glob(os.path.join(path, f"tas_fcst_cube_*_{selection_region}_1-20.nc"))+glob.glob(os.path.join(path, f"tas_fcst_cube_ALL_1-20.nc")) # fcst, members º1-20 BEST&ALL&OBS
        members_files = glob.glob(os.path.join(path, f"tas_members_fcst_cube_*_{selection_region}_1-20.nc"))+glob.glob(os.path.join(path, f"tas_members_fcst_cube_ALL_1-20.nc"))
    if "Befort" in path:
        obs_file = os.path.join(path, f"tas_obs_cube_1-20.nc")
    else:
        obs_file = os.path.join(path, f"tas_obs_cube_BEST_{selection_region}_1-20.nc")
    fcst_cubes = {f.split("/")[-1].split("_")[3]: reg_postpro(f, region) for f in fcst_files}
    # fcst_cubes.update(cwip_fcst_cubes)
    members_cubes = {f.split("/")[-1].split("_")[4]: reg_postpro(f, region) for f in members_files}
    # members_cubes.update(cwip_members_cubes)
    obs_cube = reg_postpro(obs_file, region)
    colors={"BEST": "red","BEST35": "red", "ALL": "blue", "weighted": "yellow", "ERA5": "black", "OBS": "red", "DP": "red", "DPts": "red"}
    fig, ax = plt.subplots(dpi=300)
    plt.plot(series[indices],obs_cube.data, color="black",label="OBS",zorder=101)
    min, max = [0, 0]
    for key, fcst in fcst_cubes.items():
        if key == "weighted" and not cwip:
            continue
        plt.plot(series[indices],fcst.data, color=colors[key], label=key, zorder=101)
    for key, mem in members_cubes.items():
        if "BEST" in key:
            plt.scatter(np.repeat(series[indices], len(mem.coord("member").points), axis=0),mem.data.flatten(),marker='.',s=4,alpha=0.5, zorder=100, color="red")
        else:
            for i in np.arange(mem.data.shape[1]):
                plt.plot(series[indices],mem[:,i].data, color="cornflowerblue", alpha=0.1)
            min=np.min(mem.data)
            max=np.max(mem.data)
    if rpss[subregion]["1961-2000"][method][selection_region][1] == 1:
        weight = "bold"
    else:
        weight = "normal"
    ax.text(1980, 0.85, f"1961-2000 RPSS: {rpss[subregion]['1961-2000'][method][selection_region][0]:.2f}", verticalalignment="center", horizontalalignment="center", weight=weight)
    if rpss[subregion]["1970-2000"][method][selection_region][1] == 1:
        weight = "bold"
    else:
        weight = "normal"
    ax.text(1980, 1, f"1970-2000 RPSS: {rpss[subregion]['1970-2000'][method][selection_region][0]:.2f}", verticalalignment="center", horizontalalignment="center", weight=weight)
    for ens in [method, "ALL", "ERA5"]:
        for side, period in enumerate(["1961-2000", "1970-2000"]):
            level = rpss[subregion][period][ens]
            if isinstance(level, dict):
                level = level[selection_region]
            plt.hlines(level[-1],series[indices][side*-1]-1.5,series[indices][side*-1]+1.5, color=colors[ens], alpha=0.7)
            plt.hlines(level[-2],series[indices][side*-1]-1.5,series[indices][side*-1]+1.5, color=colors[ens], alpha=0.7)
    plt.hlines(0,series[indices][0],series[indices][-1], color="grey", alpha=0.4)
    plt.fill_between([1961,1970],[min,min], [max,max], color="grey", alpha=0.2)
    plt.ylim(min,max)
    plt.xlim(series[indices][0],series[indices][-1])
    plt.ylabel("(K)")
    plt.legend()
    if len(region)==4:
        save =  f"{region[0]}_{region[1]}_{region[2]}_{region[3]}"
    else:
        save = f"{region[0]}_{region[1]}"
        
    format = {"DP": "DPfield-selection", "OBS": "OBSfield-selection", "DPts": "DPtime-selection"}
    title_sp=title.split(" ")
    title = f"{title_sp[0]} {title_sp[1]} {title_sp[-2]} {format[title_sp[2]]} {title_sp[-1]}"
    plt.title(title)
    plt.savefig(os.path.join("/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/plots/gyre",f"{method}_fcst_timeseries_{selection_region}_{save}.png"),bbox_inches='tight')
    plt.savefig(os.path.join("/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/plots/gyre",f"{method}_fcst_timeseries_{selection_region}_{save}.pdf"), format="pdf", bbox_inches='tight')
    


# def plot_2methods_regional_timeseries(path, region, title):
#     """compares the OBS, BEST, ALL forecasts for the 1-20 timeseries in an input area avgd region"""
#     "tas_fcst_cube_BEST_Glob_1-20.nc"
#     "tas_members_fcst_cube_BEST_Glob_1-20.nc"
#     selection_region = title.split(" ")[-2]
#     method = title.split(" ")[-3]    
#     obs_file = os.path.join(path, f"tas_obs_cube_BEST_{selection_region}_1-20.nc")

#     obs_cube = reg_postpro(obs_file, region)
#     colors={"BEST": "red", "ALL": "blue", "weighted": "yellow"}
#     fig, ax = plt.subplots()
#     plt.plot(series[indices],obs_cube.data, color="black",label="OBS",zorder=101)
#     min, max = [0, 0]
#     for key, fcst in fcst_cubes.items():
#         if key == "weighted":
#             continue
#         plt.plot(series[indices],fcst.data, color=colors[key], label=key, zorder=101)
#     for path in paths:
#         fcst_files = glob.glob(os.path.join(path, f"tas_fcst_cube_*_{selection_region}_1-20.nc")) # fcst, members º1-20 BEST&ALL&OBS
#         members_files = glob.glob(os.path.join(path, f"tas_members_fcst_cube_*_{selection_region}_1-20.nc"))
#         fcst_cubes = {f.split("_")[-3]: reg_postpro(f, region) for f in fcst_files}
#         members_cubes = {f.split("_")[-3]: reg_postpro(f, region) for f in members_files}
#         for key, mem in members_cubes.items():
#             if key == "BEST":
#                 plt.scatter(np.repeat(series[indices], len(mem.coord("member").points), axis=0),mem.data.flatten(),marker='.',s=4,alpha=0.5, zorder=100, color="red")
#             else:
#                 for i in np.arange(mem.data.shape[1]):
#                     plt.plot(series[indices],mem[:,i].data, color="purple", alpha=0.2)
#                 min=np.min(mem.data)
#                 max=np.max(mem.data)
#     plt.fill_between([1961,1970],[min,min], [max,max], color="grey", alpha=0.2)
#     plt.ylim(min,max)
#     plt.xlim(series[indices][0],series[indices][-1])
#     plt.legend()
#     if len(region)==4:
#         save =  f"{region[0]}_{region[1]}_{region[2]}_{region[3]}"
#     else:
#         save = f"{region[0]}_{region[1]}"

    
#     plt.title(title)
#     plt.savefig(os.path.join(path.replace("/work/", "/plots/"),f"{method}_fcst_timeseries_{selection_region}_{save}.png"),bbox_inches='tight')
    

def reg_postpro(f, region):
    if len(region) == 4:
        if isinstance(f,str):
            return area_statistics(extract_region(iris.load_cube(f), region[0], region[1], region[2], region[3]), "mean")
        else:
            return area_statistics(extract_region(f, region[0], region[1], region[2], region[3]), "mean")
        # return area_statistics(extract_region(iris.load_cube(f)[9:], region[0], region[1], region[2], region[3]), "mean")
    elif len(region) == 2:
        if isinstance(f,str):
            return extract_point(iris.load_cube(f), region[0], region[1], "nearest")
        else:
            return extract_point(f, region[0], region[1], "nearest")
        # return extract_point(iris.load_cube(f)[9:], region[0], region[1], "nearest")


def main():
    # path="/esarchive/scratch/jcos/output/recipe_light_observational_constrains_20220621_095313/work/tas_skill/collect/jja_PC/1-9y"
    # plot_regional_timeseries(path, [15,40,30,45], "Eastern Mediterranean OBS Glob 1-20") # Med
    # plot_regional_timeseries(path, [-10,15,30,45], "Western Mediterranean OBS Glob 1-20") # Med
    # plot_regional_timeseries(path, [15,40,30,45], "Eastern Mediterranean OBS NAtl 1-20") # Med
    # plot_regional_timeseries(path, [-10,15,30,45], "Western Mediterranean OBS NAtl 1-20") # Med
    # path="/esarchive/scratch/jcos/output/recipe_decadal_constrains_20220718_053727/work/tas_skill/collect/jja_PC/1-9y/"
    # plot_regional_timeseries(path, [15,40,30,45], "Eastern Mediterranean DP Glob 1-20") # Med
    # plot_regional_timeseries(path, [-10,15,30,45], "Western Mediterranean DP Glob 1-20") # Med
    # plot_regional_timeseries(path, [15,40,30,45], "Eastern Mediterranean DP NAtl 1-20") # Med
    # plot_regional_timeseries(path, [-10,15,30,45], "Western Mediterranean DP NAtl 1-20") # Med
    path="/esarchive/scratch/jcos/output/recipe_decadal_constrains_20220718_053727/work/tas_skill/collect/jja_Befort_eval-1961-2000/1-9y/"
    # plot_regional_timeseries(path, [15,40,30,45], "Eastern Mediterranean DPts Glob 1-20") # Med
    # plot_regional_timeseries(path, [-10,15,30,45], "Western Mediterranean DPts Glob 1-20") # Med
    plot_regional_timeseries(path, [15,40,30,45], "Eastern Mediterranean DPts SPG 1-20") # Med
    plot_regional_timeseries(path, [-10,15,30,45], "Western Mediterranean DPts SPG 1-20") # Med
    # plot_regional_timeseries(path, [15,40,30,45], "Eastern Mediterranean DPts NAtl 1-20") # Med
    # plot_regional_timeseries(path, [-10,15,30,45], "Western Mediterranean DPts NAtl 1-20") # Med

if __name__ == "__main__":
    main()