import numpy as np
import pytest

from giant.particles import (
    nearest_known_pdg,
    particle_mass_charge,
    particle_phys_array,
)


def test_photon_massless_neutral():
    mass, charge = particle_mass_charge(22)
    assert mass == pytest.approx(0.0)
    assert charge == pytest.approx(0.0)


def test_electron_mass_charge():
    mass, charge = particle_mass_charge(11)
    assert mass == pytest.approx(0.51099895069, rel=1e-6)
    assert charge == pytest.approx(-1.0)


def test_positron_is_charge_conjugate_of_electron():
    mass_e, charge_e = particle_mass_charge(11)
    mass_p, charge_p = particle_mass_charge(-11)
    assert mass_p == pytest.approx(mass_e)
    assert charge_p == pytest.approx(-charge_e)


def test_proton_mass_charge():
    mass, charge = particle_mass_charge(2212)
    assert mass == pytest.approx(938.27208943, rel=1e-6)
    assert charge == pytest.approx(1.0)


def test_neutrino_unmeasured_mass_treated_as_zero():
    """PDG tables store an unmeasured neutrino mass as None -- must not
    propagate a None/NaN into a physical conditioning feature."""
    mass, charge = particle_mass_charge(12)
    assert mass == pytest.approx(0.0)
    assert charge == pytest.approx(0.0)


def test_ground_state_nucleus_resolved_via_particle_package():
    """He-4 (Z=2, A=4) is a common nuclide in `particle`'s ground-state table."""
    mass, charge = particle_mass_charge(1000020040)
    assert charge == pytest.approx(2.0)
    assert mass == pytest.approx(
        4 * 931.494, rel=0.05
    )  # near A*amu, binding-energy-corrected


def test_nuclear_isomer_falls_back_to_z_a_decode():
    """An excited/isomer nuclear code (nonzero trailing digit) is absent from
    `particle`'s ground-state-only nuclide table -- confirmed necessary for
    ~32% of the nuclear codes in the multi-material dataset. Fe-56 isomer:
    Z=26, A=56, isomer level 1 -> pdgid 1000260561."""
    pdg = 1000260561
    mass, charge = particle_mass_charge(pdg)
    assert charge == pytest.approx(26.0)
    assert mass == pytest.approx(56 * 931.494, rel=1e-6)


def test_invalid_pdg_code_raises():
    with pytest.raises(ValueError):
        particle_mass_charge(999999999)


def test_particle_mass_charge_is_cached():
    particle_mass_charge.cache_clear()
    particle_mass_charge(22)
    particle_mass_charge(22)
    info = particle_mass_charge.cache_info()
    assert info.hits >= 1


def test_particle_phys_array_shape_and_dtype():
    arr = particle_phys_array(np.array([22, 11, 2212]))
    assert arr.shape == (3, 2)
    assert arr.dtype == np.float32
    np.testing.assert_allclose(arr[0], [0.0, 0.0])
    np.testing.assert_allclose(arr[2], [938.27208943, 1.0], rtol=1e-5)


# ── nearest_known_pdg (reporting-only nearest-neighbour label) ──────────────


def test_nearest_known_pdg_exact_match():
    candidates = [22, 11, -11, 2212, 2112]
    mass_e, charge_e = particle_mass_charge(11)
    result = nearest_known_pdg(np.array([mass_e]), np.array([charge_e]), candidates)
    assert result[0] == 11


def test_nearest_known_pdg_prioritises_charge_match():
    """Charge is a small conserved quantum number and should usually match
    exactly even when the queried mass is noisy/imperfect."""
    candidates = [22, 11, -11, 2212]
    # Close to electron mass but not exact, positive charge like the positron.
    result = nearest_known_pdg(np.array([0.6]), np.array([1.0]), candidates)
    assert result[0] == -11


def test_nearest_known_pdg_empty_candidates_raises():
    with pytest.raises(ValueError):
        nearest_known_pdg(np.array([1.0]), np.array([0.0]), [])


def test_nearest_known_pdg_shape():
    candidates = [22, 11, -11, 2212, 2112]
    n = 10
    result = nearest_known_pdg(
        np.random.default_rng(0).uniform(0, 1000, n),
        np.random.default_rng(1).uniform(-1, 1, n),
        candidates,
    )
    assert result.shape == (n,)
    assert set(result.tolist()) <= set(candidates)