
from abc import ABC, abstractmethod

from dataclasses import dataclass
import xarray as xr

import tinysunset as ts

@dataclass
class ScalerConfig:

    clim_type: str = None # clim | ens_mean | ens_mean_loess | none
    clim_ref_years: list = None # ex. [1981, 2010]
    std_ref_years: list = None # ex. [1981, 2010]
    ensemble_dim: str = "ensemble"
    time_dim: str = "year"
    loess_degree: int = 1
    loess_window: int = 40
    groupby_type: str = None # to allow standarization by group (e.g. by model)
    # scaler type: options are: gamma, standard
    scaler_type: str = "standard"

def parse_scaler_config(cfg: dict):

    # same as above return but only adds the parameters that exist in the cfg dict
    return ScalerConfig(**cfg)

def _check_years(ds: xr.Dataset, years: list):
    # checks that all the years in the dataset are in the years list
    #assert (ds.time.dt.year.isin(years)).all(), f"Years in dataset: {ds.time.dt.year.values} not in years list: {years}"
    pass

class BaseScaler(ABC):
    def __init__(self, config: ScalerConfig):
        self.ensemble_dim = config.ensemble_dim
        self.time_dim = config.time_dim
        self.groupby_type = config.groupby_type

    @abstractmethod
    def fit(self, da: xr.DataArray):
        pass

    @abstractmethod
    def transform(self, da: xr.DataArray):
        pass

    @abstractmethod
    def inverse_transform(self, da: xr.DataArray):
        pass

    def fit_transform(self, da: xr.DataArray):
        self.fit(da)
        return self.transform(da)