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
from esmvalcore.preprocessor import extract_region, area_statistics, extract_time

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"

        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:]
            fig = plt.figure()
            files_projections = glob.glob("/esarchive/scratch/jcos/output/recipe_decadal_constraints_befort_20231124_123725/preproc/constr_load_uninit+obs/tos/*.nc")
            files_projections.sort()
            for file_projections in files_projections:
                cube = area_statistics(extract_region(extract_time(iris.load_cube(file_projections),1970,1,1,2008,12,31), self.cfg["regions"][region][0], self.cfg["regions"][region][1], self.cfg["regions"][region][2], self.cfg["regions"][region][3]), "mean")
                years = [i.year for i in cube.coord("time").units.num2date(cube.coord("time").points)]
                plt.plot(years, cube.data, color="grey", alpha=0.15)

            for file in ref_files:
                cube = iris.load_cube(file)
                years = [i.year for i in cube.coord("time").units.num2date(cube.coord("time").points)]
                plt.plot(years, cube.data, alpha=0.5)
            
            plt.savefig(f"/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/plots/befort_decadals_{region}_one_plot.png")



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