import os
import pytest
import numpy as np

from dynasor.trajectory.mdanalysis_trajectory_reader \
    import MDAnalysisTrajectoryReader as TrajectoryReader


@pytest.fixture
def traj_fname_lammps():
    this_dir = os.path.dirname(__file__)
    traj_fname = os.path.join(
        this_dir,
        'trajectory_files/positions_and_velocities.lammpstrj')
    return traj_fname


@pytest.fixture
def traj_fname_lammps_no_velocities():
    this_dir = os.path.dirname(__file__)
    traj_fname = os.path.join(
        this_dir,
        'trajectory_files/positions.lammpstrj')
    return traj_fname


@pytest.fixture
def traj_fname_xtc():
    this_dir = os.path.dirname(__file__)
    traj_fname = os.path.join(
        this_dir,
        'trajectory_files/1frame3atoms.xtc')
    return traj_fname


@pytest.fixture
def traj_fname_trr():
    this_dir = os.path.dirname(__file__)
    traj_fname = os.path.join(
        this_dir,
        'trajectory_files/1frame3atoms.trr')
    return traj_fname


# No units from from trajectory and incorrect units from user
# Checks if error is raised
def test_no_units_from_user_or_traj(traj_fname_lammps_no_velocities):
    with pytest.raises(ValueError,
                       match='Trajectory contains no unit information and specified units not '
                             'recognized. Please check the available units.'):
        reader = TrajectoryReader(traj_fname_lammps_no_velocities, 'lammpsdump', 'm', 's')
        reader.close()


# No units from user but from trajectory
# Also checks if trr can be read and that atomic_indices is None
def test_units_only_from_traj(traj_fname_trr):
    reader = TrajectoryReader(traj_fname_trr, 'trr')
    frame = next(reader)
    reader.close()
    assert frame.n_atoms == 3
    assert frame.velocities is not None
    assert pytest.approx(frame.cell[2, 2]) == 41.5305
    assert frame.atom_types is None


# Units from user that match units from trajectory
# Also checks if xtc can be read and that atomic_indices is None
def test_units_from_user_matching_traj(traj_fname_xtc):
    reader = TrajectoryReader(traj_fname_xtc, 'xtc', 'nm', 'ps')
    frames = list(reader)
    reader.close()
    assert len(frames) == 1
    assert frames[0].n_atoms == 3
    assert frames[0].velocities is None
    assert pytest.approx(frames[0].cell[2, 2]) == 41.5305
    assert frames[0].atom_types is None


# Units from user not matching units from trajectory
# Checks if the units from the trajectory are used
def test_units_from_user_not_matching_traj(traj_fname_xtc):
    reader = TrajectoryReader(traj_fname_xtc, 'xtc', 'Angstrom', 'ps')
    frames = list(reader)
    reader.close()
    assert len(frames) == 1
    assert frames[0].n_atoms == 3
    assert frames[0].velocities is None
    assert pytest.approx(frames[0].cell[2, 2]) == 41.5305


# Units from user while no units from trajectory
# Checks if correct unit conversions happen
# Also checks if lammps with velocity can be read and that the atomic indices are read
def test_units_only_from_user(traj_fname_lammps):
    reader = TrajectoryReader(traj_fname_lammps, 'lammpsdump', 'Angstrom', 'ps')
    frame = next(reader)
    reader.close()
    assert frame.frame_index == 0
    assert frame.n_atoms == 24
    assert pytest.approx(list(frame.positions[0, :])) == [1.91468, 3.02071, 0.528818]
    assert pytest.approx(list(frame.velocities[:, 0])) == [1.05876e-6, 1.16399e-5, -1.90113e-5,
                                                           -7.75121e-6, -9.16183e-6, -4.15819e-6,
                                                           5.16741e-6, 1.47389e-5, -9.64617e-6,
                                                           1.60311e-7, 3.17508e-5, 5.01066e-6,
                                                           -4.1924e-6, -8.64071e-6, 1.56727e-6,
                                                           1.43456e-6, 2.18267e-5, -6.2517e-6,
                                                           2.45983e-6, 9.89121e-6, -8.68655e-6,
                                                           -8.36257e-7, 6.78782e-6, 2.00832e-6]
    assert frame.cell[2, 2] == 6.25
    assert isinstance(frame.atom_types, np.ndarray)
    assert np.all(frame.atom_types == np.array(['1', '2', '2', '1', '2', '2', '1', '2', '2',
                                                '1', '2', '2', '1', '2', '2', '1', '2', '2',
                                                '1', '2', '2', '1', '2', '2']))


# Checks that a closed file raises StopIteration
def test_stopiteration_of_closed_file(traj_fname_xtc):
    reader = TrajectoryReader(traj_fname_xtc, 'xtc')
    reader.close()
    with pytest.raises(StopIteration):
        next(reader)
    reader.close()
