import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy.stats import gumbel_r, genextreme
import seaborn as sns

class ExceedanceProbability:
    def __init__(self, historical_losses):
        self.losses = np.sort(historical_losses)
        self.n = len(historical_losses)
        
    def calculate_ep_curve(self):
        """Calculate exceedance probabilities for each loss level"""
        ranks = np.arange(1, self.n + 1)
        exceedance_probs = 1 - ranks / (self.n + 1)
        return self.losses, exceedance_probs
    
    def fit_extreme_value_distribution(self):
        """Fit Gumbel distribution to the tail of the loss distribution"""
        # Use top 20% of losses for extreme value fitting
        tail_losses = self.losses[int(0.8 * self.n):]
        params = genextreme.fit(tail_losses)
        return params
    
    def calculate_aal(self):
        """Calculate Average Annual Loss"""
        return np.mean(self.losses)
    
    def calculate_pml(self, return_period):
        """Calculate Probable Maximum Loss for given return period"""
        ep_curve = self.calculate_ep_curve()
        losses, probs = ep_curve
        # Interpolate to find loss at target probability
        target_prob = 1 / return_period
        pml = np.interp(target_prob, probs[::-1], losses[::-1])
        return pml
    
    def plot_ep_curve(self, save_path=None):
        """Plot the exceedance probability curve"""
        losses, probs = self.calculate_ep_curve()
        
        plt.figure(figsize=(12, 6))
        plt.semilogy(losses, probs, 'b-', label='Historical EP Curve')
        
        # Add return period labels
        return_periods = [10, 50, 100, 250, 500]
        for rp in return_periods:
            pml = self.calculate_pml(rp)
            plt.axvline(x=pml, color='r', linestyle='--', alpha=0.3)
            plt.text(pml, 0.1, f'{rp}yr', rotation=90, va='bottom')
        
        plt.xlabel('Loss Amount')
        plt.ylabel('Exceedance Probability')
        plt.title('Exceedance Probability Curve')
        plt.grid(True, which="both", ls="-", alpha=0.2)
        plt.legend()
        
        if save_path:
            plt.savefig(save_path)
        plt.show()

# Example usage with synthetic data
np.random.seed(42)
historical_losses = np.concatenate([
    np.random.lognormal(10, 1, 900),  # Small to medium losses
    np.random.lognormal(12, 1.5, 100)  # Large losses
])

ep = ExceedanceProbability(historical_losses)

# Calculate key metrics
aal = ep.calculate_aal()
pml_100yr = ep.calculate_pml(100)
pml_250yr = ep.calculate_pml(250)

print(f"Average Annual Loss: {aal:,.2f}")
print(f"100-year PML: {pml_100yr:,.2f}")
print(f"250-year PML: {pml_250yr:,.2f}")

# Plot EP curve
ep.plot_ep_curve('ep_curve.png')