import iris
import matplotlib.pyplot as plt
import os
import numpy as np
import cartopy.crs as ccrs
import matplotlib.colors as colors
import iris.plot as iplt
from matplotlib.gridspec import GridSpec
import matplotlib.path as mpath
from esmvalcore.preprocessor import extract_region
import xarray as xr
import yaml
import json
from pathlib import Path
import matplotlib.ticker as ticker

import glob
import string

def build_plot():
    path = "/esarchive/scratch/jcos/constraints_paper_figures_data_scritps/files"
    alphabet = np.arange(1,29).reshape(7,4)
    # alphabet = np.array([f"{i})" for i in string.ascii_lowercase])[:28].reshape(7,4)
    alphabet2 = np.array([np.array(list(i)+list(i)) for i in np.arange(1,2*7+1).reshape(7,2)])
    alphabet3 = np.array(list(np.arange(1,5))+list(np.arange(1,4*4+1))).reshape(5,4) 
    # alphabet2 = np.array([f"{i})" for i in string.ascii_lowercase][:10]+[f"{i})" for i in string.ascii_lowercase][:10]).reshape(4,5)
    # evals = ["1970-2000"] #, "1961-2000"]
    # evals = ["1961-2000"]
    evals = ["1970-2000", "1961-2000"]
    metrics = ["CORR", "DIFFCORR", "RESCORR", "RPSS"]
    label_format = ["ACC", r"$\Delta$ACC", "ResCor", "RPSS"]
    labels = {k: v for k, v in zip(metrics, label_format)}
    methods = ["CWIP", "OBS_NAtl-field", "OBS_Glob-field", "DP_NAtl-field", "DP_Glob-field", "DP_Gyre-time", "DP_Glob-time"]
    # methods = ["CWIP", "OBS-field_NAtl", "OBS-field_Glob", "DP-field_NAtl", "DP-field_Glob", "DP-time_NAtl", "DP-time_Glob"]
    # pearson = "PCiris"
    pearson = "PC"
    methods_dict = {i: pearson for i in methods}
    methods_dict["DP_Gyre-time"] = "Befort"
    # methods_dict["DP_NAtl-time"] = "Befort"
    methods_dict["DP_Glob-time"] = "Befort"
    reg = "Med"
    projection_period = "y1-20"

    rd = {"Med": [-10,30,50,15], "WMed": [-10,30,25,15], "EMed": [15,30,25,15], "Glob": [-179,-89,358,178], "NAm": [-120,30,45,15]}
    proj, projection, path_ext, plotextend = define_projection([rd[reg][0]-5, rd[reg][0]+rd[reg][2]+5, rd[reg][1]-5, rd[reg][1]+rd[reg][3]+5])
    top, bottom, left, right = 0.98, 0.05, 0.03, 1
    wspace, hspace = 0.0001, 0.0001

    for eval in evals:
        # ratio = 15/13
        # fig = plt.figure(figsize=(8*1.25, 10*1.25),constrained_layout=True, dpi=100)
        fig = plt.figure(figsize=(9, 9.5), dpi=300)
        # plt.gcf().subplots_adjust()
        gs = GridSpec(
            7,
            4,
            figure=fig,
            hspace=hspace,
            wspace=wspace,
            top=top,
            bottom=bottom,
            left=left,
            right=right,
        )
        cb_heights = [6/8, 4/8, 2/8, 0/8]
        for i, metric in enumerate(metrics):
            for j, method in enumerate(methods):
                if pearson == "PCiris" and "field" not in method:
                    continue
                if method == "CWIP":
                    prefix="CWIP"
                elif "OBS" in method:
                    prefix=method.split('-')[0]
                else:
                    prefix=method.split('-')[0]
                    # prefix=f"{method.split('-')[0]}_{method.split('_')[-1]}"
                if method == "DP_Glob-time":
                    pass
                file = os.path.join(path, f"{prefix}_{projection_period}_{metric}_jja_{methods_dict[method]}_eval-{eval}.nc")
                file_signi = os.path.join(path, f"{prefix}_{projection_period}_{metric}_jja_{methods_dict[method]}_eval-{eval}_signi.nc")
                ax = fig.add_subplot(gs[j, i], projection=proj)
                ax.set_facecolor("silver")
                ax.set_boundary(path_ext, transform=ccrs.PlateCarree())
                ax.set_extent(plotextend, crs=ccrs.PlateCarree())
                med_region_encloser(rd[reg])
                ax.coastlines("50m", linewidth=0.8)
                if eval == "1961-2000":
                    alphabet=alphabet2
                if pearson == "PCiris":
                    alphabet=alphabet3
                ax.text(0.01, 0.81, f"{alphabet[j,i]})", transform=ax.transAxes)
                
                cube = iris.load_cube(file)
                cube = regrid_longitude_coord(cube)
                signi = iris.load_cube(file_signi)
                signi = regrid_longitude_coord(signi)

                if i == 1:
                    bounds = np.linspace(-0.1, 0.1, 21)
                    if np.max(extract_region(cube,-10,40,30,45).data)>0.5 or np.min(extract_region(cube,-10,40,30,45).data)<-0.5:
                        bounds = np.linspace(-1, 1, 21)
                elif i == 3:
                    bounds = np.linspace(-0.45, 0.45, 13)
                    num_ticks = 3
                else:
                    bounds = np.linspace(-1, 1, 21)
                    num_ticks = 4
                cmap = plt.cm.RdYlBu_r
                norm = colors.BoundaryNorm(boundaries=bounds, ncolors=256)
                fill = iplt.pcolormesh(cube, norm=norm, coords=("longitude", "latitude"), cmap=cmap)
                # pos, neg = save_amount_of_significant_points(cube, signi)
                hatches = [4*"/",None]
                if len(signi.shape) == 3:
                    signi=signi[-1]
                if eval == "1961-2000" and method == "CWIP" and metric == "CORR":
                    signi = signi*0+1
                iplt.contourf(signi, colors="none", levels=[0,0.5,1], hatches=hatches)
                print(j)
                if j == 6:
                    ys = ax.get_position()
                    colorbar_axes = plt.gcf().add_axes([ys.x0+0.021, ys.y0-0.025, ys.width-0.042, 0.02])
                    cbar = plt.colorbar(fill, colorbar_axes, orientation="horizontal")
                    # if i != 3: 
                    # Reduce the number of ticks using a locator
                    cbar.locator = ticker.MaxNLocator(num_ticks)
                    # Update the colorbar with the modified tick locator
                    cbar.update_ticks()
                if j == 0 or (j == 1 and pearson=="PCiris"):
                    ax.text(0.5, 1, labels[metric], transform=ax.transAxes, verticalalignment="center", horizontalalignment="center")
                if i == 0:
                    m = f"{method.split('_')[-1].split('-')[0]} {method.split('_')[0]}{method.split('-')[-1]}"
                    if "CWIP" in m:
                        m="ClimWIP"
                    elif "Gyre" in m:
                        m=m.replace("Gyre", "SPG")
                    ax.text(-0.05, 0.525, m, transform=ax.transAxes, verticalalignment="center", horizontalalignment="center", rotation="vertical")
            

            # ax = fig.add_subplot(gs[:, :])
            # ax.axis("OFF")
        plt.savefig(os.path.join(path.replace("/files", "/plots"), "gyre", f"figure1_{pearson}_{eval}_{projection_period}.png"), bbox_inches='tight')
        plt.savefig(os.path.join(path.replace("/files", "/plots"), "gyre", f"figure1_{pearson}_{eval}_{projection_period}.pdf"),format="pdf",  bbox_inches='tight')

