"""Alpha-Stable distribution fitter."""

import numpy as np
import pandas as pd
from scipy.optimize import minimize
from scipy.fft import fft, ifft
from scipy.interpolate import interp1d
from scipy.stats import norm
from scipy.integrate import quad
from typing import Dict, Tuple
import warnings

from .base import DistributionFitter


class AlphaStable(DistributionFitter):
    """Alpha-Stable distribution fitter for heavy-tailed returns."""
    
    def __init__(self, returns: np.ndarray):
        """
        Initialize Alpha-Stable distribution fitter.
        
        Parameters
        ----------
        returns : np.ndarray
            Array of returns
        """
        super().__init__(returns)
        self.name = "AlphaStable"
        self._cdf_interp = None
        self._ppf_interp = None
    
    def _characteristic_function(self, t: np.ndarray, alpha: float, beta: float, 
                               c: float, mu: float) -> np.ndarray:
        """
        Characteristic function of alpha-stable distribution.
        
        Parameters
        ----------
        t : np.ndarray
            Frequency points
        alpha : float
            Stability parameter (0 < alpha <= 2)
        beta : float
            Skewness parameter (-1 <= beta <= 1)
        c : float
            Scale parameter (c > 0)
        mu : float
            Location parameter
            
        Returns
        -------
        np.ndarray
            Characteristic function values
        """
        if alpha == 1:
            # Special case: Cauchy distribution
            phi = np.exp(1j * mu * t - c * np.abs(t) * (1 + 1j * beta * np.sign(t) * np.log(np.abs(t))))
        else:
            phi = np.exp(1j * mu * t - c**alpha * np.abs(t)**alpha * 
                         (1 - 1j * beta * np.sign(t) * np.tan(np.pi * alpha / 2)))
        
        return phi
    
    def _cf_to_pdf(self, t: np.ndarray, cf_vals: np.ndarray, x_range: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        Convert characteristic function to PDF using FFT.
        
        Parameters
        ----------
        t : np.ndarray
            Frequency points
        cf_vals : np.ndarray
            Characteristic function values
        x_range : np.ndarray
            Range of x values for PDF
            
        Returns
        -------
        tuple
            (x_values, pdf_values)
        """
        # Use FFT to compute PDF from CF
        dt = t[1] - t[0]
        n = len(t)
        
        # Prepare for FFT
        cf_symmetric = np.concatenate([cf_vals, np.conj(cf_vals[-1:0:-1])])
        pdf_vals = ifft(cf_symmetric).real
        pdf_vals = pdf_vals[:n] * n / (2 * np.pi)
        
        # Corresponding x values
        dx = 2 * np.pi / (n * dt)
        x_vals = np.arange(n) * dx
        
        # Shift and scale to match desired range
        x_vals = x_vals - x_vals[-1] / 2
        
        # Interpolate to desired x_range
        pdf_interp = interp1d(x_vals, pdf_vals, kind='cubic', 
                            fill_value='extrapolate', bounds_error=False)
        pdf_final = pdf_interp(x_range)
        
        # Ensure non-negative
        pdf_final = np.maximum(pdf_final, 1e-10)
        
        return x_range, pdf_final
    
    def _stable_loglikelihood(self, params: np.ndarray) -> float:
        """
        Negative log-likelihood for alpha-stable distribution.
        
        Parameters
        ----------
        params : np.ndarray
            Array of [alpha, beta, c, mu]
            
        Returns
        -------
        float
            Negative log-likelihood
        """
        alpha, beta, c, mu = params
        
        # Parameter constraints
        if not (0.1 < alpha <= 2.0):
            return 1e10
        if not (-1 <= beta <= 1):
            return 1e10
        if not (c > 0):
            return 1e10
        
        try:
            # Evaluate PDF for all returns
            pdf_vals = self.pdf(self.returns, alpha, beta, c, mu)
            
            if np.any(pdf_vals <= 0):
                return 1e10
            
            loglik = np.sum(np.log(pdf_vals))
            
            # Add penalty for extreme parameters
            penalty = 0
            if alpha < 0.3:
                penalty += (0.3 - alpha) * 1000
            
            return -(loglik + penalty)
            
        except Exception:
            return 1e10
    
    def fit(self) -> Dict[str, float]:
        """
        Fit Alpha-Stable distribution using MLE.
        
        Returns
        -------
        dict
            Fitted parameters: alpha, beta, c, mu
        """
        # Initial parameter estimates
        mu0 = np.mean(self.returns)
        c0 = np.std(self.returns) / np.sqrt(2)  # Rough estimate
        
        # Estimate alpha from tail behavior
        # Use Hill estimator for the tail index
        sorted_returns = np.sort(np.abs(self.returns - mu0))
        tail_size = max(10, len(sorted_returns) // 20)  # Use top 5% of tail
        
        if tail_size < 10:
            alpha0 = 1.8  # Default close to normal
        else:
            # Hill estimator
            tail_data = sorted_returns[-tail_size:]
            log_ratios = np.log(tail_data / sorted_returns[-tail_size])
            hill_alpha = 1.0 / np.mean(log_ratios)
            alpha0 = min(max(hill_alpha, 0.5), 2.0)  # Constrain to valid range
        
        # Estimate beta from skewness
        empirical_skew = pd.Series(self.returns).skew()
        if abs(empirical_skew) > 0.1:
            beta0 = np.sign(empirical_skew) * min(abs(empirical_skew) * 0.5, 0.95)
        else:
            beta0 = 0.0
        
        # Initial parameter vector
        x0 = np.array([alpha0, beta0, c0, mu0])
        
        # Bounds for parameters
        bounds = [
            (0.1, 2.0),    # alpha
            (-0.999, 0.999),  # beta
            (1e-6, 10.0),  # c
            (None, None)   # mu
        ]
        
        # Optimization
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            
            result = minimize(
                self._stable_loglikelihood,
                x0,
                method='Nelder-Mead',
                options={
                    'maxiter': 500,
                    'xatol': 1e-6,
                    'fatol': 1e-6,
                    'disp': False
                }
            )
        
        alpha, beta, c, mu = result.x
        
        # Clean up parameters
        alpha = min(max(alpha, 0.1), 2.0)
        beta = min(max(beta, -0.999), 0.999)
        c = max(c, 1e-6)
        
        self.params = {
            'alpha': alpha,
            'beta': beta,
            'c': c,
            'mu': mu
        }
        
        # Calculate log-likelihood
        self.loglikelihood = -result.fun
        self.aic = -2 * self.loglikelihood + 2 * len(self.params)
        self.bic = len(self.params) * np.log(len(self.returns)) - 2 * self.loglikelihood
        
        # Build interpolation functions for faster CDF/PPF
        self._build_interpolations()
        
        self._fitted = True
        return self.params
    
    def _build_interpolations(self):
        """Build interpolation functions for CDF and PPF."""
        if not self.params:
            return
        
        # Create grid of x values
        x_min = np.percentile(self.returns, 0.1)
        x_max = np.percentile(self.returns, 99.9)
        x_range = np.linspace(x_min * 3, x_max * 3, 200)
        
        # Compute CDF values
        cdf_vals = np.zeros_like(x_range)
        for i, x in enumerate(x_range):
            cdf_vals[i], _ = quad(self.pdf, -np.inf, x)
        
        # Normalize and create interpolation
        self._cdf_interp = interp1d(x_range, cdf_vals, kind='cubic', 
                                   fill_value=(0, 1), bounds_error=False)
        
        # Create PPF interpolation (inverse of CDF)
        q_range = np.linspace(0.001, 0.999, 200)
        ppf_vals = np.zeros_like(q_range)
        
        for i, q in enumerate(q_range):
            # Use root finding to get quantile
            def f(x):
                return self._cdf_interp(x) - q
            
            try:
                from scipy.optimize import brentq
                ppf_vals[i] = brentq(f, x_range[0], x_range[-1])
            except:
                ppf_vals[i] = self.params['mu']  # Fallback
        
        self._ppf_interp = interp1d(q_range, ppf_vals, kind='cubic', 
                                   bounds_error=False, fill_value='extrapolate')
    
    def pdf(self, x: np.ndarray, alpha: float = None, beta: float = None, 
           c: float = None, mu: float = None) -> np.ndarray:
        """
        Probability density function.
        
        Parameters
        ----------
        x : np.ndarray
            Points at which to evaluate PDF
        alpha, beta, c, mu : float, optional
            Distribution parameters. If None, uses fitted parameters.
            
        Returns
        -------
        np.ndarray
            PDF values
        """
        if alpha is None:
            if not self._fitted:
                self.fit()
            alpha = self.params['alpha']
            beta = self.params['beta']
            c = self.params['c']
            mu = self.params['mu']
        
        # For special cases, use analytical formulas
        if alpha == 2.0:
            # Normal distribution
            return norm.pdf(x, loc=mu, scale=c)
        elif alpha == 1.0 and beta == 0:
            # Cauchy distribution
            return c / (np.pi * ((x - mu)**2 + c**2))
        
        # General case: use characteristic function method
        t_range = np.linspace(-50, 50, 1024)
        
        # Vectorize for efficiency
        x = np.atleast_1d(x)
        pdf_vals = np.zeros_like(x, dtype=float)
        
        for i, xi in enumerate(x):
            cf_vals = self._characteristic_function(t_range, alpha, beta, c, mu - xi)
            dt = t_range[1] - t_range[0]
            
            # Use FFT to get density
            cf_sym = np.concatenate([cf_vals, np.conj(cf_vals[-1:0:-1])])
            pdf_component = ifft(cf_sym).real
            pdf_component = pdf_component[:len(t_range)] * len(t_range) / (2 * np.pi)
            
            # The density at x=xi is the integral of exp(-i*t*xi) * phi(t)
            # This is what's computed by the FFT
            pdf_vals[i] = np.abs(pdf_component[0])
        
        return np.maximum(pdf_vals, 1e-10).reshape(np.array(x).shape)
    
    def cdf(self, x: np.ndarray, alpha: float = None, beta: float = None,
           c: float = None, mu: float = None) -> np.ndarray:
        """
        Cumulative distribution function.
        
        Parameters
        ----------
        x : np.ndarray
            Points at which to evaluate CDF
        alpha, beta, c, mu : float, optional
            Distribution parameters. If None, uses fitted parameters.
            
        Returns
        -------
        np.ndarray
            CDF values
        """
        if alpha is None and self._cdf_interp is not None:
            return self._cdf_interp(x)
        
        if alpha is None:
            if not self._fitted:
                self.fit()
            alpha = self.params['alpha']
            beta = self.params['beta']
            c = self.params['c']
            mu = self.params['mu']
        
        # For special cases, use analytical formulas
        if alpha == 2.0:
            return norm.cdf(x, loc=mu, scale=c)
        elif alpha == 1.0 and beta == 0:
            return 0.5 + np.arctan((x - mu) / c) / np.pi
        
        # Use numerical integration
        x = np.atleast_1d(x)
        cdf_vals = np.zeros_like(x)
        
        for i, xi in enumerate(x):
            cdf_vals[i], _ = quad(lambda xx: self.pdf(np.array([xx]), alpha, beta, c, mu), 
                                -np.inf, xi)
        
        return cdf_vals.reshape(np.array(x).shape)
    
    def ppf(self, q: np.ndarray, alpha: float = None, beta: float = None,
           c: float = None, mu: float = None) -> np.ndarray:
        """
        Percent point function (inverse of CDF).
        
        Parameters
        ----------
        q : np.ndarray
            Quantiles to evaluate
        alpha, beta, c, mu : float, optional
            Distribution parameters. If None, uses fitted parameters.
            
        Returns
        -------
        np.ndarray
            PPF values
        """
        if alpha is None and self._ppf_interp is not None:
            return self._ppf_interp(q)
        
        if alpha is None:
            if not self._fitted:
                self.fit()
            alpha = self.params['alpha']
            beta = self.params['beta']
            c = self.params['c']
            mu = self.params['mu']
        
        # For special cases, use analytical formulas
        if alpha == 2.0:
            return norm.ppf(q, loc=mu, scale=c)
        elif alpha == 1.0 and beta == 0:
            return mu + c * np.tan(np.pi * (q - 0.5))
        
        # Use root finding
        from scipy.optimize import brentq
        
        q = np.atleast_1d(q)
        ppf_vals = np.zeros_like(q)
        
        # Determine search range
        range_min = mu - 10 * c
        range_max = mu + 10 * c
        if alpha < 1:
            range_min = mu - 50 * c
            range_max = mu + 50 * c
        
        def f(x, target):
            return self.cdf(np.array([x]), alpha, beta, c, mu)[0] - target
        
        for i, qi in enumerate(q):
            try:
                ppf_vals[i] = brentq(f, range_min, range_max, args=(qi,))
            except:
                # Fallback approximation
                ppf_vals[i] = mu + c * np.sign(qi - 0.5) * np.abs(qi - 0.5)**(-1/alpha)
        
        return ppf_vals.reshape(np.array(q).shape)
    
    def es(self, alpha: float = 0.05) -> float:
        """
        Calculate Expected Shortfall (Conditional VaR).
        
        Parameters
        ----------
        alpha : float
            Significance level (default: 0.05)
            
        Returns
        -------
        float
            Expected Shortfall value
        """
        if not self._fitted:
            self.fit()
        
        var = self.var(alpha)
        
        # Use numerical integration for ES
        def integrand(x):
            return x * self.pdf(np.array([x]))[0]
        
        # Integrate from -∞ to VaR
        es_value, _ = quad(integrand, -np.inf, var)
        es_value = es_value / alpha
        
        return es_value
    
    def moment(self, order: int) -> float:
        """
        Calculate the theoretical moment of the distribution.
        
        Parameters
        ----------
        order : int
            Order of the moment
            
        Returns
        -------
        float
            Theoretical moment
        """
        if not self._fitted:
            self.fit()
        
        alpha = self.params['alpha']
        beta = self.params['beta']
        c = self.params['c']
        mu = self.params['mu']
        
        if order == 1:
            return mu if alpha > 1 else np.nan
        elif order == 2 and alpha == 2:
            return c**2
        else:
            return np.nan  # Higher moments may not exist