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 iris.plot as iplt
import sys
from matplotlib.gridspec import GridSpec 
import matplotlib.patches as patches

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/es-esmvaltool/common-tools/projections/diagnostics/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
        self.setup_dirs()
        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 setup_dirs(self):
        criteria_label = self.cfg["selection_criteria"]
        year_averaging = self.cfg["cstr_yavg"]
        self.cfg["out_plot_dir"] = os.path.join(self.cfg[names.PLOT_DIR],f"{criteria_label}", year_averaging)
        if not os.path.isdir(self.cfg["out_plot_dir"]):
            os.makedirs(self.cfg["out_plot_dir"])
        self.cfg["out_work_dir"] = os.path.join(self.cfg[names.WORK_DIR],f"{criteria_label}", year_averaging)
        if not os.path.isdir(self.cfg["out_work_dir"]):
            os.makedirs(self.cfg["out_work_dir"])
        
        if "_eval-" in self.cfg["selection_criteria"]:
            if self.cfg["selection_criteria"].endswith("_X"):
                attributes = self.cfg["selection_criteria"].split("_")
                criteria_label = ""
                for attr in attributes:
                    if "eval-" not in attr:
                        criteria_label += f"{attr}_"
                criteria_label = criteria_label[:-1]
            else:
                criteria_label = self.cfg["selection_criteria"].split("_eval-")[0]

        new_idata = []
        for path in self.cfg["input_files"]:
            new_path = path
            if not path.endswith('metadata.yml'):
                new_path = os.path.join(path, criteria_label, year_averaging)
            new_idata.append(new_path)
        self.cfg["input_files"] = new_idata
        print(self.cfg["input_files"])

    
    def main(self):
        # self.plot_heatmap()
        self.plot_single_line_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"
        elif "decadal" in self.cfg["recipe"]:
            method = "DP"

        N=self.cfg["subsample_size"]
        for region in self.cfg["regions"]:
            if region != "Glob" and region != "NAtl":
                continue
            fig = plt.figure(figsize=(55, 60))
            dates = np.arange(1970, 2001)
            gs = GridSpec(N, len(dates), figure=fig)
            # if region != "NAtl":
            ref_files = io.get_all_ancestor_files(self.cfg, f"{region}_reference_timeseries.nc")
            ref_files.sort()
            ref_files = ref_files[9:]
            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.sort()
            print(f"files {region}:")
            print(m_files)
            if len(m_files) == 0:
                continue
            formatter = {"RMSE": "RMSE", "PC": "Uncentered pat. corr.", "PCiris": "Centered pat. corr.", "Befort": f"{region} timeseries MAE"}
            corrs_array = []
            sds = []
            for i,m_sd in enumerate(m_files[9:]):
                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)
                selected = [k for k, v in sorted(corr.items(), key=lambda item: item[1], reverse=reverse)][:N]
                ref_ts_nc = iris.load_cube(ref_files[i])
                for j,selected_alias_id in enumerate(selected):
                    alias_files = io.get_all_ancestor_files(self.cfg, f"{region}_{selected_alias_id}_timeseries.nc")
                    alias_files.sort()
                    file = alias_files[9:][i]
                    ts = iris.load_cube(file)
                    print(selected_alias_id)
                    print(ts)
                    ax = fig.add_subplot(gs[j, i])
                    years = [i.year for i in ts.coord("time").units.num2date(ts.coord("time").points)]
                    ax.plot(years,ts.data)
                    ax.plot(years, ref_ts_nc.data, color="black")
                    ax.set_title(f"{dates[i]} {selected_alias_id.split('_')[0]}\n{selected_alias_id.split('_')[1]}\nMAE: {corr[selected_alias_id]:.3f}")
                    
            plt.tight_layout()
            print("savefig:")
            print(os.path.join("/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/plots", f"{method}_{self.cfg['selection_criteria']}_timeseries_comparison_{region}.png"))
            plt.savefig(os.path.join("/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/plots", f"{method}_{self.cfg['selection_criteria']}_timeseries_comparison_{region}.png"))


    def plot_single_line_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"
        elif "decadal" in self.cfg["recipe"]:
            method = "DP"

        N=self.cfg["subsample_size"]
        for region in self.cfg["regions"]:
            if region != "Gyre" and region != "Glob":
            # if region != "Glob" and region != "NAtl":
                continue
            # fig = plt.figure(figsize=(70, 60/35+4), dpi=300)
            fig = plt.figure(figsize=(70/31*6, 60/35+4), dpi=300)
            dates = np.arange(1970, 2001)
            gs = GridSpec(1, 6, figure=fig)
            # gs = GridSpec(1, len(dates), figure=fig)
            # if region != "NAtl":
            ref_files = io.get_all_ancestor_files(self.cfg, f"{region}_reference_timeseries.nc")
            ref_files.sort()
            ref_files = ref_files[9:]
            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.sort()
            print(f"files {region}:")
            print(m_files)
            if len(m_files) == 0:
                continue
            i_gs = 0 
            for i,m_sd in enumerate(m_files[9:]):
                logger.info(f"loading json correlation file: {m_sd}")
                if dates[i] != 1970 and dates[i] != 1971 and dates[i] != 1993 and dates[i] != 1994 and dates[i] != 1999 and dates[i] != 2000:
                    continue
                # if i < 15 or i > 18:
                #     continue
                with open(m_sd, "r") as f:
                    corr = json.load(f)
                corr = self.purge_dataset("CESM2", corr)
                selected = [k for k, v in sorted(corr.items(), key=lambda item: item[1], reverse=reverse)][:N]
                all = [k for k, v in sorted(corr.items(), key=lambda item: item[1], reverse=reverse)]
                ref_ts_nc = iris.load_cube(ref_files[i])
                ax = fig.add_subplot(gs[0, i_gs])
                if i == 0:
                    ax.set_ylabel("SST anomaly / K")
                no_selected_label = True
                no_all_label = True
                for j,selected_alias_id in enumerate(all):
                    alias_files = io.get_all_ancestor_files(self.cfg, f"{region}_{selected_alias_id}_timeseries.nc")
                    alias_files.sort()
                    file = alias_files[9:][i]
                    ts = iris.load_cube(file)
                    print(selected_alias_id)
                    print(ts)
                    years = [i.year for i in ts.coord("time").units.num2date(ts.coord("time").points)]
                    if selected_alias_id in selected:
                        color = "red"
                        if no_selected_label:
                            ax.plot(years,ts.data, color = color, alpha=0.15, label="selected members")
                            no_selected_label = False
                        else: 
                            ax.plot(years,ts.data, color = color, alpha=0.15)
                    else:
                        color = "grey" 
                        if no_all_label:
                            ax.plot(years,ts.data, color = color, alpha=0.15, label="CMIP6 members")
                            no_all_label = False
                        else:
                            ax.plot(years,ts.data, color = color, alpha=0.15)
                    if region == "Glob":
                        ax.set_ylim(-0.5,0.5)
                    elif region == "NAtl":
                        ax.set_ylim(-0.7,0.7)
                    else:
                        ax.set_ylim(-1.5,1.5)
                    # ax.plot(years,ts.data, color = color, alpha=0.15)
                    ax.set_title(f"{dates[i]}")
                ax.plot(years, ref_ts_nc.data, color="black", label="DP mean")
                plt.legend()
                i_gs += 1
                    
            plt.tight_layout()
            print("savefig:")
            print(os.path.join("/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/plots", f"{method}_{self.cfg['selection_criteria']}_timeseries_1line_comparison_{region}.png"))
            post=""
            if region == "Gyre":
                post = "/gyre"
            plt.savefig(os.path.join(f"/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/plots{post}", f"{method}_{self.cfg['selection_criteria']}_timeseries_1line_comparison_{region}.png"))
            plt.savefig(os.path.join(f"/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/plots{post}", f"{method}_{self.cfg['selection_criteria']}_timeseries_1line_comparison_{region}.pdf"), format="pdf")



    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:
                        IntermediaryPlots(config, crit, cstr_yavg, crit_collect, i, N).main()


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