import pytest

from reactor_sim.neutronics import NeutronDynamics
from reactor_sim.state import CoreState


def _core_state(
    temperature: float = 950.0,
    flux: float = 2e7,
    burnup: float = 0.01,
    power: float = 500.0,
) -> CoreState:
    return CoreState(
        fuel_temperature=temperature,
        neutron_flux=flux,
        reactivity_margin=0.0,
        power_output_mw=power,
        burnup=burnup,
    )


def test_reactivity_increases_with_rod_withdrawal():
    dynamics = NeutronDynamics()
    state = _core_state()
    rho_full_out = dynamics.reactivity(state, control_fraction=0.0)
    rho_half = dynamics.reactivity(state, control_fraction=0.5)
    assert rho_full_out > rho_half


def test_poisons_accumulate_under_power():
    dynamics = NeutronDynamics()
    state = _core_state(power=800.0, flux=1e6)
    dynamics.update_poisons(state, dt=100.0)
    dynamics.update_poisons(state, dt=100.0)
    assert state.iodine_inventory > 0.0
    assert state.xenon_inventory > 0.0


def test_xenon_penalty_caps():
    dynamics = NeutronDynamics()
    state = _core_state()
    state.xenon_inventory = 50.0
    assert dynamics.xenon_penalty(state) == 0.05
    state.xenon_inventory = 0.2
    assert dynamics.xenon_penalty(state) == pytest.approx(0.01)