def test_repex_save_and_load():

    nc_filename = tempfile.mkdtemp() + "/out.nc"

    T_min = 1.0 * unit.kelvin
    T_i = [T_min, T_min * 10., T_min * 100.]
    n_replicas = len(T_i)

    ho = testsystems.HarmonicOscillator()

    system = ho.system
    positions = ho.positions

    states = [ ThermodynamicState(system=system, temperature=T_i[i]) for i in range(n_replicas) ]

    coordinates = [positions] * n_replicas

    parameters = {"number_of_iterations":5}
    rex = replica_exchange.ReplicaExchange.create(states, coordinates, nc_filename, parameters=parameters)
    rex.run()
    
    rex.extend(5)
    
    rex = resume(nc_filename)
    rex.run()

    eq(rex.__class__.__name__, "ReplicaExchange")

def test_harmonic_oscillators_save_and_load():

    nc_filename = tempfile.mkdtemp() + "/out.nc"

    T_min = 1.0 * unit.kelvin
    T_i = [T_min, T_min * 10., T_min * 100.]
    n_replicas = len(T_i)

    ho = testsystems.HarmonicOscillator()

    system = ho.system
    positions = ho.positions

    states = [ ThermodynamicState(system=system, temperature=T_i[i]) for i in range(n_replicas) ]

    coordinates = [positions] * n_replicas

    mpicomm = dummympi.DummyMPIComm()
    parameters = {"number_of_iterations":200}
    replica_exchange = ReplicaExchange.create(states, coordinates, nc_filename, mpicomm=mpicomm, parameters=parameters)
    replica_exchange.run()
    
    replica_exchange.extend(100)

    replica_exchange = resume(nc_filename, mpicomm=mpicomm)
    eq(replica_exchange.iteration, 200)
    replica_exchange.run()

def test_unique_pairs():
    n = 10
    a = np.arange(n)
    b = np.arange(n, n+n)

    eq(md.Topology._unique_pairs(a, a).sort(), md.Topology._unique_pairs_equal(a).sort())
    eq(md.Topology._unique_pairs(a, b).sort(), md.Topology._unique_pairs_mutually_exclusive(a, b).sort())

def test_bool():
    sp = parse_selection("protein or water")
    eq(sp.source, "(atom.residue.is_protein or atom.residue.is_water)")

    sp = parse_selection("protein or water or all")
    eq(sp.source,
       "(atom.residue.is_protein or atom.residue.is_water or True)")

def _test_against_vmd(pdb):
    # this is probably not cross-platform compatible. I assume that the exact
    # path to this CHARMM topology that is included with VMD depends on
    # the install mechanism, especially for bundled mac or windows installers
    VMD_ROOT = os.path.join(os.path.dirname(os.path.realpath(VMD)), '..')
    top_paths = [os.path.join(r, f) for (r, _, fs) in os.walk(VMD_ROOT) for f in fs
         if 'top_all27_prot_lipid_na.inp' in f]
    assert len(top_paths) >= 0
    top = os.path.abspath(top_paths[0]).replace(" ", "\\ ")

    TEMPLATE = '''
package require psfgen
topology %(top)s
pdbalias residue HIS HSE
pdbalias atom ILE CD1 CD
segment U {pdb %(pdb)s}
coordpdb %(pdb)s U
guesscoord
writepdb out.pdb
writepsf out.psf
exit
    ''' % {'top': top, 'pdb' : pdb}

    with enter_temp_directory():
        with open('script.tcl', 'w') as f:
            f.write(TEMPLATE)
        os.system(' '.join([VMD, '-e', 'script.tcl', '-dispdev', 'none']))
        out_pdb = md.load('out.pdb')
        out_psf = md.load_psf('out.psf')

        # make sure the two topologies are equal
        eq(out_pdb.top, out_psf)

def test_baker_hubbard_2():
    t = md.load(get_fn('1vii_sustiva_water.pdb'))
    triplets = md.baker_hubbard(t)
    N = 1000
    rows = triplets[:, 0] * N*N + triplets[:, 1] * N + triplets[:, 2]
    # ensure that there aren't any repeat rows
    eq(len(np.unique(rows)), len(rows))

def test_parallel_tempering():
    nc_filename = tempfile.mkdtemp() + "/out.nc"

    T_min = 1.0 * unit.kelvin
    T_max = 10.0 * unit.kelvin
    n_temps = 3

    ho = testsystems.HarmonicOscillator()

    system = ho.system
    positions = ho.positions


    coordinates = [positions] * n_temps

    mpicomm = dummympi.DummyMPIComm()
    parameters = {"number_of_iterations":1000}
    replica_exchange = ParallelTempering.create(system, coordinates, nc_filename, T_min=T_min, T_max=T_max, n_temps=n_temps, mpicomm=mpicomm, parameters=parameters)
    
    eq(replica_exchange.n_replicas, n_temps)

    replica_exchange.run()

    u_permuted = replica_exchange.database.ncfile.variables["energies"][:]
    s = replica_exchange.database.ncfile.variables["states"][:]

    u = permute_energies(u_permuted, s)

    states = replica_exchange.thermodynamic_states
    u0 = np.array([[ho.reduced_potential_expectation(s0, s1) for s1 in states] for s0 in states])

    l0 = np.log(u0)  # Compare on log scale because uncertainties are proportional to values
    l1 = np.log(u.mean(0))
    eq(l0, l1, decimal=1)

