"""Neutron balance and reactivity calculations."""

from __future__ import annotations

from dataclasses import dataclass
import logging

from . import constants
from .fuel import fuel_reactivity_penalty
from .state import CoreState, clamp

LOGGER = logging.getLogger(__name__)


def temperature_feedback(temp: float) -> float:
    """Negative coefficient: higher temperature lowers reactivity."""
    reference = 900.0
    coefficient = -1.7e-5
    return coefficient * (temp - reference)


def xenon_poisoning(flux: float) -> float:
    return min(0.01, 5e-10 * flux)


@dataclass
class NeutronDynamics:
    beta_effective: float = 0.0065
    delayed_decay_const: float = 0.08  # 1/s effective precursor decay
    external_source_coupling: float = 1e-6
    shutdown_bias: float = -0.014
    iodine_yield: float = 1e-6  # inventory units per MW*s
    iodine_decay_const: float = 1.0 / 66000.0  # ~18h
    xenon_decay_const: float = 1.0 / 33000.0  # ~9h
    xenon_burnout_coeff: float = 1e-13  # per n/cm2
    xenon_reactivity_coeff: float = 0.05

    def reactivity(self, state: CoreState, control_fraction: float, rod_banks: list[float] | None = None) -> float:
        if rod_banks:
            weights = constants.CONTROL_ROD_BANK_WEIGHTS
            worth = 0.0
            total = sum(weights)
            for w, pos in zip(weights, rod_banks):
                worth += w * (1.0 - clamp(pos, 0.0, 0.95) / 0.95)
            rod_term = constants.CONTROL_ROD_WORTH * worth / total
        else:
            rod_term = constants.CONTROL_ROD_WORTH * (1.0 - control_fraction)
        rho = (
            self.shutdown_bias +
            rod_term
            + temperature_feedback(state.fuel_temperature)
            - fuel_reactivity_penalty(state.burnup)
            - self.xenon_penalty(state)
        )
        return rho

    def flux_derivative(
        self,
        state: CoreState,
        rho: float,
        delayed_source: float,
        external_source_rate: float = 0.0,
        baseline_source: float = 1e5,
    ) -> float:
        generation_time = constants.NEUTRON_LIFETIME
        beta = self.beta_effective
        source_term = self.external_source_coupling * external_source_rate
        prompt = ((rho - beta) / generation_time) * state.neutron_flux
        return prompt + delayed_source + baseline_source + source_term

    def step(
        self,
        state: CoreState,
        control_fraction: float,
        dt: float,
        external_source_rate: float = 0.0,
        rod_banks: list[float] | None = None,
    ) -> None:
        rho = self.reactivity(state, control_fraction, rod_banks)
        rho = min(rho, 0.02)
        shutdown = control_fraction >= 0.95
        if shutdown:
            rho = min(rho, -0.04)
        baseline = 0.0 if shutdown else 1e5
        source = 0.0 if shutdown else external_source_rate
        rod_positions = rod_banks if rod_banks else [control_fraction] * len(constants.CONTROL_ROD_BANK_WEIGHTS)
        self._ensure_precursors(state, len(rod_positions))
        bank_factors = self._bank_factors(rod_positions)
        bank_betas = self._bank_betas(len(bank_factors))
        delayed_source = self._delayed_source(state, bank_factors)

        d_flux = self.flux_derivative(
            state,
            rho,
            delayed_source,
            external_source_rate=source,
            baseline_source=baseline,
        )
        state.neutron_flux = max(0.0, state.neutron_flux + d_flux * dt)
        state.reactivity_margin = rho
        self._update_precursors(state, bank_factors, bank_betas, dt)
        LOGGER.debug(
            "Neutronics: rho=%.5f, flux=%.2e n/cm2/s, d_flux=%.2e",
            rho,
            state.neutron_flux,
            d_flux,
        )

    def update_poisons(self, state: CoreState, dt: float) -> None:
        prod_I = max(0.0, state.power_output_mw) * self.iodine_yield
        decay_I = state.iodine_inventory * self.iodine_decay_const
        state.iodine_inventory = max(0.0, state.iodine_inventory + (prod_I - decay_I) * dt)
        prod_Xe = decay_I
        burn_Xe = state.neutron_flux * self.xenon_burnout_coeff
        decay_Xe = state.xenon_inventory * self.xenon_decay_const
        state.xenon_inventory = max(0.0, state.xenon_inventory + (prod_Xe - decay_Xe - burn_Xe) * dt)

    def xenon_penalty(self, state: CoreState) -> float:
        """Return delta-rho penalty from xenon inventory (positive magnitude)."""
        return self._xenon_penalty(state)

    def _xenon_penalty(self, state: CoreState) -> float:
        return min(0.05, state.xenon_inventory * self.xenon_reactivity_coeff)

    def _bank_betas(self, bank_count: int) -> list[float]:
        weights = list(constants.CONTROL_ROD_BANK_WEIGHTS)
        if bank_count != len(weights):
            weights = [1.0 for _ in range(bank_count)]
        total = sum(weights) if weights else 1.0
        return [self.beta_effective * (w / total) for w in weights]

    def _bank_factors(self, positions: list[float]) -> list[float]:
        factors: list[float] = []
        for pos in positions:
            insertion = clamp(pos, 0.0, 0.95)
            factors.append(max(0.0, 1.0 - insertion / 0.95))
        return factors

    def _ensure_precursors(self, state: CoreState, bank_count: int) -> None:
        if not state.delayed_precursors or len(state.delayed_precursors) != bank_count:
            state.delayed_precursors = [0.0 for _ in range(bank_count)]

    def _delayed_source(self, state: CoreState, bank_factors: list[float]) -> float:
        decay = self.delayed_decay_const
        return sum(decay * precursor * factor for precursor, factor in zip(state.delayed_precursors, bank_factors))

    def _update_precursors(
        self, state: CoreState, bank_factors: list[float], bank_betas: list[float], dt: float
    ) -> None:
        generation_time = constants.NEUTRON_LIFETIME
        decay = self.delayed_decay_const
        new_pools: list[float] = []
        for precursor, factor, beta in zip(state.delayed_precursors, bank_factors, bank_betas):
            production = (beta / generation_time) * state.neutron_flux * factor
            loss = decay * precursor
            updated = max(0.0, precursor + (production - loss) * dt)
            new_pools.append(updated)
        state.delayed_precursors = new_pools