"""Main fit function for Pyreto distributions."""

import pandas as pd
import numpy as np
from typing import Union, Dict, Optional, List, Tuple
from pathlib import Path

from .alpha_stable import AlphaStable
from .student_t import StudentT
from .normal_inverse_gaussian import NormalInverseGaussian
from .base import DistributionComparison


def fit(
    data: Union[str, pd.DataFrame, np.ndarray, List[float]],
    column: Optional[str] = None,
    models: Optional[List[str]] = None,
    alpha_levels: List[float] = None
) -> Dict:
    """
    Fit non-Gaussian distributions to return data.
    
    This is the main entry point for the pyreto library. It fits multiple
    heavy-tailed distributions to the provided return data and calculates
    Expected Shortfall (ES) and other risk metrics.
    
    Parameters
    ----------
    data : str, pd.DataFrame, np.ndarray, or list
        - If str: path to CSV file
        - If pd.DataFrame: DataFrame containing returns
        - If np.ndarray or list: array of returns
    column : str, optional
        Column name if data is DataFrame or CSV
    models : list of str, optional
        List of models to fit. Options: 'alpha_stable', 'student_t', 'nig'
        Default: all three models
    alpha_levels : list of float, optional
        Significance levels for VaR/ES calculation
        Default: [0.05, 0.01]
    
    Returns
    -------
    dict
        Dictionary containing:
        - 'returns': The processed returns array
        - 'models': Dictionary of fitted models
        - 'comparison': DistributionComparison object
        - 'summary': DataFrame comparing all models
        - 'plots': Dictionary of plot methods
    
    Examples
    --------
    >>> import pyreto
    >>> # Fit from CSV file
    >>> result = pyreto.dist.fit('returns.csv', column='daily_returns')
    >>> 
    >>> # Fit from numpy array
    >>> import numpy as np
    >>> returns = np.random.standard_t(3, 1000) * 0.01
    >>> result = pyreto.dist.fit(returns)
    >>> 
    >>> # Show summary
    >>> print(result['models']['student_t'].summary())
    >>> 
    >>> # Plot comparisons
    >>> fig = result['comparison'].plot_comparison()
    >>> fig.show()
    """
    if alpha_levels is None:
        alpha_levels = [0.05, 0.01]
    
    # Load and process data
    returns = _load_returns(data, column)
    
    # Set default models
    if models is None:
        models = ['alpha_stable', 'student_t', 'nig']
    
    # Initialize models
    model_instances = {}
    comparison = DistributionComparison(returns)
    
    model_mapping = {
        'alpha_stable': AlphaStable,
        'student_t': StudentT,
        'nig': NormalInverseGaussian,
        'normal_inverse_gaussian': NormalInverseGaussian
    }
    
    for model_name in models:
        model_name_lower = model_name.lower()
        if model_name_lower in model_mapping:
            model_class = model_mapping[model_name_lower]
            model = model_class(returns)
            model.fit()
            model_instances[model_name] = model
            comparison.add_model(model, model_name)
        else:
            raise ValueError(f"Unknown model: {model_name}. "
                           f"Available: {list(model_mapping.keys())}")
    
    # Create summary comparison
    summary_df = comparison.compare_risk_metrics()
    
    # Helper function to print comprehensive summary
    def print_summary():
        print(f"\n{'='*70}")
        print("PYRETO TAIL RISK ANALYSIS - NON-GAUSSIAN DISTRIBUTION FIT")
        print(f"{'='*70}")
        print(f"\nData Summary:")
        print(f"  Number of observations: {len(returns)}")
        print(f"  Mean: {np.mean(returns):.6f}")
        print(f"  Std Dev: {np.std(returns):.6f}")
        print(f"  Skewness: {pd.Series(returns).skew():.6f}")
        print(f"  Kurtosis: {pd.Series(returns).kurtosis():.6f}")
        
        # Test for normality
        from scipy.stats import jarque_bera, shapiro
        try:
            jb_stat, jb_pvalue = jarque_bera(returns)
            print(f"\n  Jarque-Bera Test:")
            print(f"    Statistic: {jb_stat:.4f}")
            print(f"    P-value: {jb_pvalue:.6f}")
            print(f"    Normal? {'No' if jb_pvalue < 0.05 else 'Yes'} (α=0.05)")
        except:
            print(f"\n  Jarque-Bera Test: Could not compute")
        
        print(f"\n{'='*70}")
        print("FITTED MODELS:")
        print(f"{'='*70}")
        
        for name, model in model_instances.items():
            print(f"\n{name.upper()}:")
            for param, value in model.params.items():
                print(f"  {param}: {value:.6f}")
        
        print(f"\n{'='*70}")
        print("RISK METRICS COMPARISON:")
        print(f"{'='*70}")
        print(summary_df.round(6).to_string())
        
        print(f"\n{'='*70}")
        print("FAT TAIL EVIDENCE:")
        print(f"{'='*70}")
        
        # Calculate tail ratios compared to normal
        from scipy.stats import norm
        normal_var_5 = norm.ppf(0.05)
        normal_var_1 = norm.ppf(0.01)
        
        for name, model in model_instances.items():
            ratio_5 = model.var(0.05) / normal_var_5
            ratio_1 = model.var(0.01) / normal_var_1
            print(f"\n{name}:")
            print(f"  95% VaR is {ratio_5:.2f}x Normal VaR")
            print(f"  99% VaR is {ratio_1:.2f}x Normal VaR")
            if ratio_5 > 1.5:
                print(f"  → SIGNIFICANT FAT TAILS DETECTED")
        
        print(f"\n{'='*70}")
        print("NEXT STEPS:")
        print(f"{'='*70}")
        print("1. Use result['comparison'].plot_comparison() to visualize")
        print("2. Access individual models: result['models']['student_t']")
        print("3. Generate individual plots: model.plot_fit()")
        print(f"{'='*70}\n")
    
    # Helper function to create all plots
    def create_plots():
        return {
            'comparison': comparison.plot_comparison(),
            **{name: model.plot_fit() for name, model in model_instances.items()}
        }
    
    return {
        'returns': returns,
        'models': model_instances,
        'comparison': comparison,
        'summary': summary_df,
        'print_summary': print_summary,
        'create_plots': create_plots
    }


def _load_returns(data: Union[str, pd.DataFrame, np.ndarray, List[float]], 
                 column: Optional[str] = None) -> np.ndarray:
    """Load and validate returns data."""
    
    if isinstance(data, str):
        # Load from CSV file
        df = pd.read_csv(data)
        if column is None:
            # Try to infer column
            numeric_cols = df.select_dtypes(include=[np.number]).columns
            if len(numeric_cols) == 1:
                column = numeric_cols[0]
            else:
                raise ValueError("Multiple numeric columns found. "
                               "Please specify 'column' parameter.")
        returns = df[column].dropna().values
        
    elif isinstance(data, pd.DataFrame):
        if column is None:
            numeric_cols = data.select_dtypes(include=[np.number]).columns
            if len(numeric_cols) == 1:
                column = numeric_cols[0]
            else:
                raise ValueError("Multiple numeric columns found. "
                               "Please specify 'column' parameter.")
        returns = data[column].dropna().values
        
    elif isinstance(data, np.ndarray):
        returns = data.flatten()
        
    elif isinstance(data, list):
        returns = np.array(data)
        
    else:
        raise TypeError(f"Unsupported data type: {type(data)}")
    
    # Validate returns
    if len(returns) == 0:
        raise ValueError("No returns data provided")
    
    if not np.isfinite(returns).all():
        raise ValueError("Returns contain non-finite values (NaN, Inf)")
    
    return returns
