import os
import tempfile
import numpy as np
import pytest
from ase.build import bulk
from ase import Atoms
from ase.calculators.emt import EMT
from ase.io import write as ase_write
from phonopy import Phonopy
from phonopy.structure.atoms import PhonopyAtoms
from dynasor import ModeProjector, Trajectory
from dynasor.tools.structures import get_P_matrix, find_permutation
from dynasor.modes.project_modes import project_modes


# ── Helpers ───────────────────────────────────────────────────────────────────

def _write_traj(tmp_dir, snapshots, name):
    """Write snapshots to extxyz and return the filename (not a Trajectory).
    A fresh Trajectory must be constructed per test since dynasor iterates lazily."""
    fname = os.path.join(tmp_dir, name)
    ase_write(fname, snapshots, format='extxyz')
    return fname


def _traj(fname):
    return Trajectory(fname, trajectory_format='extxyz')


# ── Fixtures ──────────────────────────────────────────────────────────────────

@pytest.fixture(scope='module')
def tmp_dir():
    with tempfile.TemporaryDirectory() as d:
        yield d


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


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


@pytest.fixture(scope='module')
def mp(primitive, supercell):
    """ModeProjector built with EMT force constants."""
    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 s in phonon.supercells_with_displacements:
        atoms = Atoms(symbols=s.symbols, cell=s.cell,
                      scaled_positions=s.scaled_positions, pbc=True)
        atoms.calc = EMT()
        forces.append(atoms.get_forces())

    phonon.forces = forces
    phonon.produce_force_constants()

    sc = Atoms(cell=phonon.supercell.cell, numbers=phonon.supercell.numbers,
               positions=phonon.supercell.positions, pbc=True)
    p = find_permutation(sc, supercell)
    fc2 = phonon.force_constants[p][:, p]

    return ModeProjector(primitive, supercell, fc2)


@pytest.fixture(scope='module')
def fname_with_velocities(mp, tmp_dir):
    """5-snapshot trajectory file with random displacements and velocities."""
    rng = np.random.default_rng(42)
    N = len(mp.supercell.to_ase())
    snapshots = []
    for _ in range(5):
        atoms = mp.supercell.to_ase().copy()
        atoms.positions += rng.normal(scale=0.01, size=(N, 3))
        atoms.set_velocities(rng.normal(scale=1.0, size=(N, 3)))
        snapshots.append(atoms)
    return _write_traj(tmp_dir, snapshots, name='traj_with_velocities.extxyz')


@pytest.fixture(scope='module')
def fname_no_velocities(mp, tmp_dir):
    """5-snapshot trajectory file with random displacements but no velocities."""
    rng = np.random.default_rng(42)
    N = len(mp.supercell.to_ase())
    snapshots = []
    for _ in range(5):
        atoms = mp.supercell.to_ase().copy()
        atoms.positions += rng.normal(scale=0.01, size=(N, 3))
        snapshots.append(atoms)
    return _write_traj(tmp_dir, snapshots, name='traj_no_velocities.extxyz')


# ── Regression test ───────────────────────────────────────────────────────────
def test_project_modes_no_zeros_as_bug(mp, fname_with_velocities):
    """Regression: np.zeros_as (invalid) was used instead of np.zeros_like
    in dynasor's project_modes. This test FAILS with unfixed dynasor and
    PASSES once the bug is fixed."""
    Q, P = project_modes(_traj(fname_with_velocities), mp.eigenmodes, mp.supercell.to_ase())


# ── Shape / dtype tests ───────────────────────────────────────────────────────
def test_project_modes_output_shapes(mp, fname_with_velocities):
    """Q and P should have shape (n_frames, n_qpoints, n_modes) and match each other."""
    Q, P = project_modes(_traj(fname_with_velocities), mp.eigenmodes, mp.supercell.to_ase())
    assert Q.ndim == 3, "Q should be 3D (n_frames, n_qpoints, n_modes)"
    assert Q.shape == P.shape, "Q and P should have the same shape"


def test_project_modes_output_complex(mp, fname_with_velocities):
    """Q and P should be complex-valued."""
    Q, P = project_modes(_traj(fname_with_velocities), mp.eigenmodes, mp.supercell.to_ase())
    assert np.iscomplexobj(Q), "Q should be complex"
    assert np.iscomplexobj(P), "P should be complex"


def test_project_modes_output_finite(mp, fname_with_velocities):
    """Q and P should contain no NaN or Inf."""
    Q, P = project_modes(_traj(fname_with_velocities), mp.eigenmodes, mp.supercell.to_ase())
    assert np.all(np.isfinite(Q)), "Q contains NaN or Inf"
    assert np.all(np.isfinite(P)), "P contains NaN or Inf"


# ── No-velocity tests ─────────────────────────────────────────────────────────
def test_project_modes_no_velocities_P_is_zero(mp, fname_no_velocities):
    """Without velocities in the trajectory, P should be zero."""
    Q, P = project_modes(_traj(fname_no_velocities), mp.eigenmodes, mp.supercell.to_ase())
    assert np.allclose(P, 0.0), "P should be zero when trajectory has no velocities"


def test_project_modes_no_velocities_Q_nonzero(mp, fname_no_velocities):
    """Without velocities, Q should still be computed from displacements."""
    Q, P = project_modes(_traj(fname_no_velocities), mp.eigenmodes, mp.supercell.to_ase())
    assert np.any(np.abs(Q) > 1e-10), "Q should be nonzero for nonzero displacements"


# ── Consistency test ──────────────────────────────────────────────────────────
def test_project_modes_consistent_with_mp_get_Q_get_P(mp, tmp_dir):
    """project_modes and mp.get_Q/get_P should give the same result
    for the same snapshot."""
    rng = np.random.default_rng(0)
    N = len(mp.supercell.to_ase())
    u = rng.normal(scale=0.01, size=(N, 3))
    v = rng.normal(scale=1.0,  size=(N, 3))

    snapshots = []
    for _ in range(5):
        atoms = mp.supercell.to_ase().copy()
        atoms.positions += u
        atoms.set_array('vel', v)
        snapshots.append(atoms)
    fname = _write_traj(tmp_dir, snapshots, name='traj_consistency.extxyz')

    Q_pm, P_pm = project_modes(_traj(fname), mp.eigenmodes, mp.supercell.to_ase())

    mp.set_u(u)
    mp.set_v(v)
    Q_mp = mp.get_Q()
    P_mp = mp.get_P()

    assert np.allclose(Q_pm[0], Q_mp, atol=1e-10), "Q mismatch between project_modes and mp.get_Q"
    assert np.allclose(P_pm[0], P_mp, atol=1e-10), "P mismatch between project_modes and mp.get_P"
