import iris

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 random

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_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
        var = self.cfg['input_files'][0].split("/")[-3].split("_")[0]
        recipe_name = self.cfg['input_files'][0].split("/")[-6]
        if "observation" in self.cfg["recipe"]:
            method = "OBS"
        elif "decadal" in self.cfg["recipe"]:
            method = "DP"

        for region in self.cfg["regions"]:
            if region != "Glob" and region != "NAtl":
                    continue
            for pearson_ratio in ["numerator", "denominator"]:
                print(region)
                print(pearson_ratio)
                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"{pearson_ratio}_{region}_*.json")
                m_files.sort()
                cfg_copy = self.cfg.copy()
                cfg_copy["input_files"] = [i.replace("analysis", "PC") for i in cfg_copy["input_files"]]
                m_files_selection = io.get_all_ancestor_files(cfg_copy, f"{region}_*.json")
                m_files_selection.sort()
                if len(m_files) == 0:
                    continue
                formatter = {"RMSE": "RMSE", "PC": "Uncentered pat. corr.", "PCiris": "Centered pat. corr.", "Befort": f"{region} timeseries MAE", "analysis": f"uncentered Pearson coeficient {pearson_ratio}", "analysis_centered": f"centered Pearson coeficient {pearson_ratio}"}
                N=self.cfg["subsample_size"]
                corrs_array = []
                sds = []
                for i,m_sd in enumerate(m_files_selection):
                    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")])
                    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):
                            seldict[d] = 1
                    selection[i] = list(seldict.values())

                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:
                        values = np.zeros((len(m_files),len(list(corr.keys()))))
                        if pearson_ratio == "numerator":
                            values_num = 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")])
                    
                    selected = [k for k, v in sorted(corr.items(), key=lambda item: item[1], reverse=reverse)][:N]
                    
                    corrs = [corr[member] for member in members]
                    values[i] = corrs
                    if pearson_ratio == "numerator":
                        values_num[i] = corrs
                
                    # corrs.sort(reverse=reverse)
                    # corrs_array.append(corrs)
                    # sds.append(int(m_sd.split(".json")[0].split("_")[-1]))

                print("selection")
                print(selection)
                plt.figure(figsize=(10,5))
                all_values = values[9:15].flatten()
                sel_values = all_values[np.where(selection[9:15].flatten()==1)[0]]
                non_values = all_values[np.where(selection[9:15].flatten()==0)[0]]
                # plt.boxplot(1,all_values, color="black")
                # plt.boxplot(1.2,sel_values, color="red")
                plt.boxplot([non_values], positions=[1], widths=0.5,vert=False, patch_artist=True, 
                            boxprops=dict(facecolor='white', color='black'), 
                            capprops=dict(color='black'), whiskerprops=dict(color='black'), 
                            flierprops=dict(markeredgecolor='black'), medianprops=dict(color='black'), labels=["unselected\nmembers"])
                plt.boxplot([sel_values], positions=[1.5], widths=0.5,vert=False, patch_artist=True, 
                            boxprops=dict(facecolor="white", color='red'), 
                            capprops=dict(color='red'), whiskerprops=dict(color='red'), 
                            flierprops=dict(markeredgecolor='red'), medianprops=dict(color='red'), labels=["selected\nmembers"])

                plt.title(f"Uncentered Pearson correlation coefficint {pearson_ratio} in the {region} DPfield-selection method\nfor the constraints performed in start dates 1970-1975 (over-selection of ECS)")
                if "numerator" in pearson_ratio:
                    plt.xlabel(f"uncentered pearson numerator / K2")
                else:
                    plt.xlabel(f"uncentered pearson denominator / K")
                plt.tight_layout()
                plt.savefig(os.path.join("/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/plots", f"{method}_selection_{self.cfg['selection_criteria']}_heatmapscatter_best{self.cfg['subsample_size']}_{region}_{pearson_ratio}_new_regrid.png"))
                # for sel, val in zip(selection[9:15].flatten(), values[9:15].flatten()):
                #     if sel == 1:
                #         color="red"
                #     else:
                #         color ="black"
                #     plt.scatter(random.uniform(-0.5,0.5), val, alpha=0.5, color=color)
            #     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)
            #     vmin = np.min(values.transpose()[:,9:])
            #     vmax = np.max(values.transpose()[:,9:])
            #     fill = plt.pcolormesh(dates[9:], members, values.transpose()[:,9:], cmap=plt.cm.plasma, vmin=vmin, vmax=vmax, 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']} through\n"
            #             f"{method}-sel. {self.cfg['cstr_yavg']} {formatter[self.cfg['selection_criteria']]}\n"
            #             f"in region: {region.split('_signi')[0]}\n", fontsize="x-large", pad=-10)
            #     plt.tight_layout()
            #     minimal_region = region
            #     if "corr" in minimal_region:
            #         minimal_region=minimal_region.partition(".")[-1].rpartition("_signi_corr-mask")[0]
            #     if "decadal" in self.cfg["out_plot_dir"]:
            #         method = "DP"
            #         rect = patches.Rectangle((1969.6, canesm[0]+ymin), 30.8, canesm[-1]-canesm[0]-ymin*2, linewidth=3, edgecolor='r', facecolor='none')
            #         ax.add_patch(rect)
            #         ax.text(-0.2, 1.02, "A)", transform=ax.transAxes, fontsize="xx-large")
            #     else:
            #         method = "OBS"
            #         ax.text(-0.2, 1.02, "B)", transform=ax.transAxes, fontsize="xx-large")
            #     print(os.path.join("/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/plots", f"{method}_selection_{self.cfg['selection_criteria']}_heatmap_best{self.cfg['subsample_size']}_{minimal_region}_{pearson_ratio}_new_regrid.png"))
            #     plt.savefig(os.path.join("/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/plots", f"{method}_selection_{self.cfg['selection_criteria']}_heatmap_best{self.cfg['subsample_size']}_{minimal_region}_{pearson_ratio}_new_regrid.png"))
            
            # plt.figure()
            # pearson = values_num/values
            # for sel, val in zip(selection[9:15].flatten(), pearson[9:15].flatten()):
            #     if sel == 1:
            #         color="red"
            #     else:
            #         color ="black"
            #     plt.scatter(random.uniform(-0.5,0.5), val, alpha=0.5, color=color)
            # plt.savefig(os.path.join("/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/plots", f"{method}_selection_{self.cfg['selection_criteria']}_heatmapscatter_best{self.cfg['subsample_size']}_{region}_pearsonratio_new_regrid.png"))


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