#!/usr/bin/env python

# Tim Fenn (fenn@stanford.edu)
# pymol isoprobability ellipsoid function
# requires MMTK and numeric python
# see: http://starship.python.net/crew/hinsen/MMTK/

from math import *
from sys import *

from Numeric import *
from LinearAlgebra import *

from Scientific.IO.PDB import *

from pymol.opengl.gl import *
from pymol.opengl.glu import *
from pymol.opengl.glut import *
from pymol.callback import Callback
from pymol import cmd

# Table 6.1 from ORTEP-III manual
pcrit = [
  [0.01, 0.3389], [0.02, 0.4299], [0.03, 0.4951], [0.04, 0.5479], [0.05, 0.5932],
  [0.06, 0.6334], [0.07, 0.6699], [0.08, 0.7035], [0.09, 0.7349], [0.10, 0.7644],
  [0.11, 0.7924], [0.12, 0.8192], [0.13, 0.8447], [0.14, 0.8694], [0.15, 0.8932],
  [0.16, 0.9162], [0.17, 0.9386], [0.18, 0.9605], [0.19, 0.9818], [0.20, 1.0026],
  [0.21, 1.0230], [0.22, 1.0430], [0.23, 1.0627], [0.24, 1.0821], [0.25, 1.1012],
  [0.26, 1.1200], [0.27, 1.1386], [0.28, 1.1570], [0.29, 1.1751], [0.30, 1.1932],
  [0.31, 1.2110], [0.32, 1.2288], [0.33, 1.2464], [0.34, 1.2638], [0.35, 1.2812],
  [0.36, 1.2985], [0.37, 1.3158], [0.38, 1.3330], [0.39, 1.3501], [0.40, 1.3672],
  [0.41, 1.3842], [0.42, 1.4013], [0.43, 1.4183], [0.44, 1.4354], [0.45, 1.4524],
  [0.46, 1.4695], [0.47, 1.4866], [0.48, 1.5037], [0.49, 1.5209], [0.50, 1.5382],
  [0.51, 1.5555], [0.52, 1.5729], [0.53, 1.5904], [0.54, 1.6080], [0.55, 1.6257],
  [0.56, 1.6436], [0.57, 1.6616], [0.58, 1.6797], [0.59, 1.6980], [0.60, 1.7164],
  [0.61, 1.7351], [0.62, 1.7540], [0.63, 1.7730], [0.64, 1.7924], [0.65, 1.8119],
  [0.66, 1.8318], [0.67, 1.8519], [0.68, 1.8724], [0.69, 1.8932], [0.70, 1.9144],
  [0.71, 1.9360], [0.72, 1.9580], [0.73, 1.9804], [0.74, 2.0034], [0.75, 2.0269],
  [0.76, 2.0510], [0.77, 2.0757], [0.78, 2.1012], [0.79, 2.1274], [0.80, 2.1544],
  [0.81, 2.1824], [0.82, 2.2114], [0.83, 2.2416], [0.84, 2.2730], [0.85, 2.3059],
  [0.86, 2.3404], [0.87, 2.3767], [0.88, 2.4153], [0.89, 2.4563], [0.90, 2.5003],
  [0.91, 2.5478], [0.92, 2.5997], [0.93, 2.6571], [0.94, 2.7216], [0.95, 2.7955],
  [0.96, 2.8829], [0.97, 2.9912], [0.98, 3.1365], [0.99, 3.3682], [1.00, 6.0000]]

class anisouCallback(Callback):

    # setup variables, blah
    def __init__(self, pdbfile="", crit=50, transparency=0.0, segments=10, color=[1.0,1.0,1.0], drawprin=False):
        try:
            self.model = Structure(pdbfile)
        except IOError, detail:
            print detail
            self.model = False

        if not 1 <= crit <= 100:
            crit = 50
        self.crit = pcrit[crit - 1][1]

        if not 0.0 <= transparency <= 1.0:
            transparency = 0.0
        self.transparency = transparency

        if segments < 3:
            segments = 5
        self.segments = segments
        
        self.color = color
        self.drawprin = drawprin

    def __call__(self):

        if not self.model:
            return

        for residue in self.model.residues:
            for atom in residue:
                pos = atom['position']

                # if no aniso U available,
                # use temperature factor to generate Ueq
                try:
                    anisou = (array) (atom['u'])
                except KeyError:
                    ueq = atom['temperature_factor'] / (8.0 * pi * pi)
                    anisou = (array) ([[ueq, 0.0, 0.0], [0.0, ueq, 0.0], [0.0, 0.0, ueq]])

                eval, evec = eigenvectors(anisou)

                # put matrix into GL format
                gl_mat = [evec[0][0], evec[0][1], evec[0][2], 0.0,
                          evec[1][0], evec[1][1], evec[1][2], 0.0,
                          evec[2][0], evec[2][1], evec[2][2], 0.0,
                          0.0, 0.0, 0.0, 1.0]

                glColor4d(self.color[0], self.color[1], self.color[2], 1.0);

                glPushMatrix()
                # translate to atom location, apply transform
                glTranslatef(pos.x(), pos.y(), pos.z())
                glMultMatrixd(gl_mat)
                glScalef(self.crit * sqrt(eval[0]),
                         self.crit * sqrt(eval[1]),
                         self.crit * sqrt(eval[2]))

                # start drawin stuff
                if self.drawprin:
                    # principal axes
                    qobj = gluNewQuadric()
                    gluCylinder(qobj, 0.04, 0.04, 1.0, self.segments, self.segments)
                    glRotatef(180.0, 0.0, 1.0, 0.0)
                    gluCylinder(qobj, 0.04, 0.04, 1.0, self.segments, self.segments)

                    # for some reason, this doesn't work on my machine.
                    # gluDisk(qobj, 0.0, 1.0, self.segments, self.segments)

                    glPushMatrix()
                    glRotatef(90.0, 1.0, 0.0, 0.0)
                    gluCylinder(qobj, 0.04, 0.04, 1.0, self.segments, self.segments)
                    glRotatef(180.0, 0.0, 1.0, 0.0)
                    gluCylinder(qobj, 0.04, 0.04, 1.0, self.segments, self.segments)
                    # gluDisk(qobj, 0.0, 1.0, self.segments, self.segments)
                    glPopMatrix()

                    glPushMatrix()
                    glRotatef(90.0, 0.0, 1.0, 0.0)
                    gluCylinder(qobj, 0.04, 0.04, 1.0, self.segments, self.segments)
                    glRotatef(180.0, 0.0, 1.0, 0.0)
                    gluCylinder(qobj, 0.04, 0.04, 1.0, self.segments, self.segments)
                    # gluDisk(qobj, 0.0, 1.0, self.segments, self.segments)
                    glPopMatrix()


                glColor4d(self.color[0], self.color[1], self.color[2], 1.0 - self.transparency);

                # draw sphere last (to avoid transparency issues w/OpenGL)
                glutSolidSphere(1.0, self.segments, self.segments)
                glPopMatrix()

#cmd.load_callback(anisouCallback(pdbfile='phe.pdb'), 'anisou')