def test_compute_distances2_nimage_exclude():
    n0 = 10
    n1 = 5
    nimage = 1
    cell = get_system_water32().cell
    pos0 = np.random.normal(0, 10, (n0, 3))
    pos1 = np.random.normal(0, 10, (n1, 3))
    factor = (1+2*nimage)**3
    exclude = np.array([[1, 0], [2, 1]])
    output = np.zeros(factor*n0*n1 - 2, float)
    cell.compute_distances(output, pos0, pos1, pairs=exclude, nimage=nimage)
    counter = 0
    ex_counter = 0
    for r0 in xrange(-1, 2):
        for r1 in xrange(-1, 2):
            for r2 in xrange(-1, 2):
                for i0 in xrange(n0):
                    for i1 in xrange(n1):
                        if (r0 == 0) and (r1 == 0) and (r2 == 0):
                            if ex_counter < len(exclude) and i0 == exclude[ex_counter,0] and i1 == exclude[ex_counter, 1]:
                                ex_counter += 1
                                continue
                        delta = pos0[i0] - pos1[i1]
                        cell.mic(delta)
                        cell.add_vec(delta, np.array([r0, r1, r2]))
                        assert abs(output[counter] - np.linalg.norm(delta)) < 1e-10
                        counter += 1
    assert ex_counter == 2

def test_generator_water32_cross():
    system = get_system_water32()
    fn_pars = pkg_resources.resource_filename(__name__, '../../data/test/parameters_water_cross.txt')
    ff = ForceField.generate(system, fn_pars)
    assert len(ff.parts) == 1
    assert isinstance(ff.parts[0], ForcePartValence)
    part_valence = ff.parts[0]
    assert part_valence.dlist.ndelta == 64
    for i, j in system.bonds:
        row0 = part_valence.dlist.lookup.get((i, j))
        row1 = part_valence.dlist.lookup.get((j, i))
        assert row0 is not None or row1 is not None
    assert part_valence.iclist.nic == 96
    vlist = part_valence.vlist
    iclist = vlist.iclist
    assert vlist.nv == 96
    for irow in range(vlist.nv):
        row = vlist.vtab[irow]
        assert row['kind'] == 3
        ic0 = iclist.ictab[row['ic0']]
        ic1 = iclist.ictab[row['ic1']]
        if ic0['kind'] == 0 and ic1['kind'] == 0:
            np.testing.assert_allclose(row['par0'], 2.0000000000e+01*(kjmol/angstrom**2), rtol=0.0, atol=1e-10)
            np.testing.assert_allclose(row['par1'], 0.9470000000e+00*angstrom, rtol=0.0, atol=1e-10)
            np.testing.assert_allclose(row['par2'], 0.9470000000e+00*angstrom, rtol=0.0, atol=1e-10)
        elif ic0['kind'] == 0 and ic1['kind'] == 2:
            np.testing.assert_allclose(row['par0'], 1.0000000000e+01*(kjmol/angstrom*rad), rtol=0.0, atol=1e-10)
            np.testing.assert_allclose(row['par1'], 0.9470000000e+00*angstrom, rtol=0.0, atol=1e-10)
            np.testing.assert_allclose(row['par2'], 1.0500000000e+02*deg, rtol=0.0, atol=1e-10)
        else:
            raise AssertionError('ICs in Cross term should be Bond-Bond or Bond-BendAngle')

def test_topology_water32():
    system = get_system_water32()
    assert system.bonds[0,0] == 0
    assert system.bonds[0,1] == 1
    assert system.bonds[1,0] == 0
    assert system.bonds[1,1] == 2
    assert system.bonds[2,0] == 3
    assert system.bonds[2,1] == 4
    assert system.bonds[3,0] == 3
    assert system.bonds[3,1] == 5
    for i in range(system.natom):
        if system.numbers[i] == 8:
            assert len(system.neighs1[i]) == 2
            n0, n1 = system.neighs1[i]
            assert system.numbers[n0] == 1
            assert system.numbers[n1] == 1
            assert len(system.neighs2[i]) == 0
            assert len(system.neighs3[i]) == 0
        elif system.numbers[i] == 1:
            assert len(system.neighs1[i]) == 1
            n, = system.neighs1[i]
            assert system.numbers[n] == 8
            assert len(system.neighs2[i]) == 1
            n, = system.neighs2[i]
            assert system.numbers[n] == 1
            assert len(system.neighs3[i]) == 0

