import matplotlib.pyplot as plt
import os
import logging
import iris
import matplotlib.path as mpath
import cartopy.crs as ccrs
import matplotlib.colors as colors
import numpy as np
import iris.plot as iplt
from esmvaltool.diag_scripts.shared import run_diagnostic, names, io
from esmvalcore.preprocessor import (
    area_statistics,
    extract_region,
)
from matplotlib.gridspec import GridSpec
import json
import glob
import string

def compute():
    # path = "/home/jcos/Desktop/paper#2/constraints_intercomparison_paper/files"
    path = "/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/files"
    alphabet = np.array([f"{i})" for i in string.ascii_lowercase])[:14].reshape(7,2)
    regions = ["WMed", "EMed"]
    methods = ["CWIP", "OBS_NAtl", "OBS_Glob", "DP_NAtl", "DP_Glob", "DPtime_Gyre", "DPtime_Glob"]
    bold=[[False, False],[False, True],[False, True],[True, False],[True, True],[False, True],[False, False]]
    signis=[[-0.04, -0.10],[-0.08, 0.08],[-0.04, 0.20],[0.15, 0.05],[0.26, 0.21],[0.06, 0.06],[-0.18, -0.12]]
    bold_corr=[[True, False],[False, True],[False, True],[False, False],[False, False],[False, False],[False, True]]
    signis_corr=[[-0.60, 0.03],[0.25, 0.39],[0.34, 0.82],[0.12, -0.30],[0.28, -0.17],[0.16, 0.29],[0.42, 0.50]]
    fig = plt.figure(figsize=(8, 16/5*6),constrained_layout=True, dpi=100)
    # plt.gcf().subplots_adjust()
    top, bottom, left, right = 0.9, 0.05, 0.2, 0.99
    wspace, hspace = 0.15, 0.1
    gs = GridSpec(
        7,
        2,
        figure=fig,
        hspace=hspace,
        wspace=wspace,
        top=top,
        bottom=bottom,
        left=left,
        right=right,
    )

    for j, region in enumerate(regions):
        for i, method in enumerate(methods):
            file = os.path.join(path, f"{method}_y1-20_{region}_CDF.json")
            with open(file, "r") as f:
                cdf_dict = json.load(f)

            ax = fig.add_subplot(gs[i, j])
            cstr = cdf_dict["CSTR"]
            raw = cdf_dict["RAW"]
            plt.plot(cstr[0], cstr[1], color="red", label="constr. CMIP6")
            plt.plot(raw[0], raw[1], color="blue", label="raw CMIP6")
            plt.ylim(0,100.5)

            xmax = np.max([cstr[0][-1], raw[0][-1]])
            xmin = np.min([cstr[0][0], raw[0][0]])
            ax.hlines(y=100*0.1,xmin=xmin,xmax=xmax,color="grey", alpha=0.5, clip_on=False, zorder=100)
            ax.hlines(y=100*0.25,xmin=xmin,xmax=xmax,color="grey", alpha=0.5, clip_on=False, zorder=100)
            ax.hlines(y=100*0.5,xmin=xmin,xmax=xmax,color="grey", alpha=0.5, clip_on=False, zorder=100)
            ax.hlines(y=100*0.75,xmin=xmin,xmax=xmax,color="grey", alpha=0.5, clip_on=False, zorder=100)
            ax.hlines(y=100*0.9,xmin=xmin,xmax=xmax,color="grey", alpha=0.5, clip_on=False, zorder=100)
            if i != 0:
                ax.hlines(y=100*-0.14,xmin=cstr[0][0],xmax=cstr[0][-1],color="red", clip_on=False, zorder=100)
                ax.hlines(y=100*-0.16,xmin=raw[0][0],xmax=raw[0][-1],color="blue", clip_on=False, zorder=100)
            plt.xlim(xmin,xmax)
            plt.yticks([0,10,25,50,75,90,100])
            if i == len(methods)-1:
                plt.xlabel("Temperature anomaly (K)", labelpad=10)
            if j == 0:
                plt.ylabel("Cumulative probability (%)", labelpad=-1, fontsize="medium")
            if j == 0:
                m = f"{method.split('_')[-1]} {method.split('_')[0]}"
                if m.endswith("OBS") or m.endswith("DP"):
                    m += "field"
                if "CWIP" in m:
                    m=m.split(" ")[0].replace("CWIP", "ClimWIP")
                if "Gyre" in m:
                    m=m.replace("Gyre", "SPG")
                ax.text(-0.25, 0.5, m, transform=ax.transAxes, rotation="vertical", verticalalignment="center", horizontalalignment="center", fontsize="large")
            if i == 0:
                ax.text(0.5, 1.1, region, transform=ax.transAxes, verticalalignment="center", horizontalalignment="center", fontsize="large")

            ax.text(0.03, 0.925, alphabet[i,j], transform=ax.transAxes)
            if bold[i][j]:
                ax.text(0.67, 0.53, r"RPSS: $\bf{%.2f}$" % signis[i][j], transform=ax.transAxes)
            else:
                ax.text(0.67, 0.53, 'RPSS: %.2f' % signis[i][j], transform=ax.transAxes)
            if bold_corr[i][j]:
                ax.text(0.625, 0.43, r"ResCor: $\bf{%.2f}$" % signis_corr[i][j], transform=ax.transAxes)
            else:
                ax.text(0.625, 0.43, 'ResCor: %.2f' % signis_corr[i][j], transform=ax.transAxes)
            if i == 0 and j == 0:
                plt.legend(fontsize="small", loc="lower right").set_zorder(101)
    plt.savefig(os.path.join(path.replace("/files", "/plots"), "gyre", "figure2.png"), bbox_inches='tight')
    plt.savefig(os.path.join(path.replace("/files", "/plots"), "gyre", "figure2.pdf"),format="pdf",  bbox_inches='tight')


        # best_ind_50 = np.argmin(np.absolute(best_increments-50))
        # best_ind_10 = np.argmin(np.absolute(best_increments-10))
        # best_ind_90 = np.argmin(np.absolute(best_increments-90))
        # tas_proj = np.array(best_tas)[[best_ind_10, best_ind_50, best_ind_90]]
        # plt.title(f"{region} area averaged 2015-2036 warming distributions\n ref. 1981-2010")
        # plt.savefig(os.path.join(self.cfg[names.PLOT_DIR], f"1-20_{region}_{sel_region}_CDF.png"), bbox_inches='tight')
        # with open(os.path.join(self.cfg[names.PLOT_DIR], f"1-20_{region}_{sel_region}_CDF.txt"), "w") as f:
        #     np.savetxt(f, tas_proj)


