# Wave branches of RNA at frequencies corresponding to branch size.
# This is to to add some life to the RNA molecule.
#
# runscript waggle.py <mol> <pairing-file> <angle> <speed> <frames>

mid = int(arguments[0])
from chimera import openModels as om
m = om.list(id = mid)[0]

pairing_file = arguments[1]
angle, speed = [float(a) for a in arguments[2:4]]
frames = int(arguments[4])

def stem_ends(pairing_file):

    from RNALayout import rna_layout
    stems = rna_layout.read_base_pairs(pairing_file)
    pairs = []
    for s,e,l in stems:
        pairs.append((s,e))
        if l > 1:
            pairs.append((s+l-1,e-l+1))
    return pairs

def rotate_branch(coords, pivot1, pivot2, r1, r2, angle):

    p1, p2 = coords[pivot1], coords[pivot2]
    import Matrix as M
    tf = M.rotation_transform(p2-p1, angle, p1)
    M.transform_points(coords[r1:r2,:], tf)

class branch_wiggler:

    def __init__(self, mol, pairs, angle, speed):

        self.mol = mol
        self.angle = angle
        self.speed = float(speed)

        cpairs = []
        for s, e in pairs:
            r1, r2 = mol.findResidue(s), mol.findResidue(e)
            if r1 is None or r2 is None:
                continue
            a1, a2 = r1.findAtom('P'), r2.findAtom("O3'")
            if a1 is None or a2 is None:
                continue
            c1 = min(a.coordIndex for a in r1.atoms)
            c2 = max(a.coordIndex for a in r2.atoms) + 1
            i1, i2 = a1.coordIndex, a2.coordIndex
            w = e-s+1
            cpairs.append((c1,c2,i1,i2,w))
        self.cpairs = cpairs

        from numpy import array, float32
        self.coords = array([p.data() for p in mol.activeCoordSet.coords()],
                            float32)

    def rotate_branches(self, frame):

        s = self.speed
        ar = self.angle
        xyz = self.coords
        from math import sin, pi
        for c1,c2,i1,i2,w in self.cpairs:
            a = sin(pi * frame * s/w) * ar
            ap = (sin(pi * (frame-1) * s/w) * ar) if frame > 0 else 0
            rotate_branch(xyz, i1, i2, c1, c2, a-ap)

        from chimera import Point
        for a in self.mol.atoms:
            a.setCoord(Point(*xyz[a.coordIndex]))
        print 'frame', frame, 'rotated', len(self.cpairs)

def wiggle(mol, pairs, angle, speed, frames):

    bw = branch_wiggler(mol, pairs, angle, speed)
    params = [None, None]
    def new_frame_cb(t, unused, f, bw=bw, frames=frames, params=params):
        h, fstart = params
        if fstart is None:
            params[1] = fstart = f
        t = f - fstart
        if t >= frames:
            triggers.deleteHandler('new frame', h)
        else:
            bw.rotate_branches(t)
    from chimera import triggers
    h = triggers.addHandler('new frame', new_frame_cb, frames)
    params[0] = h

pairs = stem_ends(pairing_file)
wiggle(m, pairs, angle, speed, frames)