import xarray as xr
import tinysunset as ts
from tinysunset.scalers.base import ScalerConfig
import os
import pickle

class CompositeScaler:
    def __init__(
        self,
        config: ScalerConfig,
        path: str = None
    ):

        self.config = config   
        self.scalers = {}

        self.path = path

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

        for var in ds.data_vars:
            
            _SCALERS = {"pr": "gamma", "else": "standard"}
            #_SCALERS = {"pr": "standard", "else": "standard"}
            _var_scaler = _SCALERS.get(var, _SCALERS["else"])

            if _var_scaler == "gamma":
                self.scalers[var] = ts.scalers.GammaScaler(self.config)
            elif _var_scaler == "standard":
                self.scalers[var] = ts.scalers.StandardScaler(self.config)
            else:
                raise ValueError(f"Scaler type {_var_scaler} not recognized")
        
            self.scalers[var].fit(ds[var])
            
    def transform(self, ds: xr.Dataset):
        transformed_data = {}
        for var, scaler in self.scalers.items():
            if var in ds:
                transformed_data[var] = scaler.transform(ds[var])
        return xr.Dataset(transformed_data)

    def inverse_transform(self, ds: xr.Dataset):
        inverse_transformed_data = {}
        for var, scaler in self.scalers.items():
            if var in ds:
                inverse_transformed_data[var] = scaler.inverse_transform(ds[var])

        return xr.Dataset(inverse_transformed_data)

    def fit_transform(self, ds: xr.Dataset):
        self.fit(ds)
        return self.transform(ds)

    def dump(self, path: str, overwrite: bool = False):
        # dump all the scalers per var in a unique diretory containing all the scalers
        os.makedirs(path, exist_ok=True)
        for var, scaler in self.scalers.items():
            pickle.dump(scaler, open(os.path.join(path, f"{var}.pkl"), "wb"))

    def load(self):
        # look for all the scalers in the directory and load them as vars
        for file in os.listdir(self.path):
            
            var = file.split(".")[0]
            _filepath = os.path.join(self.path, file)
            with open(_filepath, "rb") as f:
                self.scalers[var] = pickle.load(f)
            
    #def load(self, path: str):
    #    for var, scaler in self.scalers.items():
    #        with open(path, "rb") as f:
    #            self.scalers[var] = scaler.load(f)
    
    