"""Unit tests for NeutronScatteringLengths, XRayFormFactors, and
ElectronScatteringFactors, verified against tabulated reference values
from the primary sources cited in each class's docstring.
"""
import numpy as np
import pytest
from numpy import pi

from dynasor.post_processing import (
    NeutronScatteringLengths,
    XRayFormFactors,
    ElectronScatteringFactors,
)


# ---------------------------------------------------------------------------
# NeutronScatteringLengths
# Source: NIST database / Neutron News 3, 26 (1992)
# Units: fm
# ---------------------------------------------------------------------------

class TestNeutronScatteringLengths:

    def test_natural_hydrogen_coherent_length(self):
        """b_coh for natural H equals the abundance-weighted average.

        1H: abundance=99.985%, b_coh=-3.7406 fm
        2H: abundance=0.015%,  b_coh= 6.671  fm
        3H: abundance=0.0,     b_coh= 4.792  fm
        """
        expected = (0.99985 * (-3.7406) + 0.00015 * 6.671 + 0.0 * 4.792)
        nsl = NeutronScatteringLengths(['H'])
        assert np.isclose(nsl.get_weight_coh('H').real, expected, rtol=1e-5)

    def test_natural_hydrogen_incoherent_length(self):
        """b_inc for natural H equals the abundance-weighted average.

        1H: abundance=99.985%, b_inc=25.274 fm
        2H: abundance=0.015%,  b_inc= 4.04  fm
        """
        expected = 0.99985 * 25.274 + 0.00015 * 4.04
        nsl = NeutronScatteringLengths(['H'])
        assert np.isclose(nsl.get_weight_incoh('H').real, expected, rtol=1e-5)

    def test_natural_oxygen_coherent_length(self):
        """b_coh for natural O equals the abundance-weighted average.

        16O: abundance=99.762%, b_coh=5.803 fm
        17O: abundance=0.038%,  b_coh=5.78  fm
        18O: abundance=0.2%,    b_coh=5.84  fm
        """
        expected = 0.99762 * 5.803 + 0.00038 * 5.78 + 0.002 * 5.84
        nsl = NeutronScatteringLengths(['O'])
        assert np.isclose(nsl.get_weight_coh('O').real, expected, rtol=1e-5)

    def test_natural_oxygen_incoherent_length_nearly_zero(self):
        """b_inc for natural O is nearly zero (only 17O contributes)."""
        nsl = NeutronScatteringLengths(['O'])
        assert abs(nsl.get_weight_incoh('O')) < 0.01

    def test_pure_protium_gives_tabulated_b_coh(self):
        """Selecting pure 1H returns the tabulated value b_coh = -3.7406 fm."""
        nsl = NeutronScatteringLengths(['H'], abundances={'H': {1: 1.0}})
        assert np.isclose(nsl.get_weight_coh('H').real, -3.7406, rtol=1e-4)

    def test_pure_deuterium_gives_tabulated_b_coh(self):
        """Selecting pure 2H (D) returns the tabulated value b_coh = 6.671 fm."""
        nsl = NeutronScatteringLengths(['H'], abundances={'H': {2: 1.0}})
        assert np.isclose(nsl.get_weight_coh('H').real, 6.671, rtol=1e-4)

    def test_pure_deuterium_gives_tabulated_b_inc(self):
        """Selecting pure 2H (D) returns the tabulated value b_inc = 4.04 fm."""
        nsl = NeutronScatteringLengths(['H'], abundances={'H': {2: 1.0}})
        assert np.isclose(nsl.get_weight_incoh('H').real, 4.04, rtol=1e-4)

    def test_multiple_species_loaded_independently(self):
        """Loading H and O together returns the same values as loading each separately."""
        nsl_both = NeutronScatteringLengths(['H', 'O'])
        nsl_H = NeutronScatteringLengths(['H'])
        nsl_O = NeutronScatteringLengths(['O'])
        assert np.isclose(nsl_both.get_weight_coh('H'), nsl_H.get_weight_coh('H'))
        assert np.isclose(nsl_both.get_weight_coh('O'), nsl_O.get_weight_coh('O'))

    def test_supports_incoherent_is_true(self):
        nsl = NeutronScatteringLengths(['O'])
        assert nsl.supports_incoherent

    def test_supports_currents_is_false(self):
        nsl = NeutronScatteringLengths(['O'])
        assert not nsl.supports_currents

    def test_unknown_species_raises_value_error(self):
        with pytest.raises(ValueError, match='requested species Xx'):
            NeutronScatteringLengths(['Xx'])

    def test_invalid_isotope_in_abundances_raises_value_error(self):
        with pytest.raises(ValueError):
            NeutronScatteringLengths(['H'], abundances={'H': {99: 1.0}})

    def test_abundances_not_summing_to_one_raises_value_error(self):
        with pytest.raises(ValueError):
            NeutronScatteringLengths(['H'], abundances={'H': {1: 0.5}})

    def test_parameters_dataframe_contains_expected_columns(self):
        nsl = NeutronScatteringLengths(['H', 'O'])
        cols = nsl.parameters.columns.tolist()
        for col in ['species', 'isotope', 'abundance', 'b_coh', 'b_inc']:
            assert col in cols

    def test_abundances_property_sums_to_one_per_species(self):
        nsl = NeutronScatteringLengths(['H', 'O'])
        for species, isotope_abundances in nsl.abundances.items():
            assert np.isclose(sum(isotope_abundances.values()), 1.0)


