class mvMovie:    
    def __init__(self):
        """Permanently stores what should appear in the movie."""
        self.movie=[]
        self.backmovie=[]        

mv = mvMovie()

# import some python modules
from pymol import cmd
import string
from math import *
from copy import deepcopy

# ---------------------------------------------------------
# internal stuff

def parseFrames(frames):
    """Returns a tuple of first and last frame from a string like
    1-10;
    10;
    99-200;
    5;
    """
    firstFrame=1
    lastFrame=1

    if string.find(frames,"-")>=0:
        # range info given
        t=string.split(frames,"-")
        firstFrame=int(t[0])
        lastFrame=int(t[1])
    else:
        # no range specified
        firstFrame=int(frames)
        lastFrame=firstFrame
    return (firstFrame,lastFrame)

# -------------------------------------
# here come the commands

def mvClear():
    """Deletes the movie."""
    mv.movie=[]
    mv.backmovie=[]
    cmd.mclear()
    cmd.frame(1)

def mvRot(frames="1",axis="z",totalAngle=360):
    """
    mvRot(frames,axis,totalAngle) - rotates the environment over the given
       frame range, the rotation angle summing up to the totalAngle given.
       """    
    frameRange = parseFrames(frames)
    nFrames = frameRange[1]-frameRange[0]+1
    angleIncrement = float(totalAngle)/(1.0*nFrames)

    i=frameRange[0]
    while i<=frameRange[1]:
        mv.movie.append((i,"turn %s,%f"%(axis,angleIncrement)))
        i+=1

def mvCxRot(frames="1",X=0,Y=0,Z=0):
    """
    mvCxRot(frames,Xangle,Yangle,Zangle) - Rotates the scene gradually
       over the three axes, corresponding to Euler angle rotations;
       For the end scene, consider rotation first around z-axis, then
       around y and finally around x.
       """
    frameRange = parseFrames(frames)
    nFrames = frameRange[1]-frameRange[0]+1
    Xinc = float(X)/(1.0*nFrames)
    Yinc = float(Y)/(1.0*nFrames)
    Zinc = float(Z)/(1.0*nFrames)

    i=frameRange[0]
    j=1
    while i<=frameRange[1]:
        mv.movie.append((i, "cmd.turn('x',%f);cmd.turn('y',%f);cmd.turn('z',%f);cmd.turn('y',%f);cmd.turn('x',%f)" %
                         (-1.0*(j-1)*Xinc, -1.0*(j-1)*Yinc, Zinc, j*Yinc, j*Xinc)))
        i+=1
        j+=1
        
def mvXrotYtoZturn(frames="1",X=0,YtoZ=0):
    """
    Actually this is another attempt to get these complex rotations going a
       bit smoothly, now trying to keep the thing rotating, while turning around
       the z-axis.
       """
    frameRange = parseFrames(frames)
    nFrames = frameRange[1]-frameRange[0]+1
    Xinc = float(X)/(1.0*nFrames)
    YtoZinc = float(YtoZ)/(1.0*nFrames)

    i=frameRange[0]
    j=1
    while i<=frameRange[1]:
        mv.movie.append((i, "cmd.turn('x',%f);cmd.turn('y',%f);cmd.turn('z',%f);cmd.turn('y',%f);cmd.turn('x',%f)" %
                         (-1.0*(j-1)*Xinc, -1.0*(j-1)*Yinc, Zinc, j*Yinc, j*Xinc)))
        i+=1
        j+=1
           
def mvMove(frames="1",axis="x",totalDistance="0"):
    """
    mvMove(frames,axis,totalDistance) - moves the environment over the given
       frame range, the moved distance summing up to the totalDistance given.
       """    
    frameRange = parseFrames(frames)
    nFrames = frameRange[1]-frameRange[0]+1
    distanceIncrement = float(totalDistance)/(1.0*nFrames)

    i=frameRange[0]
    while i<=frameRange[1]:
        mv.movie.append((i,"move %s,%f"%(axis,distanceIncrement)))
        i+=1

           
