import sys
import os
import string
import Numeric
import Geometry
import copy

GUAAtoms = [ "N2", "O6", "C6", "C5", "N7", "C8", "N9", "C4", "N3", "C2", "N1" ]

def rotmat(phi, theta, psi, tx=0, ty=0, tz=0):

    s1 = Numeric.sin(phi)
    s2 = Numeric.sin(theta)
    s3 = Numeric.sin(psi)
    c1 = Numeric.cos(phi)
    c2 = Numeric.cos(theta)
    c3 = Numeric.cos(psi)

    newmat = Numeric.array([
        [ c2*c3, s2*s1*c3 - c1*s3, s2*c1*c3 + s1*s3, 0],
        [ c2*s3, s2*s1*s3 + c1*c3, s2*c1*s3 - s1*c3, 0],
        [-s2, c2*s1, c2*c1, 0],
        [tx, ty, tz, 1]])

    return newmat

def matrix_apply(crd, mat):
    o = [crd[0], crd[1], crd[2], 1.0]
    n = [0.0, 0.0, 0.0, 0.0]

    for i in range(4):
	for j in range(4):
	    n[i] = n[i] + o[j] * mat[j][i]
	    
    new = [None, None, None]
    
    for i in range(3):
	new[i] = (n[i]/n[3])
    return new



class PDBRecord:
    def __init__(self, type=None, anum=None, atom=None, residue=None, chain=None, rnum=None):
	self.type = type
	self.anum = anum
	self.atom = atom
	self.residue = residue
	self.chain = chain
	self.rnum = rnum



class PDB:
    def __init__(self, filename = None, crds = None, records = None, connect = None):
	if crds == None:
	    crds = []
	if records == None:
	    records = []
	self.records = records
	self.crds = crds
	self.connect = connect

	if filename != None:
	    sys.stderr.write("Reading in new PDB %s\n" % filename)
	    self.Read(filename)


## note: we read the anum here, which doesn't necessarily correspond to the actual
## record number.
    def Read(self, filename):
	pdbfile = open(filename)
	sys.stderr.write("Opened '%s' for reading as PDB\n" % filename)

	while(1):
	    line = pdbfile.readline()

	    if line == '':
		break

	    if line[0:4] == 'ATOM' or line[0:6] == 'HETATM':
		type = string.strip(line[0:6])
		anum = string.atoi(string.strip(line[7:11]))
		atom = string.strip(line[12:17])
		residue = string.strip(line[17:20])
		chain = string.strip(line[21:22])
		rnum = string.atoi(string.strip(line[23:26]))
		x = string.atof(string.strip(line[31:38]))
		y = string.atof(string.strip(line[39:46]))
		z = string.atof(string.strip(line[47:54]))

		self.records.append(PDBRecord(type, anum, atom, residue, chain, rnum))
		self.crds.append((x,y,z))
		
	
	self.crds = Numeric.array(self.crds)

    def Write(self, filename):
	self.pdbfd=open(filename,"w")
	    
	for i in range(len(self.records)):
	    record = self.records[i]
	    self.pdbfd.write("%-6s%5d %-4s%c%-4s%c%4d%c   %8.3f%8.3f%8.3f\n" % \
			     ( record.type, record.anum, record.atom,' ', record.residue,' ', record.rnum, ' ', 
			       self.crds[i][0], self.crds[i][1], self.crds[i][2]))

	self.pdbfd.write("TER\n")

	if self.connect != None:
	    for i in self.connect:
		self.pdbfd.write("CONECT")
		for j in i:
		    self.pdbfd.write("%5d" % (j+1))
		self.pdbfd.write("\n")

	self.pdbfd.write("END\n")
	self.pdbfd.close()


    def Rotate(self,alpha, beta, gamma, tx, ty, tz):
	r = rotmat(alpha, beta, gamma, tx, ty, tz)
	self.RotateMatrix(r)

    def RotateMatrix(self,r):
	for i in range(len(self.crds)):
	    self.crds[i] = matrix_apply(self.crds[i], r)

    def Center(self):
	center = Numeric.add.reduce(self.crds) / len(self.crds)
	self.crds = Numeric.subtract(self.crds, center)

    def Print(self):
	for i in self.records:
	    print i.type, i.anum, i.atom, i.residue, i.chain, i.rnum

    def ReturnAnum(self, atom, rnum):
	for i in range(len(self.records)):
	    record = self.records[i]
	    if record.atom == atom and record.rnum == rnum:
 		return i
	sys.stderr.write("Unable to find atom '%s' in residue %d\n" % (atom, rnum))
	return None

    def CrdByName(self, atom, res):
	x = self.ReturnAnum(atom, res)
	if x == None:
	    return None
	return self.crds[x]

    def FixResNum(self):
	rnum = 1
	for i in range(len(self.records)):
	    newrnum = rnum
	    if i < len(self.records)-1 and \
	       self.records[i+1].rnum != self.records[i].rnum:
		rnum = rnum + 1
	    self.records[i].rnum = newrnum
	    self.records[i].anum = i

    def Copy(self):
	crds = Numeric.array(self.crds)
	records = copy.copy(self.records)
	return PDB(None, crds, records)


    def Append(self, struct2):
	records = self.records + copy.copy(struct2.records)
	crds = Numeric.concatenate((self.crds, struct2.crds))

	return PDB(None, crds, records)