
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

import logging
logger = logging.getLogger(__name__)

import tinysunset as ts

from tinysunset.scalers import ScalerConfig, StandardScaler, GammaScaler

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

@dataclass
class VerificationConfig:
    
    metrics: list = None 
    verification_years: list = None
    plots_layout: list = None
    scaler_hcst: ScalerConfig = None
    scaler_obs: ScalerConfig = 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)

def _scaler_func(
    conf: ScalerConfig,
    var_name: str
    ):
    if var_name == "pr":
        return GammaScaler(conf)
    else:
        return StandardScaler(conf)
    

class Verification:
    
    def __init__(
        self,
        verification_config: VerificationConfig
        ):

        self.conf = verification_config
        
    def compute(
        self,
        obs: xr.DataArray,
        hcst: xr.DataArray,
        ref_hcst: xr.DataArray = None,
        scaler_ref_hcst: ScalerConfig = None
    ) -> None:

        assert obs.name == hcst.name, f"Variables are different {obs.name} {hcst.name}"
        #_scaler_func = alp

        # Apply scalers
        # -------------------------------------------------------------------

        if self.conf.scaler_hcst is not None:
            hcst_scaler = _scaler_func(self.conf.scaler_hcst, hcst.name)
            hcst = hcst_scaler.fit_transform(hcst)
            
        if scaler_ref_hcst is not None and ref_hcst is not None:
            ref_hcst_scaler = _scaler_func(scaler_ref_hcst, ref_hcst.name)
            ref_hcst = ref_hcst_scaler.fit_transform(ref_hcst)

        if self.conf.scaler_obs is not None:
                
            obs_scaler = _scaler_func(self.conf.scaler_obs, obs.name)
            obs = obs_scaler.fit_transform(obs)
            obs = obs.isel(month=0) ## TODO: hardcoded
    
        # Sel Verification years only
        # -------------------------------------------------------------------
        
        if self.conf.verification_years is not None:
        
            obs = obs.sel(
                time=obs.time.isin(self.conf.verification_years)
                )
            fcst = hcst.sel(
                time=hcst.time.isin(self.conf.verification_years)
                )
            ## TODO: some forecst don't have time dimension 
            #if ref_hcst is not None:
            #    ref_hcst = ref_hcst.sel(
            #        time=ref_hcst.time.isin(self.conf.verification_years)
            #        )
            
        _error_msg = f"Verification years are not the same in obs and fcst {obs.time.values} {fcst.time.values}"
        assert np.all(obs.time.values == fcst.time.values), _error_msg

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

        # if lat lon in dims transpose to have lat lon as last dimensions
        
        if "lat" in obs.dims and "lon" in obs.dims:
            obs = obs.transpose(..., "lat", "lon")
            fcst = fcst.transpose(..., "lat", "lon")
            if ref_hcst is not None:
                ref_hcst = ref_hcst.transpose(..., "lat", "lon")

        # Compute hcst metrics
        # -------------------------------------------------------------------
        hcst_metrics = ts.verification.comp_metrics(
            fcst, 
            obs.mean("ensemble"), 
            self.conf.metrics
            )    
        
        # Climatological forecast and metrics against climatological forecast
        # -------------------------------------------------------------------
        if ref_hcst 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
            # if time does not exist in obs
            if "time" not in ref_hcst.dims:
                ref_hcst = ref_hcst.expand_dims(
                    dim={"time": obs.time}
                    )
            else:
                ref_hcst = ref_hcst.sel(
                    time=obs.time.values
                    )
            # expand ensemble mean to have 1 value on ensemble dimension
            #clim_fcst = clim_fcst.drop(["year","month", "ensemble_id"])

            ref_hcst_metrics = ts.verification.comp_metrics(
                ref_hcst, 
                obs.mean("ensemble"), 
                self.conf.metrics
                )
            
            # compute skill scores from fcst_metrics and clim_fcst_metrics
            # TODO: expand this to other benchmarks such as persistence
            hcst_metrics = ts.verification.skill_scores(
                hcst_metrics,
                ref_hcst_metrics
                )
        else:
            ref_hcst_metrics = None
            
        return hcst_metrics, ref_hcst_metrics
            
    @staticmethod
    def plot_verification(
        verification: xr.Dataset,
        plots_layout: list,
        title: str = None,
        outfile: Path = None,
        mask = None
        ) -> None:

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

        gs = gridspec.GridSpec(
            len(plots_layout),
            max([len(row) for row in 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(plots_layout):
            for j, col in enumerate(row):
                
                # plots metrics
                if col not in ["fcst", "fcst_ensmean", "obs", "fcst_ensmedian"]:
                
                    _, _ = ts.verification.plot_metric(
                        verification[col],
                        ax=fig.add_subplot(gs[i, j], projection=proj),
                        **plts_kwargs
                    )

                else:
                    pass
                    #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()