def test_gpos_vtens_mm3quartic_water32():
    system = get_system_water32()
    part = ForcePartValence(system)
    for i, j in system.bonds:
        part.add_term(MM3Quartic(1.5, 2.0+0.01*i, Bond(i, j)))
    check_gpos_part(system, part)
    check_vtens_part(system, part)

def test_gpos_vtens_morse_water32():
    system = get_system_water32()
    part = ForcePartValence(system)
    for i, j in system.bonds:
        part.add_term(Morse(0.3, 1.7, 2.0, Bond(i, j)))
    check_gpos_part(system, part)
    check_vtens_part(system, part)

def test_ewald_gpos_vtens_reci_water32():
    system = get_system_water32()
    dielectric = 1.4
    for alpha in 0.05, 0.1, 0.2:
        part_ewald_reci = ForcePartEwaldReciprocal(system, alpha, gcut=alpha / 0.75, dielectric=dielectric)
        check_gpos_part(system, part_ewald_reci)
        check_vtens_part(system, part_ewald_reci)

def test_ewald_gpos_vtens_corr_water32():
    system = get_system_water32()
    scalings = Scalings(system, 0.0, 0.0, 0.5)
    for alpha in 0.05, 0.1, 0.2:
        part_ewald_corr = ForcePartEwaldCorrection(system, alpha, scalings, dielectric=0.8)
        check_gpos_part(system, part_ewald_corr)
        check_vtens_part(system, part_ewald_corr)

def get_part_water32_14A_eidip():
    # Initialize system, nlist and scaling
    system = get_system_water32()
    #Reset charges
    system.charges *= 0.0
    #Set dipoles to random values
    dipoles = np.random.rand( system.natom, 3 )
    #Set polarizations to infinity (no energy to create dipoles)
    poltens_i = np.tile( np.diag([0.0,0.0,0.0]) , np.array([system.natom, 1]) )
    nlist = NeighborList(system)

    scalings = Scalings(system, 1.0, 1.0, 1.0)
    # Create the pair_pot and part_pair
    rcut = 14*angstrom
    alpha = 5.5/rcut
    pair_pot = PairPotEIDip(system.charges, dipoles, poltens_i, alpha, rcut)
    part_pair = ForcePartPair(system, nlist, scalings, pair_pot)
    part_pair.nlist.update()
    # The pair function
    def pair_fn(i, j, d, delta):
        energy = 0.0
        #Dipole-Dipole (only term for this test)
        fac1 = erfc(alpha*d) + 2.0*alpha*d/np.sqrt(np.pi)*np.exp(-alpha**2*d**2)
        fac2 = 3.0*erfc(alpha*d) + 4.0*alpha**3*d**3/np.sqrt(np.pi)*np.exp(-alpha**2*d**2) \
                + 6.0*alpha*d/np.sqrt(np.pi)*np.exp(-alpha**2*d**2)
        energy += np.dot( pair_pot.dipoles[i,:] , pair_pot.dipoles[j,:] )*fac1/d**3 - \
                         1.0*np.dot(pair_pot.dipoles[i,:],delta)*np.dot(delta,pair_pot.dipoles[j,:])*fac2/d**5
        return energy
    return system, nlist, scalings, part_pair, pair_fn