def weighted_quantiles(
    self, values, sample_weight, quantiles=[0.05, 0.25, 0.5, 0.75, 0.95]
):
    """ Very close to numpy.percentile, but supports weights.
    param values: numpy.array with data
    param sample_weight: array-like of the same length as `array`
    param quantiles: array-like with many quantiles needed
    return: numpy.array with computed quantiles.
    values_N > 2
    """
    values = np.ma.array(values)
    if len(values) != values.count():
        # There are masked values, don't compute.
        return np.ma.masked_array(values, mask=True)
    quantiles = np.array(quantiles)
    if sample_weight is None:
        sample_weight = np.ones(len(values))
    sample_weight = np.array(sample_weight)
    assert np.all(quantiles >= 0) and np.all(
        quantiles <= 1
    ), "quantiles should be in [0, 1]"

    # make it sum one
    sample_weight = np.array(sample_weight)
    sample_weight /= np.sum(sample_weight)

    sorter = np.argsort(values)
    values = values[sorter]
    sample_weight = sample_weight[sorter]

    weighted_quantiles = np.ma.cumsum(sample_weight) - 0.5 * sample_weight
    weighted_quantiles -= weighted_quantiles[0]
    weighted_quantiles /= weighted_quantiles[-1]

    return np.interp(quantiles, weighted_quantiles, values)

def main():
   compute()

if __name__ == "__main__":
    main()