"""

This module contains a class for weighted Gaussian families.

"""
import numpy as np
import warnings
from scipy.stats import norm as normal_dbn
from dataclasses import dataclass
from typing import Optional, List, Callable

from .discrete_family import discrete_family

@dataclass
class TruncatedGaussian(object):
    """
    Represents a truncated Gaussian distribution for post-selection inference.

    Attributes
    ----------
    estimate : float
        The observed value or estimate.
    sigma : float
        Standard deviation of the underlying Gaussian.
    smoothing_sigma : float
        Standard deviation for the smoothing noise.
    lower_bound : float
        Lower truncation bound.
    upper_bound : float
        Upper truncation bound.
    noisy_estimate : float
        Noisy version of the estimate (for randomized selection).
    factor : float
        Scaling factor for the distribution (default 1).
    """
    
    estimate: float
    sigma: float
    smoothing_sigma: float
    lower_bound: float
    upper_bound: float
    noisy_estimate: float
    factor: float = 1.

    def weight(self,
               x):
        t = x * self.factor
        return (normal_dbn.cdf((self.upper_bound - t) / self.smoothing_sigma)
                - normal_dbn.cdf((self.lower_bound - t) / self.smoothing_sigma))
    
@dataclass
class WeightedGaussianFamily(object):
    """
    Represents a family of weighted Gaussian distributions for selective inference.

    Attributes
    ----------
    estimate : float
        The observed value or estimate.
    sigma : float
        Standard deviation of the underlying Gaussian.
    weight_fns : list of callables
        List of weight functions to apply to the distribution.
    num_sd : float, optional
        Number of standard deviations to use for grid (default 10).
    num_grid : int, optional
        Number of grid points (default 4000).
    use_sample : bool, optional
        Whether to use a sample-based approximation (default False).
    seed : int, optional
        Random seed for reproducibility (default 0).
    """

    estimate: float
    sigma: float
    weight_fns: list
    num_sd: Optional[float] = 10
    num_grid: Optional[float] = 4000
    use_sample: Optional[bool] = False
    seed: Optional[int] = 0

    def _get_family(self,
                    basept=None):
   
        if basept is None:
            basept = self.estimate
            
        self._rng = np.random.default_rng(self.seed)
        
        if not self.use_sample:
            grid = np.linspace(basept - self.num_sd * self.sigma,
                               basept + self.num_sd * self.sigma, self.num_grid)

            log_weight = np.sum([np.log(np.clip(w(grid), 1e-16,1)) for w in self.weight_fns], 0)
            log_weight -= log_weight.max() + 10
            weight = np.exp(log_weight)

            weight *= normal_dbn.pdf((grid - basept) / self.sigma)

            return discrete_family(grid, weight)

        else:
            sample = self._rng.standard_normal(self.num_grid) * self.sigma + basept

            log_weight = np.sum([np.log(w(sample)) for w in self.weight_fns])
            log_weight -= log_weight.max() + 10
            weight = np.exp(log_weight)

            return discrete_family(sample, weight)

    def pvalue(self,
               null_value=0,
               alternative='twosided',
               basept=None):

        if basept is None:
            basept = null_value

        if alternative not in ['twosided', 'greater', 'less']:
            raise ValueError("alternative should be one of ['twosided', 'greater', 'less']")
        _family = self._get_family(basept=basept)
        tilt = (null_value - basept) / self.sigma**2
        if alternative in ['less', 'twosided']:
            _cdf = _family.cdf(tilt, x=self.estimate)
            if alternative == 'less':
                pvalue = _cdf
            else:
                pvalue = 2 * min(_cdf, 1 - _cdf)
        else:
            pvalue = _family.ccdf(tilt, x=self.estimate)
        return pvalue

    def interval(self,
                 basept=None,
                 level=0.9):
                 
        if basept is None:
            basept = self.estimate
            
        _family = self._get_family(basept=basept)

        L, U = _family.equal_tailed_interval(self.estimate,
                                             alpha=1-level)
        L *= self.sigma**2; L += basept
        U *= self.sigma**2; U += basept

        return L, U

    def MLE(self,
            basept=None):

        if basept is None:
            basept = self.estimate
            
        _family = self._get_family(basept=basept)

        mle, _, _ = _family.MLE(self.estimate)
        mle *= self.sigma**2; mle += basept
        return mle