def mvCmd(frames="1",command=""):
    """
    mvCmd(frames,command) - executes a command in all frames specified.
       """    
    frameRange = parseFrames(frames)
    nFrames = frameRange[1]-frameRange[0]+1

    i=frameRange[0]
    while i<=frameRange[1]:
        mv.movie.append((i,command))
        i+=1

def mvSet(frames="1",variable="",startValue="0.0",endValue="1.0",selection=''):
    """
    mvSet(frames,variable,startValue,endValue,selection) - lets a variable go through a gradient in the specified frame range. Great for fading effects!
       """    
    frameRange = parseFrames(frames)
    nFrames = frameRange[1]-frameRange[0]+1

    stV=float(startValue)
    endV=float(endValue)
    increment = (endV-stV)/(1.0*nFrames)

    i=frameRange[0]
    j=0
    while i<=frameRange[1]:
        mv.movie.append((i,"set %s,%f,%s"%(variable,stV+j*increment,selection)))
        j+=1
        i+=1

def mvSinrot(frames="1",axis="z",totalAngle=360):
    """
    mvSinrot(frames,axis,totalAngle) - rotates the environment over the
       given frame range, the angles summing up to the totalAngle given.
       The incremental steps will be calculated with the sinus function to
       make the rotation smoother.
       """    
    frameRange = parseFrames(frames)
    nFrames = frameRange[1]-frameRange[0]+1    
    angle = float(totalAngle)
    arcIncrement=pi/(1.0*nFrames)
    
    i=frameRange[0]
    j=1
    prev=1.0
    while i<=frameRange[1]:
        arc=cos(j*arcIncrement)
        # print arc,"  *   ",abs(arc-prev)*0.5
        angleIncrement=angle*abs(arc-prev)*0.5
        mv.movie.append((i,"turn %s,%f"%(axis,angleIncrement)))
        prev=arc
        j+=1
        i+=1        

def mvSinmove(frames="1",axis="x",totalDistance="0"):
    """
    mvSinmove(frames,axis,totalDistance) - moves the environment over the given
       frame range, the moved distance summing up to the totalDistance given.
       The incremental steps will be calculated with the sinus function to
       make the movement smoother.
       """    
    frameRange = parseFrames(frames)
    nFrames = frameRange[1]-frameRange[0]+1    
    dist = float(totalDistance)
    arcIncrement=pi/(1.0*nFrames)
    
    i=frameRange[0]
    j=1
    prev=1.0
    while i<=frameRange[1]:
        arc=cos(j*arcIncrement)
        # print arc,"  *   ",abs(arc-prev)*0.5
        distanceIncrement=dist*abs(arc-prev)*0.5
        mv.movie.append((i,"move %s,%f"%(axis,distanceIncrement)))
        prev=arc
        j+=1
        i+=1        
    
def mvSinset(frames="1",variable="",startValue="0.0",endValue="1.0",selection=''):
    """
    mvSet(frames,variable,startValue,endValue) - lets a variable go through
    a gradient in the specified frame range. Great for fading effects!
    The incremental steps will be calculated with the sinus function to
    make the movement smoother.
    """        
    frameRange = parseFrames(frames)
    nFrames = frameRange[1]-frameRange[0]+1    
    arcIncrement=pi/(1.0*nFrames)
    stV=float(startValue)
    endV=float(endValue)
    
    i=frameRange[0]
    j=1
    prev=1.0
    sum=0.0
    while i<=frameRange[1]:
        arc=cos(j*arcIncrement)
        increment=(endV-stV)*abs(arc-prev)*0.5
        sum+=increment
        mv.movie.append((i,"set %s,%f,%s"%(variable,stV+sum,selection)))
        prev=arc
        j+=1
        i+=1        

def mvClip(frames="1",mode="near",distance="0"):
    frameRange = parseFrames(frames)
    nFrames = frameRange[1]-frameRange[0]+1
 
    dist=float(distance)
    increment = (dist)/(1.0*nFrames)

    i=frameRange[0]
    j=0
    while i<=frameRange[1]:
        mv.movie.append((i,"clip %s,%f"%(mode,increment)))
        i+=1

