
from dataclasses import dataclass
from pathlib import Path

import xarray as xr
import numpy as np

import cartopy.crs as ccrs
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
from matplotlib.colors import LinearSegmentedColormap

import logging
logger = logging.getLogger(__name__)

import tinysunset as ts

CBAR_KWARGS = {
    "shrink": 0.8, 
    "pad": 0.05,
    "orientation": "horizontal", #"label": None
    }

@dataclass
class VerificationConfig:
    
    metrics: list = None 
    comp_clim_fcst: bool = False
    verification_years: list = None
    plots_layout: list = None

def parse_verification_config(
    cfg: dict
    ) -> VerificationConfig:

    cfg['verification_years'] = list(
        range(cfg['reference_years'][0], cfg['reference_years'][1]+1)
        )
    # drop reference_years from cfg
    cfg.pop("reference_years")

    return VerificationConfig(**cfg)
    
class Verification:
    
    def __init__(
        self,
        obs: xr.Dataset,
        fcst: xr.Dataset,
        verification_config: VerificationConfig
        ):

        
        self.metrics = verification_config.metrics
        self.comp_clim_fcst = verification_config.comp_clim_fcst
        self.verification_years = verification_config.verification_years
        self.plots_layout = verification_config.plots_layout

        if self.verification_years is not None:
        
            self.obs = obs.sel(
                time=obs.time.isin(self.verification_years)
                )
            self.fcst = fcst.sel(
                time=fcst.time.isin(self.verification_years)
                ) 
        else:
            
            self.obs = obs
            self.fcst = fcst

        _error_msg = f"Verification years are not the same in obs and fcst {self.obs.time.values} {self.fcst.time.values}"
        assert np.all(self.obs.time.values == self.fcst.time.values), _error_msg

        logger.debug(f"Verification years: {self.obs.time.values}")
        self.fcst["time"] = self.obs.time.values

        # if lat lon in dims transpose to have lat lon as last dimensions
        
        if "lat" in self.obs.dims and "lon" in self.obs.dims:
            self.obs = self.obs.transpose(..., "lat", "lon")
            self.fcst = self.fcst.transpose(..., "lat", "lon")
        
    def compute(
        self,
        clim_fcst: xr.Dataset = None
    ) -> None:

        self.clim_fcst = clim_fcst

        # Fcst metrics
        # -------------------------------------------------------------------
        self.fcst_metrics = ts.verification.comp_metrics(
            self.fcst, 
            self.obs.mean("ensemble"), 
            self.metrics
            )    
        
        # Climatological forecast and metrics against climatological forecast
        # -------------------------------------------------------------------
        
        if self.clim_fcst is not None:
        
            # check that self.obs contains all the ref_years
            ## TODO:
            #assert all([y in self.obs.time.dt.year.values for y in self.reference_years])
            #self.clim_fcst = self.obs.sel(
            #    time=self.obs.time.dt.year.isin(self.reference_years))
            
            # expand clim fcst to have 30 values on time dimension
            # add new dim time
            self.clim_fcst = self.clim_fcst.expand_dims(
                dim={"time": self.obs.time}
                )
            
            # expand ensemble mean to have 1 value on ensemble dimension
            #clim_fcst = clim_fcst.drop(["year","month", "ensemble_id"])

            self.clim_fcst_metrics = ts.verification.comp_metrics(
                self.clim_fcst, 
                self.obs.mean("ensemble"), 
                self.metrics
                )
            
            # compute skill scores from fcst_metrics and clim_fcst_metrics
            # TODO: expand this to other benchmarks such as persistence
            self.fcst_metrics = ts.verification.skill_scores(
                self.fcst_metrics,
                self.clim_fcst_metrics
                )

    def plot_verification(
        self, 
        title: str = None,
        outfile: Path = None,
        mask = None
        ) -> None:

        _ncols = len(self.plots_layout[0])
        _nrows = len(self.plots_layout)
        
        fig = plt.figure(figsize=(10*_ncols, 10*_nrows), constrained_layout=True, facecolor="white")

        gs = gridspec.GridSpec(
            len(self.plots_layout),
            max([len(row) for row in self.plots_layout]), 
            figure=fig
            )

        proj = ccrs.Robinson(central_longitude=0)

        plts_kwargs = dict(
            mask = mask,
            projection = proj
            )

        # PLOTS EVERY METRIC ACCORDING TO THE LAYOUT
        for i, row in enumerate(self.plots_layout):
            for j, col in enumerate(row):
                
                # plots metrics
                if col not in ["fcst", "fcst_ensmean", "obs", "fcst_ensmedian"]:
                
                    _, _ = ts.verification.plot_metric(
                        self.fcst_metrics[col],
                        ax=fig.add_subplot(gs[i, j], projection=proj),
                        **plts_kwargs
                    )

                # plts fcst, fcst_ensmean, obs examples
                else:

                    if col == "fcst":
                        _data = self.fcst.isel(time=0, ensemble=0)
                        sub_title = "FCST"
                    elif col == "fcst_ensmean":
                        _data = self.fcst.isel(time=0).mean("ensemble")
                        sub_title = "FCST ENSMEAN"
                    elif col == "fcst_ensmedian":
                        _data = self.fcst.isel(time=0).median("ensemble")
                        sub_title = "FCST ENSMEAN"
                    elif col == "obs": 
                        _data = self.obs.isel(time=0,ensemble=0)
                        sub_title = "OBS"
                    else:
                        raise ValueError(f"Unknown metric {col}")

                
                    kwgs = dict(
                        cmap="bwr", vmin=-3, vmax=3, levels=21,
                        robust=True,
                        cbar_kwargs= CBAR_KWARGS,
                        **plts_kwargs
                    )

                    ax, im = ts.plots.plot_array_to_map(
                        _data, 
                        ax=fig.add_subplot(gs[i, j], projection=proj), 
                        **kwgs
                    )
                    
                    plt.title(sub_title, fontsize=40, fontweight="bold")

        if title is not None:
            plt.suptitle(
                title, fontsize=30, fontweight="bold",
                horizontalalignment="right", y=0.6, x=0.6
                )

        if outfile is not None:
            plt.savefig(f"{outfile}")
            plt.close()