"""Tests for Pyreto distributions."""

import numpy as np
import pandas as pd
import pytest
import sys
import os

# Add src to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'src'))

from pyreto.distributions import StudentT, AlphaStable, NormalInverseGaussian
from pyreto.distributions.fit import fit, _load_returns


class TestDistributionBase:
    """Test base distribution functionality."""
    
    def setup_method(self):
        """Set up test data."""
        np.random.seed(42)
        # Generate heavy-tailed returns
        self.returns = np.random.standard_t(5, 1000) * 0.01
        
    def test_student_t_fit(self):
        """Test Student's T distribution fitting."""
        model = StudentT(self.returns)
        params = model.fit()
        
        assert 'df' in params
        assert 'loc' in params
        assert 'scale' in params
        assert 1 < params['df'] < 10  # Reasonable range
        assert model._fitted is True
        
    def test_student_t_risk_metrics(self):
        """Test Student's T VaR and ES calculation."""
        model = StudentT(self.returns)
        model.fit()
        
        var_5 = model.var(0.05)
        es_5 = model.es(0.05)
        var_1 = model.var(0.01)
        es_1 = model.es(0.01)
        
        assert var_5 < 0  # Should be negative (loss)
        assert es_5 < var_5  # ES should be more extreme than VaR
        assert var_1 < var_5  # 1% VaR should be more extreme than 5% VaR
        assert es_1 < es_5   # 1% ES should be more extreme than 5% ES
        
    def test_student_t_pdf_cdf_ppf(self):
        """Test Student's T PDF, CDF, and PPF."""
        model = StudentT(self.returns)
        model.fit()
        
        x_test = np.array([-0.02, -0.01, 0, 0.01, 0.02])
        
        pdf = model.pdf(x_test)
        cdf = model.cdf(x_test)
        ppf = model.ppf(np.array([0.05, 0.5, 0.95]))
        
        assert all(pdf > 0)
        assert all((cdf >= 0) & (cdf <= 1))
        assert len(ppf) == 3
        assert ppf[0] < ppf[1] < ppf[2]  # Monotonic increasing
        
    def test_student_t_moments(self):
        """Test Student's T moment calculation."""
        model = StudentT(self.returns)
        model.fit()
        
        mean = model.moment(1)
        # For df > 2, variance should be finite
        if model.params['df'] > 2:
            var = model.moment(2)
            assert var > 0
            
    def test_student_t_summary(self):
        """Test Student's T summary method."""
        model = StudentT(self.returns)
        model.fit()
        
        summary = model.summary()
        assert isinstance(summary, str)
        assert 'StudentT' in summary
        assert 'VaR' in summary
        assert 'ES' in summary
        
    def test_nig_fit(self):
        """Test Normal-Inverse Gaussian distribution fitting."""
        # NIG can be sensitive, use slightly larger sample
        returns_large = np.random.standard_t(5, 2000) * 0.01
        model = NormalInverseGaussian(returns_large)
        params = model.fit()
        
        assert 'alpha' in params
        assert 'beta' in params
        assert 'delta' in params
        assert 'mu' in params
        assert params['alpha'] > 0
        assert params['delta'] > 0
        assert abs(params['beta']) < params['alpha']  # NIG constraint
        assert model._fitted is True
        
    def test_nig_risk_metrics(self):
        """Test NIG VaR and ES calculation."""
        returns_large = np.random.standard_t(5, 2000) * 0.01
        model = NormalInverseGaussian(returns_large)
        model.fit()
        
        var_5 = model.var(0.05)
        es_5 = model.es(0.05)
        
        assert var_5 < 0
        assert es_5 < var_5
        assert np.isfinite(var_5)
        assert np.isfinite(es_5)
        
    def test_alpha_stable_fit(self):
        """Test Alpha-Stable distribution fitting."""
        # Alpha-stable needs more data for stable estimation
        returns_large = np.random.standard_t(3, 3000) * 0.01
        model = AlphaStable(returns_large)
        params = model.fit()
        
        assert 'alpha' in params
        assert 'beta' in params
        assert 'c' in params
        assert 'mu' in params
        assert 0.1 < params['alpha'] <= 2.0
        assert -1 <= params['beta'] <= 1
        assert params['c'] > 0
        assert model._fitted is True
        
    def test_alpha_stable_risk_metrics(self):
        """Test Alpha-Stable VaR and ES calculation."""
        returns_large = np.random.standard_t(3, 3000) * 0.01
        model = AlphaStable(returns_large)
        model.fit()
        
        var_5 = model.var(0.05)
        es_5 = model.es(0.05)
        
        assert var_5 < 0
        assert es_5 < var_5
        assert np.isfinite(var_5)
        # ES might be infinite for very heavy tails (alpha < 1)
        if model.params['alpha'] > 1:
            assert np.isfinite(es_5)
            
    def test_plot_methods(self):
        """Test plotting methods work without errors."""
        import matplotlib
        matplotlib.use('Agg')  # Non-interactive backend
        
        model = StudentT(self.returns)
        model.fit()
        
        fig = model.plot_fit()
        assert fig is not None
        assert len(fig.axes) == 2


