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
from dynasor.modes.project_modes import project_modes
from dynasor.trajectory.trajectory_frame import TrajectoryFrame


@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


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])


# ------------------------------------------------------------------ shared fixture ---

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


# --------------------------------------------------------- constructor validation ---


def test_mode_projector_primitive_larger_than_supercell(primitive, supercell, fc2):
    """Swapping primitive and supercell (primitive > supercell) must raise."""
    with pytest.raises(ValueError, match='Primitive cell larger than supercell'):
        ModeProjector(supercell, primitive, np.zeros((len(primitive), len(primitive), 3, 3)))


def test_mode_projector_fc2_wrong_n(primitive, supercell):
    """Force constants with incompatible N must raise."""
    wrong_fc2 = np.zeros((10, 10, 3, 3))
    with pytest.raises(ValueError, match='not compatible with supercell size'):
        ModeProjector(primitive, supercell, wrong_fc2)


def test_mode_projector_fc2_wrong_shape(primitive, supercell):
    """Force constants with wrong trailing dimensions must raise."""
    N = len(supercell)
    wrong_fc2 = np.zeros((N, N, 3))
    with pytest.raises(ValueError, match='shape should be'):
        ModeProjector(primitive, supercell, wrong_fc2)


# --------------------------------------------------------- __getitem__ bounds ---

def test_mode_projector_getitem_out_of_bounds(mp):
    with pytest.raises(IndexError):
        mp[1000]


def test_qpoint_band_out_of_bounds(mp):
    qp = mp[0]
    with pytest.raises(IndexError):
        qp[1000]


# ----------------------------------------------- harmonic force accessors ---

def test_get_F_harmonic(mp):
    """get_F_harmonic returns -omega^2 * Q."""
    np.random.seed(5)
    N = mp.supercell.n_atoms
    mp.set_u(np.random.normal(scale=0.01, size=(N, 3)))

    F_harm = mp.get_F_harmonic()
    assert F_harm.shape == mp.get_Q().shape
    assert np.allclose(F_harm, -mp._w2 * mp.get_Q())


def test_get_f_harmonic(mp):
    """get_f_harmonic returns real forces consistent with get_atoms(harmonic_forces=True)."""
    np.random.seed(6)
    N = mp.supercell.n_atoms
    mp.set_u(np.random.normal(scale=0.01, size=(N, 3)))

    f_harm = mp.get_f_harmonic()
    assert f_harm.shape == (N, 3)
    assert f_harm.dtype == float

    # Must match the forces attached to the atoms object when harmonic_forces=True
    atoms = mp.get_atoms(harmonic_forces=True)
    assert np.allclose(f_harm, atoms.get_forces())


def test_get_atoms_harmonic_forces(mp):
    """get_atoms with harmonic_forces=True uses the harmonic force field."""
    np.random.seed(7)
    N = mp.supercell.n_atoms
    mp.set_u(np.random.normal(scale=0.01, size=(N, 3)))

    atoms = mp.get_atoms(harmonic_forces=True)
    assert np.allclose(atoms.get_forces(), mp.get_f_harmonic())


# -------------------------------------------------- set_Q/P/F symmetry guard ---

def test_set_Q_violates_conjugate_symmetry(mp):
    """set_Q must reject Q that does not satisfy Q(-q) = Q(q)*."""
    Q = mp.get_Q()
    Q[0, 0] = 1j  # imaginary part violates Q=Q* for a self-conjugate q-point
    with pytest.raises(ValueError, match='Q\\(-q\\)'):
        mp.set_Q(Q)


def test_set_P_violates_conjugate_symmetry(mp):
    P = mp.get_P()
    P[0, 0] = 1j
    with pytest.raises(ValueError, match='P\\(-q\\)'):
        mp.set_P(P)


def test_set_F_violates_conjugate_symmetry(mp):
    F = mp.get_F()
    F[0, 0] = 1j
    with pytest.raises(ValueError, match='F\\(-q\\)'):
        mp.set_F(F)


