import numpy as np
import pytest
from ase.build import bulk
from dynasor.tools.acfs import compute_acf, fermi_dirac, smoothing_function
from dynasor.tools.acfs import psd_from_acf, psd_from_time_signal, psd_from_acf_2d
from dynasor.tools.structures import align_structure, _align_a_onto_xy
from dynasor.tools.structures import get_offset_index


@pytest.fixture
def signal_complex():
    np.random.seed(42)
    N = 1000
    x = np.random.random((N, ))
    y = np.random.random((N, ))
    Z_t = x + 1j * y
    return Z_t


def test_numpy_scipy_acf(signal_complex):
    delta_t = 0.25
    t1, acf1 = compute_acf(signal_complex, delta_t=delta_t, method='numpy')
    t2, acf2 = compute_acf(signal_complex, delta_t=delta_t, method='scipy')
    assert np.allclose(t1, t2)
    assert np.allclose(acf1, acf2)


def test_compute_acf_invalid_arguments(signal_complex):
    with pytest.raises(ValueError):
        t1, acf1 = compute_acf(signal_complex, method='asd')


def test_fermi_dirac_time_function(signal_complex):
    signal = signal_complex.real
    time = np.arange(0, 1000, 1)
    t_0 = 500
    t_width = 20
    f = fermi_dirac(time, t_0, t_width)
    signal_damped = f * signal
    assert np.isclose(signal_damped[0], signal[0])
    assert np.isclose(signal_damped[-1], 0.0)


def test_smoothing_function():
    data = np.array([1, 2, 2.5, 4, 5.5, 6.5, 7, 8, 7, 2])

    res1 = smoothing_function(data, window_size=1, window_type='boxcar')
    assert np.allclose(res1, data)

    res2 = smoothing_function(data, window_size=2, window_type='boxcar')
    assert len(res2) == len(data)
    assert np.allclose(res2[0], 1)
    assert np.allclose(res2[1], 1.5)
    assert np.allclose(res2[2], 2.25)

    res3 = smoothing_function(data, window_size=3, window_type='boxcar')
    assert len(res3) == len(data)
    assert np.allclose(res3[0], 1.5)
    assert np.allclose(res3[1], 5.5/3)
    assert np.allclose(res3[2], 8.5/3)
    assert np.allclose(res3[-1], 4.5)


def test_align_structure():

    # randomly rotated Al FCC cubic cell
    atoms = bulk('Al', 'fcc', cubic=True).repeat(2)
    ref_cell = atoms.cell[:].copy()
    atoms.rotate(42, [1, 2, 5], rotate_cell=True)
    print(atoms.cell[:])
    align_structure(atoms)
    assert np.isclose(atoms.cell[0, 1], 0)
    assert np.isclose(atoms.cell[0, 2], 0)
    assert np.isclose(atoms.cell[1, 2], 0)
    assert np.allclose(atoms.cell, ref_cell)
    print(atoms.cell[:])

    # BCC primtive cell
    atoms = bulk('Ti', 'bcc', a=3.3).repeat(2)
    print(atoms.cell[:])
    align_structure(atoms)
    assert np.isclose(atoms.cell[0, 1], 0)
    assert np.isclose(atoms.cell[0, 2], 0)
    assert np.isclose(atoms.cell[1, 2], 0)
    print(atoms.cell[:])

    # strange cell which fails with very low tolerance
    cell = np.array([[6.66725037e+01, -7.15537500e-02, 3.71540563e-08],
                     [1.40460691e-01,  6.63955753e+01, 5.54413709e-02],
                     [5.91136223e-02, -3.04975220e-02, 4.73122339e+01]])
    atoms = bulk('Al', 'fcc', cubic=True).repeat(2)
    atoms.set_cell(cell, scale_atoms=True)
    align_structure(atoms, atol=1e-5)


