import os
import numpy as np
import pytest

from dynasor.trajectory.ase_trajectory_reader import ASETrajectoryReader


@pytest.fixture
def traj_filename():
    this_dir = os.path.dirname(__file__)
    filename_traj = os.path.join(this_dir, 'trajectory_files/dump.traj')
    return filename_traj


@pytest.fixture
def traj_filename_with_velocities():
    this_dir = os.path.dirname(__file__)
    filename_traj_with_velocities = os.path.join(
        this_dir, 'trajectory_files/dump_with_velocities.traj'
    )
    return filename_traj_with_velocities


@pytest.fixture
def first_frame():
    """first frame data on the five first atoms"""
    x_target5 = np.array(
        [
            [35.49864674, 35.46639586, 0.17918569],
            [3.01637108, 3.02819058, 2.89811028],
            [3.33355666, 3.05802449, 0.00510028],
            [2.38218237, -0.02001685, 2.92525985],
            [35.51266599, 3.65396702, 2.90811710],
        ]
    )
    v_target5 = np.array(
        [
            [0.00110336, 0.00111412, 0.00155735],
            [-0.00100465, 0.00085346, -0.00125993],
            [0.00139634, 0.00269955, -0.00236018],
            [0.00059981, -0.00173836, -0.00372708],
            [0.00247676, 0.00041072, -0.00134462],
        ]
    )
    return x_target5, v_target5


def test_open_no_velocities(traj_filename):
    reader = ASETrajectoryReader(traj_filename)
    reader.close()


def test_open_with_velocities(traj_filename_with_velocities):
    reader = ASETrajectoryReader(traj_filename_with_velocities)
    reader.close()


def test_read_frames(traj_filename):
    # dump with position
    reader = ASETrajectoryReader(traj_filename)
    frames = list(reader)
    assert len(frames) == 4
    for it, frame in enumerate(frames):
        assert it == frame.frame_index
        assert frame.positions.shape == (1080, 3)
    reader.close()


def test_read_frames_with_velocities(traj_filename_with_velocities):
    # dump with position and velocities
    reader = ASETrajectoryReader(traj_filename_with_velocities)
    frames = list(reader)
    assert len(frames) == 4
    for it, frame in enumerate(frames):
        assert it == frame.frame_index
        assert frame.positions.shape == (1080, 3)
        assert frame.velocities.shape == (1080, 3)
    reader.close()


def test_first_frame_contents(traj_filename, first_frame):
    reader = ASETrajectoryReader(traj_filename)
    frame = next(reader)

    assert frame.n_atoms == 1080
    assert frame.frame_index == 0

    # first five poisitions
    x_target5, _ = first_frame
    x_read5 = frame.positions[0:5]
    assert np.allclose(x_read5, x_target5)

    # cell
    cell_target = 35.60125893 * np.eye(3)
    assert np.allclose(frame.cell, cell_target)
    reader.close()


def test_first_frame_contents_with_velocities(
    traj_filename_with_velocities, first_frame
):
    reader = ASETrajectoryReader(traj_filename_with_velocities)
    frame = next(reader)
    assert frame.n_atoms == 1080
    assert frame.frame_index == 0

    # first five velocities
    _, v_target5 = first_frame
    v_target5 = np.array(
        [
            [0.00110336, 0.00111412, 0.00155735],
            [-0.00100465, 0.00085346, -0.00125993],
            [0.00139634, 0.00269955, -0.00236018],
            [0.00059981, -0.00173836, -0.00372708],
            [0.00247676, 0.00041072, -0.00134462],
        ]
    )

    v_read5 = frame.velocities[0:5]
    assert np.allclose(v_read5, v_target5)
    reader.close()


def test_first_frame_contents_with_units(traj_filename_with_velocities, first_frame):
    reader = ASETrajectoryReader(
        traj_filename_with_velocities, length_unit='nm', time_unit='ps'
    )
    frame = next(reader)
    assert frame.n_atoms == 1080
    assert frame.frame_index == 0

    # first five velocities
    x_target5, v_target5 = first_frame

    x_read5 = frame.positions[0:5]
    v_read5 = frame.velocities[0:5]
    assert np.allclose(x_read5, 10 * x_target5)
    assert np.allclose(v_read5, 1 / 100 * v_target5)

    # cell
    cell_target = 10 * 35.60125893 * np.eye(3)
    assert np.allclose(frame.cell, cell_target)
    reader.close()
