import os, sys
from dataclasses import dataclass

import xarray as xr
import numpy as np

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)

def parse_scaler_config(cfg: dict):

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

    # Normalization setups
    # ----------------------------

    # 1. Compute a single mean and std value for each point for the whole dataset

    # 2. Compute mean based on ensemble mean and std based on ensemble both time and ensemble dim

    # 3. Compute mean based on ensemble mean and de-trend using loess 
    #    (this can be applied to obs ts directly too when ensemble size = 1) 

    def __init__(
        self, 
        scaler_config: ScalerConfig
        ):

        self.ensemble_dim = scaler_config.ensemble_dim
        self.time_dim = scaler_config.time_dim
        self.clim_type = scaler_config.clim_type
        self.std_ref_years = scaler_config.std_ref_years
        self.clim_ref_years = scaler_config.clim_ref_years 
        self.loess_degree = scaler_config.loess_degree
        self.loess_window = scaler_config.loess_window
        self.groupby_type = scaler_config.groupby_type

    def fit(self, ds: xr.Dataset):
        
        assert self.ensemble_dim in list(ds.dims), f"{self.ensemble_dim} not in ds dims: {ds.dims}"
        assert self.time_dim in list(ds.dims), f"{self.time_dim} not in ds dims: {ds.dims}"
    
        if self.std_ref_years is not None:
            _selection = {self.time_dim: slice(str(self.std_ref_years[0]), str(self.std_ref_years[-1]))}
            self.std = ds.sel(_selection)

            if self.groupby_type is not None:
                self.std = self.std.groupby(self.groupby_type)

            self.std = self.std.std([self.time_dim, self.ensemble_dim])

        else:

            self.std = 1 # = No standardization
    
        if self.clim_type == "clim":

            assert self.clim_ref_years is not None, "clim_ref_years must be specified"
            _selection = {self.time_dim: slice(str(self.clim_ref_years[0]), str(self.clim_ref_years[-1]))}
            self.mean = ds.sel(_selection)
            
            if self.groupby_type is not None:
                self.mean = self.mean.groupby(self.groupby_type)
            
            self.mean = self.mean.mean([self.time_dim, self.ensemble_dim])

        else:
            if self.groupby_type is not None:
                self.mean = ds.groupby(self.groupby_type).mean([self.ensemble_dim])
            else:
                self.mean = ds.mean([self.ensemble_dim])

            if self.clim_type == "ens_mean_loess":

                self.mean = xr.apply_ufunc(
                    lambda x: ts.tools.loess_ts(x, window=self.loess_window, degree=self.loess_degree),
                    self.mean,
                    input_core_dims = [[self.time_dim]],
                    output_core_dims=[[self.time_dim]],
                    vectorize=True,
                    dask='parallelized',
                    output_dtypes=[float]
                    )

    def transform(self, ds: xr.Dataset):

        if self.groupby_type is not None:
            return (ds.groupby(self.groupby_type) - self.mean).groupby(self.groupby_type) / self.std
        else:
            return (ds - self.mean) / self.std
            
    def fit_transform(self, ds: xr.Dataset):
            
        self.fit(ds)

        return self.transform(ds)

    def inverse_transform(self, ds: xr.Dataset):

        if self.groupby_type is not None:
            return ds.groupby(self.groupby_type) * self.std + self.mean
        else:
            return ds * self.std + self.mean