"""
Tests from MDTraj adapted to forces.
"""
import numpy as np
import tempfile
import os
import mdtraj as md
from bgmol.tpl.hdf5 import HDF5TrajectoryFile
from mdtraj.testing import eq
import pytest

try:
    from bgmol.util.importing import import_openmm
    _, units, _ = import_openmm()
    HAVE_UNITS = True
except ImportError:
    HAVE_UNITS = False

needs_units = pytest.mark.skipif(not HAVE_UNITS, reason='requires simtk.units')


fd, temp = tempfile.mkstemp(suffix='.h5')


def teardown_module(module):
    """remove the temporary file created by tests in this file
    this gets automatically called by pytest"""
    os.close(fd)
    os.unlink(temp)


def test_write_coordinates():
    coordinates = np.random.randn(4, 10,3)
    with HDF5TrajectoryFile(temp, 'w') as f:
        f.write(coordinates)

    with HDF5TrajectoryFile(temp) as f:
        assert eq(f.root.coordinates[:], coordinates)
        assert eq(str(f.root.coordinates.attrs['units']), 'nanometers')


def test_write_coordinates_reshape():
    coordinates = np.random.randn(10,3)
    with HDF5TrajectoryFile(temp, 'w') as f:
        f.write(coordinates)

    with HDF5TrajectoryFile(temp) as f:
        assert eq(f.root.coordinates[:], coordinates.reshape(1,10,3))
        assert eq(str(f.root.coordinates.attrs['units']), 'nanometers')


def test_write_multiple():
    coordinates = np.random.randn(4, 10,3)
    with HDF5TrajectoryFile(temp, 'w') as f:
        f.write(coordinates)
        f.write(coordinates)

    with HDF5TrajectoryFile(temp) as f:
        assert eq(f.root.coordinates[:], np.vstack((coordinates, coordinates)))


def test_write_inconsistent():
    coordinates = np.random.randn(4, 10,3)
    with HDF5TrajectoryFile(temp, 'w') as f:
        f.write(coordinates)
        # since the first frames we saved didn't contain velocities, we
        # can't save more velocities
        with pytest.raises(ValueError):
            f.write(coordinates, velocities=coordinates)


def test_write_inconsistent_2():
    coordinates = np.random.randn(4, 10,3)
    with HDF5TrajectoryFile(temp, 'w') as f:
        f.write(coordinates, velocities=coordinates)
        # we're saving a deficient set of data, since before we wrote
        # more information.
        with pytest.raises(ValueError):
            f.write(coordinates)


@needs_units
def test_write_units():
    # simtk.units are automatically converted into MD units for storage on disk
    coordinates = units.Quantity(np.random.randn(4, 10,3), units.angstroms)
    velocities = units.Quantity(np.random.randn(4, 10,3), units.angstroms/units.year)

    with HDF5TrajectoryFile(temp, 'w') as f:
        f.write(coordinates, velocities=velocities)

    with HDF5TrajectoryFile(temp) as f:
        assert eq(f.root.coordinates[:], coordinates.value_in_unit(units.nanometers))
        assert eq(str(f.root.coordinates.attrs['units']), 'nanometers')

        assert eq(f.root.velocities[:], velocities.value_in_unit(units.nanometers/units.picosecond))
        assert eq(str(f.root.velocities.attrs['units']), 'nanometers/picosecond')


def test_write_units2():
    from mdtraj.utils import unit
    coordinates = unit.quantity.Quantity(np.random.randn(4, 10,3),
                    unit.unit_definitions.angstroms)
    velocities = unit.quantity.Quantity(np.random.randn(4, 10,3),
                    unit.unit_definitions.angstroms/unit.unit_definitions.year)

    with HDF5TrajectoryFile(temp, 'w') as f:
        f.write(coordinates, velocities=velocities)

    with HDF5TrajectoryFile(temp) as f:
        assert eq(f.root.coordinates[:], coordinates.value_in_unit(unit.unit_definitions.nanometers))
        assert eq(str(f.root.coordinates.attrs['units']), 'nanometers')

        assert eq(f.root.velocities[:], velocities.value_in_unit(unit.unit_definitions.nanometers/unit.unit_definitions.picosecond))
        assert eq(str(f.root.velocities.attrs['units']), 'nanometers/picosecond')


@needs_units
def test_write_units_mismatch():
    velocities = units.Quantity(np.random.randn(4, 10,3), units.angstroms/units.picosecond)

    with HDF5TrajectoryFile(temp, 'w') as f:
        # if you try to write coordinates that are unitted and not
        # in the correct units, we find that
        with pytest.raises(TypeError):
            f.write(coordinates=velocities)


def test_topology(get_fn):
    top = md.load_pdb(get_fn('native.pdb')).topology

    with HDF5TrajectoryFile(temp, 'w') as f:
        f.topology = top

    with HDF5TrajectoryFile(temp) as f:
        assert f.topology == top


def test_constraints():
    c = np.array([(1,2,3.5)], dtype=np.dtype([('atom1', np.int32), ('atom2', np.int32), ('distance', np.float32)]))

    with HDF5TrajectoryFile(temp, 'w') as f:
        f.constraints = c

    with HDF5TrajectoryFile(temp) as f:
        assert eq(f.constraints, c)