# ---------------------------------------------------------------------------
# XRayFormFactors
# Source: Waasmaier & Kirfel, Acta Cryst. A51, 416 (1995)
# f0(q) = sum_i a_i * exp(-b_i * s^2) + c,  s = q / (4*pi)
# ---------------------------------------------------------------------------

class TestXRayFormFactors:

    def test_oxygen_f0_at_q_zero_equals_atomic_number(self):
        """f0(q=0) = sum(a_i) + c ≈ Z for all elements.

        For O (Z=8): sum of a_i + c = 7.9997 (within parametrization accuracy).
        """
        ff = XRayFormFactors(['O'])
        # q→0 limit: use a very small q to avoid division issues
        f0 = ff.get_weight_coh('O', 1e-6)
        assert np.isclose(f0, 8.0, atol=0.01)

    def test_iron_f0_at_q_zero_equals_atomic_number(self):
        """f0(q=0) ≈ Z=26 for Fe."""
        ff = XRayFormFactors(['Fe'])
        f0 = ff.get_weight_coh('Fe', 1e-6)
        assert np.isclose(f0, 26.0, atol=0.02)

    def test_oxygen_f0_at_s_equals_one_matches_manual_calculation(self):
        """f0(q=4*pi) for O matches manual sum of Waasmaier-1995 coefficients.

        At s=1 (q=4*pi rad/Å), using tabulated O coefficients:
        a=[2.960427, 2.508818, 0.637853, 0.722838, 1.142756]
        b=[14.182259, 5.936858, 0.112726, 34.958481, 0.390240]
        c=0.027014
        """
        a = [2.960427, 2.508818, 0.637853, 0.722838, 1.142756]
        b = [14.182259, 5.936858, 0.112726, 34.958481, 0.390240]
        c = 0.027014
        s2 = 1.0  # s = q/(4*pi) = 1 when q = 4*pi
        expected = c + sum(ai * np.exp(-bi * s2) for ai, bi in zip(a, b))

        ff = XRayFormFactors(['O'])
        assert np.isclose(ff.get_weight_coh('O', 4 * pi), expected, rtol=1e-6)

    def test_f0_decreases_with_q(self):
        """The form factor f0 is a monotonically decreasing function of |q|."""
        ff = XRayFormFactors(['O'])
        q_values = np.linspace(0.1, 20.0, 20)
        f_values = [ff.get_weight_coh('O', q) for q in q_values]
        assert all(f_values[i] > f_values[i + 1] for i in range(len(f_values) - 1))

    def test_itc_2006_parametrization_loads_and_gives_reasonable_values(self):
        """The 'itc-2006' parametrization gives f0(q→0) ≈ Z."""
        ff = XRayFormFactors(['O', 'Fe'], source='itc-2006')
        assert np.isclose(ff.get_weight_coh('O', 1e-6), 8.0, atol=0.05)
        assert np.isclose(ff.get_weight_coh('Fe', 1e-6), 26.0, atol=0.05)

    def test_hydrogen_returns_zero_with_warning(self):
        """H has no parametrization; class should warn and return ~0."""
        with pytest.warns(UserWarning, match='H'):
            ff = XRayFormFactors(['H'])
        assert np.isclose(ff.get_weight_coh('H', 5.0), 0.0)

    def test_unknown_species_raises_value_error(self):
        with pytest.raises(ValueError, match='Missing'):
            XRayFormFactors(['Xx'])

    def test_supports_currents_is_false(self):
        ff = XRayFormFactors(['O'])
        assert not ff.supports_currents

    def test_supports_incoherent_is_false(self):
        ff = XRayFormFactors(['O'])
        assert not ff.supports_incoherent

    def test_parameters_dataframe_contains_expected_columns(self):
        ff = XRayFormFactors(['O', 'Fe'])
        cols = ff.parameters.columns.tolist()
        for col in ['a1', 'b1', 'c']:
            assert col in cols

    def test_unknown_source_raises_value_error(self):
        with pytest.raises(ValueError, match='Unknown'):
            XRayFormFactors(['O'], source='unknown-source')


# ---------------------------------------------------------------------------
# ElectronScatteringFactors
# Source: Peng et al., Acta Cryst. A52, 257 (1996); Peng, Acta Cryst. A54, 481 (1998)
# f(q) = sum_i a_i * exp(-b_i * s^2),  s = q / (4*pi)
# Ions additionally have a Coulomb offset: 0.023934 * charge / s^2
# ---------------------------------------------------------------------------

