"""SIR epidemiological model (discrete-time stochastic) for CrisisWorld."""

from __future__ import annotations

from math import exp

import numpy as np

from ..models import RegionState
from ._internal import EpiParams


def _advance_single_region(
    region: RegionState,
    params: EpiParams,
    rng: np.random.Generator,
) -> RegionState:
    """Intra-region SIR step. Returns updated region (new object)."""
    if region.infected == 0:
        return region

    pop = region.population
    s = pop - region.infected - region.recovered - region.deceased

    if s <= 0:
        # No susceptible left -- only recovery and death
        new_infections = 0
    else:
        effective_beta = params.beta * 0.5 if region.restricted else params.beta
        p_infect = 1.0 - exp(-effective_beta * region.infected / pop)
        p_infect = max(0.0, min(1.0, p_infect))
        new_infections = int(rng.binomial(s, p_infect))

    new_recoveries = int(rng.binomial(region.infected, params.gamma))
    new_deaths = int(rng.binomial(region.infected, params.mu))

    # Clamp: recoveries + deaths cannot exceed current infected
    total_outflow = new_recoveries + new_deaths
    if total_outflow > region.infected:
        scale = region.infected / total_outflow
        new_recoveries = int(new_recoveries * scale)
        new_deaths = region.infected - new_recoveries

    infected = max(0, region.infected + new_infections - new_recoveries - new_deaths)
    recovered = max(0, region.recovered + new_recoveries)
    deceased = max(0, region.deceased + new_deaths)

    # Final safety: ensure I+R+D <= population
    total = infected + recovered + deceased
    if total > pop:
        excess = total - pop
        infected = max(0, infected - excess)

    return region.model_copy(
        update={
            "infected": infected,
            "recovered": recovered,
            "deceased": deceased,
        }
    )


def _compute_spillovers(
    regions: tuple[RegionState, ...],
    adjacency: dict[str, list[str]],
    params: EpiParams,
    rng: np.random.Generator,
) -> dict[str, int]:
    """Compute inter-region spillover infections.

    Returns a mapping of region_id -> total incoming spillover count.
    """
    by_id = {r.region_id: r for r in regions}
    incoming: dict[str, int] = {r.region_id: 0 for r in regions}

    for region in regions:
        if region.infected == 0:
            continue
        neighbors = adjacency.get(region.region_id, [])
        for neighbor_id in neighbors:
            spillover = int(rng.binomial(region.infected, params.inter_region_spread))
            if region.restricted:
                spillover = int(spillover * 0.1)
            # Cap at target susceptible count
            target = by_id.get(neighbor_id)
            if target is None:
                continue
            target_s = target.population - target.infected - target.recovered - target.deceased
            spillover = min(spillover, max(0, target_s))
            incoming[neighbor_id] += spillover

    return incoming


def advance_epi_state(
    regions: tuple[RegionState, ...],
    adjacency: dict[str, list[str]],
    params: EpiParams,
    rng: np.random.Generator,
) -> tuple[RegionState, ...]:
    """Advance all regions by one SIR time step.

    1. Intra-region dynamics (infection, recovery, death).
    2. Inter-region spillover from adjacent infected regions.

    All stochastic draws use the provided ``rng`` for reproducibility.
    """
    # Early exit: nothing to do if no region is infected
    if all(r.infected == 0 for r in regions):
        return regions

    # Phase 1: intra-region dynamics
    updated = tuple(_advance_single_region(r, params, rng) for r in regions)

    # Phase 2: inter-region spillover (computed from *original* regions)
    incoming = _compute_spillovers(regions, adjacency, params, rng)

    # Apply spillover to updated regions
    result: list[RegionState] = []
    for r in updated:
        spill = incoming.get(r.region_id, 0)
        if spill <= 0:
            result.append(r)
            continue

        # Cap spillover at susceptible in the *updated* region
        s = r.population - r.infected - r.recovered - r.deceased
        spill = min(spill, max(0, s))
        new_infected = r.infected + spill

        # Final safety clamp
        total = new_infected + r.recovered + r.deceased
        if total > r.population:
            new_infected = r.population - r.recovered - r.deceased

        result.append(r.model_copy(update={"infected": max(0, new_infected)}))

    return tuple(result)