def regrid_longitude_coord(cube):
    """Sorts the longitudes of the cubes from 0/360 degrees to -180/180"""
    coord = cube.coord("longitude")
    lon_extent = iris.coords.CoordExtent(coord, -180., 180., True, False)
    cube = cube.intersection(lon_extent)
    return cube

def save_amount_of_significant_points(cube, signi):
    cube = extract_region(cube,-10,40,30,45)
    signi = extract_region(signi,-10,40,30,45)
    positives = len(np.ma.where((signi.data==1)& (cube.data>0))[0])
    negatives = len(np.ma.where((signi.data==1)& (cube.data<0))[0])
    cube_key = (f"{self.cfg['selection_criteria']}.{self.cfg['cstr_yavg']}."
                f"{cube.attributes['variable']}.{cube.attributes['fcst_years']}y.{cube.attributes['region']}")

    json_file = os.path.join(Path(self.cfg["work_dir"]).parent.absolute(), "significant_mediterranean_skill_points.json")
    if not os.path.isfile(json_file):
        signi_dict = {}
    else:
        with open(json_file, "r") as f:
            signi_dict = json.load(f)

    signi_dict[cube_key] = {"positives": positives, "negatives": negatives}
    with open(json_file, "w") as f:
            json.dump(signi_dict, f)
    return positives, negatives


