import os
import numpy as np
import pytest
from ase import Atoms
from ase.md.verlet import VelocityVerlet
from ase.md.langevin import Langevin
from ase.calculators.emt import EMT
from ase.calculators.singlepoint import SinglePointCalculator
import taco_format


@pytest.fixture
def temp_dir(tmpdir):
    return tmpdir


def create_trajectory(num_frames=10):
    """Create a simple MD trajectory using ASE's EMT calculator"""
    # Create a copper atom
    atoms = Atoms("Cu", positions=[[0.0, 0.0, 0.0]])
    atoms.calc = EMT()

    # Set initial velocity
    atoms.set_velocities([[0.01, 0.01, 0.01]])

    # Run molecular dynamics
    dyn = VelocityVerlet(atoms, timestep=1.0)

    traj = []
    for _ in range(num_frames):
        dyn.run(steps=5)  # Run 5 steps before saving

        # Clone the atoms object to save current state
        frame = atoms.copy()

        # Create a calculator with the current energy and forces
        calc = SinglePointCalculator(
            frame, energy=atoms.get_potential_energy(), forces=atoms.get_forces()
        )
        frame.calc = calc

        traj.append(frame)

    return traj


def test_write_read_trajectory(temp_dir):
    """Test writing and reading a trajectory"""
    filename = os.path.join(temp_dir, "trajectory.taco")

    # Create a test trajectory
    traj = create_trajectory(10)

    # Write to TACO file
    taco_format.write(filename, traj)

    # Read from TACO file
    traj_read = taco_format.read(filename)

    # Check that we got the right number of frames
    assert len(traj_read) == len(traj)

    # Check each frame
    for i in range(len(traj)):
        print(f"Comparing frame {i}")
        np.testing.assert_allclose(
            traj_read[i].get_positions(), traj[i].get_positions(), atol=1e-3
        )
        np.testing.assert_allclose(
            traj_read[i].get_velocities(), traj[i].get_velocities(), atol=1e-3
        )
        np.testing.assert_allclose(
            traj_read[i].get_forces(), traj[i].get_forces(), atol=1e-3
        )
        assert (
            abs(traj_read[i].get_potential_energy() - traj[i].get_potential_energy())
            < 1e-6
        )


def test_read_frame_range(temp_dir):
    """Test reading a range of frames from a trajectory"""
    filename = os.path.join(temp_dir, "range_trajectory.taco")

    # Create a test trajectory
    traj = create_trajectory(20)

    # Write to TACO file
    taco_format.write(filename, traj)

    # Read a subset of frames
    start_frame = 5
    end_frame = 15
    traj_subset = taco_format.read_frame_range(filename, start_frame, end_frame)

    # Check that we got the right number of frames
    assert len(traj_subset) == end_frame - start_frame

    # Check that the frames match
    for i, frame in enumerate(traj_subset):
        orig_idx = i + start_frame
        np.testing.assert_allclose(
            frame.get_positions(), traj[orig_idx].get_positions(), atol=1e-3
        )
        np.testing.assert_allclose(frame.get_forces(), traj[orig_idx].get_forces())


def test_read_specific_frames(temp_dir):
    """Test reading specific frames from a trajectory"""
    filename = os.path.join(temp_dir, "specific_trajectory.taco")

    # Create a test trajectory
    traj = create_trajectory(20)

    # Write to TACO file
    taco_format.write(filename, traj)

    # Read specific frames
    frame_indices = [0, 5, 10, 15]
    traj_subset = taco_format.read_frames(filename, frame_indices)

    # Check that we got the right number of frames
    assert len(traj_subset) == len(frame_indices)

    # Check that the frames match
    for i, idx in enumerate(frame_indices):
        np.testing.assert_allclose(
            traj_subset[i].get_positions(), traj[idx].get_positions(), atol=1e-3
        )
        np.testing.assert_allclose(traj_subset[i].get_forces(), traj[idx].get_forces())


def test_read_single_frame_from_trajectory(temp_dir):
    """Test reading a single frame from a trajectory"""
    filename = os.path.join(temp_dir, "single_from_traj.taco")

    # Create a test trajectory
    traj = create_trajectory(10)

    # Write to TACO file
    taco_format.write(filename, traj)

    # Read a specific frame
    frame_idx = 5
    frames = taco_format.read(filename, frame_index=frame_idx)
    assert len(frames) == 1
    frame = frames[0]

    # Check that the frame matches
    np.testing.assert_allclose(
        frame.get_positions(), traj[frame_idx].get_positions(), atol=1e-3
    )
    np.testing.assert_allclose(frame.get_forces(), traj[frame_idx].get_forces())


def test_full_frame_interval(temp_dir):
    """Test writing with different full_frame_interval values"""
    filename1 = os.path.join(temp_dir, "interval_10.taco")
    filename2 = os.path.join(temp_dir, "interval_5.taco")

    # Create a test trajectory
    traj = create_trajectory(20)

    # Write with different full_frame_interval values
    taco_format.write(filename1, traj, full_frame_interval=10)
    taco_format.write(filename2, traj, full_frame_interval=5)

    # Read back
    traj1 = taco_format.read(filename1)
    traj2 = taco_format.read(filename2)

    # Check that both have the correct number of frames
    assert len(traj1) == len(traj)
    assert len(traj2) == len(traj)

    # Check that a sample frame matches
    np.testing.assert_allclose(traj1[10].get_positions(), traj[10].get_positions())
    np.testing.assert_allclose(traj2[10].get_positions(), traj[10].get_positions())
