import os
import numpy as np
import pytest

from dynasor.trajectory.extxyz_trajectory_reader import ExtxyzTrajectoryReader


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


@pytest.fixture
def traj_filename_with_velocities():
    this_dir = os.path.dirname(__file__)
    filename_xyz_with_velocities = os.path.join(
        this_dir, 'trajectory_files/dump_with_velocities.xyz')
    return filename_xyz_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 = ExtxyzTrajectoryReader(traj_filename)
    reader.close()


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


def test_read_frames(traj_filename):
    # dump with position
    reader = ExtxyzTrajectoryReader(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)


def test_read_frames_with_velocities(traj_filename_with_velocities):
    # dump with position and velocities
    reader = ExtxyzTrajectoryReader(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 = ExtxyzTrajectoryReader(traj_filename)
    frame = next(reader)
    reader.close()

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

    # first five poisitions
    x_target5, v_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)


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

    # first five velocities
    x_target5, 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)


def test_first_frame_contents_with_units(traj_filename_with_velocities, first_frame):
    reader = ExtxyzTrajectoryReader(traj_filename_with_velocities, length_unit='nm', time_unit='ps')
    frame = next(reader)
    reader.close()
    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)