def region_to_square(region, dimension):
    """Definition of the region polygon"""
    if dimension == "latitude":
        return [
            region["start_latitude"],
            region["start_latitude"],
            region["end_latitude"],
            region["end_latitude"],
            region["start_latitude"],
        ]
    elif dimension == "longitude":
        return [
            region["start_longitude"],
            region["end_longitude"],
            region["end_longitude"],
            region["start_longitude"],
            region["start_longitude"],
        ]
    else:
        return "dimension unknown"

def define_projection(region):
    """
    projection definition to get LambertConformal borders
    """
    region = {
        "start_longitude": region[0],
        "end_longitude": region[1],
        "start_latitude": region[2],
        "end_latitude": region[3],
    }
    projection = "LambertConformal"
    plotextend = [
        region["start_longitude"],
        region["end_longitude"],
        region["start_latitude"],
        region["end_latitude"],
    ]
    if projection == "LambertConformal":
        # plotextend has to be a little larger so everything is on there
        plotextend = [
            plotextend[0] - 1.0,
            plotextend[1] + 1.0,
            plotextend[2] - 1.0,
            plotextend[3] + 1.0,
        ]
        # path to cut out is exact though
        lons = region_to_square(region, "longitude")
        lats = region_to_square(region, "latitude")
        path_ext = [[lon, lat] for lon, lat in zip(lons, lats)]
        path_ext = mpath.Path(path_ext).interpolated(20)
    # South Hemisfere
    if region["start_latitude"] <= 0 and region["end_latitude"] <= 0:
        proj = ccrs.LambertConformal(
            central_longitude=np.sum(plotextend[:2]) / 2.0,
            central_latitude=np.sum(plotextend[2:]) / 2.0,
            cutoff=+30,
            standard_parallels=(-33, -45),
        )
    # North Hemisphere
    else:
        proj = ccrs.LambertConformal(
            central_longitude=np.sum(plotextend[:2]) / 2.0,
            central_latitude=np.sum(plotextend[2:]) / 2.0,
        )
    return proj, projection, path_ext, plotextend

def med_region_encloser(region):
    """
    function that draws a dashed box over the Med
    region as defined in the setup.yml file
    as "REGION".
    """
    reg=[region[0], region[0]+region[2], region[1], region[1]+region[3]]
    lon_steps = 11
    lat_steps = 9
    lon_pad = 0.5
    lat_pad = 0
    box_x_points = np.concatenate(
        (
            np.ones(lat_steps) * reg[0] + lon_pad,
            np.linspace(reg[0], reg
                        [1], lon_steps)[1:-1],
            np.ones(lat_steps) * reg[1] - lon_pad,
            np.linspace(reg[1], reg
                        [0], lon_steps)[1:-1],
            [reg[0] + lon_pad],
        )
    )
    box_y_points = np.concatenate(
        (
            np.linspace(reg[2], reg[3], lat_steps),
            np.ones(lat_steps) * reg[3] - lat_pad,
            np.linspace(reg[3], reg[2], lat_steps),
            np.ones(lat_steps) * reg[2] + lat_pad,
            [reg[2] + lat_pad],
        )
    )
    plt.plot(
        box_x_points,
        box_y_points,
        linewidth=2.1,
        color="black",
        linestyle="dashed",
        transform=ccrs.PlateCarree(),
        zorder=5
    )

def main():
    build_plot()


if __name__ == "__main__":
    main()
