import iris

"""
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

"""

import matplotlib.pyplot as plt
import glob
import os
import warnings
import yaml
import json
import logging
import numpy as np
from dask import array as da
import esmvaltool.diag_scripts.shared as e
from esmvaltool.diag_scripts.shared import names, extract_variables, io
from datetime import datetime
import sys
import matplotlib.patches as patches
import matplotlib.colors as mcolors
import glob

current = os.path.dirname(os.path.realpath(__file__))
parent_directory = os.path.dirname(current)
sys.path.append(parent_directory)

logger = logging.getLogger(__name__)
FORMAT = "%(asctime)-15s %(message)s"
logging.basicConfig(format=FORMAT)


startTime = datetime.now()
warnings.filterwarnings("ignore")
config_loops_file = "/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/config_loops.yml"
with open(config_loops_file, "r") as f:
    config_loops = yaml.full_load(f)


class IntermediaryPlots(object):
    def __init__(self, config, crit, cstr_yavg, crit_collect, i, N):
        self.cfg = config
        self.cfg["earl_anomaly"] = False
        self.cfg["selection_criteria"] = crit
        self.cfg["collect_criteria"] = crit_collect
        self.cfg["cstr_yavg"] = cstr_yavg
        self.cfg["subsample_size"] = N
        if i != 0:
            keys = list(self.cfg["regions"].keys())
            [self.cfg["regions"].pop(key, None) for key in keys if "corr" not in key]

    
    def main(self):
        self.plot_heatmap()
        # self.plot_correlation_timeseries()

    def plot_heatmap(self, ):
        reverse = True if "RMSE" not in self.cfg["selection_criteria"] and "Befort" not in self.cfg["selection_criteria"] else False
        if "observation" in self.cfg["recipe"]:
            method = "OBS"
            subfolders="tos*"
        elif "decadal" in self.cfg["recipe"]:
            method = "DP"
            subfolders = "constr_s*/s*"

        highECS = ["CESM2","CanESM5","CanESM5-CanOE","E3SM-1-0","HadGEM3-GC31-LL","HadGEM3-GC31-MM","UKESM1-0-LL"] # over 5ºC ECS 
        midECS = ["NESM3","IPSL-CM6A-LR","CNRM-CM6-1","CNRM-CM6-1-HR","CESM2-WACCM","ACCESS-CM2", "CNRM-ESM2-1", "EC-Earth3"] # https://www.ipcc.ch/report/ar6/wg1/figures/chapter-1/figure-1-16
        colors = ['white', 'orange', 'red', 'black']
        bounds = [0, 0.5, 0.75, 0.99, 1]
        work_dir = {"DP": "/esarchive/scratch/jcos/output/recipe_decadal_constrains_20220718_053727/work", 
                    "OBS": "/esarchive/scratch/jcos/output/recipe_light_observational_constrains_20220621_095313/work"}

        # Create a colormap using ListedColormap
        cmap = mcolors.ListedColormap(colors)
        norm = mcolors.BoundaryNorm(bounds, cmap.N)

        cmap_text = plt.cm.RdBu_r
        # Define the normalization range
        norm_text = mcolors.Normalize(vmin=-0.3, vmax=0.7)


        for region in self.cfg["regions"]:
            # if region != "NAtl":
            if region != "Gyre" and region != "Glob":
            # if region != "Gyre":
            # if region != "Glob" and region != "NAtl":
                continue
            print(region)
            if "corr" in region:
                region = self.cfg["regions"][region].replace("CSTR-YRS",f"{self.cfg['cstr_yavg']}").replace("CRIT-CLCT",f"{self.cfg['collect_criteria']}")
                region = region.split("OBS_")[-1].split(".nc")[0]
            # m_files = io.get_all_ancestor_files(self.cfg, f"{region}_*.json")
            m_files = glob.glob(f"{work_dir[method]}/{subfolders}/{self.cfg['selection_criteria']}/{self.cfg['cstr_yavg']}/{region}_*.json")
            m_files.sort()
            if len(m_files) == 0:
                continue
            formatter = {"RMSE": "RMSE", "PC": "Uncentered pat. corr.", "PCiris": "Centered pat. corr.", "Befort": f"{region} timeseries MAE"}
            if region == "Gyre":
                region = "SPG"
            N=self.cfg["subsample_size"]
            corrs_array = []
            sds = []
            hECS_ratio_list = []
            submethod = "field"
            if "Befort" in self.cfg["selection_criteria"]:
                submethod = "time"
            for i,m_sd in enumerate(m_files):
                logger.info(f"loading json correlation file: {m_sd}")
                with open(m_sd, "r") as f:
                    corr = json.load(f)
                    corrs = list(corr.values())
                corr = self.purge_dataset("CESM2", corr)
                if i == 0:
                    selection = np.zeros((len(m_files),len(list(corr.keys()))))
                    members = list(corr.keys())
                    members = sorted([m for m in corr if not m.startswith("OBS")])
                    hECS_ratio = len([m for m in members if m.split("_")[0] in highECS or m.split("_")[0] in midECS])/len(members)
                selected = [k for k, v in sorted(corr.items(), key=lambda item: item[1], reverse=reverse)][:N]
                seldict = {m: 0 for m in members}
                for d in seldict:
                    if (d in selected) and d.split("_")[0] in highECS:
                        seldict[d] = 0.80
                    if (d in selected) and d.split("_")[0] in midECS:
                        seldict[d] = 0.60
                    if (d in selected) and (d.split("_")[0] not in midECS) and (d.split("_")[0] not in highECS):
                        seldict[d] = 0.995
                selection[i] = list(seldict.values())
                hECS_ratio_list.append(len([m for m in selected if m.split("_")[0] in highECS or m.split("_")[0] in midECS])/len(selected)-hECS_ratio)
                corrs.sort(reverse=reverse)
                corrs_array.append(corrs[:N])
                sds.append(int(m_sd.split(".json")[0].split("_")[-1]))

            corrs_array = np.array(corrs_array)[9:]
            canesm = [i for i, m in enumerate(members) if "CanESM5" in m]
            members = [m.replace("_", " ") for m in members]
            dates = np.arange(1961, 2001)
            fig, ax = plt.subplots(figsize=(12,28/210*len(members)), dpi=150)
            fill = plt.pcolormesh(dates[9:], members, selection.transpose()[:,9:], cmap=cmap, norm=norm, edgecolors="black")
            ymin, ymax = plt.ylim()
            # plt.vlines(1969.5, ymin, ymax, color = "red", linewidth=2)
            # plt.colorbar(fill, location="right", orientation="vertical", shrink=0.3, cmap=plt.cm.Greys, label="Weight")
            plt.title(f"Selection of the best {self.cfg['subsample_size']} CMIP6 members; method:"
                       f"{region.split('_signi')[0]} {method}{submethod}-sel."
                       , fontsize="x-large", pad=40)
            plt.tight_layout()
            minimal_region = region
            for i, (year, percentage) in enumerate(zip(dates[9:], hECS_ratio_list[9:])):
                # color = plt.cm.RdBu_r((percentage - 0) / 0.7)  # Normalizing between -0.3 and 0.4 for RdBu_r colormap
                color = cmap_text(norm_text(percentage)) # Normalizing between -0.3 and 0.4 for RdBu_r colormap
                plt.text(year, len(members) + 0.42, f"{percentage:.1f}", ha='center', color=color, backgroundcolor="black", rotation=90, fontsize="large")
            plt.text(1966.65, len(members) + 0.36, f"high ECS members\nover-representation", ha='center', color="white", backgroundcolor="black", fontsize="large")
            if "corr" in minimal_region:
                minimal_region=minimal_region.partition(".")[-1].rpartition("_signi_corr-mask")[0]
            if "decadal" in self.cfg["recipe"]:
                method = "DP"
                if minimal_region == "Glob" and "Befort" not in self.cfg['selection_criteria']:
                    rect = patches.Rectangle((1969.6, canesm[0]+ymin), 30.8, canesm[-1]-canesm[0]-ymin*2, linewidth=3, edgecolor='blue', facecolor='none')
                    ax.add_patch(rect)
                ax.text(-0.22, 1.022, "A)", transform=ax.transAxes, fontsize="xx-large")
            else:
                method = "OBS"
                ax.text(-0.22, 1.022, "B)", transform=ax.transAxes, fontsize="xx-large")
            
            post=""
            if region == "SPG":
                post = "/gyre"
            plt.savefig(os.path.join(f"/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/plots{post}", f"{method}_selection_{self.cfg['selection_criteria']}_heatmap_best{self.cfg['subsample_size']}_{minimal_region}_new_regrid.png"))
            plt.savefig(os.path.join(f"/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/plots{post}", f"{method}_selection_{self.cfg['selection_criteria']}_heatmap_best{self.cfg['subsample_size']}_{minimal_region}_new_regrid.pdf"), format="pdf", bbox_inches='tight')

            fig, ax = plt.subplots(figsize=(12,3), dpi=150)
            plt.plot(dates[9:], np.median(corrs_array, axis=1), color="black", label="median")
            plt.plot(dates[9:], np.max(corrs_array, axis=1), linestyle="--", color="black", label="max/min")
            plt.plot(dates[9:], np.min(corrs_array, axis=1), linestyle="--", color="black")
            plt.title(f"Constraining metric values for the best {N} members; method: {region} {method}{submethod}-selection", fontsize="large")
            if submethod == "time":
                plt.ylabel("mean absolute error / K")
                plt.ylim(0,0.25)
            else:
                plt.ylabel("uncentered pearson correlation / 1")
                plt.ylim(0,1)
            plt.legend(fontsize="large")
            plt.savefig(os.path.join(f"/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/plots{post}", f"{method}_selection_{self.cfg['selection_criteria']}_value-series_best{self.cfg['subsample_size']}_{region}_new_regrid.png"), bbox_inches='tight')
            plt.savefig(os.path.join(f"/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/plots{post}", f"{method}_selection_{self.cfg['selection_criteria']}_value-series_best{self.cfg['subsample_size']}_{region}_new_regrid.pdf"), format="pdf", bbox_inches='tight')


    def purge_dataset(self, datasets, dictionary):
        """
        Purge a list of datasets from the dictionary
        datasets: List of strings
        dictionary: Keys are datasets or aliases
        """
        if not isinstance(datasets, list):
            datasets = [datasets]
        keys_list = list(dictionary.keys())
        purgeable_keys = []
        for dataset in datasets:
            full_name_dataset_keys=[k for k in keys_list if dataset == k]
            part_name_dataset_keys=[k for k in keys_list if dataset in k.split("_")]
            dataset_keys=full_name_dataset_keys+part_name_dataset_keys
            purgeable_keys += dataset_keys
        return {k: v for k, v in dictionary.items() if k not in purgeable_keys}
    
def main():
    print("----------- MAIN ----------")
    for crit in config_loops["criteria_cstr"]:
        for cstr_yavg in config_loops["year_averages"]:
            for N in config_loops["subsample_size"]:
                for i, crit_collect in enumerate(config_loops["criteria_collect"]):
                    # with e.run_diagnostic() as config:
                    config = {"regions": {"Gyre":[-50,-20,45,60],"Glob":[-180,180,-90,90]}}
                    # config = {"regions": {"Glob":[-180,180,-90,90], "NAtl":[-80,0,0,60]},}
                    config["recipe"] = config_loops["recipe"]
                    IntermediaryPlots(config, crit, cstr_yavg, crit_collect, i, N).main()


if __name__ == "__main__":
    main()
    print(f" TIME SPENT: {datetime.now() - startTime}")
