"""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

LOGGER = logging.getLogger(__name__)


def temperature_feedback(temp: float) -> float:
    """Negative coefficient: higher temperature lowers reactivity."""
    reference = 900.0
    coefficient = -1.5e-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_neutron_fraction: float = 0.0008
    external_source_coupling: float = 1e-6

    def reactivity(self, state: CoreState, control_fraction: float) -> float:
        rho = (
            0.02 * (1.0 - control_fraction)
            + temperature_feedback(state.fuel_temperature)
            - fuel_reactivity_penalty(state.burnup)
            - xenon_poisoning(state.neutron_flux)
        )
        return rho

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

    def step(self, state: CoreState, control_fraction: float, dt: float, external_source_rate: float = 0.0) -> None:
        rho = self.reactivity(state, control_fraction)
        d_flux = self.flux_derivative(state, rho, external_source_rate)
        state.neutron_flux = max(0.0, state.neutron_flux + d_flux * dt)
        state.reactivity_margin = rho
        LOGGER.debug(
            "Neutronics: rho=%.5f, flux=%.2e n/cm2/s, d_flux=%.2e",
            rho,
            state.neutron_flux,
            d_flux,
        )