import unittest
import IMP
import IMP.core
import IMP.algebra
import IMP.atom
import IMP.rmf
import math
import IMP.test

class GenericTest(IMP.test.TestCase):
    def setup_protein():
        # TMH sequences
        seq=[]
        seq.append(("L","F","A","Y","I","V","F","L","I","T","M","F","I","N","V","S","I","L","I"))
        seq.append(("G","I","E","R","F","A","K","I","A","M","P","T","L","F","I","L","A","V","F","L","V","I"))
        # residue index
        TMH= [[24,42], [75,96]]
        return seq,TMH
    def create_representation(m,seq,tmh):
        all=IMP.atom.Hierarchy.setup_particle(IMP.Particle(m))
        def generate_tm(id,seq,tmh):
            pm=IMP.Particle(m)
            tm=IMP.atom.Molecule.setup_particle(pm)
            tm.set_name("TM"+str(id))
            cc=IMP.atom.Chain.setup_particle(pm, chr(65+id))
            nres=len(seq)
            atoms=[]
            for i in range(nres):
                x=2.3*math.cos(math.radians(100.0)*float(i))
                y=2.3*math.sin(math.radians(100.0)*float(i))
                z=1.51*(float(i)-float((nres-1))/2.0)
                # set up residue
                p=IMP.Particle(m)
                r=IMP.atom.Residue.setup_particle(p, IMP.atom.get_residue_type(seq[i]), i+tmh[0])
                rd=IMP.core.XYZR.setup_particle(p, IMP.algebra.Sphere3D(IMP.algebra.Vector3D(x,y,z),2.273))
                # set up atom
                p1=IMP.Particle(m)
                a=IMP.atom.Atom.setup_particle(p1, IMP.atom.AT_CA)
                ad=IMP.core.XYZR.setup_particle(p1, IMP.algebra.Sphere3D(IMP.algebra.Vector3D(x,y,z),2.273))
                r.add_child(a)
                tm.add_child(r)
                atoms.append(ad)
            all.add_child(tm)
            rb=IMP.Particle(m)
            rb=IMP.atom.create_rigid_body(atoms,"TM"+str(id))

        for i in range(len(seq)):
            generate_tm(i,seq[i],tmh[i])
        return all

    def test_perturbed(self):
        """Test creating a handing of rigid bodies"""
        m=IMP.Model()
        print "reading protein data"
        (seq,TMH)=self.setup_protein()
        print "creating representation"
        protein=self.create_representation(m,seq,TMH)

        # open file for writing
        rh = IMP.rmf.RootHandle("traj.hdf5", True)
        # adding hierarchies
        for hs in protein.get_children():
            IMP.rmf.add_hierarchy(rh, hs)

        print "sampling"
for steps in range(100):
    mc.optimize(100)
    print steps, m.evaluate(False)
    for hs in protein.get_children():
        IMP.rmf.save_frame(rh, steps+1, hs)

# close file
del rh

print "reading trajectory"
# open file for reading
rh = IMP.rmf.RootHandle("traj.hdf5", False)
# setting hierarchies
IMP.rmf.set_hierarchies(rh, protein.get_children())

for steps in range(100):
    for hs in protein.get_children():
        IMP.rmf.load_frame(rh, steps+1, hs)
    print steps, m.evaluate(False)   

if __name__ == '__main__':
    unittest.main()