"""Fuel behavior, burnup, and decay heat modeling."""

from __future__ import annotations

from dataclasses import dataclass, field
import logging

from . import constants
from .atomic import Atom, AtomicPhysics, FissionEvent, make_atom
from .state import CoreState

LOGGER = logging.getLogger(__name__)


def fuel_reactivity_penalty(burnup: float) -> float:
    """Simplistic model that penalizes reactivity as burnup increases."""
    # Burnup is 0-1, penalty grows quadratically to mimic depletion.
    return 0.4 * burnup**2


def decay_heat_fraction(burnup: float) -> float:
    """Return remaining decay heat fraction relative to nominal power."""
    return min(0.07 + 0.2 * burnup, 0.15)


@dataclass
class FuelAssembly:
    enrichment: float  # fraction U-235
    mass_kg: float
    fissile_atom: Atom = field(default_factory=lambda: make_atom(92, 143))
    atomic_physics: AtomicPhysics = field(default_factory=AtomicPhysics)
    decay_activity_base: float = 1e16  # nominal decay events per second at burnup 0

    def available_energy_j(self, state: CoreState) -> float:
        fraction_remaining = max(0.0, 1.0 - state.burnup)
        return self.mass_kg * constants.FUEL_ENERGY_DENSITY * fraction_remaining

    def simulate_electron_hit(self) -> FissionEvent:
        return self.atomic_physics.electron_induced_fission(self.fissile_atom)

    def prompt_energy_rate(self, flux: float, control_fraction: float) -> tuple[float, FissionEvent]:
        """Compute MW thermal from prompt fission by sampling atomic physics."""
        event = self.simulate_electron_hit()
        effective_flux = max(0.0, flux * max(0.0, 1.0 - control_fraction))
        atoms = self.mass_kg / self.fissile_atom.atomic_mass_kg
        event_rate = max(
            0.0,
            effective_flux * constants.ELECTRON_FISSION_CROSS_SECTION * atoms * self.enrichment,
        )
        power_watts = event_rate * event.energy_mev * constants.MEV_TO_J
        power_mw = power_watts / constants.MEGAWATT
        LOGGER.debug(
            "Prompt fission products %s-%d + %s-%d yielding %.2f MW",
            event.products[0].symbol,
            event.products[0].mass_number,
            event.products[1].symbol,
            event.products[1].mass_number,
            power_mw,
        )
        return max(0.0, power_mw), event_rate, event

    def decay_reaction_effects(
        self, state: CoreState
    ) -> tuple[float, float, dict[str, float], dict[str, float]]:
        """Return (power MW, neutron source rate, product rates, particle rates)."""
        # Scale activity with burnup (fission products) and a small flux contribution.
        activity = self.decay_activity_base * (0.05 + state.burnup) * (1.0 + 0.00001 * state.neutron_flux)
        activity = max(0.0, min(activity, 1e20))

        reactions = [
            self.atomic_physics.alpha_decay(self.fissile_atom),
            self.atomic_physics.beta_minus_decay(self.fissile_atom),
            self.atomic_physics.beta_plus_decay(self.fissile_atom),
            self.atomic_physics.electron_capture(self.fissile_atom),
            self.atomic_physics.gamma_emission(self.fissile_atom),
            self.atomic_physics.spontaneous_fission(self.fissile_atom),
        ]
        weights = [1.0, 1.2, 0.5, 0.4, 1.6, 0.05]
        total_weight = sum(weights)

        power_watts = 0.0
        neutron_source = 0.0
        product_rates: dict[str, float] = {}
        particle_rates: dict[str, float] = {}

        for event, weight in zip(reactions, weights):
            rate = activity * (weight / total_weight)
            power_watts += rate * event.energy_mev * constants.MEV_TO_J
            for atom in event.products:
                product_rates[atom.symbol] = product_rates.get(atom.symbol, 0.0) + rate
            for particle in event.emitted_particles:
                particle_rates[particle] = particle_rates.get(particle, 0.0) + rate
                if particle == "n":
                    neutron_source += rate

        power_mw = power_watts / constants.MEGAWATT
        capped_power = min(power_mw, 0.1 * max(state.power_output_mw, 1.0))
        LOGGER.debug(
            "Decay reactions contribute %.2f MW (raw %.2f MW) with neutron source %.2e 1/s",
            capped_power,
            power_mw,
            neutron_source,
        )
        return max(0.0, capped_power), neutron_source, product_rates, particle_rates

    def decay_reaction_power(self, state: CoreState) -> float:
        """Compatibility shim: return only power contribution."""
        power_mw, _, _, _ = self.decay_reaction_effects(state)
        return power_mw