class TestFitFunction:
    """Test the main fit function."""
    
    def setup_method(self):
        """Set up test data."""
        np.random.seed(42)
        self.returns = np.random.standard_t(5, 1000) * 0.01
        
    def test_fit_with_array(self):
        """Test fit function with numpy array."""
        result = fit(self.returns)
        
        assert 'models' in result
        assert 'comparison' in result
        assert 'summary' in result
        assert 'student_t' in result['models']
        assert 'alpha_stable' in result['models']
        assert 'nig' in result['models']
        
    def test_fit_with_list(self):
        """Test fit function with list."""
        returns_list = self.returns.tolist()
        result = fit(returns_list)
        
        assert result is not None
        assert len(result['models']) == 3
        
    def test_fit_with_dataframe(self):
        """Test fit function with DataFrame."""
        df = pd.DataFrame({'returns': self.returns})
        result = fit(df, column='returns')
        
        assert result is not None
        assert np.array_equal(result['returns'], self.returns)
        
    def test_fit_select_models(self):
        """Test fit function with model selection."""
        result = fit(self.returns, models=['student_t', 'nig'])
        
        assert 'student_t' in result['models']
        assert 'nig' in result['models']
        assert 'alpha_stable' not in result['models']
        
    def test_fit_custom_alpha(self):
        """Test fit function with custom alpha levels."""
        result = fit(self.returns, alpha_levels=[0.1, 0.05, 0.01])
        
        summary = result['summary']
        assert 'VaR_10%' in summary.columns
        assert 'VaR_1%' in summary.columns
        
    def test_load_returns_csv(self, tmp_path):
        """Test loading returns from CSV."""
        # Create temporary CSV
        csv_path = tmp_path / "returns.csv"
        df = pd.DataFrame({'returns': self.returns})
        df.to_csv(csv_path, index=False)
        
        result = fit(str(csv_path), column='returns')
        assert result is not None
        assert len(result['models']) == 3


class TestDistributionComparison:
    """Test distribution comparison functionality."""
    
    def setup_method(self):
        """Set up test data."""
        np.random.seed(42)
        self.returns = np.random.standard_t(5, 2000) * 0.01
        
    def test_comparison_risk_metrics(self):
        """Test risk metrics comparison."""
        from pyreto.distributions.base import DistributionComparison
        
        comparison = DistributionComparison(self.returns)
        
        model1 = StudentT(self.returns)
        model1.fit()
        comparison.add_model(model1, 'student_t')
        
        model2 = NormalInverseGaussian(self.returns[:2000])  # Use larger sample
        model2.fit()
        comparison.add_model(model2, 'nig')
        
        metrics_df = comparison.compare_risk_metrics()
        
        assert isinstance(metrics_df, pd.DataFrame)
        assert 'VaR_5%' in metrics_df.columns
        assert 'ES_5%' in metrics_df.columns
        assert len(metrics_df) == 2


class TestEdgeCases:
    """Test edge cases and error handling."""
    
    def test_load_returns_errors(self):
        """Test error handling in _load_returns."""
        # Empty data
        with pytest.raises(ValueError):
            _load_returns(np.array([]))
        
        # Invalid type
        with pytest.raises(TypeError):
            _load_returns(123)
        
        # DataFrame without column
        df = pd.DataFrame({'a': [1, 2, 3], 'b': [4, 5, 6]})
        with pytest.raises(ValueError):
            _load_returns(df)
    
    def test_invalid_model_name(self):
        """Test invalid model name."""
        returns = np.random.standard_t(5, 100) * 0.01
        
        with pytest.raises(ValueError):
            fit(returns, models=['invalid_model'])
    
    def test_nan_handling(self):
        """Test handling of NaN values."""
        returns_valid = np.random.standard_t(5, 100) * 0.01
        returns_with_nan = np.concatenate([returns_valid, [np.nan]])
        
        with pytest.raises(ValueError):
            _load_returns(returns_with_nan)


if __name__ == "__main__":
    # Run tests
    pytest.main([__file__, "-v"])