# ------------------------------------------------------------ write / read ---

def test_write_and_read(mp, tmp_path):
    """ModeProjector survives a round-trip through pickle via write/read."""
    path = tmp_path / 'mp.pkl'
    mp.write(str(path))
    mp2 = ModeProjector.read(str(path))

    assert np.allclose(mp2.q_reduced, mp.q_reduced)
    assert np.allclose(mp2.omegas, mp.omegas)
    assert np.allclose(mp2._X, mp._X)


# ------------------------------------------------- update_from_atoms edge cases ---

def test_update_from_atoms_large_displacement_warns(mp):
    """Displacements > 2 Å trigger a warning about possible permutation error."""
    N = mp.supercell.n_atoms
    atoms = mp.supercell.to_ase()
    atoms.positions[0, 0] += 2.5  # exceeds the 2 Å threshold
    atoms.set_velocities(np.zeros((N, 3)))

    with pytest.warns(UserWarning, match='Displacements larger than 2'):
        mp.update_from_atoms(atoms)


def test_update_from_atoms_forces_in_arrays(mp):
    """Forces stored in atoms.arrays (no calculator) are read correctly."""
    N = mp.supercell.n_atoms
    atoms = mp.supercell.to_ase()
    atoms.set_velocities(np.zeros((N, 3)))
    np.random.seed(8)
    forces = np.random.normal(size=(N, 3))
    atoms.arrays['forces'] = forces  # no calculator attached

    mp.update_from_atoms(atoms)
    assert np.allclose(mp.get_f(), forces)


def test_update_from_atoms_no_forces_uses_zeros(mp):
    """When no calculator and no forces array are present, forces default to zero."""
    N = mp.supercell.n_atoms
    atoms = mp.supercell.to_ase()
    atoms.set_velocities(np.zeros((N, 3)))
    # No calculator, no 'forces' in arrays

    mp.update_from_atoms(atoms)
    assert np.allclose(mp.get_f(), 0)


# --------------------------------------------------- Band object ---

def test_band_str_and_repr(mp):
    """Band __str__ and __repr__ execute without error."""
    np.random.seed(9)
    N = mp.supercell.n_atoms
    mp.set_u(np.random.normal(scale=0.01, size=(N, 3)))
    mp.set_v(np.random.normal(scale=0.1, size=(N, 3)))
    mp.set_f(np.random.normal(scale=0.01, size=(N, 3)))

    band = mp[0][0]
    s = str(band)
    assert 'Band' in s
    assert str(band) == repr(band)


def test_band_properties(mp):
    """All Band property accessors return values with correct shapes/types."""
    Np = mp.primitive.n_atoms
    N = mp.supercell.n_atoms

    band = mp[0][0]

    assert isinstance(band.index, int)
    assert band.index == 0
    assert isinstance(band.omega, (float, np.floating))
    assert band.polarization.shape == (Np, 3)
    assert band.eigenmode.shape == (N, 3)
    assert isinstance(band.potential_energy, (float, np.floating))
    assert isinstance(band.kinetic_energy, (float, np.floating))
    assert np.isscalar(band.virial_energy) or band.virial_energy.ndim == 0
    assert band.qpoint.index == 0


# --------------------------------------- ComplexCoordinate accessors and setters ---

def _first_non_gamma_q(mp):
    """Return (q_index, band_index) for the first non-Gamma q-point."""
    for i, q in enumerate(mp.q_reduced):
        if not np.allclose(q, 0):
            return i, 0
    raise RuntimeError('No non-Gamma q-point found')


