import os, sys
from dataclasses import dataclass

import xarray as xr
import numpy as np
import scipy.stats as stats

import tinysunset as ts
from tinysunset.scalers.base import ScalerConfig, BaseScaler

class GammaScaler(BaseScaler):
    def __init__(self, config: ScalerConfig):
        super().__init__(config)

    def _fit_gamma(self, values):
        #try:
        #    params = stats.gamma.fit(x, floc=0)
        #    return params[0:2]
        #except:
        #    return (np.nan, np.nan)
       #def get_gamma_params(values):
        means = np.nanmean(values, axis=0)
        log_means = np.log(means)
        logs = np.log(values)
        mean_logs = np.nanmean(logs, axis=0)
        a = log_means - mean_logs
        alphas = (1 + np.sqrt(1 + 4 * a / 3)) / (4 * a)
        betas = means / alphas

        return (alphas, betas)
     
        
        
    def fit(self, da: xr.DataArray):

        # clip negative values to 0
        da = da.where(da > 0, 0)

        self.mean, self.var = xr.apply_ufunc(
            lambda x: self._fit_gamma(x.flatten()),
            da,
            input_core_dims=[[self.time_dim, self.ensemble_dim]],
            output_core_dims=[[],[]],
            vectorize=True,
            output_dtypes=[None, None],
            dask='parallelized',
        )

        da = self.transform(da)
        self.inverse_mean, self.inverse_var = da.mean([self.time_dim, self.ensemble_dim]), da.var([self.time_dim, self.ensemble_dim])


    def transform(self, da: xr.DataArray):
        return xr.apply_ufunc(
            lambda data, mean, var: stats.norm.ppf(stats.gamma.cdf(data, a=mean, scale=var)),
            #lambda data, params: stats.norm.ppf(stats.gamma.cdf(data, **params)),
            da, self.mean, self.var,
            input_core_dims=[[self.time_dim, self.ensemble_dim], [], []],
            output_core_dims=[[self.time_dim, self.ensemble_dim]],
            vectorize=True,
            dask='parallelized',
            output_dtypes=[None]
        )

    def inverse_transform(self, da: xr.DataArray):
        return xr.apply_ufunc(
            lambda data, mean, var: stats.gamma.ppf(stats.norm.cdf(data), a=mean, scale=var),
            da, self.mean, self.var,
            input_core_dims=[[self.time_dim, self.ensemble_dim], [], []],
            output_core_dims=[[self.time_dim, self.ensemble_dim]],
            vectorize=True,
            dask='parallelized',
            output_dtypes=[float]
        )