def mvGradCol(frames="1",selection="",tmpc="",startR="1.0",startG="1.0",startB="1.0",endR="0.0",endG="0.0",endB="0.0"):
    """
    mvGradCol(frames,selection,tmpc,R1,G1,B1,R2,G2,B2) - changes colour
    gradually from (R1,G1,B1) to (R1,G1,B1) through the specified frame range
    and on the specified selection.
       """
    frameRange = parseFrames(frames)
    nFrames = frameRange[1]-frameRange[0]+1

    stRV=float(startR)
    endRV=float(endR)
    incR = (endRV-stRV)/(1.0*nFrames)
    
    stGV=float(startG)
    endGV=float(endG)
    incG = (endGV-stGV)/(1.0*nFrames)
    
    stBV=float(startB)
    endBV=float(endB)
    incB = (endBV-stBV)/(1.0*nFrames)

    i=frameRange[0]
    j=0
    while i<=frameRange[1]:
        mv.movie.append((i,"set_color %s,[ %f,%f,%f ]"%(tmpc,stRV+j*incR,stGV+j*incG,stBV+j*incB)))
        mv.movie.append((i,"color %s, %s"%(tmpc,selection)))
        j+=1
        i+=1

def mvBG(frames='1',R1=0.0,B1=0.0,G1=0.0,R2=0.0,G2=0.0,B2=0.0,):
    """
    mvBG lets the background colour fade from RGB1 to RGB2. Isn't that a nice feature ;)
    """
    frameRange = parseFrames(frames)
    nFrames = frameRange[1]-frameRange[0]+1

    R1=float(R1)
    R2=float(R2)
    Ri = (R2-R1)/(1.0*nFrames)
    
    G1=float(G1)
    G2=float(G2)
    Gi = (G2-G1)/(1.0*nFrames)
    
    B1=float(B1)
    B2=float(B2)
    Bi = (B2-B1)/(1.0*nFrames)
    
    i=frameRange[0]
    j=0
    while i<=frameRange[1]:
        mv.movie.append((i,"cmd.set('bg_rgb',[%f,%f,%f])" %(R1+j*Ri,G1+j*Gi,B1+j*Bi)))
        j+=1
        i+=1

def mvSinBG(frames='1',R1=0.0,B1=0.0,G1=0.0,R2=0.0,G2=0.0,B2=0.0,):
    """
    This does essentially the same as the one above, but uses the gradual
    transition as the other mvSin* commands. Even better! :) (Why would you
    ever need this?)
    """
    frameRange = parseFrames(frames)
    nFrames = frameRange[1]-frameRange[0]+1    
    arcIncrement=pi/(1.0*nFrames)

    R1=float(R1)
    R2=float(R2)
    Ri = (R2-R1)/(1.0*nFrames)
    
    G1=float(G1)
    G2=float(G2)
    Gi = (G2-G1)/(1.0*nFrames)
    
    B1=float(B1)
    B2=float(B2)
    Bi = (B2-B1)/(1.0*nFrames)
    
    i=frameRange[0]
    j=1
    prev=1.0
    Rt=0.0
    Gt=0.0
    Bt=0.0
    while i<=frameRange[1]:
        arc=cos(j*arcIncrement)
        Ri=(R2-R1)*abs(arc-prev)*0.5
        Rt+=Ri
        Gi=(G2-G1)*abs(arc-prev)*0.5
        Gt+=Gi
        Bi=(B2-B1)*abs(arc-prev)*0.5
        Bt+=Bi
        mv.movie.append((i,"cmd.set('bg_rgb',[%f,%f,%f])"%(Rt,Gt,Bt)))
        prev=arc
        j+=1
        i+=1        