def test_complex_coordinate_accessors(mp):
    """Q, P, F complex coordinates expose amplitude, phase, and complex value."""
    np.random.seed(10)
    N = mp.supercell.n_atoms
    mp.set_u(np.random.normal(scale=0.01, size=(N, 3)))
    mp.set_v(np.random.normal(scale=0.1, size=(N, 3)))
    mp.set_f(np.random.normal(scale=0.01, size=(N, 3)))

    q_idx, s_idx = _first_non_gamma_q(mp)
    band = mp[q_idx][s_idx]

    for cc in (band.Q, band.P, band.F):
        assert isinstance(complex(cc), complex)
        assert isinstance(cc.complex, complex)
        assert isinstance(cc.amplitude, (float, np.floating))
        assert isinstance(cc.phase, float)

    # str of Q, P, F subclasses
    assert 'Q' in str(band.Q)
    assert 'P' in str(band.P)
    assert 'F' in str(band.F)


def test_complex_coordinate_amplitude_setter(mp):
    """Setting amplitude scales the complex value while preserving the phase."""
    np.random.seed(11)
    N = mp.supercell.n_atoms
    mp.set_u(np.random.normal(scale=0.01, size=(N, 3)))

    q_idx, s_idx = _first_non_gamma_q(mp)
    band = mp[q_idx][s_idx]
    old_phase = band.Q.phase
    band.Q.amplitude = 0.5

    assert np.isclose(band.Q.amplitude, 0.5)
    assert np.isclose(np.exp(1j * band.Q.phase), np.exp(1j * old_phase))


def test_complex_coordinate_phase_setter(mp):
    """Setting phase rotates the complex value while preserving the amplitude."""
    np.random.seed(12)
    N = mp.supercell.n_atoms
    mp.set_u(np.random.normal(scale=0.01, size=(N, 3)))

    q_idx, s_idx = _first_non_gamma_q(mp)
    band = mp[q_idx][s_idx]
    old_amplitude = band.Q.amplitude
    band.Q.phase = 0.0

    assert np.isclose(band.Q.amplitude, old_amplitude)
    assert np.isclose(band.Q.phase, 0.0)


def test_complex_coordinate_complex_setter(mp):
    """Setting the complex value directly updates Q and respects Q(-q)=Q(q)*."""
    q_idx, s_idx = _first_non_gamma_q(mp)

    # All q-points in a 2×2×2 Al supercell are self-conjugate; Q must be real
    band = mp[q_idx][s_idx]
    band.Q.complex = 0.3  # purely real value is valid for a self-conjugate q-point

    assert np.isclose(band.Q.complex, 0.3)
    # Conjugate symmetry must still hold globally
    assert np.allclose(mp.get_Q().conj(), mp.get_Q()[mp.q_minus])


# ----------------------------------------------- QPoint wavenumber / wavelength ---

def test_qpoint_gamma_wavelength_is_infinite(mp):
    """The Gamma point (q=0) has infinite wavelength."""
    gamma_idx = [tuple(q) for q in mp.q_reduced].index((0.0, 0.0, 0.0))
    qp = mp[gamma_idx]
    assert qp.wavenumber == pytest.approx(0.0)
    assert qp.wavelength == np.inf


def test_qpoint_nonzero_wavelength(mp):
    """A non-Gamma q-point has a finite, positive wavelength."""
    # Find a non-Gamma point
    for i, q in enumerate(mp.q_reduced):
        if not np.allclose(q, 0):
            qp = mp[i]
            break
    assert qp.wavenumber > 0
    assert np.isfinite(qp.wavelength)
    assert np.isclose(qp.wavelength, 2 * np.pi / qp.wavenumber)


# ----------------------------------------------- atoms.scaled_positions ---

def test_atoms_scaled_positions(mp):
    """Prim and Supercell expose scaled_positions (spos alias)."""
    spos = mp.primitive.scaled_positions
    assert spos.shape == (mp.primitive.n_atoms, 3)
    assert spos.dtype == float


# ------------------------------------------------- constructor edge cases ---


