import numpy as np
import pytest

from giant.data.loader import TopNMap
from giant.particles import (
    decode_embedding_nearest,
    decode_topn_class,
    invert_dense_map,
    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_skips_unresolvable_candidate():
    """One unresolvable code in the candidate set (e.g. a training-vocab
    entry giant.particles can't decode) must not crash the lookup -- it's
    simply excluded from the nearest-neighbour candidate pool."""
    candidates = [22, 11, -11, 999999999]
    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_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)


# ── invert_dense_map ─────────────────────────────────────────────────────


def test_invert_dense_map_round_trips():
    pdg_map = {22: 0, 11: 1, -11: 2, 2212: 3}
    inv = invert_dense_map(pdg_map)
    for pdg, idx in pdg_map.items():
        assert inv[idx] == pdg


# ── decode_topn_class ────────────────────────────────────────────────────


def _topn_fixture():
    # n_classes=4: photon/electron/positron get their own class (0,1,2),
    # everything else (proton, neutron) falls into "other" (class 3).
    class_map = {22: 0, 11: 1, -11: 2, 2212: 3, 2112: 3}
    other_members = {2212: 7, 2112: 3}
    return TopNMap(class_map=class_map, other_members=other_members), 4


def test_decode_topn_class_known_classes_are_exact():
    topn_map, n_classes = _topn_fixture()
    out = decode_topn_class(np.array([0, 1, 2]), topn_map, n_classes)
    np.testing.assert_array_equal(out, [22, 11, -11])


def test_decode_topn_class_other_modal_picks_most_frequent():
    topn_map, n_classes = _topn_fixture()
    out = decode_topn_class(np.array([3, 3]), topn_map, n_classes, other_policy="modal")
    assert (out == 2212).all()  # count 7 > 3


def test_decode_topn_class_other_drop_returns_zero_sentinel():
    topn_map, n_classes = _topn_fixture()
    out = decode_topn_class(np.array([3]), topn_map, n_classes, other_policy="drop")
    assert out[0] == 0


def test_decode_topn_class_other_sample_stays_within_members():
    topn_map, n_classes = _topn_fixture()
    rng = np.random.default_rng(0)
    out = decode_topn_class(
        np.full(50, 3), topn_map, n_classes, other_policy="sample", rng=rng
    )
    assert set(out.tolist()) <= {2212, 2112}


def test_decode_topn_class_unknown_other_policy_raises():
    topn_map, n_classes = _topn_fixture()
    with pytest.raises(ValueError):
        decode_topn_class(np.array([3]), topn_map, n_classes, other_policy="bogus")


def test_decode_topn_class_empty_other_members_raises():
    class_map = {22: 0, 11: 1}
    topn_map = TopNMap(class_map=class_map, other_members={})
    with pytest.raises(ValueError):
        decode_topn_class(np.array([1]), topn_map, 2, other_policy="sample")


def test_decode_topn_class_preserves_shape():
    topn_map, n_classes = _topn_fixture()
    idx = np.array([[0, 1], [2, 3]])
    out = decode_topn_class(idx, topn_map, n_classes, other_policy="modal")
    assert out.shape == (2, 2)


# ── decode_embedding_nearest ─────────────────────────────────────────────


def test_decode_embedding_nearest_exact_row_recovers_pdg():
    emb_weight = np.array([[1.0, 0.0], [0.0, 1.0], [-1.0, -1.0]])
    idx_to_pdg = {0: 22, 1: 11, 2: 2212}
    vectors = np.array([[0.0, 1.0], [-1.0, -1.0]])  # exact rows 1, 2
    pdg, dist = decode_embedding_nearest(vectors, emb_weight, idx_to_pdg)
    np.testing.assert_array_equal(pdg, [11, 2212])
    np.testing.assert_allclose(dist, [0.0, 0.0], atol=1e-8)


def test_decode_embedding_nearest_off_manifold_snaps_to_closest_row():
    emb_weight = np.array([[1.0, 0.0], [0.0, 1.0]])
    idx_to_pdg = {0: 22, 1: 11}
    vectors = np.array([[0.9, 0.2]])  # closer to row 0
    pdg, dist = decode_embedding_nearest(vectors, emb_weight, idx_to_pdg)
    assert pdg[0] == 22
    assert dist[0] > 0.0


def test_decode_embedding_nearest_preserves_leading_shape():
    emb_weight = np.array([[1.0, 0.0], [0.0, 1.0]])
    idx_to_pdg = {0: 22, 1: 11}
    vectors = np.random.default_rng(0).standard_normal((3, 4, 2))
    pdg, dist = decode_embedding_nearest(vectors, emb_weight, idx_to_pdg)
    assert pdg.shape == (3, 4)
    assert dist.shape == (3, 4)