def test_constraints2():
    c = np.array([(1,2,3.5)], dtype=np.dtype([('atom1', np.int32), ('atom2', np.int32), ('distance', np.float32)]))

    with HDF5TrajectoryFile(temp, 'w') as f:
        f.constraints = c
        f.constraints = c

    with HDF5TrajectoryFile(temp) as f:
        assert eq(f.constraints, c)


def test_read_0():
    coordinates = np.random.randn(4, 10,3)
    with HDF5TrajectoryFile(temp, 'w') as f:
        f.write(coordinates, alchemicalLambda=np.array([1,2,3,4]))

    with HDF5TrajectoryFile(temp) as f:
        got = f.read()
        assert eq(got.coordinates, coordinates)
        assert eq(got.velocities, None)
        assert eq(got.alchemicalLambda, np.array([1,2,3,4]))


@needs_units
def test_read_1():
    coordinates = units.Quantity(np.random.randn(4, 10,3), units.angstroms)
    velocities = units.Quantity(np.random.randn(4, 10,3), units.angstroms/units.years)
    forces = units.Quantity(np.random.randn(4, 10,3), units.kilocalorie_per_mole/units.angstroms)

    with HDF5TrajectoryFile(temp, 'w') as f:
        f.write(coordinates, velocities=velocities, forces=forces)

    with HDF5TrajectoryFile(temp) as f:
        got = f.read()
        assert eq(got.coordinates, coordinates.value_in_unit(units.nanometers))
        assert eq(got.velocities, velocities.value_in_unit(units.nanometers/units.picoseconds))
        assert eq(got.forces, forces.value_in_unit(units.kilojoules_per_mole/units.nanometers), decimal=5)


def test_read_slice_0():
    coordinates = np.random.randn(4, 10,3)
    with HDF5TrajectoryFile(temp, 'w') as f:
        f.write(coordinates, alchemicalLambda=np.array([1,2,3,4]))

    with HDF5TrajectoryFile(temp) as f:
        got = f.read(n_frames=2)
        assert eq(got.coordinates, coordinates[:2])
        assert eq(got.velocities, None)
        assert eq(got.forces, None)
        assert eq(got.alchemicalLambda, np.array([1,2]))


def test_read_slice_1():
    coordinates = np.random.randn(4, 10,3)
    with HDF5TrajectoryFile(temp, 'w') as f:
        f.write(coordinates)

    with HDF5TrajectoryFile(temp) as f:
        got = f.read(n_frames=2)
        assert eq(got.coordinates, coordinates[:2])
        assert eq(got.forces, None)
        assert eq(got.velocities, None)

        got = f.read(n_frames=2)
        assert eq(got.coordinates, coordinates[2:])
        assert eq(got.forces, None)
        assert eq(got.velocities, None)


def test_read_slice_2():
    coordinates = np.random.randn(4, 10,3)
    with HDF5TrajectoryFile(temp, 'w') as f:
        f.write(coordinates, alchemicalLambda=np.arange(4))

    with HDF5TrajectoryFile(temp) as f:
        got = f.read(atom_indices=np.array([0,1]))
        assert eq(got.coordinates, coordinates[:, [0,1], :])
        assert eq(got.alchemicalLambda, np.arange(4))


def test_read_slice_3():
    coordinates = np.random.randn(4, 10,3)
    with HDF5TrajectoryFile(temp, 'w') as f:
        f.write(coordinates, alchemicalLambda=np.arange(4))

    with HDF5TrajectoryFile(temp) as f:
        got = f.read(stride=2, atom_indices=np.array([0,1]))
        assert eq(got.coordinates, coordinates[::2, [0,1], :])
        assert eq(got.alchemicalLambda, np.arange(4)[::2])


def test_do_overwrite():
    with open(temp, 'w') as f:
        f.write('a')

    with HDF5TrajectoryFile(temp, 'w', force_overwrite=True) as f:
        f.write(np.random.randn(10,5,3))


def test_vsite_elements(get_fn):
    #  Test case for issue #265
    pdb_filename = get_fn('GG-tip4pew.pdb')
    trj = md.load(pdb_filename)
    trj.save_hdf5(temp)

    trj2 = md.load(temp, top=pdb_filename)


def test_dont_overwrite():
    with open(temp, 'w') as f:
        f.write('a')
    with pytest.raises(IOError):
        with HDF5TrajectoryFile(temp, 'w', force_overwrite=False) as f:
            f.write(np.random.randn(10,5,3))


def test_attributes():
    constraints = np.zeros(10, dtype=[('atom1', np.int32), ('atom2', np.int32), ('distance', np.float32)])
    with HDF5TrajectoryFile(temp, 'w') as f:
        f.title = 'mytitle'
        f.reference = 'myreference'
        f.forcefield = 'amber99'
        f.randomState = 'sdf'
        f.application = 'openmm'
        f.constraints = constraints

    with HDF5TrajectoryFile(temp) as g:
        eq(g.title, 'mytitle')
        eq(g.reference, 'myreference')
        eq(g.forcefield, 'amber99')
        eq(g.randomState, 'sdf')
        eq(g.application, 'openmm')
        eq(g.constraints, constraints)


def test_append():
    x1 = np.random.randn(10,5,3)
    x2 = np.random.randn(8,5,3)
    with HDF5TrajectoryFile(temp, 'w') as f:
        f.write(x1)
    with HDF5TrajectoryFile(temp, 'a') as f:
        f.write(x2)

    with HDF5TrajectoryFile(temp) as f:
        eq(f.root.coordinates[:], np.concatenate((x1,x2)))