def test_mode_projector_nonmultiple_supercell_raises(primitive):
    """A supercell whose size is not a multiple of the primitive must raise."""
    # 3-atom supercell vs 1-atom primitive → 3/1=3, OK; use 2-atom fake primitive
    prim2 = Atoms('Al2', positions=[[0, 0, 0], [2, 0, 0]], cell=[4, 4, 4], pbc=True)
    sc3 = Atoms('Al3', positions=[[0, 0, 0], [2, 0, 0], [4, 0, 0]], cell=[6, 4, 4], pbc=True)
    fc2 = np.zeros((3, 3, 3, 3))
    with pytest.raises(ValueError, match='not multiple of primitive size'):
        ModeProjector(prim2, sc3, fc2)


def test_mode_projector_call_invalid_qpoint_raises(mp):
    """Calling mp with a q-point not in the mesh must raise ValueError."""
    with pytest.raises(ValueError, match='qpoint not compatible'):
        mp((0.123, 0.456, 0.789))


# ------------------------------------------------- project_modes ---

class FakeTraj:
    """Minimal trajectory stub for testing project_modes."""

    def __init__(self, frames, n_atoms, has_velocities=True):
        self._atomic_indices = {'Al': list(range(n_atoms))}
        self.n_atoms = n_atoms
        self.has_velocities = has_velocities
        self._frames = frames

    def __iter__(self):
        return iter(self._frames)


@pytest.fixture
def fake_traj_with_velocities(supercell):
    N = len(supercell)
    atomic_indices = {'Al': list(range(N))}
    np.random.seed(42)
    n_frames = 3
    frames = [
        TrajectoryFrame(
            atomic_indices, i,
            supercell.positions + np.random.normal(scale=0.01, size=(N, 3)),
            np.random.normal(size=(N, 3)),
        )
        for i in range(n_frames)
    ]
    return FakeTraj(frames, N, has_velocities=True)


@pytest.fixture
def fake_traj_no_velocities(supercell):
    N = len(supercell)
    atomic_indices = {'Al': list(range(N))}
    np.random.seed(43)
    frames = [
        TrajectoryFrame(
            atomic_indices, 0,
            supercell.positions + np.random.normal(scale=0.01, size=(N, 3)),
            velocities=None,
        )
    ]
    return FakeTraj(frames, N, has_velocities=False)


def test_project_modes_with_velocities(mp, supercell, fake_traj_with_velocities):
    """project_modes returns Q and P arrays of the correct shape."""
    modes = mp.eigenmodes  # shape: (Nk, Ns, N, 3)
    Q_traj, P_traj = project_modes(fake_traj_with_velocities, modes, supercell)

    Nk = len(mp.q_reduced)
    Ns = mp.primitive.n_atoms * 3
    n_frames = len(fake_traj_with_velocities._frames)
    assert Q_traj.shape == (n_frames, Nk, Ns)
    assert P_traj.shape == (n_frames, Nk, Ns)


def test_project_modes_no_velocities(mp, supercell, fake_traj_no_velocities):
    """project_modes with frames lacking velocities sets P to zero."""
    modes = mp.eigenmodes
    Q_traj, P_traj = project_modes(fake_traj_no_velocities, modes, supercell)
    assert np.allclose(P_traj, 0)


def test_project_modes_modes_n_atoms_mismatch_raises(supercell):
    """project_modes raises if modes[-2] != traj.n_atoms."""
    N = len(supercell)
    traj = FakeTraj([], N)
    modes = np.zeros((1, N + 1, 3))  # modes has wrong atom count
    with pytest.raises(ValueError, match='Second dim in modes'):
        project_modes(traj, modes, supercell)


def test_project_modes_supercell_n_atoms_mismatch_raises(supercell):
    """project_modes raises if traj.n_atoms != len(ideal_supercell)."""
    N = len(supercell)
    traj = FakeTraj([], N + 1)
    modes = np.zeros((1, N + 1, 3))
    with pytest.raises(ValueError, match='same number of atoms'):
        project_modes(traj, modes, supercell)
