"""
Cholera Risk Assessment - Simulation Model

This module contains the core Monte Carlo simulation logic for evaluating 
transboundary cholera infection risk. The model multiplies six probability 
components representing steps in the disease transmission chain:

    P_final = P1 × P2 × P3 × P4 × P5 × P6

Where:
    P1  = Probability of Cholera outbreak in Bangladesh (Poisson process)
    P2  = Probability of Cholera infection in humans (Beta distribution)
    P3  = False Negative clinical examination in Bangladesh (1 - Sensitivity)
    P4  = False Negative clinical examination in India (1 - Sensitivity)
    P5  = Exposure risk due to unsafe water and sanitation access
    P6  = Probability of mortality (case fatality rate)
"""

import numpy as np
from scipy.stats import spearmanr
from typing import Dict, Any, Optional
from dataclasses import dataclass, field


@dataclass
class SimulationConfig:
    """Configuration parameters for the cholera risk simulation.
    
    Attributes:
        n_samples: Number of Monte Carlo samples to generate.
        seed: Random seed for reproducibility.
        t: Time period (years) for the Poisson process.
        total_outbreaks: Total cholera outbreaks observed.
        observation_years: Number of years of observation data.
        observation_months_per_year: Months per year in the observation window.
        beta_alpha: Alpha parameter for the Beta distribution (P2).
        beta_beta: Beta parameter for the Beta distribution (P2).
        sensitivity_bangladesh_low: Lower bound of CholKit RDT sensitivity (Bangladesh).
        sensitivity_bangladesh_high: Upper bound of CholKit RDT sensitivity (Bangladesh).
        sensitivity_india_low: Lower bound of CholKit RDT sensitivity (India).
        sensitivity_india_high: Upper bound of CholKit RDT sensitivity (India).
        unsafe_water_fraction: Fraction of population with unsafe water access.
        unsafe_sanitation_fraction: Fraction of population with unsafe sanitation.
        case_fatality_rate: Cholera case fatality rate.
    """
    n_samples: int = 10000
    seed: int = 42
    t: float = 1.0
    total_outbreaks: int = 2779
    observation_years: int = 22
    observation_months_per_year: int = 12
    beta_alpha: int = 1605
    beta_beta: int = 24618
    sensitivity_bangladesh_low: float = 0.549
    sensitivity_bangladesh_high: float = 0.906
    sensitivity_india_low: float = 0.884
    sensitivity_india_high: float = 0.999
    unsafe_water_fraction: float = 0.11
    unsafe_sanitation_fraction: float = 0.01
    case_fatality_rate: float = 0.03

    @property
    def lambda_rate(self) -> float:
        """Compute the Poisson rate parameter λ from observed data."""
        return self.total_outbreaks / (self.observation_years * self.observation_months_per_year)

    @property
    def P5(self) -> float:
        """Compute exposure risk from water and sanitation."""
        return self.unsafe_water_fraction * self.unsafe_sanitation_fraction

    @property
    def P6(self) -> float:
        """Return case fatality rate."""
        return self.case_fatality_rate


