# This script loads a protein structure in PDB format and then it creates a nurbs curve following its backbone ribbon.
# Written by Andres Colubri, May 2005.

from PL import *
from Blender import Curve, Object, Scene, Window
from Blender import Mathutils
from Blender.Mathutils import *

global cp_center
global cp_side1
global cp_side2
global cp_up

global normal_vector
global flip_test_vector
global ribbon_width
global ribbon_height

global curve_center
global curve_side1
global curve_side2
global curve_up

def init_curves():
    global cp_center
    global cp_side1
    global cp_side2
    global cp_up

    global curve_center
    global curve_side1
    global curve_side2
    global curve_up

    curve_center.appendNurb([cp_center[2][0], cp_center[2][1], cp_center[2][2], 100])
    curve_side1.appendNurb([cp_side1[2][0], cp_side1[2][1], cp_side1[2][2], 100])
    curve_side2.appendNurb([cp_side2[2][0], cp_side2[2][1], cp_side2[2][2], 100])
    curve_up.appendNurb([cp_up[2][0], cp_up[2][1], cp_up[2][2], 100])

def add_point_to_curves():
    global cp_center
    global cp_side1
    global cp_side2
    global cp_up

    global curve_center
    global curve_side1
    global curve_side2
    global curve_up

    curve_center[0].append([cp_center[3][0], cp_center[3][1], cp_center[3][2], 100])
    curve_side1[0].append([cp_side1[3][0], cp_side1[3][1], cp_side1[3][2], 100])
    curve_side2[0].append([cp_side2[3][0], cp_side2[3][1], cp_side2[3][2], 100])  
    curve_up[0].append([cp_up[3][0], cp_up[3][1], cp_up[3][2], 100])

def init_cp():
    global cp_center
    global cp_side1
    global cp_side2
    global cp_up

    cp_center = [Vector([0.0, 0.0, 0.0])] * 4
    cp_side1 = [Vector([0.0, 0.0, 0.0])] * 4
    cp_side2 = [Vector([0.0, 0.0, 0.0])] * 4
    cp_up = [Vector([0.0, 0.0, 0.0])] * 4

def init_all():
    global normal_vector
    global flip_test_vector
    global ribbon_width
    global ribbon_height

    init_cp()
    normal_vector = Vector([0.0, 0.0, 0.0])
    flip_test_vector = Vector([0.0, 0.0, 0.0])
    ribbon_width = 1.0
    ribbon_height = 1.0

def first_res(n, i):
    return plFirstAA(n) == i

def last_res(n, i):
    return plLastAA(n) == i

def shift_control_points():
    global cp_center
    global cp_side1
    global cp_side2
    global cp_up

    cp_center[0] = cp_center[1]
    cp_side1[0] = cp_side1[1]
    cp_side2[0] = cp_side2[1]
    cp_up[0] = cp_up[1]

    cp_center[1] = cp_center[2]
    cp_side1[1] = cp_side1[2]
    cp_side2[1] = cp_side2[2]
    cp_up[1] = cp_up[2]

    cp_center[2] = cp_center[3]
    cp_side1[2] = cp_side1[3]
    cp_side2[2] = cp_side2[3]
    cp_up[2] = cp_up[3]

#Adds a new control point to the arrays cp_center, cp_side1 and cp_side2
def add_control_points(n, i):
    global cp_center
    global cp_side1
    global cp_side2
    global cp_up

    global normal_vector
    global flip_test_vector
    global ribbon_width
    global ribbon_height

    # Getting atom coordinates.
    v = new_doubleArray(3)
    plAtomPos(n, i, plAT(n, i, MC_CA), v)
    ca0 = Vector([float(doubleArray_getitem(v, 0)),
                  float(doubleArray_getitem(v, 1)),
                  float(doubleArray_getitem(v, 2))])
    plAtomPos(n, i, plAT(n, i, MC_O), v)
    oxi = Vector([float(doubleArray_getitem(v, 0)),
                  float(doubleArray_getitem(v, 1)),
                  float(doubleArray_getitem(v, 2))])
    plAtomPos(n, i + 1, plAT(n, i + 1, MC_CA), v)
    ca1 = Vector([float(doubleArray_getitem(v, 0)),
                  float(doubleArray_getitem(v, 1)),
                  float(doubleArray_getitem(v, 2))])
    delete_doubleArray(v)

    a = ca1 - ca0
    b = oxi - ca0

    # Vector normal to the peptide plane (pointing outside in the case of the
    # alpha helix.
    c = CrossVecs(a, b)

    # Vector contained in the peptide plane (perpendicular to its direction).
    d = CrossVecs(c, a)

    # Normalizing vectors.
    c.normalize()
    d.normalize()

    # Flipping test (to avoid self crossing in the strands).
    if 90.0 < AngleBetweenVecs(flip_test_vector, d):
        # Flip detected. The plane vector is inverted.
        d = -1.0 * d
        c = -1.0 * c

    # The central control point is constructed.
    cp_center[3] = 0.5 * (ca0 + ca1)

    # The control points for the side ribbons are constructed.
    cp_side1[3] = cp_center[3] + ribbon_width * d
    cp_side2[3] = cp_center[3] - ribbon_width * d
    cp_up[3] = cp_center[3] + ribbon_height *  c;

    normal_vector = c
    # Saving the plane vector (for the flipping test in the next call).
    flip_test_vector = d

