from __future__ import annotations

import numpy as np
from numpy.testing import assert_allclose
from pytest import approx
from scipy import stats

from pymatgen.core.spectrum import Spectrum
from pymatgen.util.testing import MatSciTest


class TestSpectrum(MatSciTest):
    def setup_method(self):
        rng = np.random.default_rng()
        self.spec1 = Spectrum(np.arange(0, 10, 0.1), rng.standard_normal(100))
        self.spec2 = Spectrum(np.arange(0, 10, 0.1), rng.standard_normal(100))

        self.multi_spec1 = Spectrum(np.arange(0, 10, 0.1), rng.standard_normal((100, 2)))
        self.multi_spec2 = Spectrum(np.arange(0, 10, 0.1), rng.standard_normal((100, 2)))

    def test_normalize(self):
        self.spec1.normalize(mode="max")
        assert np.max(self.spec1.y) == approx(1)
        self.spec1.normalize(mode="sum")
        assert np.sum(self.spec1.y) == approx(1)

        self.multi_spec1.normalize(mode="sum")
        assert np.sum(self.multi_spec1.y[:, 0]) == approx(1)
        assert np.sum(self.multi_spec1.y[:, 1]) == approx(1)

        # XRD style mode.
        self.spec1.normalize(mode="max", value=100)
        assert np.max(self.spec1.y) == approx(100)

    def test_operators(self):
        scaled_spect = 3 * self.spec1 + self.spec2
        assert_allclose(scaled_spect.y, 3 * self.spec1.y + self.spec2.y)
        assert self.spec1.get_interpolated_value(0.05) == approx((self.spec1.y[0] + self.spec1.y[1]) / 2)

        scaled_spect = self.spec1 - self.spec2
        assert_allclose(scaled_spect.y, self.spec1.y - self.spec2.y)

        scaled_spect = self.spec1 / 3
        assert_allclose(scaled_spect.y, self.spec1.y / 3)

        scaled_spect = 3 * self.multi_spec1 + self.multi_spec2
        assert_allclose(scaled_spect.y, 3 * self.multi_spec1.y + self.multi_spec2.y)
        assert_allclose(
            self.multi_spec1.get_interpolated_value(0.05),
            (self.multi_spec1.y[0, :] + self.multi_spec1.y[1, :]) / 2,
        )

    def test_smear(self):
        y = np.zeros(100)
        y[25] = 1
        y[50] = 1
        y[75] = 1
        spec = Spectrum(np.linspace(-10, 10, 100), y)
        spec.smear(0.3)
        assert not np.allclose(y, spec.y)
        assert sum(y) == approx(sum(spec.y))

        # Test direct callable use of smearing.
        spec2 = Spectrum(np.linspace(-10, 10, 100), y)
        spec2.smear(0, func=lambda x: stats.norm.pdf(x, scale=0.3))
        assert_allclose(spec.y, spec2.y)

        spec = Spectrum(np.linspace(-10, 10, 100), y)
        spec.smear(0.3, func="lorentzian")
        assert not np.allclose(y, spec.y)
        assert sum(y) == approx(sum(spec.y))

        y = np.array(self.multi_spec1.y)
        self.multi_spec1.smear(0.2)
        assert not np.allclose(y, self.multi_spec1.y)
        assert_allclose(np.sum(y, axis=0), np.sum(self.multi_spec1.y, axis=0))

    def test_str(self):
        expected = "Spectrum\nx: [0.  0.1 0.2 0.3 0.4 0.5 0.6 0.7 0.8 0.9 1.  1.1 "
        assert str(self.spec1).startswith(expected)
        assert str(self.multi_spec1).startswith(expected)

    def test_copy(self):
        spec1copy = self.spec1.copy()
        spec1copy.y[0] += 1
        assert spec1copy.y[0] != self.spec1.y[0]
        assert spec1copy.y[1] == self.spec1.y[1]