class CholeraSimulator:
    """Monte Carlo simulator for transboundary cholera risk assessment.
    
    This class encapsulates the entire simulation pipeline:
    1. Generate random samples from appropriate distributions
    2. Compute the compound probability of risk
    3. Calculate summary statistics  
    4. Perform sensitivity analysis using Spearman rank correlation
    
    Example:
        >>> sim = CholeraSimulator(n_samples=10000, seed=42)
        >>> results = sim.run_simulation()
        >>> print(f"Expected Risk: {results['summary']['expected_prob']:.6e}")
        >>> print(f"90% CI: [{results['summary']['ci_5']:.2e}, {results['summary']['ci_95']:.2e}]")
    """

    def __init__(self, n_samples: int = 10000, seed: int = 42, 
                 config: Optional[SimulationConfig] = None):
        """Initialize the simulator.
        
        Args:
            n_samples: Number of Monte Carlo samples. Overrides config if provided.
            seed: Random seed for reproducibility. Overrides config if provided.
            config: Optional SimulationConfig with custom parameters.
                    If None, default epidemiological parameters are used.
        """
        if config is not None:
            self.config = config
            self.config.n_samples = n_samples
            self.config.seed = seed
        else:
            self.config = SimulationConfig(n_samples=n_samples, seed=seed)

        self._results: Optional[Dict[str, Any]] = None

    def run_simulation(self) -> Dict[str, Any]:
        """Execute the Monte Carlo simulation.
        
        Returns:
            Dictionary containing:
                - 'samples': Array of final risk probability samples (n_samples,)
                - 'summary': Dict with expected_prob, ci_5, ci_95, min, max
                - 'inputs': Dict with P1, P2, P3, P4, P5, P6 arrays/values
                - 'sensitivity': Dict with Spearman correlations for stochastic inputs
                - 'config': The SimulationConfig used
        """
        np.random.seed(self.config.seed)
        cfg = self.config

        # ─── Step 1: Outbreak probability (deterministic from Poisson process) ───
        P1 = 1.0 - np.exp(-cfg.t * cfg.lambda_rate)

        # ─── Step 2: Infection rate in humans (stochastic - Beta distribution) ───
        P2 = np.random.beta(cfg.beta_alpha, cfg.beta_beta, cfg.n_samples)

        # ─── Step 3: False Negative in Bangladesh (stochastic - Uniform) ───
        P3_sensitivity = np.random.uniform(
            cfg.sensitivity_bangladesh_low, 
            cfg.sensitivity_bangladesh_high, 
            cfg.n_samples
        )
        P3 = 1.0 - P3_sensitivity  # False Negative = 1 - Sensitivity

        # ─── Step 4: False Negative in India (stochastic - Uniform) ───
        P4_sensitivity = np.random.uniform(
            cfg.sensitivity_india_low, 
            cfg.sensitivity_india_high, 
            cfg.n_samples
        )
        P4 = 1.0 - P4_sensitivity  # False Negative = 1 - Sensitivity

        # ─── Step 5: Exposure risk (deterministic) ───
        P5 = cfg.P5

        # ─── Step 6: Mortality (deterministic) ───
        P6 = cfg.P6

        # ─── Monte Carlo compound probability ───
        P_final = P1 * P2 * P3 * P4 * P5 * P6

        # ─── Summary statistics ───
        summary = {
            "expected_prob": float(np.mean(P_final)),
            "median": float(np.median(P_final)),
            "std_dev": float(np.std(P_final)),
            "ci_5": float(np.percentile(P_final, 5.0)),
            "ci_95": float(np.percentile(P_final, 95.0)),
            "ci_25": float(np.percentile(P_final, 25.0)),
            "ci_75": float(np.percentile(P_final, 75.0)),
            "min": float(np.min(P_final)),
            "max": float(np.max(P_final)),
        }

        # ─── Sensitivity analysis (Spearman rank correlation) ───
        sensitivity = self._sensitivity_analysis(P_final, P2, P3, P4)

        # ─── Assemble results ───
        self._results = {
            "samples": P_final,
            "summary": summary,
            "inputs": {
                "P1": P1,
                "P2": P2,
                "P3": P3,
                "P4": P4,
                "P5": P5,
                "P6": P6,
            },
            "input_means": {
                "P1": float(P1),
                "P2": float(np.mean(P2)),
                "P3": float(np.mean(P3)),
                "P4": float(np.mean(P4)),
                "P5": float(P5),
                "P6": float(P6),
            },
            "sensitivity": sensitivity,
            "config": cfg,
            "lambda_rate": cfg.lambda_rate,
        }

        return self._results

    @staticmethod
    def _sensitivity_analysis(P_final: np.ndarray, 
                               P2: np.ndarray, 
                               P3: np.ndarray, 
                               P4: np.ndarray) -> Dict[str, float]:
        """Perform sensitivity analysis using Spearman's rank correlation.
        
        Measures how strongly each stochastic input variable influences 
        the final risk output.
        
        Args:
            P_final: Array of final compound risk probabilities.
            P2: Array of infection rate samples (Beta distribution).
            P3: Array of false negative rate samples (Bangladesh).
            P4: Array of false negative rate samples (India).
            
        Returns:
            Dictionary mapping factor names to their Spearman ρ values.
        """
        factors = {
            "P2 (Infection Rate)": P2,
            "P3 (FN Bangladesh)": P3,
            "P4 (FN India)": P4,
        }

        correlations = {}
        for name, values in factors.items():
            rho, p_value = spearmanr(values, P_final)
            correlations[name] = {
                "rho": float(rho),
                "p_value": float(p_value),
            }

        return correlations

    @property
    def results(self) -> Optional[Dict[str, Any]]:
        """Access the last simulation results (None if not yet run)."""
        return self._results

    def __repr__(self) -> str:
        status = "run" if self._results else "not run"
        return (f"CholeraSimulator(n_samples={self.config.n_samples}, "
                f"seed={self.config.seed}, status={status})")


# ═══════════════════════════════════════════════════════════════════
# Convenience function for quick one-shot usage
# ═══════════════════════════════════════════════════════════════════

def quick_risk_assessment(n_samples: int = 10000, seed: int = 42) -> Dict[str, Any]:
    """Run a cholera risk assessment with default parameters.
    
    This is a convenience function for quick, one-off simulations.
    For repeated use or customized parameters, use CholeraSimulator directly.
    
    Args:
        n_samples: Number of Monte Carlo samples.
        seed: Random seed for reproducibility.
        
    Returns:
        Dictionary of simulation results (same format as CholeraSimulator.run_simulation).
    
    Example:
        >>> results = quick_risk_assessment(n_samples=50000)
        >>> print(f"Risk: {results['summary']['expected_prob']:.6e}")
    """
    sim = CholeraSimulator(n_samples=n_samples, seed=seed)
    return sim.run_simulation()
