"""Base class for distribution fitters."""

from abc import ABC, abstractmethod
from typing import Dict, Tuple, Optional, Any
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import norm


class DistributionFitter(ABC):
    """Abstract base class for distribution fitting."""
    
    def __init__(self, returns: np.ndarray):
        """
        Initialize with returns data.
        
        Parameters
        ----------
        returns : np.ndarray
            Array of returns (daily returns expected)
        """
        self.returns = np.array(returns)
        self.params: Optional[Dict[str, float]] = None
        self._fitted = False
        
    @abstractmethod
    def fit(self) -> Dict[str, float]:
        """
        Fit the distribution to the returns.
        
        Returns
        -------
        dict
            Dictionary of fitted parameters
        """
        pass
    
    @abstractmethod
    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
        """
        pass
    
    @abstractmethod
    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
        """
        pass
    
    @abstractmethod
    def es(self, alpha: float = 0.05) -> float:
        """
        Calculate Expected Shortfall (Conditional VaR).
        
        Parameters
        ----------
        alpha : float
            Significance level (default: 0.05 for 95% confidence)
            
        Returns
        -------
        float
            Expected Shortfall value
        """
        pass
    
    def var(self, alpha: float = 0.05) -> float:
        """
        Calculate Value at Risk.
        
        Parameters
        ----------
        alpha : float
            Significance level (default: 0.05 for 95% confidence)
            
        Returns
        -------
        float
            VaR value
        """
        if not self._fitted:
            self.fit()
        
        # VaR is the quantile at alpha level
        return self.ppf(alpha)
    
    @abstractmethod
    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
        """
        pass
    
    def plot_fit(self, bins: int = 50, figsize: Tuple[int, int] = (12, 8)) -> plt.Figure:
        """
        Plot the fitted distribution against histogram of data.
        
        Parameters
        ----------
        bins : int
            Number of histogram bins
        figsize : tuple
            Figure size (width, height)
            
        Returns
        -------
        plt.Figure
            Matplotlib figure object
        """
        if not self._fitted:
            self.fit()
        
        fig, (ax1, ax2) = plt.subplots(2, 1, figsize=figsize)
        
        # Plot 1: Histogram vs PDF
        ax1.hist(self.returns, bins=bins, density=True, alpha=0.6, 
                 label='Empirical Returns', color='steelblue')
        
        x_range = np.linspace(self.returns.min(), self.returns.max(), 1000)
        pdf_fitted = self.pdf(x_range)
        ax1.plot(x_range, pdf_fitted, 'r-', lw=2, 
                label=f'{self.__class__.__name__} PDF')
        
        # Add Normal distribution for comparison
        mu, sigma = norm.fit(self.returns)
        pdf_normal = norm.pdf(x_range, mu, sigma)
        ax1.plot(x_range, pdf_normal, 'g--', lw=2, label='Normal PDF')
        
        ax1.set_xlabel('Returns')
        ax1.set_ylabel('Density')
        ax1.set_title(f'Distribution Fit: {self.__class__.__name__} vs Normal')
        ax1.legend()
        ax1.grid(True, alpha=0.3)
        
        # Plot 2: QQ-Plot (focus on tails)
        sorted_returns = np.sort(self.returns)
        n = len(sorted_returns)
        theoretical_quantiles = norm.ppf((np.arange(1, n + 1) - 0.5) / n)
        
        ax2.scatter(theoretical_quantiles, sorted_returns, 
                   alpha=0.6, s=20, color='steelblue')
        
        # Add reference line
        z = np.polyfit(theoretical_quantiles, sorted_returns, 1)
        p = np.poly1d(z)
        ax2.plot(theoretical_quantiles, p(theoretical_quantiles), 
                "r--", alpha=0.8, lw=2)
        
        ax2.set_xlabel('Normal Theoretical Quantiles')
        ax2.set_ylabel('Sample Quantiles')
        ax2.set_title('QQ-Plot (Deviation from Normal = Fat Tails)')
        ax2.grid(True, alpha=0.3)
        
        plt.tight_layout()
        return fig
    
    def summary(self) -> str:
        """
        Generate a summary of the fitted distribution.
        
        Returns
        -------
        str
            Formatted summary string
        """
        if not self._fitted:
            self.fit()
        
        summary_lines = [
            f"{'='*60}",
            f"{self.__class__.__name__} Distribution Fit Summary",
            f"{'='*60}",
            "",
            "Parameters:",
        ]
        
        for param, value in self.params.items():
            summary_lines.append(f"  {param}: {value:.6f}")
        
        summary_lines.extend([
            "",
            "Risk Metrics:",
            f"  VaR (5%): {self.var(0.05):.6f}",
            f"  ES (5%): {self.es(0.05):.6f}",
            f"  VaR (1%): {self.var(0.01):.6f}",
            f"  ES (1%): {self.es(0.01):.6f}",
            "",
            "Comparison with Normal:",
            f"  95% VaR Ratio: {self.var(0.05) / norm.ppf(0.05):.3f}x",
            f"  99% VaR Ratio: {self.var(0.01) / norm.ppf(0.01):.3f}x",
        ])
        
        if hasattr(self, 'aic'):
            summary_lines.append(f"  AIC: {self.aic:.2f}")
        if hasattr(self, 'bic'):
            summary_lines.append(f"  BIC: {self.bic:.2f}")
        
        summary_lines.append(f"{'='*60}")
        
        return "\n".join(summary_lines)


class DistributionComparison:
    """Utility class to compare multiple distributions."""
    
    def __init__(self, returns: np.ndarray):
        """
        Initialize with returns data.
        
        Parameters
        ----------
        returns : np.ndarray
            Array of returns
        """
        self.returns = np.array(returns)
        self.fitted_models: Dict[str, DistributionFitter] = {}
    
    def add_model(self, model: DistributionFitter, name: str):
        """
        Add a fitted model to comparison.
        
        Parameters
        ----------
        model : DistributionFitter
            Fitted distribution model
        name : str
            Name for the model
        """
        if not model._fitted:
            model.fit()
        self.fitted_models[name] = model
    
    def compare_risk_metrics(self) -> pd.DataFrame:
        """
        Compare risk metrics across all models.
        
        Returns
        -------
        pd.DataFrame
            DataFrame with risk metrics comparison
        """
        metrics = []
        
        for name, model in self.fitted_models.items():
            metrics.append({
                'Model': name,
                'VaR_5%': model.var(0.05),
                'ES_5%': model.es(0.05),
                'VaR_1%': model.var(0.01),
                'ES_1%': model.es(0.01),
            })
        
        return pd.DataFrame(metrics).set_index('Model')
    
    def plot_comparison(self, figsize: Tuple[int, int] = (14, 10)) -> plt.Figure:
        """
        Plot comparison of all fitted distributions.
        
        Parameters
        ----------
        figsize : tuple
            Figure size
            
        Returns
        -------
        plt.Figure
            Matplotlib figure object
        """
        fig, axes = plt.subplots(2, 2, figsize=figsize)
        axes = axes.flatten()
        
        # Plot 1: PDF comparison
        ax = axes[0]
        x_range = np.linspace(self.returns.min(), self.returns.max(), 1000)
        
        for name, model in self.fitted_models.items():
            ax.plot(x_range, model.pdf(x_range), lw=2, label=name)
        
        # Add histogram
        ax.hist(self.returns, bins=50, density=True, alpha=0.3, 
                label='Empirical', color='gray')
        
        ax.set_xlabel('Returns')
        ax.set_ylabel('Density')
        ax.set_title('PDF Comparison')
        ax.legend()
        ax.grid(True, alpha=0.3)
        
        # Plot 2: Tail comparison (left tail)
        ax = axes[1]
        x_tail = np.linspace(self.returns.min(), np.percentile(self.returns, 5), 500)
        
        for name, model in self.fitted_models.items():
            ax.plot(x_tail, model.pdf(x_tail), lw=2, label=f'{name} PDF')
        
        ax.set_xlabel('Returns')
        ax.set_ylabel('Density')
        ax.set_title('Left Tail Detail')
        ax.legend()
        ax.grid(True, alpha=0.3)
        
        # Plot 3: Risk metrics
        ax = axes[2]
        metrics_df = self.compare_risk_metrics()
        metrics_df.plot(kind='bar', ax=ax)
        ax.set_title('Risk Metrics Comparison')
        ax.set_ylabel('Value')
        ax.set_xticklabels(ax.get_xticklabels(), rotation=45)
        ax.grid(True, alpha=0.3)
        
        # Plot 4: VaR/ES ratios
        ax = axes[3]
        normal_var_5 = norm.ppf(0.05)
        normal_var_1 = norm.ppf(0.01)
        
        ratios = []
        model_names = []
        
        for name, model in self.fitted_models.items():
            ratios.extend([
                model.var(0.05) / normal_var_5,
                model.var(0.01) / normal_var_1
            ])
            model_names.extend([f'{name}_5%', f'{name}_1%'])
        
        ax.bar(range(len(ratios)), ratios)
        ax.set_xticks(range(len(ratios)))
        ax.set_xticklabels(model_names, rotation=45)
        ax.axhline(y=1, color='r', linestyle='--', label='Normal Ratio = 1')
        ax.set_ylabel('Ratio to Normal')
        ax.set_title('Fat Tail Severity (VaR Ratio)')
        ax.legend()
        ax.grid(True, alpha=0.3)
        
        plt.tight_layout()
        return fig
