"""Normal-Inverse Gaussian (NIG) distribution fitter."""

import numpy as np
import pandas as pd
from scipy.optimize import minimize
from scipy.special import kv as bessel_k  # Modified Bessel function of 2nd kind
from typing import Dict

from .base import DistributionFitter


class NormalInverseGaussian(DistributionFitter):
    """Normal-Inverse Gaussian distribution fitter for heavy-tailed returns."""
    
    def __init__(self, returns: np.ndarray):
        """
        Initialize NIG distribution fitter.
        
        Parameters
        ----------
        returns : np.ndarray
            Array of returns
        """
        super().__init__(returns)
        self.name = "NormalInverseGaussian"
    
    def _nig_pdf(self, x: np.ndarray, alpha: float, beta: float, 
                 delta: float, mu: float) -> np.ndarray:
        """
        Normal-Inverse Gaussian PDF.
        
        Parameters
        ----------
        x : np.ndarray
            Points to evaluate
        alpha : float
            Tail heaviness parameter (alpha > 0)
        beta : float
            Asymmetry parameter (|beta| < alpha)
        delta : float
            Scale parameter (delta > 0)
        mu : float
            Location parameter
            
        Returns
        -------
        np.ndarray
            PDF values
        """
        gamma = np.sqrt(alpha**2 - beta**2)
        
        z = (x - mu) / delta
        factor = alpha * delta * np.sqrt(1 + z**2)
        
        numerator = alpha * delta * bessel_k(1, factor)
        denominator = np.pi * np.sqrt(1 + z**2)
        
        pdf = (numerator / denominator) * np.exp(delta * gamma + beta * (x - mu))
        
        return pdf
    
    def _nig_loglikelihood(self, params: np.ndarray) -> float:
        """
        Negative log-likelihood for NIG distribution.
        
        Parameters
        ----------
        params : np.ndarray
            Array of [alpha, beta, delta, mu]
            
        Returns
        -------
        float
            Negative log-likelihood
        """
        alpha, beta, delta, mu = params
        
        # Parameter constraints
        if alpha <= 0 or delta <= 0 or abs(beta) >= alpha:
            return 1e10  # Large penalty
        
        try:
            pdf_vals = self._nig_pdf(self.returns, alpha, beta, delta, mu)
            if np.any(pdf_vals <= 0):
                return 1e10
            
            loglik = np.sum(np.log(pdf_vals))
            
            # Add penalties for extreme parameters
            penalty = 0
            if alpha > 10:
                penalty += (alpha - 10) * 100
            if delta > 10:
                penalty += (delta - 10) * 100
                
            return -(loglik - penalty)  # Negative for minimization
            
        except:
            return 1e10
    
    def fit(self) -> Dict[str, float]:
        """
        Fit NIG distribution using MLE.
        
        Returns
        -------
        dict
            Fitted parameters: alpha, beta, delta, mu
        """
        # Initial parameter estimates
        mu0 = np.mean(self.returns)
        sigma0 = np.std(self.returns)
        skew0 = pd.Series(self.returns).skew()
        
        # Heuristic initial values
        delta0 = sigma0 * 0.8
        alpha0 = 1.0 / (sigma0 * 0.5)
        
        if abs(skew0) > 0.1:
            beta0 = skew0 * alpha0 * 0.5
            # Ensure |beta| < alpha
            beta0 = np.sign(beta0) * min(abs(beta0), alpha0 * 0.95)
        else:
            beta0 = 0.0
        
        # Initial parameter vector
        x0 = np.array([alpha0, beta0, delta0, mu0])
        
        # Bounds: alpha>0, |beta|<alpha, delta>0, mu unrestricted
        bounds = [
            (1e-4, 50),      # alpha
            (-49, 49),       # beta (will be adjusted in constraint)
            (1e-4, 50),      # delta
            (None, None)     # mu
        ]
        
        # Constraint: |beta| < alpha
        def constraint(params):
            alpha, beta, delta, mu = params
            return alpha - abs(beta) - 1e-6
        
        constraints = [{'type': 'ineq', 'fun': constraint}]
        
        # Optimization
        result = minimize(
            self._nig_loglikelihood,
            x0,
            method='SLSQP',
            bounds=bounds,
            constraints=constraints,
            options={
                'ftol': 1e-8,
                'maxiter': 1000,
                'disp': False
            }
        )
        
        if not result.success:
            # Try different starting values
            x0_alt = np.array([alpha0 * 2, beta0 * 0.5, delta0 * 1.2, mu0])
            result_alt = minimize(
                self._nig_loglikelihood,
                x0_alt,
                method='SLSQP',
                bounds=bounds,
                constraints=constraints,
                options={'ftol': 1e-8, 'maxiter': 1000, 'disp': False}
            )
            
            if result_alt.fun < result.fun:
                result = result_alt
        
        alpha, beta, delta, mu = result.x
        
        self.params = {
            'alpha': alpha,
            'beta': beta,
            'delta': delta,
            'mu': mu,
            'gamma': np.sqrt(alpha**2 - beta**2)  # Convenience parameter
        }
        
        # Calculate log-likelihood, AIC, BIC
        self.loglikelihood = -result.fun
        self.aic = 2 * len(self.params) - 2 * self.loglikelihood
        self.bic = len(self.params) * np.log(len(self.returns)) - 2 * self.loglikelihood
        
        self._fitted = True
        return self.params
    
    def pdf(self, x: np.ndarray) -> np.ndarray:
        """
        Probability density function.
        
        Parameters
        ----------
        x : np.ndarray
            Points at which to evaluate PDF
            
        Returns
        -------
        np.ndarray
            PDF values
        """
        if not self._fitted:
            self.fit()
        
        return self._nig_pdf(x, **{k: self.params[k] 
                                 for k in ['alpha', 'beta', 'delta', 'mu']})
    
    def cdf(self, x: np.ndarray) -> np.ndarray:
        """
        Cumulative distribution function.
        
        Parameters
        ----------
        x : np.ndarray
            Points at which to evaluate CDF
            
        Returns
        -------
        np.ndarray
            CDF values
        """
        if not self._fitted:
            self.fit()
        
        # Use numerical integration for CDF
        from scipy.integrate import quad
        
        cdf_vals = np.zeros_like(x)
        for i, xi in enumerate(x):
            cdf_vals[i], _ = quad(self.pdf, -np.inf, xi)
        
        return cdf_vals
    
    def ppf(self, q: np.ndarray) -> np.ndarray:
        """
        Percent point function (inverse of CDF).
        
        Parameters
        ----------
        q : np.ndarray
            Quantiles to evaluate
            
        Returns
        -------
        np.ndarray
            PPF values
        """
        if not self._fitted:
            self.fit()
        
        # Use numerical root finding
        from scipy.optimize import brentq
        
        def cdf_func(x, target):
            return self.cdf(np.array([x]))[0] - target
        
        ppf_vals = np.zeros_like(q)
        
        # Find approximate range for root finding
        x_min = self.returns.min() * 3
        x_max = self.returns.max() * 3
        
        for i, qi in enumerate(q):
            try:
                ppf_vals[i] = brentq(cdf_func, x_min, x_max, args=(qi,))
            except:
                # Fallback to approximation
                ppf_vals[i] = self.params['mu'] + self.params['delta'] * np.tan(np.pi * (qi - 0.5))
        
        return ppf_vals
    
    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
        from scipy.integrate import quad
        
        def integrand(x):
            return x * self.pdf(x)
        
        # 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']
        delta = self.params['delta']
        mu = self.params['mu']
        
        if order == 1:
            return mu + delta * beta / np.sqrt(alpha**2 - beta**2)
        elif order == 2 and alpha > 2:
            gamma = np.sqrt(alpha**2 - beta**2)
            return delta / gamma + delta * beta**2 / (alpha**2 - beta**2)**(3/2)
        else:
            return np.nan