class TestElectronScatteringFactors:

    def test_oxygen_f_at_s_equals_one_matches_manual_calculation(self):
        """f(q=4*pi) for O matches manual sum of Peng-1996 (smax6) coefficients.

        Tabulated O coefficients (smax6):
        a=[0.0365, 0.1729, 0.5805, 0.8814, 0.3121]
        b=[0.0652, 0.6184, 2.9449, 9.6298, 28.2194]
        """
        a = [0.0365, 0.1729, 0.5805, 0.8814, 0.3121]
        b = [0.0652, 0.6184, 2.9449, 9.6298, 28.2194]
        s2 = 1.0  # s=1 when q=4*pi
        expected = sum(ai * np.exp(-bi * s2) for ai, bi in zip(a, b))

        esf = ElectronScatteringFactors(['O'])
        assert np.isclose(esf.get_weight_coh('O', 4 * pi), expected, rtol=1e-5)

    def test_iron_f_at_s_equals_one_matches_manual_calculation(self):
        """f(q=4*pi) for Fe matches manual sum of Peng-1996 (smax6) coefficients.

        Tabulated Fe coefficients (smax6):
        a=[0.1929, 0.8239, 1.8689, 2.3694, 1.9060]
        b=[0.1087, 1.0806, 4.7637, 22.850, 76.7309]
        """
        a = [0.1929, 0.8239, 1.8689, 2.3694, 1.9060]
        b = [0.1087, 1.0806, 4.7637, 22.850, 76.7309]
        s2 = 1.0
        expected = sum(ai * np.exp(-bi * s2) for ai, bi in zip(a, b))

        esf = ElectronScatteringFactors(['Fe'])
        assert np.isclose(esf.get_weight_coh('Fe', 4 * pi), expected, rtol=1e-5)

    def test_ionic_species_o2minus_includes_coulomb_offset(self):
        """f(q=4*pi) for O2- includes the ionic Coulomb offset 0.023934*charge/s^2.

        Tabulated O2- coefficients (Peng-1998):
        a=[0.0421, 0.2100, 0.8520, 1.8200, 1.1700]
        b=[0.0609, 0.5590, 2.9600, 11.500, 37.700]
        charge = -2
        """
        a = [0.0421, 0.2100, 0.8520, 1.8200, 1.1700]
        b = [0.0609, 0.5590, 2.9600, 11.500, 37.700]
        charge = -2
        s2 = 1.0
        expected = 0.023934 * charge / s2 + sum(ai * np.exp(-bi * s2) for ai, bi in zip(a, b))

        esf = ElectronScatteringFactors(['O2-'])
        assert np.isclose(esf.get_weight_coh('O2-', 4 * pi), expected, rtol=1e-5)

    def test_elemental_and_ionic_species_can_be_loaded_together(self):
        """Elemental and ionic species can be instantiated in the same object."""
        esf = ElectronScatteringFactors(['O', 'Fe', 'O2-'])
        assert esf.get_weight_coh('O', 4 * pi) > 0
        assert esf.get_weight_coh('Fe', 4 * pi) > 0
        assert esf.get_weight_coh('O2-', 4 * pi) > 0

    def test_smax2_parametrization_gives_reasonable_values(self):
        """The 'smax2' parametrization gives reasonable values within its range."""
        esf = ElectronScatteringFactors(['O'], parametrization='smax2')
        f = esf.get_weight_coh('O', 4 * pi)  # s=1, within smax2 range
        assert 0 < f < 5

    def test_f_decreases_with_q(self):
        """The scattering factor f is a decreasing function of |q| for neutral atoms."""
        esf = ElectronScatteringFactors(['O'])
        q_values = np.linspace(1.0, 20.0, 20)
        f_values = [esf.get_weight_coh('O', q) for q in q_values]
        assert all(f_values[i] > f_values[i + 1] for i in range(len(f_values) - 1))

    def test_unknown_species_raises_value_error(self):
        with pytest.raises(ValueError, match='Missing'):
            ElectronScatteringFactors(['Xx'])

    def test_unknown_parametrization_raises_value_error(self):
        with pytest.raises(ValueError, match='Unknown'):
            ElectronScatteringFactors(['O'], parametrization='unknown')

    def test_supports_currents_is_false(self):
        esf = ElectronScatteringFactors(['O'])
        assert not esf.supports_currents

    def test_supports_incoherent_is_false(self):
        esf = ElectronScatteringFactors(['O'])
        assert not esf.supports_incoherent

    def test_parameters_dataframe_contains_expected_columns(self):
        esf = ElectronScatteringFactors(['O', 'Fe'])
        cols = esf.parameters.columns.tolist()
        for col in ['a1', 'b1', 'a5', 'b5']:
            assert col in cols