def test_pair_pot_ei_water32_dielectric():
    #Using a relative permittivity epsilon (!=1) should give the same results
    #as epsilon==1 with all charges scaled by 1.0/sqrt(epsilon)
    # Initialize system, nlist, scaling, ...
    system = get_system_water32()
    nlist = NeighborList(system)
    scalings = Scalings(system, 0.0, 0.5, 1.0)
    rcut = 14*angstrom
    alpha = 5.5/rcut
    dielectric = 1.44
    #Compute energy with epsilon 1 and scaled charges
    pair_pot = PairPotEI(system.charges/np.sqrt(dielectric), alpha, rcut)
    part_pair = ForcePartPair(system, nlist, scalings, pair_pot)
    ff = ForceField(system, [part_pair], nlist)
    ff.update_pos(system.pos)
    gpos0 = np.zeros(system.pos.shape, float)
    vtens0 = np.zeros((3, 3), float)
    energy0 = ff.compute(gpos0, vtens0)
    #Compute energy with epsilon=dielctric and original charges
    pair_pot = PairPotEI(system.charges, alpha, rcut, dielectric=dielectric)
    part_pair = ForcePartPair(system, nlist, scalings, pair_pot)
    ff = ForceField(system, [part_pair], nlist)
    ff.update_pos(system.pos)
    gpos1 = np.zeros(system.pos.shape, float)
    vtens1 = np.zeros((3, 3), float)
    energy1 = ff.compute(gpos1, vtens1)
    assert np.abs(energy0-energy1) < 1.0e-10
    assert np.all(np.abs(gpos0-gpos1)) < 1.0e-10
    assert np.all(np.abs(vtens0-vtens1) < 1.0e-10 )

def get_part_water32_9A_grimme():
    # Initialize system, nlist and scaling
    system = get_system_water32()
    nlist = NeighborList(system)
    scalings = Scalings(system)
    # Initialize parameters
    r0_table = {1: 1.001*angstrom, 8: 1.342*angstrom}
    c6_table = {1: 0.14*1e-3*kjmol*nanometer**6, 8: 0.70*1e-3*kjmol*nanometer**6}
    r0s = np.zeros(96, float)
    c6s = np.zeros(96, float)
    for i in xrange(system.natom):
        r0s[i] = r0_table[system.numbers[i]]
        c6s[i] = c6_table[system.numbers[i]]
    # Create the pair_pot and part_pair
    rcut = 9*angstrom
    pair_pot = PairPotGrimme(r0s, c6s, rcut, Hammer(1.0))
    assert abs(pair_pot.r0 - r0s).max() == 0.0
    assert abs(pair_pot.c6 - c6s).max() == 0.0
    part_pair = ForcePartPair(system, nlist, scalings, pair_pot)
    # Create a pair function:
    def pair_fn(i, j, d, delta):
        r0 = (r0s[i]+r0s[j])
        c6 = np.sqrt(c6s[i]*c6s[j])
        if d<rcut:
            return -1.1/(1.0 + np.exp(-20.0*(d/r0-1.0)))*c6/d**6*np.exp(1.0/(d-rcut))
        else:
            return 0.0
    return system, nlist, scalings, part_pair, pair_fn

def get_part_water32_9A_mm3():
    # Initialize system, nlist and scaling
    system = get_system_water32()
    nlist = NeighborList(system)
    scalings = Scalings(system)
    # Initialize parameters
    sigma_table  = {1: 1.62*angstrom, 8: 1.82*angstrom}
    epsilon_table = {1: 0.020*kcalmol, 8: 0.059*kcalmol}
    sigmas = np.zeros(96, float)
    epsilons = np.zeros(96, float)
    onlypaulis = np.zeros(96, np.int32)
    for i in xrange(system.natom):
        sigmas[i] = sigma_table[system.numbers[i]]
        epsilons[i] = epsilon_table[system.numbers[i]]
    # Create the pair_pot and part_pair
    rcut = 9*angstrom
    pair_pot = PairPotMM3(sigmas, epsilons, onlypaulis, rcut, Hammer(1.0))
    assert abs(pair_pot.sigmas - sigmas).max() == 0.0
    assert abs(pair_pot.epsilons - epsilons).max() == 0.0
    part_pair = ForcePartPair(system, nlist, scalings, pair_pot)
    # Create a pair function:
    def pair_fn(i, j, d, delta):
        sigma = sigmas[i]+sigmas[j]
        epsilon = np.sqrt(epsilons[i]*epsilons[j])
        x = (sigma/d)
        if d<rcut:
            return epsilon*(1.84e5*np.exp(-12.0/x)-2.25*x**6)*np.exp(1.0/(d-rcut))
        else:
            return 0.0
    return system, nlist, scalings, part_pair, pair_fn