def test_align_structure_error_cases():

    # a nearly along z-axis — norm_a_xy < atol
    atoms = bulk('Al', 'fcc', cubic=True)
    cell = atoms.cell.array.copy()
    cell[0] = [0.0, 0.0, 4.05]
    atoms.set_cell(cell, scale_atoms=True)
    with pytest.raises(AssertionError, match='nearly parallel to the z-axis'):
        _align_a_onto_xy(atoms, atol=1e-5)

    # a has near-zero length — norm_a < atol
    atoms = bulk('Al', 'fcc', cubic=True)
    cell = atoms.cell.array.copy()
    cell[0] = [1e-10, 0.0, 0.0]
    atoms.set_cell(cell, scale_atoms=True)
    with pytest.raises(AssertionError, match='near-zero length'):
        _align_a_onto_xy(atoms, atol=1e-5)

    # a[2] < 0 — exercises the sign = +1 rotation branch
    atoms = bulk('Al', 'fcc', cubic=True).repeat(2)
    atoms.rotate(30, [0, 1, 0], rotate_cell=True)  # tilts a so that a[2] < 0
    assert atoms.cell[0, 2] < 0
    align_structure(atoms)
    assert np.isclose(atoms.cell[0, 1], 0)
    assert np.isclose(atoms.cell[0, 2], 0)
    assert np.isclose(atoms.cell[1, 2], 0)


def test_index_offset():
    prim = bulk('Ti', 'hcp')
    prim.set_array('basis_index', np.array([0, 1]))
    supercell = prim.repeat((4, 3, 2))

    # check indices
    offset, index = get_offset_index(prim, supercell)
    assert np.allclose(index, supercell.get_array('basis_index'))

    # fail if positions dont match
    supercell.positions[0, 0] += 0.03
    with pytest.raises(ValueError):
        offset, index = get_offset_index(prim, supercell)

    # works again if using higher tolerenaces
    offset, index = get_offset_index(prim, supercell, tol=0.05)
    assert np.allclose(index, supercell.get_array('basis_index'))


def test_psd_from_acf():
    acf_vals = [np.array([1.0, 0.5, 0.25]), np.array([1.0, 0.5, 0.25, 0.18])]
    for acf in acf_vals:
        dt = 2.0

        freqs, psd = psd_from_acf(acf, dt=dt, even=True)

        assert np.isrealobj(freqs)
        assert np.isrealobj(psd)

        signal = np.hstack((acf, acf[:0:-1]))
        freqs_expected = 2 * np.pi * np.fft.rfftfreq(len(signal), dt)
        psd_expected = (dt * np.fft.rfft(signal)).real

        assert np.allclose(freqs, freqs_expected)
        assert np.allclose(psd, psd_expected)


def test_psd_from_time_signal_real_matches_manual_rfft():
    x = np.array([0.5, 1.5, -0.5, 2.0, 1.0])
    dt = 0.25

    freqs, psd = psd_from_time_signal(x, dt=dt, complex=False)

    fft = np.fft.rfft(x)
    freqs_expected = 2 * np.pi * np.fft.rfftfreq(len(x), dt)
    psd_expected = dt * np.abs(fft) ** 2 / len(x)

    assert np.allclose(freqs, freqs_expected)
    assert np.allclose(psd, psd_expected)


def test_psd_from_time_signal_complex_matches_manual_fft():
    x = np.array([1.0 + 0.0j, 0.0 + 1.0j, -1.0 + 0.5j, 0.5 - 0.25j, 2.0 + 1.0j])
    dt = 0.25

    freqs, psd = psd_from_time_signal(x, dt=dt, complex=True)

    fft = np.fft.fft(x)
    freqs_expected = 2 * np.pi * np.fft.fftfreq(len(x), dt)
    psd_expected = dt * np.abs(fft) ** 2 / len(x)

    assert np.allclose(freqs, freqs_expected)
    assert np.allclose(psd, psd_expected)


def test_psd_from_acf_2d():
    acf_2d_vals = [
        np.array([[1.0, 0.5, 0.25],
                  [1.0, 0.4, 0.1]]),
        np.array([[1.0, 0.5, 0.25, 0.18],
                  [1.0, 0.3, 0.2, 0.05]])
    ]
    dt = 2.0

    for acf_2d in acf_2d_vals:
        for even in [True, False]:
            freqs_2d, psd_2d = psd_from_acf_2d(acf_2d, dt=dt, even=even)

            psd_1d = []
            for acf in acf_2d:
                freqs_1d, psd = psd_from_acf(acf, dt=dt, even=even)
                psd_1d.append(psd)

            psd_1d = np.array(psd_1d)

            assert np.allclose(freqs_2d, freqs_1d)
            assert np.allclose(psd_2d, psd_1d)


def test_psd_from_acf_zero_signal():
    acf = np.zeros(10)
    freqs, psd = psd_from_acf(acf, dt=1.0, even=True)
    assert np.allclose(psd, 0)


def test_psd_from_acf_2d_zero_signal():
    acf = np.zeros((5, 10))
    freqs, psd = psd_from_acf_2d(acf, dt=1.0, even=True)
    assert np.allclose(psd, 0)
