from MMTK import *
from MMTK.Proteins import Protein
from MMTK.PDB import PDBConfiguration
from MMTK.Trajectory import Trajectory, TrajectoryOutput
from MMTK.DCD import DCDReader
import numpy
import math

##numpy.set_printoptions(precision=8)
numpy.set_printoptions(suppress=True)


## setting up refrence and projecting 1SX4
refuniverse = InfiniteUniverse()
conf = PDBConfiguration('1SX4_trimer_noADP_noH.pdb.1')
chains = conf.createPeptideChains(model='no_hydrogens')
refuniverse.addObject(Protein(chains))

helixD1=refuniverse[0][0][96:108]
helixD2=refuniverse[0][2][96:108]

pointA=refuniverse[0][0][96:108]
pointB=refuniverse[0][2][96:108]

cm, mi = helixD1.centerAndMomentOfInertia()
ev, axes = mi.diagonalization()
N = axes[numpy.argmin(ev)].asVector().normal()
print N

cmA = pointA.centerOfMass()
cmB = pointB.centerOfMass()
R = cmA-cmB
print R

Uref = R - (R*N)*N
newX = Uref.normal()
newY = N.cross(newX).normal()

print newX
print newY

## with matrix object
#Ainv = numpy.matrix([ newX, newY, N ])
#Ainv = numpy.transpose(Ainv)
#AA = [a*Ainv.I for a in Ainv]

## with array object
Ainv = numpy.array([ newX, newY, N ]).T
A = numpy.linalg.inv(Ainv)
AA = numpy.dot(A, Ainv)

## testing
numpy.dot(A, newX)
numpy.dot(A, newY)
numpy.dot(A, N)

# reference projection
Pref = numpy.dot(A, Uref)
print Pref
print Uref


## projecting 1XCK
refuniverse = InfiniteUniverse()
conf = PDBConfiguration('1XCK_chainALM_noWAT_noH_fitted.pdb')
chains = conf.createPeptideChains(model='no_hydrogens')
refuniverse.addObject(Protein(chains))

helixD1=refuniverse[0][0][96:108]
helixD2=refuniverse[0][2][96:108]

pointA=refuniverse[0][0][96:108]
pointB=refuniverse[0][2][96:108]

cmA = pointA.centerOfMass()
cmB = pointB.centerOfMass()
R = cmA-cmB

U = R - (R*N)*N
Pxck = numpy.dot(A, U)
print cmA
print cmB
print U


## beging with MD trajectories
trjs = []
trjs.append("trimertraj_0.nc")
trjs.append("trimertraj_1.nc")

t = Trajectory(None, trjs[0])
universe=t.universe
           
i=0
for trj in trjs:
    i+=1
    proj = []
    t = Trajectory(None, trj)
    universe=t.universe

    for step in t:
        universe.setConfiguration(step['configuration'])
        pointB=universe[0][2][105:108]
        cmB = pointB.centerOfMass()
        R = cmA-cmB
        U = R - (R*N)*N
        P=numpy.dot(A, U)
        #print P
        proj.append(P)
        
    file=open('projected_%s.dat' % str(i), 'w')
    file.write("%s %s %s \n" % (str(Pref[0]), str(Pref[1]), str(Pref[2])))
    file.write("%s %s %s \n" % (str(Pxck[0]), str(Pxck[1]), str(Pxck[2])))
    for v in proj:
        file.write("%s %s %s \n" % (str(v[0]), str(v[1]), str(v[2])))
        
    file.close()