def get_part_water32_9A_lj():
    # Initialize system, nlist and scaling
    system = get_system_water32()
    nlist = NeighborList(system)
    scalings = Scalings(system)
    # Initialize parameters
    rminhalf_table = {1: 0.2245*angstrom, 8: 1.7682*angstrom}
    epsilon_table = {1: -0.0460*kcalmol, 8: -0.1521*kcalmol}
    sigmas = np.zeros(96, float)
    epsilons = np.zeros(96, float)
    for i in xrange(system.natom):
        sigmas[i] = rminhalf_table[system.numbers[i]]*(2.0)**(5.0/6.0)
        epsilons[i] = epsilon_table[system.numbers[i]]
    # Create the pair_pot and part_pair
    rcut = 9*angstrom
    pair_pot = PairPotLJ(sigmas, epsilons, rcut, Hammer(1.0))
    assert abs(pair_pot.sigmas - sigmas).max() == 0.0
    assert abs(pair_pot.epsilons - epsilons).max() == 0.0
    part_pair = ForcePartPair(system, nlist, scalings, pair_pot)
    # Create a pair function:
    def pair_fn(i, j, d, delta):
        sigma = 0.5*(sigmas[i]+sigmas[j])
        epsilon = np.sqrt(epsilons[i]*epsilons[j])
        x = (sigma/d)**6
        return 4*epsilon*(x*(x-1))*np.exp(1.0/(d-rcut))
    return system, nlist, scalings, part_pair, pair_fn

def test_vlist_cross_water32():
    system = get_system_water32()
    part = ForcePartValence(system)
    for j in range(system.natom):
        if len(system.neighs1[j])==2:
            i, k = system.neighs1[j]
            part.add_term(Cross(
                    1.2,
                    1.7 + 0.01*i,
                    1.9 + 0.01*k,
                    Bond(i, j),
                    Bond(j, k),
            ))
    energy = part.compute()
    check_energy = 0.0
    for j in range(system.natom):
        if len(system.neighs1[j])==2:
            i, k = system.neighs1[j]
            delta0 = system.pos[j] - system.pos[i]
            delta1 = system.pos[k] - system.pos[j]
            system.cell.mic(delta0)
            system.cell.mic(delta1)
            bond0 = np.linalg.norm(delta0)
            bond1 = np.linalg.norm(delta1)
            check_energy += 1.2*(bond0 - 1.7 - 0.01*i)*(bond1 - 1.9 - 0.01*k)
    assert abs(energy - check_energy) < 1e-8

def test_ewald_vtens_neut_water32():
    # fake water model, negative oxygens and neutral hydrogens
    system = get_system_water32()
    system.charges -= 0.1
    for alpha in 0.05, 0.1, 0.2:
        part_ewald_neut = ForcePartEwaldNeutralizing(system, alpha)
        check_vtens_part(system, part_ewald_neut)

def test_gpos_vtens_bond_fues_water32():
    system = get_system_water32()
    part = ForcePartValence(system)
    for i, j in system.bonds:
        part.add_term(Fues(0.3, 1.7, Bond(i, j)))
    check_gpos_part(system, part)
    check_vtens_part(system, part)