def mvMorph(frames='1',source='',target='',filename='morph_script.pml'):
    """
    Atoms from source selection one will be _alter_ed to be on positions from
    target atoms at end time. This is very time consuming! It writes and uses
    an external script which is read every frame, since the list of commands
    will be very long usually.
    """
    frameRange = parseFrames(frames)
    nFrames = 1.0*(frameRange[1]-frameRange[0]+1)

    sourceobj = cmd.get_model(source)
    targetobj = cmd.get_model(target)
    natoms = len(sourceobj.atom)
    if not len(targetobj.atom) == natoms:
        print "Something wrong, target selection is of different size as source!"
        return

    if filename[-4:] == '.pml':
        filename = filename[:-4]
    filename = filename + repr(frameRange[0]) + '.pml'
    fout = open(filename, 'a')

    for i in range(natoms):
        r1 = sourceobj.atom[i].coord
        r2 = targetobj.atom[i].coord
        if r1 != r2:
            shift = [ (r2[0]-r1[0])/(nFrames - 1),
                      (r2[1]-r1[1])/(nFrames - 1),
                      (r2[2]-r1[2])/(nFrames - 1) ]
            atomid = 'id %d' % sourceobj.atom[i].id
            if shift[0]:
                fout.write('alter_state 1,%s,x=x+ %f\n' % (source + ' and ' + atomid, shift[0]))
            if shift[1]:
                fout.write('alter_state 1,%s,y=y+ %f\n' % (source + ' and ' + atomid, shift[1]))
            if shift[2]:
                fout.write('alter_state 1,%s,z=z+ %f\n' % (source + ' and ' + atomid, shift[2]))

    i=frameRange[0]
    j=0
    while i<=frameRange[1]:
        mv.movie.append((i,"cmd.do('@%s')"%filename))
        j+=1
        i+=1


def mvDel(frames="1",S1="all",S2="all",Rstart=50):
    frameRange = parseFrames(frames)
    nFrames = frameRange[1]-frameRange[0]+1
 
    Rstart=float(Rstart)
    Rinc = (Rstart)/(1.0*nFrames)

    i=frameRange[0]
    j=0
    while i<=frameRange[1]:
        mv.movie.append((i,"remove %s and not (%s within %f of %s)"%(S1,S1,Rstart - j*Rinc,S2)))
        i+=1
        j+=1


def mvForward():
    """Creates the movie and plays it."""
    # find out the frame range
    nFrames=1
    for m in mv.movie:
        if m[0]>nFrames:
            nFrames=m[0]
            
    # Specify movie length
    print "creating movie with %i frames."%(nFrames)
    cmd.mset("1 x%i"%(nFrames))

    # create empty frame-2do-lists
    do=["zero frame is unused"]
    for i in range(nFrames):
        do.append("")

    # push all movie commands to the 2do-list
    for m in mv.movie:
        do[m[0]]+=m[1]+";"

    # now let action happen in the frames
    i=1
    while i<=nFrames:
        cmd.mdo(i,do[i])
        i+=1                

    # start the movie
    cmd.mplay()

def mvBackward():
    # not implemented yet
    return

# ----------------------------------------------

def fwd(amount):
    for i in range(int(amount)):
        mvMovie.append("move z,5")
        backmovie.append("move z,-5")

# -----------------------------------------------------------------

# define a set of new PyMOL commands for creation of movies
# this part will be executed first upon the 'run movie.py' command from within PyMOL.
cmd.extend('movie',mvForward)
# cmd.extend('movie_back',mvBackward)
cmd.extend('mvClear',mvClear)

cmd.extend('mvRot',mvRot)
cmd.extend('mvMove',mvMove)
cmd.extend('mvCmd',mvCmd)
cmd.extend('mvSet',mvSet)

cmd.extend('mvSinrot',mvSinrot)
cmd.extend('mvSinmove',mvSinmove)
cmd.extend('mvSinset',mvSinset)

# TAW additions
cmd.extend('mvGradCol',mvGradCol)
cmd.extend('mvClip',mvClip)
cmd.extend('mvBG',mvBG)
cmd.extend('mvSinBG',mvSinBG)
cmd.extend('mvMorph',mvMorph)
cmd.extend('mvCxRot',mvCxRot)
cmd.extend('mvDel',mvDel)