"""Tests for exp_data.spectrum."""

import numpy as np
import pytest

from exp_data.spectrum import ExpSpectrum


class TestExpSpectrumBasic:
    """Basic ExpSpectrum construction and properties."""

    def test_construct_with_errors(self):
        spectrum = ExpSpectrum(
            energies=np.array([0.0, 1.0, 2.0]),
            counts=np.array([10.0, 20.0, 30.0]),
            errors=np.array([3.16, 4.47, 5.48]),
        )
        assert spectrum.n_channels == 3
        np.testing.assert_array_equal(spectrum.energies, [0.0, 1.0, 2.0])

    def test_automatic_poisson_errors(self):
        spectrum = ExpSpectrum(
            energies=np.array([0.0, 1.0, 2.0]),
            counts=np.array([10.0, 20.0, 30.0]),
        )
        expected = np.sqrt([10.0, 20.0, 30.0])
        np.testing.assert_allclose(spectrum.errors, expected)

    def test_zero_counts_get_zero_error(self):
        """Zero counts produce zero error (not sqrt(1))."""
        spectrum = ExpSpectrum(
            energies=np.array([0.0, 1.0]),
            counts=np.array([0.0, 100.0]),
        )
        assert spectrum.errors[0] == 0.0  # sqrt(0) = 0

    def test_shape_mismatch_raises(self):
        with pytest.raises(AssertionError):
            ExpSpectrum(
                energies=np.array([0.0, 1.0]),
                counts=np.array([10.0, 20.0, 30.0]),
            )

    def test_negative_energies_raises(self):
        with pytest.raises(AssertionError):
            ExpSpectrum(
                energies=np.array([-1.0, 0.0]),
                counts=np.array([10.0, 20.0]),
            )

    def test_negative_counts_allowed_for_difference_spectra(self):
        """Negative counts are allowed — they arise in cumulative subtraction.

        Difference spectra can have negative counts in channels where the true
        count is zero and Poisson noise causes the earlier cumulative to slightly
        exceed the later one.
        """
        spectrum = ExpSpectrum(
            energies=np.array([0.0, 1.0]),
            counts=np.array([10.0, -5.0]),
        )
        # Should not raise; errors use abs(counts) for negative values
        assert spectrum.errors[1] == np.sqrt(5.0)

    def test_total_counts(self):
        spectrum = ExpSpectrum(
            energies=np.array([0.0, 1.0, 2.0]),
            counts=np.array([10.0, 20.0, 30.0]),
        )
        assert spectrum.total_counts == 60.0

    def test_energy_range(self):
        spectrum = ExpSpectrum(
            energies=np.array([0.5, 1.5, 2.5]),
            counts=np.array([10.0, 20.0, 30.0]),
        )
        assert spectrum.energy_range == (0.5, 2.5)

    def test_normalize(self):
        spectrum = ExpSpectrum(
            energies=np.array([0.0, 1.0, 2.0]),
            counts=np.array([10.0, 20.0, 30.0]),
        )
        norm = spectrum.normalize()
        integral = np.trapezoid(norm, spectrum.energies)
        assert np.isclose(integral, 1.0, rtol=1e-10)


class TestExpSpectrumDeadTime:
    """Dead-time correction on ExpSpectrum."""

    def test_no_correction_when_dead_time_zero(self):
        spectrum = ExpSpectrum(
            energies=np.array([0.0, 1.0]),
            counts=np.array([10.0, 20.0]),
        )
        corrected = spectrum.apply_dead_time_correction()
        np.testing.assert_array_equal(corrected.counts, spectrum.counts)

    def test_correction_applied(self):
        spectrum = ExpSpectrum(
            energies=np.array([0.0, 1.0]),
            counts=np.array([10.0, 20.0]),
            dead_time=0.1,
        )
        corrected = spectrum.apply_dead_time_correction()
        expected = spectrum.counts / 0.9
        np.testing.assert_allclose(corrected.counts, expected, rtol=1e-10)
        assert corrected.dead_time == 0.0

    def test_live_time_corrected_rate(self):
        spectrum = ExpSpectrum(
            energies=np.array([0.0, 1.0]),
            counts=np.array([10.0, 20.0]),
            dead_time=0.1,
            live_time=10.0,
        )
        rate = spectrum.live_time_corrected_rate
        assert rate is not None
        # Total counts / (live_time * (1 - dead_time))
        expected = 30.0 / (10.0 * 0.9)
        assert np.isclose(rate, expected)


class TestExpSpectrumSerialization:
    """Dict round-trip for ExpSpectrum."""

    def test_to_from_dict(self):
        spectrum = ExpSpectrum(
            energies=np.array([0.0, 1.0, 2.0]),
            counts=np.array([10.0, 20.0, 30.0]),
            dead_time=0.05,
            live_time=100.0,
            source="Tc99",
            run_id="run_042",
            date="2025-01-15",
            metadata={"detector": "MMC", "temperature": "25 mK"},
        )

        recovered = ExpSpectrum.from_dict(spectrum.to_dict())

        np.testing.assert_allclose(recovered.energies, spectrum.energies)
        np.testing.assert_allclose(recovered.counts, spectrum.counts)
        assert recovered.dead_time == spectrum.dead_time
        assert recovered.live_time == spectrum.live_time
        assert recovered.source == spectrum.source
        assert recovered.run_id == spectrum.run_id
        assert recovered.date == spectrum.date
        assert recovered.metadata == spectrum.metadata