def test_generator_water32_qmdffrep():
    system = get_system_water32()
    print(system.ffatypes)
    print(system.ffatype_ids)
    fn_pars = pkg_resources.resource_filename(__name__, '../../data/test/parameters_fake_qmdffrep.txt')
    ff = ForceField.generate(system, fn_pars)
    assert len(ff.parts) == 1
    qmdffrep = ff.part_pair_qmdffrep
    # check parameters
    A_cross = qmdffrep.pair_pot.amp_cross
    assert A_cross.shape == (2,2)
    print(A_cross)
    assert abs(A_cross[0,0] - 3.2490000000e+01) < 1e-10
    assert abs(A_cross[0,1] - 1.3395000000e+01) < 1e-10
    assert abs(A_cross[1,0] - 1.3395000000e+01) < 1e-10
    assert abs(A_cross[1,1] - 5.5225000000e+00) < 1e-10

    B_cross = qmdffrep.pair_pot.b_cross
    assert B_cross.shape == (2,2)
    assert abs(B_cross[0,0] - 4.08560961303e+00) < 1e-10
    assert abs(B_cross[0,1] - 5.00924416592e+00) < 1e-10
    assert abs(B_cross[1,0] - 5.00924416592e+00) < 1e-10
    assert abs(B_cross[1,1] - 6.34212100945e+00) < 1e-10

    # check scalings
    scalings = qmdffrep.scalings
    assert abs(scalings.scale1 - 0.0) < 1e-10
    assert abs(scalings.scale2 - 0.0) < 1e-10
    assert abs(scalings.scale3 - 0.5) < 1e-10
    assert abs(scalings.scale4 - 0.5) < 1e-10

def test_gpos_vtens_polyfour_water32():
    system = get_system_water32()
    part = ForcePartValence(system)
    for i, j in system.bonds:
        part.add_term(PolyFour([-0.5, 0.3, -0.16, 0.09], Bond(i, j)))
    check_gpos_part(system, part)
    check_vtens_part(system, part)

def test_compute_distances1_nimage_exclude():
    natom = 10
    nimage = 1
    cell = get_system_water32().cell
    pos = np.random.normal(0, 10, (natom, 3))
    exclude = np.array([[1, 0], [2, 1]])
    factor = (1+2*nimage)**3
    output = np.zeros(factor*(natom*(natom-1))//2 - 2, float)
    cell.compute_distances(output, pos, pairs=exclude, nimage=nimage)
    counter = 0
    ex_counter = 0
    for r0 in range(-1, 2):
        for r1 in range(-1, 2):
            for r2 in range(-1, 2):
                for i0 in range(natom):
                    for i1 in range(i0):
                        if (r0 == 0) and (r1 == 0) and (r2 == 0):
                            if ex_counter < len(exclude) and i0 == exclude[ex_counter,0] and i1 == exclude[ex_counter, 1]:
                                ex_counter += 1
                                continue
                        delta = pos[i0] - pos[i1]
                        cell.mic(delta)
                        cell.add_vec(delta, np.array([r0, r1, r2]))
                        assert abs(output[counter] - np.linalg.norm(delta)) < 1e-10
                        counter += 1
    assert ex_counter == 2

def test_add_part():
    system = get_system_water32()
    fn_pars = pkg_resources.resource_filename(__name__, '../../data/test/parameters_water_bondharm.txt')
    ff = ForceField.generate(system, fn_pars)
    part_press = ForcePartPressure(system, 1e-3)
    ff.add_part(part_press)
    assert part_press in ff.parts
    assert ff.part_press is part_press
    assert ff.compute() == ff.part_valence.energy + ff.part_press.energy

def test_clean_momenta_3d():
    sys = get_system_water32()
    sys.set_standard_masses()
    masses = sys.masses
    vel = get_random_vel(300, False, masses)
    com_mom = np.dot(masses, vel)
    assert abs(com_mom).max() > 1e-10
    clean_momenta(sys.pos, vel, masses, sys.cell)
    com_mom = np.dot(masses, vel)
    assert abs(com_mom).max() < 1e-10