from __future__ import print_function
import IMP
import IMP.test
import sys
import IMP.em
import IMP.multifit
import os
try:
    from time import process_time  # needs python 3.3 or later
except ImportError:
    from time import clock as process_time


class Tests(IMP.test.TestCase):

    """Tests for sampled density maps"""

    def METHOD_NAME(self):
        """initialize IMP environment create particles"""
        IMP.test.TestCase.METHOD_NAME(self)
        # init IMP model ( the environment)
        self.mdl = IMP.Model()
        self.sel = IMP.atom.CAlphaPDBSelector()
        self.mhs = IMP.atom.Hierarchies()
        self.mhs.append(
            IMP.atom.read_pdb(self.get_input_file_name("1z5s_A.pdb"),
                              self.mdl, self.sel))
        self.mhs.append(
            IMP.atom.read_pdb(self.get_input_file_name("1z5s_C1.pdb"),
                              self.mdl, self.sel))
        self.voxel_size = 1.
        self.rbs = []
        for mh in self.mhs:
            IMP.atom.add_radii(mh)
            IMP.multifit.add_surface_index(mh, self.voxel_size)
            self.rbs.append(IMP.atom.create_rigid_body(mh))
        restraints = []
        # set the restraint
        hub = IMP.core.HarmonicUpperBound(0, 1)
        sdps = IMP.core.SphereDistancePairScore(hub)
        rdps = IMP.core.RigidBodyDistancePairScore(sdps,
                                           IMP.multifit.RigidLeavesRefiner())
        lsc = IMP.container.ListSingletonContainer(self.mdl, self.rbs)
        self.c_r = IMP.core.ConnectivityRestraint(rdps, lsc)

        restraints.append(self.c_r)
        self.sf = IMP.core.RestraintsScoringFunction(restraints)
        print("going to evaluate 2")
        self.sf.evaluate(False)
        self.wev_r = IMP.multifit.create_weighted_excluded_volume_restraint(
            self.rbs[0], self.rbs[1])

        restraints.append(self.wev_r)
        self.sf = IMP.core.RestraintsScoringFunction(restraints)
        print("going to evaluate 1")
        self.sf.evaluate(False)
        print("end setup")

    def test_weighted_excluded_volume_restraint(self):
        """Check that weighted excluded volume restraint works"""
        # IMP.set_log_level(IMP.VERBOSE)
        rotations = [[0.960739, 0.177613, -0.196201, 0.0833023],
                     [0.98373, -0.0268444, -0.115434, -0.135015],
                     [0.995413, 0.0545123, -0.0635521, 0.0462946],
                     [0.99739, 0.0503421, -0.000272958, -0.0517479],
                     [0.953478, 0.148336, -0.198021, 0.17223],
                     [0.994239, 0.0570374, -0.0140131, 0.089658],
                     [0.98401, 0.148801, -0.00937242, 0.0974403],
                     [0.864194, -0.251682, 0.325705, -0.28938],
                     [0.952222, -0.155235, -0.253447, -0.0702893],
                     [0.881024, -0.405223, 0.154964, -0.188619]]
        translations = [IMP.algebra.Vector3D(7.37269, -32.1143, -68.123),
                        IMP.algebra.Vector3D(-26.968, 30.7037, -18.9437),
                        IMP.algebra.Vector3D(-0.990538, -13.8685, -18.359),
                        IMP.algebra.Vector3D(-4.79513, 3.45006, -9.39788),
                        IMP.algebra.Vector3D(18.2439, -40.6508, -62.0047),
                        IMP.algebra.Vector3D(10.7605, -26.9859, -11.0509),
                        IMP.algebra.Vector3D(13.1573, -34.8041, -26.6549),
                        IMP.algebra.Vector3D(46.1903, 87.4569, 58.9899),
                        IMP.algebra.Vector3D(-32.0694, 24.0887, -33.828),
                        IMP.algebra.Vector3D(20.1398, 111.715, 60.5263)]
        # move chain A and calculate weighted excluded volume
        ps1 = IMP.core.get_leaves(self.mhs[0])
        IMP.set_log_level(IMP.SILENT)  # VERBOSE)
        for i in range(10):
            t = IMP.algebra.Transformation3D(
                IMP.algebra.Rotation3D(
                    rotations[i][0],
                    rotations[i][1],
                    rotations[i][2],
                    rotations[i][3]),
                translations[i])
            xyz = IMP.core.XYZ(self.rbs[0].get_particle())
            xyz.set_coordinates(t.get_transformed(xyz.get_coordinates()))
            # check that when the proteins are not connected (self.c_r>0) the excluded volume
            # restraint is bigger than 0
            start = process_time()
            # to make sure the coordinates were transformed
            self.mdl.update()
            end = process_time()
            print("Time elapsed for PairRestraint evaluation = ", end - start, "seconds")
            conn_r = self.c_r.evaluate(False)
            w_exc_vol_r = self.wev_r.evaluate(False)
            self.assertTrue(
                ((conn_r == 0.)and(w_exc_vol_r > 0.))or(
                    (conn_r > 0.)and(w_exc_vol_r == 0.)),
                "inconsistency between connectivity and excluded volume restraint")
            # print "connectivity "+str(conn_r) + " | excluded-volume " +
            # str(w_exc_vol_r)
            xyz.set_coordinates(
                t.get_inverse().get_transformed(xyz.get_coordinates()))
if __name__ == '__main__':
    IMP.test.main()