def construct_control_points(n, i):
    global cp_center
    global cp_side1
    global cp_side2
    global cp_up

    global normal_vector
    global flip_test_vector
    global ribbon_width
    global ribbon_height

    if first_res(n, i):
        # The control points 2 and 3 are created.
        normal_vector = Vector([0.0, 0.0, 0.0])
        flip_test_vector = Vector([0.0, 0.0, 0.0])
        add_control_points(n, i)
        cp_center[2] = cp_center[3]
        cp_side1[2] = cp_side1[3]
        cp_side2[2] = cp_side2[3]
        cp_up[2] = cp_up[3]
        add_control_points(n, i + 1)

        # We still need the two first control points.
        # Moving backwards along the cp_center[2] - cp_center[3] direction.
        cp_center[1] = 2.0 * cp_center[2] - cp_center[3]
        cp_side1[1] = cp_center[1] + ribbon_width * flip_test_vector
        cp_side2[1] = cp_center[1] - ribbon_width * flip_test_vector
        cp_up[1] = cp_center[1] + ribbon_height * normal_vector

        cp_center[0] = 2.0 * cp_center[1] - cp_center[2]
        cp_side1[0] = cp_center[0] + ribbon_width * flip_test_vector
        cp_side2[0] = cp_center[0] - ribbon_width * flip_test_vector
        cp_up[0] = cp_center[0] + ribbon_height * normal_vector

        init_curves()
    else:
        shift_control_points()
        if last_res(n, i) or last_res(n, i + 1):
            # Moving forward along the cp_center[1] - cp_center[2] direction.
            cp_center[3] = 2.0 * cp_center[2] - cp_center[1]
            cp_side1[3] = cp_center[3] + ribbon_width * flip_test_vector
            cp_side2[3] = cp_center[3] - ribbon_width * flip_test_vector
            cp_up[3] = cp_center[3] + ribbon_height * normal_vector
        else:
            add_control_points(n, i + 1)

    add_point_to_curves()

def main():
    global curve_center
    global curve_side1
    global curve_side2
    global curve_up

    pdb_fn = '/home/andres/Art/S3 project/pdb/1AMM.pdb'
    pdb_plg = '/usr/lib/libPDBread.so' 
    plLoadStrFromFile(pdb_fn, pdb_plg, 0, 0, 0, 0)

    init_all()

    # Creating curves. 
    curve_center = Curve.New()
    curve_side1 = Curve.New()
    curve_side2 = Curve.New()
    curve_up = Curve.New()

    # Creating Blender objects and linking them to the curves. 
    curve_center_obj = Object.New('Curve', 'ribbon_curve')
    curve_center_obj.link(curve_center)
    curve_side1_obj = Object.New('Curve', 'guide_mesh_0')
    curve_side1_obj.link(curve_side1)
    curve_side2_obj = Object.New('Curve', 'guide_mesh_1')
    curve_side2_obj.link(curve_side2)
    curve_up_obj = Object.New('Curve', 'camera_guide')
    curve_up_obj.link(curve_up)

    # Linking objects to the scene. 
    current_scene = Scene.getCurrent()
    current_scene.link(curve_center_obj)
    current_scene.link(curve_side1_obj)
    current_scene.link(curve_side2_obj)
    current_scene.link(curve_up_obj)

    for n in range(plFirstID(), plLastID() + 1):
        for i in range(plFirstAA(n), plLastAA(n) + 1):
            construct_control_points(n, i)

    curve_center.update()
    curve_side1.update()
    curve_side2.update()
    curve_up.update()

    Window.Redraw()
main()