def test_double_set_thermodynamic_states():
    nc_filename = tempfile.mkdtemp() + "/out.nc"

    T_min = 1.0 * unit.kelvin
    T_max = 10.0 * unit.kelvin
    n_temps = 3

    ho = testsystems.HarmonicOscillator()

    system = ho.system
    positions = ho.positions

    coordinates = [positions] * n_temps

    mpicomm = dummympi.DummyMPIComm()
    parameters = {"number_of_iterations": 3}
    replica_exchange = ParallelTempering.create(
        system,
        coordinates,
        nc_filename,
        T_min=T_min,
        T_max=T_max,
        n_temps=n_temps,
        mpicomm=mpicomm,
        parameters=parameters,
    )

    eq(replica_exchange.n_replicas, n_temps)

    replica_exchange.run()

    states = replica_exchange.thermodynamic_states
    replica_exchange.database.thermodynamic_states = states

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

    with LH5TrajectoryFile(temp) as f:
       eq(f.read(), coordinates.reshape(1,10,3), decimal=3)

def test_hrex_save_and_load(mpicomm):

    nc_filename = tempfile.mkdtemp() + "/out.nc"

    temperature = 1 * unit.kelvin

    K0 = 100.0  # Units are automatically added by the testsystem
    K = [K0, K0 * 10., K0 * 1.]
    powers = [2., 2., 4.]
    n_replicas = len(K)

    oscillators = [testsystems.PowerOscillator(b=powers[i]) for i in range(n_replicas)]

    systems = [ho.system for ho in oscillators]
    positions = [ho.positions for ho in oscillators]

    state = ThermodynamicState(system=systems[0], temperature=temperature)

    parameters = {"number_of_iterations":200}
    replica_exchange = hamiltonian_exchange.HamiltonianExchange.create(state, systems, positions, nc_filename, mpicomm=mpicomm, parameters=parameters)
    replica_exchange.run()

    replica_exchange.extend(100)
    
    replica_exchange = resume(nc_filename, mpicomm=mpicomm)
    eq(replica_exchange.iteration, 200)
    replica_exchange.run()

def test_parse_ligand_filename():
    molecule_name = "sustiva"
    input_filename = utils.get_data_filename("chemicals/sustiva/sustiva.mol2")
    name, ext = utils.parse_ligand_filename(input_filename)

    eq(name, "sustiva")
    eq(ext, ".mol2")

def test_power_oscillators(mpicomm):

    nc_filename = tempfile.mkdtemp() + "/out.nc"

    temperature = 1 * unit.kelvin

    K0 = 100.0  # Units are automatically added by the testsystem
    K = [K0, K0 * 10., K0 * 1.]
    powers = [2., 2., 4.]
    n_replicas = len(K)

    oscillators = [testsystems.PowerOscillator(b=powers[i]) for i in range(n_replicas)]

    systems = [ho.system for ho in oscillators]
    positions = [ho.positions for ho in oscillators]

    state = ThermodynamicState(system=systems[0], temperature=temperature)

    parameters = {"number_of_iterations":2000}
    replica_exchange = hamiltonian_exchange.HamiltonianExchange.create(state, systems, positions, nc_filename, mpicomm=mpicomm, parameters=parameters)
    replica_exchange.run()

    u_permuted = replica_exchange.database.ncfile.variables["energies"][:]
    s = replica_exchange.database.ncfile.variables["states"][:]
    u = permute_energies(u_permuted, s)

    beta = (state.temperature * kB) ** -1.
    u0 = np.array([[testsystems.PowerOscillator.reduced_potential(beta, ho2.K, ho2.b, ho.K, ho.b) for ho in oscillators] for ho2 in oscillators])

    l = np.log(u.mean(0))
    l0 = np.log(u0)

    eq(l0, l, decimal=1)

def test_inertia():
    assert eq(order.compute_inertia_tensor(TRAJ1),
              order._compute_inertia_tensor_slow(TRAJ1))
    assert eq(order.compute_inertia_tensor(TRAJ2),
              order._compute_inertia_tensor_slow(TRAJ2))
    assert eq(order.compute_inertia_tensor(TRAJ3),
              order._compute_inertia_tensor_slow(TRAJ3))

def test_parallel_tempering_save_and_load():

    nc_filename = tempfile.mkdtemp() + "/out.nc"

    T_min = 1.0 * unit.kelvin
    T_max = 10.0 * unit.kelvin
    n_temps = 3

    ho = testsystems.HarmonicOscillator()

    system = ho.system
    positions = ho.positions


    coordinates = [positions] * n_temps
    
    parameters = {"number_of_iterations":5}
    rex = parallel_tempering.ParallelTempering.create(system, coordinates, nc_filename, T_min=T_min, T_max=T_max, n_temps=n_temps, parameters=parameters)
    rex.run()
    
    rex.extend(5)
    
    rex = resume(nc_filename)
    rex.run()
    
    eq(rex.__class__.__name__, "ParallelTempering")

