from __future__ import annotations

from dataclasses import dataclass
from typing import Dict, Iterable, List, Optional, Tuple

import numpy as np
import pandas as pd

from .sexual_debut import SexualDebutModel


@dataclass(slots=True)
class Campaign:
    age_from: int
    age_to: int
    coverage: float


def age_coverage(
    ages: Iterable[int],
    routine_coverage: float,
    vaccine_efficacy_before: float,
    vaccine_efficacy_after: float,
    campaigns: Optional[List[Campaign]],
    sexual_debut_model: SexualDebutModel,
    age_at_vaccination: int,
    protection_duration_years: Optional[int] = None,
) -> pd.DataFrame:
    """
    Calculate age-specific vaccination coverage.
    
    Args:
        ages: Iterable of ages to calculate coverage for
        routine_coverage: Routine vaccination coverage rate
        vaccine_efficacy_before: Efficacy for those not yet sexually active
        vaccine_efficacy_after: Efficacy for those already sexually active
        campaigns: Optional list of catch-up campaigns
        sexual_debut_model: Model for sexual debut probabilities
        age_at_vaccination: Target age for routine vaccination
        protection_duration_years: Years of protection (None = lifelong)
        
    Returns:
        DataFrame with 'age' and 'coverage' columns
    """
    ages_arr = np.asarray(list(ages), dtype=int)
    cov = np.zeros_like(ages_arr, dtype=float)

    # Routine at target age
    idx = np.where(ages_arr == age_at_vaccination)[0]
    if idx.size:
        p_debut = sexual_debut_model.cdf(age_at_vaccination)
        base_cov = (
            routine_coverage * vaccine_efficacy_before * (1.0 - p_debut)
            + routine_coverage * vaccine_efficacy_after * p_debut
        )
        cov[idx[0]] = base_cov
        
        # Apply coverage to all ages after vaccination, with waning if specified
        for i, age in enumerate(ages_arr):
            if age > age_at_vaccination:
                years_since_vac = age - age_at_vaccination
                
                # Calculate waning multiplier (1.0 if unlimited protection)
                if protection_duration_years is None:
                    waning_multiplier = 1.0  # Lifelong protection
                else:
                    waning_multiplier = calculate_waning_multiplier(
                        years_since_vac, protection_duration_years
                    )
                
                # Base coverage remains the same as at vaccination age
                # (coverage is set at vaccination and persists, modulated by waning)
                cov[i] = base_cov * waning_multiplier

    # Campaigns
    if campaigns:
        for camp in campaigns:
            mask = (ages_arr >= camp.age_from) & (ages_arr <= camp.age_to)
            if not np.any(mask):
                continue
            # Use debut at midpoint age
            mid_age = float((camp.age_from + camp.age_to) / 2.0)
            p_debut = sexual_debut_model.cdf(mid_age)
            base_camp_cov = (
                camp.coverage * vaccine_efficacy_before * (1.0 - p_debut)
                + camp.coverage * vaccine_efficacy_after * p_debut
            )
            
            # Apply waning for campaign coverage
            if protection_duration_years is not None:
                for i, age in enumerate(ages_arr):
                    if mask[i]:
                        years_since_vac = max(0, age - mid_age)
                        waning_multiplier = calculate_waning_multiplier(
                            years_since_vac, protection_duration_years
                        )
                        cov[i] += base_camp_cov * waning_multiplier
            else:
                cov[mask] += base_camp_cov

    cov = np.clip(cov, 0.0, 1.0)
    return pd.DataFrame({"age": ages_arr, "coverage": cov})


def calculate_waning_multiplier(
    years_since_vaccination: float, protection_duration_years: int
) -> float:
    """
    Calculate the waning multiplier based on years since vaccination.
    
    Uses a step function model:
    - Full protection (multiplier = 1.0) for years < protection_duration_years
    - No protection (multiplier = 0.0) for years >= protection_duration_years
    
    Args:
        years_since_vaccination: Years elapsed since vaccination
        protection_duration_years: Duration of vaccine protection in years
        
    Returns:
        Waning multiplier (0.0 to 1.0)
    """
    if years_since_vaccination < protection_duration_years:
        return 1.0
    else:
        return 0.0