import numpy as np
import pytest

from ase.build import bulk
from ase.units import fs

from dynasor.tools.structures import get_P_matrix
from phonopy import Phonopy
from phonopy.structure.atoms import PhonopyAtoms
from ase.calculators.emt import EMT
from ase import Atoms
from dynasor.tools.structures import find_permutation

from dynasor import ModeProjector


@pytest.fixture
def primitive():
    return bulk('Al', cubic=False)


@pytest.fixture
def supercell():
    repeat = 2
    unit = bulk('Al', cubic=True)
    return unit.repeat(repeat)


@pytest.fixture
def fc2(primitive, supercell):

    P = get_P_matrix(primitive.cell, supercell.cell)

    primitive_phonopy = PhonopyAtoms(numbers=primitive.numbers, cell=primitive.cell,
                                     positions=primitive.positions)
    phonon = Phonopy(primitive_phonopy, P.T)
    phonon.generate_displacements()

    forces = []
    for structure_ph in phonon.supercells_with_displacements:
        structure_ase = Atoms(symbols=structure_ph.symbols,
                              cell=structure_ph.cell,
                              scaled_positions=structure_ph.scaled_positions,
                              pbc=True)
        structure_ase.calc = EMT()
        forces.append(structure_ase.get_forces())

    phonon.forces = forces
    phonon.produce_force_constants()
    fc2_phonopy = phonon.force_constants

    supercell_phonopy = Atoms(
            cell=phonon.supercell.cell,
            numbers=phonon.supercell.numbers,
            positions=phonon.supercell.positions,
            pbc=True)

    p = find_permutation(supercell_phonopy, supercell)

    fc2 = fc2_phonopy[p][:, p]

    return fc2


@pytest.mark.filterwarnings('ignore:`product` is deprecated as of NumPy 1.25.0')
def test_modes(primitive, supercell, fc2):

    mp = ModeProjector(primitive, supercell, fc2)
    print(mp)

    print(repr(mp))

    qp = mp((0, 0, 1/2))

    qp = mp[1]

    list(mp)

    Nk = len(mp.q_reduced)
    Ns = len(primitive) * 3
    Nn = len(primitive)
    N = len(supercell)

    assert mp.q_cartesian.shape == (Nk, 3)
    assert mp.q_reduced.shape == (Nk, 3)
    assert mp.omegas.shape == (Nk, Ns)
    assert mp.polarizations.shape == (Nk, Ns, Nn, 3)
    assert mp.eigenmodes.shape == (Nk, Ns, N, 3)
    assert mp.kinetic_energies.shape == (Nk, Ns)
    assert mp.potential_energies.shape == (Nk, Ns)
    assert mp.virial_energies.shape == (Nk, Ns)

    assert mp.get_u().shape == (N, 3)
    assert mp.get_v().shape == (N, 3)
    assert mp.get_f().shape == (N, 3)

    assert mp.get_Q().shape == (Nk, Ns)
    assert mp.get_P().shape == (Nk, Ns)
    assert mp.get_F().shape == (Nk, Ns)

    u = np.random.normal(scale=0.1, size=(N, 3))
    v = np.random.normal(size=(N, 3))
    f = np.random.normal(size=(N, 3))

    mp.set_u(u)
    mp.set_v(v)
    mp.set_f(f)

    Q = mp.get_Q()
    P = mp.get_P()
    F = mp.get_F()

    mp.set_Q(Q)
    mp.set_P(P)
    mp.set_F(F)

    assert np.allclose(mp.get_u(), u)
    assert np.allclose(mp.get_v(), v)
    assert np.allclose(mp.get_f(), f)

    atoms = mp.supercell.to_ase()

    atoms.positions += u
    atoms.set_velocities(v / fs)
    atoms.arrays['forces'] = f

    mp.set_Q(0)
    mp.set_P(0)
    mp.set_F(0)

    mp.update_from_atoms(atoms)

    assert np.allclose(mp.get_u(), u)
    assert np.allclose(mp.get_v(), v)
    assert np.allclose(mp.get_f(), f)

    atoms = mp.get_atoms()

    assert np.allclose(atoms.get_positions() - mp.supercell.positions, u)
    assert np.allclose(atoms.get_velocities() * fs, v)
    assert np.allclose(atoms.get_forces(), f)
    assert np.allclose(atoms.get_potential_energy(), mp.potential_energies.sum())

    assert np.isclose(1/2 * np.einsum('ijab,ia,jb', fc2, u, u), mp.potential_energies.sum())
    assert np.isclose(atoms.get_kinetic_energy(), mp.kinetic_energies.sum())
    assert np.isclose(-1/2 * np.sum(u * f), mp.virial_energies.sum())

    # Test qpoints
    qp = mp[0]

    print(qp)

    repr(qp)

    assert np.allclose(qp.q_reduced, mp.q_reduced[qp.index])
    assert np.allclose(qp.q_cartesian, mp.q_cartesian[qp.index])
    assert np.allclose(qp.omegas, mp.omegas[qp.index])
    assert np.allclose(qp.polarizations, mp.polarizations[qp.index])
    assert np.allclose(qp.eigenmodes, mp.eigenmodes[qp.index])
    assert np.allclose(qp.potential_energies, mp.potential_energies[qp.index])
    assert np.allclose(qp.kinetic_energies, mp.kinetic_energies[qp.index])
    assert np.allclose(qp.virial_energies, mp.virial_energies[qp.index])