def test_select_atom_indices():
    top = md.load(get_fn("native.pdb")).topology

    yield lambda: eq(top.select_atom_indices("alpha"), np.array([8]))
    yield lambda: eq(top.select_atom_indices("minimal"), np.array([4, 5, 6, 8, 10, 14, 15, 16, 18]))

    assert_raises(ValueError, lambda: top.select_atom_indices("sdfsdfsdf"))

def test_atom_indices_1():
    atom_indices = np.arange(10)
    top = md.load(get_fn("native.pdb"))
    t0 = md.load(get_fn("frame0.mdcrd"), top=top)
    t1 = md.load(get_fn("frame0.mdcrd"), top=top, atom_indices=atom_indices)

    eq(t0.xyz[:, atom_indices], t1.xyz)

def test_fitghmm():
    with tempdir():
        RawPositionsFeaturizer().save('featurizer.pickl')
        shell('hmsm fit-ghmm --featurizer featurizer.pickl  --n-init 10  '
                  ' --n-states 4 --dir %s --ext h5 --top %s' % (
                      DATADIR, os.path.join(DATADIR, 'Trajectory0.h5')))
        shell('hmsm inspect -i hmms.jsonlines --details')
        shell('hmsm sample-ghmm --no-match-vars -i hmms.jsonlines --lag-time 1 --n-state 4 '
              '--featurizer featurizer.pickl --dir %s --ext h5 --top %s' % (
                  DATADIR, os.path.join(DATADIR, 'Trajectory0.h5')))
        shell('hmsm means-ghmm -i hmms.jsonlines --lag-time 1 --n-state 4 '
              '--featurizer featurizer.pickl --dir %s --ext h5 --top %s' % (
                  DATADIR, os.path.join(DATADIR, 'Trajectory0.h5')))
        shell('hmsm structures means.csv --ext pdb --prefix means --top %s' %  os.path.join(DATADIR, 'Trajectory0.h5'))
        
        samples_csv = pd.read_csv('samples.csv', skiprows=1)
        means_csv = pd.read_csv('means.csv', skiprows=1)
        
        model = next(iterobjects('hmms.jsonlines'))
        means_pdb = md.load(glob('means-*.pdb'))

    means = np.array(sorted(model['means'], key=lambda e: e[0]))
    print('true\n', HMM.means_)
    print('learned\n', means)

    eq(HMM.means_, means, decimal=0)

    means_pdb_xyz = np.array(sorted(means_pdb.xyz.reshape(4, 3), key=lambda e: e[0]))
    eq(means_pdb_xyz, np.array(sorted(model['means'], key=lambda e:e[0])), decimal=0)

def test_read_1():
    with MDCRDTrajectoryFile(get_fn("frame0.mdcrd"), n_atoms=22) as f:
        xyz, _ = f.read()
    with MDCRDTrajectoryFile(get_fn("frame0.mdcrd"), n_atoms=22) as f:
        xyz3, _ = f.read(stride=3)

    eq(xyz[::3], xyz3)

def test_hbonds_against_dssp():
    t = md.load(get_fn('2EQQ.pdb'))[0]
    pdb = os.path.join(tmpdir, 'f.pdb')
    dssp = os.path.join(tmpdir, 'f.pdb.dssp')
    t.save(pdb)

    cmd = ['mkdssp', '-i', pdb, '-o', dssp]
    subprocess.check_output(' '.join(cmd), shell=True)
    energy = scipy.sparse.lil_matrix((t.n_residues, t.n_residues))

    # read the dssp N-H-->O column from the output file
    with open(dssp) as f:
        # skip the lines until the description of each residue's hbonds
        while not f.readline().startswith('  #  RESIDUE AA STRUCTURE'):
            continue

        for i, line in enumerate(f):
            line = line.rstrip()
            offset0, e0 = map(float, line[39:50].split(','))
            offset1, e1 = map(float, line[61:72].split(','))
            if e0 <= -0.5:
                energy[int(i+offset0), i] = e0
            if e1 <= -0.5:
                energy[int(i+offset1), i] = e1

    dssp = energy.todense()
    ours = md.geometry.hbond.kabsch_sander(t)[0].todense()

    # There is tricky issues with the rounding right at the -0.5 cutoff,
    # so lets just check for equality with DSSP at -0.6 or less
    eq((dssp < -0.6), (ours < -0.6))
    eq(dssp[dssp < -0.6], ours[ours < -0.6], decimal=1)

def test_slice():
    t = md.load(fn)
    yield lambda: eq((t[0:5] + t[5:10]).xyz, t[0:10].xyz)
    yield lambda: eq((t[0:5] + t[5:10]).time, t[0:10].time)
    yield lambda: eq((t[0:5] + t[5:10]).unitcell_vectors, t[0:10].unitcell_vectors)
    yield lambda: eq((t[0:5] + t[5:10]).unitcell_lengths, t[0:10].unitcell_lengths)
    yield lambda: eq((t[0:5] + t[5:10]).unitcell_angles, t[0:10].unitcell_angles)