import unittest
import sys
mv = None
ct = 0
totalCt = 4

#define the 'showwarning'  variable that is used in the code returned by maa.getSourceCode()
showwarning = False

class CustomAnimations_Tests(unittest.TestCase):

    def setUp(self):
        """Create DejaVu Viewer
        """
        global mv
        if mv is None:
            from Pmv.moleculeViewer import MoleculeViewer
            mv = MoleculeViewer(customizer = './.empty', logMode = 'no',
                                withShell=0, trapExceptions=False)
                        #mv.setUserPreference(('warningMsgFormat', 'printed'), log = 0)
            mv.setUserPreference(('trapExceptions', '0'), log = 0)
            mv.loadModule('dejaVuCommands', 'ViewerFramework')
            mv.browseCommands("fileCommands",commands=['readMolecule',],package= 'Pmv')
            mv.browseCommands("bondsCommands", commands=["buildBondsByDistance",],package= "Pmv")
            mv.setOnAddObjectCommands(['buildBondsByDistance','displayLines'], log=0)
            mv.browseCommands('displayCommands',package='Pmv')
            mv.browseCommands('colorCommands',package='Pmv')
            mv.browseCommands('secondaryStructureCommands',package='Pmv')
            import os
            print os.getcwd()
            mv.readMolecule('../../Tests/Data/1crn.pdb')
        self.mv = mv 


    def tearDown(self):
        """
        clean-up
        """
        global ct, totalCt
        #delete any molecules left due to errors
        #for m in self.mv.Mols:
        #    self.mv.deleteMol(m)
        ct = ct + 1
        if ct==totalCt:
            print 'destroying mv'
            self.mv.Exit(0)
            del self.mv        



    def test_PartialFadeMolGeomsMAA_in(self):
        """ Test creation of PartialFadeMolGeom (in) and playing different frames of the maa"""
        from Pmv.scenarioInterface.animations import PartialFadeMolGeomsMAA
        import numpy
        self.mv.displayCPK("1crn: :THR2,ILE7;")
        cpk = self.mv.GUI.VIEWER.FindObjectByName("root|1crn|cpk")
        assert cpk.name == "cpk"
        #cpk.Set(visible = False)
        # test fade in
        nds = "1crn: :THR2,ILE7;"
        maa1 = PartialFadeMolGeomsMAA(cpk, nodes = nds,
                   pmv=self.mv, kfpos = [0, 40], fade = "in", destValue = 1.0)
        #check we can reproduce the maa from it's sourcecode:
        maa2 =  None
        # we need to define 'pmv' and 'viewer' variables since they are used in
        # the code returned by getSourceCode()
        pmv = self.mv
        viewer = self.mv.GUI.VIEWER
        maasrc = maa1.getSourceCode("maa2")
        exec(maasrc)
        assert maa2 != None
        for maa in [maa1, maa2]:
            actors = maa.actors
            self.assertEqual(len(actors), 3)
            initval = cpk.materials[1028].prop[1][:,3]
            # play 0, 20, 40 frames of the maa and check the animated attribute(opacity)
            # values
            maa.setValuesAt(0)
            self.mv.GUI.VIEWER.OneRedraw()
            self.assertEqual(cpk.visible, True)
            val = cpk.materials[1028].prop[1][:,3]
            #print "val1:", val
            self.assertEqual(sum(val)/len(val), 0)
            #self.assertEqual(len(val), 15)
            maa.setValuesAt(20)
            self.mv.GUI.VIEWER.OneRedraw()
            val = cpk.materials[1028].prop[1][:,3]
            #print "val2:", val
            self.assertEqual(sum(val)/len(val), 0.5)
            #self.assertEqual(len(val), 15)
            maa.setValuesAt(40)
            self.mv.GUI.VIEWER.OneRedraw()
            val = cpk.materials[1028].prop[1][:,3]
            #print "val3:", val
            self.assertEqual(sum(val)/len(val), 1.0)
            #self.assertEqual(len(val), 15)

            cpk.Set(opacity = 0.5, inheritMaterial= 0, transparent='implicit')
            val = cpk.materials[1028].prop[1][:,3]
            self.assertEqual(sum(val)/len(val), 0.5)
            # run maa and check that the afterAnimation_cb set the animated attribute to
            # the original values        
            maa.run()
            self.assertEqual(cpk.visible, True)
            val = cpk.materials[1028].prop[1][:,3]
            #print "val4:", val
            self.assertEqual(numpy.alltrue(numpy.equal(val, initval)), True)


    def test_PartialFadeMolGeomsMAA_out(self):
        """ Test creation of PartialFadeMolGeom (out) and playing different frames of the maa"""        
        from Pmv.scenarioInterface.animations import PartialFadeMolGeomsMAA
        import numpy
        self.mv.displayCPK("1crn", log=0, cpkRad=0.0, quality=0, only=False, negate=False, scaleFactor=1.0)
        cpk = self.mv.GUI.VIEWER.FindObjectByName("root|1crn|cpk")
        cpk.Set(visible = True)
        assert cpk.name == "cpk"
        
        # test fade out
        nds = "1crn: :THR2,ILE7;"
        maa1 = PartialFadeMolGeomsMAA(cpk, nodes = nds,
                   pmv=self.mv, kfpos=[0, 20], fade = "out", destValue = 0.0)
        #check we can reproduce the maa from it's sourcecode:
        maa2 =  None
        # we need to define 'pmv' and 'viewer' variables since they are used in
        # the code returned by getSourceCode()
        pmv = self.mv
        viewer = self.mv.GUI.VIEWER
        maasrc = maa1.getSourceCode("maa2")
        exec(maasrc)
        assert maa2 != None
        for maa in [maa1, maa2]:
            actors = maa.actors
            self.assertEqual(len(actors), 3)
            initval = cpk.materials[1028].prop[1][:,3]
            # play 0, 10, 20 frames of the maa and check the animated attribute(opacity)
            # values:
            maa.setValuesAt(0)
            self.mv.GUI.VIEWER.OneRedraw()
            val = cpk.materials[1028].prop[1][:,3]
            #print "val1:", val
            self.assertEqual(sum(val)/len(val), 1.0)

            maa.setValuesAt(10)
            self.mv.GUI.VIEWER.OneRedraw()
            val = cpk.materials[1028].prop[1][:,3]
            #print "val2:", val
            # current opacity array should contain 1.0s and 0.5s (for spheres that are fading out).
            # To check that there are 15 elements in the array that == 0.5 and the rest of them
            # == 1.0, replace all 0.5 in the val array with 0.0 and compute the sum of
            # all array elements
            val1 = sum(numpy.choose(numpy.equal(val, 0.5), (val, 0.)))

            self.assertEqual(val1 , len(val)-15)

            maa.setValuesAt(20)
            self.mv.GUI.VIEWER.OneRedraw()
            val = cpk.materials[1028].prop[1][:,3]
            # now there should be 15 array entries == 0.0 and the rest == 1.0
            #print "val3:", val
            self.assertEqual(sum(val) , len(val)-15)
            val3 = val
            # run maa 
            maa.run()
            val = cpk.materials[1028].prop[1][:,3]
            #print 'val', val
            #print 'initval', initval
            #self.assertEqual(numpy.alltrue(numpy.equal(val, initval)), True)
            self.assertEqual(numpy.alltrue(numpy.equal(val, val3)), True)


    def test_PmvColorObjectMAA_0(self):
        """Test creation of  PmvColorObjectMAA object. Check values of animated parameters
        of the object at different frames"""
        from Pmv.scenarioInterface.animations import PmvColorObjectMAA, colorCmds
        import numpy
        cpk = self.mv.GUI.VIEWER.FindObjectByName("root|1crn|cpk")
        self.mv.displayCPK("1crn", log=0, cpkRad=0.0, quality=0, only=False, negate=False, scaleFactor=1.0)
        cpk.Set(visible = False)
        viewer = self.mv.GUI.VIEWER
        pmv = self.mv
        for color in colorCmds.keys():
            # create maa for all available color types:
            if color == 'color by second.struct.': continue
            initval = cpk.materials[1028].prop[1][:, :3]
            #print "initval:", initval[:10]
            maa1 = PmvColorObjectMAA(cpk, nbFrames=20, pmv=pmv,
                                    colortype=color)
            #check we can reproduce the maa from it's sourcecode:
            maa2 =  None
            maasrc = maa1.getSourceCode("maa2")
            exec(maasrc)
            assert maa2 != None
            for i, maa in enumerate([maa1, maa2]):
                #print i, "color type", color, maa.name
                self.assertEqual(len(maa.actors), 3)
                # check the values of object's colors at frames 0, 10, 20  
                maa.setValuesAt(0)
                self.assertEqual(cpk.visible, True)
                val0 = cpk.materials[1028].prop[1][:, :3]
                maa.setValuesAt(10)
                viewer.OneRedraw()
                val1 = cpk.materials[1028].prop[1][:, :3]
                #assert  numpy.alltrue(numpy.equal(val0, val1)) == False
                maa.setValuesAt(20)
                viewer.OneRedraw()
                val2 = cpk.materials[1028].prop[1][:, :3]
                #assert  numpy.alltrue(numpy.equal(val2, val1)) == False
                maa.afterAnimation_cb()
                #check that we get original values 
                #self.assertEqual(cpk.visible, False)
                val = cpk.materials[1028].prop[1][:, :3]
                #self.assertEqual(numpy.alltrue(numpy.equal(val, initval)), True)
                #maa.run()
            #print "--------------------------------"


    def test_PmvColorObjectMAA_1(self):
        """ Test creation of  PmvColorObjectMAA for a 'complex' object -
        parent geometry that has 'children'
        """
        from Pmv.scenarioInterface.animations import PmvColorObjectMAA

        cpk = self.mv.GUI.VIEWER.FindObjectByName("root|1crn|cpk")
        cpk.Set(visible = False)
        self.mv.displayExtrudedSS("1crn", negate=False, only=False, log=0)
        ss = self.mv.GUI.VIEWER.FindObjectByName("root|1crn|secondarystructure")
        assert ss !=  None
        
        maa1 = PmvColorObjectMAA(ss, nbFrames=20, pmv=self.mv,
                                      colortype="color by second.struct.")
        
        #check we can reproduce the maa from it's sourcecode:
        maa2 =  None
        # we need to define 'pmv' and 'viewer' variables since they are used in
        # the code returned by getSourceCode()
        pmv = self.mv
        viewer = self.mv.GUI.VIEWER
        maasrc = maa1.getSourceCode("maa2")
        exec(maasrc)
        assert maa2 != None
        for maa in [maa1, maa2]:
            #print maa
            sschildren = ss.children
            actorsDict = {}
            actors = maa.actors
            for actor in actors:
                obj = actor.object
                if obj in sschildren:
                    oname = obj.name
                    if not actorsDict.get(oname):
                        actorsDict[oname] = []
                    if actor.name.find("visible") > 0 or actor.name.find("color"):
                        actorsDict[oname].append(actor.name)
            #print "len actors:", len(actors)
            self.assertEqual(len(actorsDict), len(sschildren))
            for item in actorsDict.values():
                self.assertEqual(len(item), 2)
            #print actorsDict
            maa.run()
        
