import os,pickle
from string import lstrip,strip, rstrip
from re import split
from math import  log, sin, cos, sqrt, pi, fabs, radians,degrees,ceil
from libmath import dihedral_vect
import numpy as Npy

from copy import copy 
from libio import MIO
_array = Npy.array

pi2 = pi * 2.0
def getVecLen(vec):
    return  sqrt(vec[0]*vec[0]+vec[1]*vec[1]+vec[2]*vec[2])

def getNorm(vect):
    #return sqrt(vect[0]*vect[0] + vect[1]*vect[1] + vect[2]*vect[2])
    return getVecLen(vect)

class MVector:
    def normVect(self, vect):
        """ return a normalized vector)"""
        return vect/sqrt(vect[0]*vect[0] + vect[1]*vect[1] + vect[2]*vect[2])

    def getVecLen(self, vec):
        return  sqrt(vec[0]*vec[0]+vec[1]*vec[1]+vec[2]*vec[2])

    def getVectSub(self, ca1, ca2):
        return (ca2[0]-ca1[0], ca2[1]-ca1[1], ca2[2]-ca1[2])

    def getVectAdd(self, ca1, ca2):
        return (ca2[0]+ca1[0], ca2[1]+ca1[1], ca2[2]+ca1[2])

    def getVectMulbyVal(self, vec, val):
        return Npy.array(vec)*val 

    def getVectNorm(self, vec):
        return sqrt(vec[0]*vec[0] + vec[1]*vec[1] + vec[2] * vec[2])

    def getVecInnerProduct(self, v1, v2):
        return (v1[0]*v2[0] + v1[1]*v2[1] + v1[2]*v2[2])

    def rot_along_X(self, angle, vect):
        vect_axle = [1, 0, 0]
        return self.arb_rot(vect, vect_axle, angle)

    def rot_along_Y(self, angle, vect):
        vect_axle = [0, 1, 0]
        return self.arb_rot(vect, vect_axle, angle)

    def rot_along_Z(self, angle, vect):
        vect_axle = [0, 0, 1]
        return self.arb_rot(vect, vect_axle, angle)

    def getDistSquare(self, ca, cb):
        dx = ca[0] - cb[0]
        dy = ca[1] - cb[1]
        dz = ca[2] - cb[2]
        dr2 = dx * dx + dy*dy + dz*dz 
        return dr2 


    def getDist(self, ca, cb):
        dx = ca[0] - cb[0]
        dy = ca[1] - cb[1]
        dz = ca[2] - cb[2]
        dr2 = dx * dx + dy*dy + dz*dz 
        dr = sqrt(dr2)
        return dr 


    def arb_rot(self, vect, vect_axis, gamma):
        vect = self.normVect(vect)
        #vect_axis = normVect(vect_axis)
        x,y,z  = vect[0], vect[1],vect[2]
        v0,v1,v2 = vect_axis[0],vect_axis[1],vect_axis[2]
        v1_2 = v1 * v1 
        v2_2 = v2 * v2 
        lyz_2 = v1_2 + v2_2
        cos_gamma = cos(gamma)
        sin_gamma = sin(gamma)
        #Ry = nmatrix([[lyz,    0.,    -vect_axis[0] ],
                        #[0,      1.,         0 ],
                        #[vect_axis[0],  0,  lyz]])

        #Rz = nmatrix([[ cos_gamma,     -sin_gamma,         0],
                        #[ sin_gamma,      cos_gamma,         0],
                        #[0,           0,         1   ]])
        if lyz_2 == 0.0:
            #Rgamma = inv(Ry)*Rz*Ry
            rot_vector = _array([x, y*cos_gamma - z * sin_gamma, z * cos_gamma + y * sin_gamma])
        else:
            #Rx = npy.array([[ 1,           0,          0],
                            #[0,       vect_axis[2]/lyz,   -vect_axis[1]/lyz],
                            #[0,       vect_axis[1]/lyz,   vect_axis[2]/lyz ]])
            #Rgamma = inv(Rx)*inv(Ry)*Rz*Ry*Rx;
            v0x = v0 * x 
            v1y = v1 * y
            v2z = v2 * z 
            vv = v0x + v1y + v2z 
            vx = v0*vv -(-lyz_2*x +v0*v1y +v0*v2z)*cos_gamma + \
               (-v2*y + v1*z) * sin_gamma
            vy = v1 * vv - (v1*v0x + (-1.0 + v1_2)*y + v1*v2z)*cos_gamma + \
               (v2*x - v0*z) * sin_gamma 
            vz = v2*v0x + v2*v1y + v2_2 * z + \
               (-v2*v0x + (1.0 - v2_2)*z -v2*v1y) * cos_gamma + \
               (-v1*x +v0*y)*sin_gamma 
            rot_vector = _array([vx,vy,vz])
        return  self.normVect(rot_vector)




    def getAtomDist(self, atom1, atom2):
        return self.getDist(atom1, atom2)
    
    
class Atom:
    def __init__(self, **rest):
        self.setAtomId(**rest)
        self.setAtomName(**rest)
        self.setResid(**rest)
        self.setResName(**rest)
        self.setXYZ(**rest)
        self.setSegid(**rest)
        self.setMisc(**rest)
        self.setChainID(**rest)

    def setChainID(self, **rest):
        chain_id = rest.get('chain_id', " ")
        if chain_id in [" ", ""]:
            chain_id = self.segid
        self.chain_id = chain_id
        #self.chain_id = rest.get('chain_id'," ")

    def setAtomName(self, **rest):
        self.name = rest.get("atomname", "" ) 

    def setResName(self, **rest):
        self.resname = rest.get("resname", "")

    def setXYZ(self, **rest):
        if 'xyz' in rest:
            xyz = rest.get("xyz")
            self.x = xyz[0]
            self.y = xyz[1]
            self.z = xyz[2]
            #self.xyz = Vector(xyz)
            self.axyz = _array(xyz)

    def getXYZ(self):
        #return #self.axyz
        return _array([self.x, self.y, self.z])

    def upDateXYZ(self):
        self.new_axyz = _array([self.x, self.y, self.z])

    def setX(self, **rest):
        self.x = rest.get('x',9999.)

    def setY(self, **rest):
        self.y = rest.get('y',9999.)

    def setZ(self, **restt):
        self.z = rest.get('z',9999.)

    def setAtomId(self, **rest):
        self.atomid = rest.get('atomid', 9999)

    def setResid(self, **rest):
        self.resid = rest.get('resid', 999)

    def setSegid(self, **rest):
        self.segid = rest.get('segid', "")

    def setMisc(self, **rest):
        self.misc = rest.get('misc', " ")


class RESIDUE:
    def __init__(self, name="", chain_id=" ", resid="", segid="", atom=None, res_sid=0):
        self.atom_objs = {}
        self.atom_list = []
        self.num_atoms = 0
        self.atom_names = []
        if atom is None:
            self.name = name
            self.chain_id = chain_id 
            self.resid = resid
            self.segid = segid
        else:
            self.name = atom.resname
            self.resid = atom.resid
            self.chain_id = atom.chain_id
            self.segid = atom.segid 
        self.res_sid = res_sid

    
    def AddAtom(self, atom):
        self.atom_list.append(atom)
        setattr(self, atom.name, atom.axyz)
        self.atom_names.append(atom.name)
        #self.atom_objs[atom.name] = atom 
        #setattr(atom.name, atom)
        #residue.setProp(atomname=atomname, resid=resid, xyz=[x,y,z])
        self.num_atoms += 1

    def getAtomVect(self, atom_name, vb=0):
        if atom_name in self.atom_names:
            #print "# get atom_name: ", atom_name
            return getattr(self, atom_name)
        else:
            if vb:
                print "# no atom: %s in residue: %4d %s" % (atom_name, self.resid, self.name)
                print "# res has atoms: %s " %  self.atom_names
                raise " atom %s could not find in residue " %  atom             
            else:
                return None 

    def prn(self):
        print self.__dict__

    def setProp(self, atomname="N", resid=1, xyz=[0,0,0]):
        setattr(self, atomname, Vector(xyz))
        self.atoms.append((atomname,resid))
        self.atom_names.append(atomname)

    def __str__(self):
        return "resid= %4d  resname= %4s %s " % (self.resid, self.name )

    

    def getBondVect(self, atom_name1, atom_name2):
        """ input: (atom_name1,  atom_name2) 
            return  vector:   atom_name1 --> atom_name2 """
        vec1 = self.getAtomVect(atom_name1)
        vec2 = self.getAtomVect(atom_name2)
        vec = vec2 - vec1
        return vec 
    
    
    def getNH(self):
        if "N" in self.atom_names:
            if "HN" in self.atom_names:
                return self.getBondVect("HN", "N")
            elif "HT1" in self.atom_names:
                return self.getBondVect("HT1", "N")
            elif "H" in self.atom_names:
                return self.getBondVect("H", "N")
        print"# No NH vector "
        print "# resid= %3d resname= %4s " % (self.resid, self.name)
        return None
            

    def getBondLen(self, *args):
        """ input: (atom_name1,  atom_name2) 
            return  vector lenght :   atom_name1 --> atom_name2  """
        vec = self.getBondVec(*args)
        vec = vec2 - vec1
        vec_len = self.getVecLen(vec)
        return vec_len 

    
    def getCenter(self):
        """ get residue center  """
        res_c = _array([0., 0., 0.])
        for _atom in self.atom_list:
            res_c += _atom.axyz
        res_c = res_c / self.num_atoms
        #print "# residue %3d  num_atom= %2d center:  %8.2f, %8.2f, %8.2f" % (self.resid, self.num_atoms, res_c[0], res_c[1], res_c[2])
        return  res_c 

    
    def getSAXCenter(self):
        """ get residue center without proton  """
        res_c = _array([0., 0., 0.])
        cnt = 0
        for _atom in self.atom_list:
            if _atom.name[0] != "H":
                res_c += _atom.axyz
                cnt += 1.0
        res_c = res_c / cnt
        #print "# residue %3d  num_atom= %2d center:  %8.2f, %8.2f, %8.2f" % (self.resid, self.num_atoms, res_c[0], res_c[1], res_c[2])
        return  res_c 

    
    def getRadius(self):
        """ spherical approximation"""
        res_c = self.getCenter()
        res_sax_c = self.getSAXCenter()
        radius_sax = 0.0
        radius = 0.0 
        for _atom in self.atom_list:
            if _atom.name[0] != "H":
                vec = _atom.axyz - res_sax_c
                vec_len = getVecLen(vec)
                if vec_len > radius_sax: 
                    radius_sax  = vec_len 
                    atom_name_sax = _atom.name 

        for _atom in self.atom_list:
            vec = _atom.axyz - res_c
            vec_len = getVecLen(vec)    
            if vec_len > radius: 
                radius  = vec_len 
                atom_name = _atom.name 

        print "# resid= %4d res_name= %4s  atom_name= %4s radius= %8.3f  sax: atom_name= %4s radius= %8.3f dif= %8.4f" %  \
              (self.resid, self.name, atom_name, radius, atom_name_sax, radius_sax, radius-radius_sax)
        
        
        return radius, radius_sax
    
    
    def getResXYZ(self, list=1):
        res_xyz = []   
        for _atom in self.atom_names:
            res_xyz.append(getattr(self, _atom))
        if list: 
            return res_xyz
        else:
            return _array(res_xyz)
    
    def getResXYZ_woH(self, list=1):
        """ remove Hydrogen atoms """
        res_xyz = []   
        
        for _atom in self.atom_names:
            _atom = _atom.lstrip()
            if _atom[0] == "H": continue 
            res_xyz.append(getattr(self, _atom))
            print "# %4s : %4s " % (self.name, _atom)
        if list: 
            return res_xyz
        else:
            return _array(res_xyz)
    
    def getResBBXYZ(self, list=1):
        """ amino acids back bone atoms' coordinattes """
        res_xyz = []   
        for _atom in ["N","CA","C","O"]:
            #print "# %3d %4s : %4s " % (self.resid, self.name, _atom)
            res_xyz.append(getattr(self, _atom))
            
        if list: 
            return res_xyz
        else:
            return _array(res_xyz)
        
        
    def setResSid(self, iresid):
        """ set residue unique serial index in memeory to solve the same residue number in different chains """
        self.res_sid = iresid
        
    
    def getDihedraAngle_by_Atoms(self, dih_atoms ):
        """ return dihedral angle by 4 atoms of dihedral angle"""
        atom1, atom2,atom3,atom4 = dih_atoms  
        return degrees(dihedral_vect(self.getAtomVect(atom1), self.getAtomVect(atom2), self.getAtomVect(atom3), self.getAtomVect(atom4)))
    
         
class PDB(MIO, MVector): # pass
    iresid = 0
    # exclude ficticious residues from PDB files generated by Xplor or other softwares
    resid_cutoff = 495
    def __init__(self,  **rest):
        self.pdb_head_lns = []
        self.pdb_atom_lns = []
        self.chain_ids = []
        self.res_sid = None
        self.all_residues = {}
        self.pdbfn2pdb(**rest)
        
    
    
    def synResids(self):
        """ synchonize resid with resid from pdb file """
        old_residues = self.residues
        new_residues = {}
        new_resids = []
        for resid in self.resids:
            _res = old_residues[resid]
            resid_pdb = _res.resid
            new_resids.append(resid_pdb)
            new_residues[resid_pdb] = _res
        self.residues = new_residues
        self.resids = new_resids
        
    def getSegResXYZ(self, seg):
        """ return residues' atom coordinates in a segment(resid_start, resid_end) """
        if len(seg)> 2:
            seg_resids = seg 
        elif len(seg) == 2:
            seg_resids = range(seg[0], seg[1]+1)
        else:
            print "# Error: getSegResXYZ: only one or no resid in segment "
            raise " no input "
        seg_xyz = []
        for resid in seg_resids:
            seg_xyz += self.residues[resid].getResXYZ(resid)
        return _array(seg_xyz)
    
    
    def getSegResXYZ_woH(self, seg):
        """ return residues' atom coordinates in a segment(resid_start, resid_end) """
        if len(seg) == 2:
            seg_resids = range(seg[0], seg[1]+1)        
        else:
            seg_resids = seg
        seg_xyz = []
        for resid in seg_resids:
            seg_xyz += self.residues[resid].getResXYZ_woH(resid)
        return _array(seg_xyz)
    
    
    def getSegBBResXYZ(self, seg):
        """ return residues' atom coordinates in a segment(resid_start, resid_end) """
        if len(seg) == 2:
            seg_resids = range(seg[0], seg[1]+1)        
        else:
            seg_resids = seg
        #print seg_resids
        seg_xyz = []
        for resid in seg_resids:
            seg_xyz += self.residues[resid].getResBBXYZ(resid)
        return _array(seg_xyz)
            
    
    def pdbfn2pdb(self, fn_pdb=None, **rest):
        if fn_pdb:
            self.readPDB(fn_pdb, **rest)
            self.parsePDB()
            self.fn_pdb = fn_pdb
        #else:
            #self.err_exit("no input pdb filename: %s " % fn_pdb )
            
    
    def readPDB(self, fp, vb=1):
        if vb: print "# reading %s " % fp
        self.pdb_lns = self.readTxtFile(fp) 

    def prnSeq(self, cols=2):
        for resid in self.resids:
            ln = " %4d %4s" % (resid, self.residues[resid].name)
            print ln 
     

    def addNewResidue(self, residue):
        self.res_sid += 1
        residue.setResSid(self.res_sid)
        self.residues[self.res_sid] = residue
      
        
    
    def setUniqResids(self, bs_resid=None):
        """ adjust residue resids in chains so that each residue has a uniq resid
            bs_resid: chain b start resid 
        """
        if hasattr(self, 'UNIQRES') and self.UNIQRES:
            print "# uniqres, nothing done"
            return 
        else:
            self.UNIQRES = 1
            
        if len(self.chain_ids) == 1:
            self.residues = self.all_residues[self.chain_ids[0]]
            resids = self.residues.keys()
            resids.sort()
            self.resids = resids 
        elif len(self.chain_ids)==2:
            chain_A_residues = self.all_residues[self.chain_ids[0]]
            chain_A_resids = chain_A_residues.keys()
            chain_A_resids.sort()
            chain_B_residues = self.all_residues[self.chain_ids[1]]
            chain_B_resids = chain_B_residues.keys()
            chain_B_resids.sort()
            residues = chain_A_residues
            last_a_resid = chain_A_resids[-1]
            if bs_resid is None:
                b_resid = int(ceil(last_a_resid/100.0)*100.0)
                b_start_resid = b_resid + 1
            else:
                b_start_resid = bs_resid
            print "# Chain B: start_resid: ", b_start_resid    
            self.b_start_resid = b_start_resid 
            for resid in chain_B_resids:
                residues[b_start_resid] = chain_B_residues[resid]
                b_start_resid += 1
            self.residues = residues
            resids = residues.keys()
            resids.sort()
            self.resids = resids
        else:
            print "# --------------------------------------"
            print "# only allow two chains in the pdb file."
            print "#_______________________________________"
            sys.exit(0)

    
    def parseAtom(self, ln):
        """ parse ATOM text line in PDB format 
        1-6    Record name "ATOM  " or "HETATM"
        7-11   Atom serial number                   (see note i)
        13-14   Chemical symbol (right justified)  )
        15      Remoteness indicator               ) (see note ii)
        16      Branch designator                  )
        17      Alternate location indicator         (see note iii)
        18-20   Residue name                         (see note iv)
        21      Reserved                   )
        22      Chain identifier           )
                                           )         (see note v)
        23-26   Residue sequence number    )
        27      Code for inserting residue )
        31-38   X   )
        39-46   Y   ) Orthogonal Angstrom coordinates
        47-54   Z   )
        55-60   Occupancy
        61-66   Isotropic B-factor
        73-76   Segment identifier, left justified (used by XPLOR)
        77-78   Element symbol, right justified )
                                                )    (see note vi)
        79-80   Charge on atom                  )
        """
        ln = lstrip(ln)
        atom_dt = {}
        atom_dt['atomid'] = int(ln[6:11])
        atom_dt['atomname'] = strip(ln[12:16])
        atom_dt['resname'] = ln[17:20]
        atom_dt['chain_id'] = ln[21]
        atom_dt['resid'] = int(ln[22:26])
        x = float(ln[30:38])
        y = float(ln[38:46])
        z = float(ln[46:54])
        atom_dt['xyz'] = (x,y,z)
        atom_dt['misc'] = ln[54:72]
        atom_dt['segid'] = rstrip(ln[72:76])
        return atom_dt        
    
    
    def parsePDB(self, resid_cutoff=495):
        """ parse pdb datafile to get residues, sequence, cartesian coordinates
        resid upper limit is 495, all resids larger thatn 495 will be not readed 
        """
        self.pdb_atom_lns = []
        self.residues = {}
        pre_resid = -999
        body = 0
        residue  = None
        self.pdb_start_resid = None 
        for ln in self.pdb_lns:
            if ln[0:4] == "REMA" and body == 0:
                self.pdb_head_lns.append(ln)
            elif ln[0:4] == "ATOM": 
                body = 1
                atom_dt = self.parseAtom(ln)
                atom = Atom(**atom_dt)
                resid = atom_dt['resid']
                if pre_resid != resid: #new residue 
                    if residue:
                        if self.pdb_start_resid is None:
                            self.pdb_start_resid = residue.resid
                            self.res_sid = self.pdb_start_resid - 1 
                        self.addNewResidue(residue)
                    if resid > resid_cutoff: 
                        residue = None 
                        break
                    residue = RESIDUE(atom=atom)
                    pre_resid = resid
                residue.AddAtom(atom)
                self.pdb_atom_lns.append(ln)
        if residue: self.addNewResidue(residue)  
        self.resids = self.residues.keys()
        self.resids.sort()
        
    def getPdbCode(self):
        fn_pdb = os.path.basename(self.fn_pdb)
        pdb_code, fn_ext = split("\.", fn_pdb)
        return pdb_code 
    
    
    def getSequence(self):
        self.setUniqResids()
        seq_dt = {}
        for resid in self.resids:
            seq_dt[resid] = self.residues[resid].name 
        return seq_dt 
    
    
    def prnChains(self):
        """ print out protein chains information """
        for i in range(len(self.chain_ids)):
            chain_id = self.chain_ids[i]
            print "# chain %s  Num_res= %3d " % (chain_id, len(self.residues[chain_id]))
            #print self.residues[self.chain_ids[0]].keys()

    
    def getNumRes(self):
        """ return the number of residues in the pdb file """
        return self.num_res 


    def prnResName(self):
        print "# pdb: cutinase " 
        for resid in self.resids:
            print "%4d %8s " % (resid, self.residues[resid].name)

 
        
    def prnAAtomXYZ(self, res_first, res_end, vb=1):
        # array index starts from 0, but the amino residues from 1.
        # frist resudue's N atom is not included in the pdb file generated by gen_pdb function
        resxyz, atom_names = self.getResXYZ(res_first, res_end)
        xyz_lns = []
        for i in range(len(resxyz)):
            #print "%3d %3s %8.4f %8.4f %8.4f" % (i+1, atom_names[i],xyz[i].x(), xyz[i].y(), xyz[i].z())
            ln =  "%3d %3s %8.3f %8.3f %8.3f" % (atom_names[i][0],atom_names[i][1], resxyz[i,0], resxyz[i,1], resxyz[i,2])
            xyz_lns.append(ln)
            if vb: print ln 
        return resxyz, xyz_lns
    


    def writePDB(self, fn_pdb, lns=None, vb=1, resid_shift=0, bresid_shift=None,  resids=None, resid_cutoff=500,  save=1):
        """ PDB format
            Cols.  
        6   1-6    Record name "ATOM  " or "HETATM"
        5   7-11   Atom serial number                   (see note i)
        2   13-14   Chemical symbol (right justified)  )
        1   15      Remoteness indicator               ) (see note ii)
        1   16      Branch designator                  )
        1   17      Alternate location indicator         (see note iii)
        4   18-20   Residue name                         (see note iv)
        1   21      Reserved                   )
        1   22      Chain identifier           )         (see note v)
        4   23-26   Residue sequence number    )
        1   27      Code for inserting residue )
        8   31-38   X   )
        8   39-46   Y   ) Orthogonal Angstrom coordinates
        8   47-54   Z   )
            55-60   Occupancy
            61-66   Isotropic B-factor
            73-76   Segment identifier, left justified (used by XPLOR)
            77-78   Element symbol, right justified )     (see note vi)
            79-80   Charge on atom                  )
        """
        fmt_pdb = '%-6s%5d %-4s %-4s%1s%4d%1s%3s%8.3f%8.3f%8.3f%18s%-4s'
        cnt = 0
        if lns is None: lns = []
        if resids is None: resids = self.resids
        for i in resids:
            _amino = self.residues[i]
            chain_id = _amino.chain_id
            #print "chain_id= %s " % chain_id 
            #if resid_shift and resid_shift <0:
                
            _resid = _amino.resid + resid_shift
            #print "# _resid: ", _resid
            
            if bresid_shift and chain_id=='B':
                _resid += bresid_shift
            if _resid > resid_cutoff: continue
            #print _amino.atom_list 
            for _atom in _amino.atom_list:
                atom_name = _atom.name 
                if len(atom_name) <= 3: atom_name= " %s" % atom_name
                ln =  fmt_pdb  % ('ATOM', _atom.atomid, atom_name, _atom.resname,chain_id, _resid, "","", _atom.x, _atom.y, _atom.z, _atom.misc, _atom.segid)
                lns.append(ln)
                #print ln 
                
        if vb: print "# pdb was saved in < %s >: " %  fn_pdb
        lns = self.pdb_head_lns + lns + ["END"]
        if save: self.saveLines(lns, fn_pdb, vb=vb)
        return lns 
    

    def setPDBXYZ(self): 
        pdb_res_iatom_ids = []
        pdb_res_eatom_ids = []
        cnt  = 0    
        pdb_xyz = []
        for resid in self.resids:
            pdb_res_iatom_ids.append(cnt)
            for _atom in self.residues[resid].atom_list:
                pdb_xyz.append(copy(_atom.axyz))
                cnt += 1
            pdb_res_eatom_ids.append(cnt-1)
        #print pdb_res_eatom_ids
        #print pdb_res_iatom_ids
        self.pdb_xyz_arr = _array(pdb_xyz)
        self.pres_ia_ids_arr = _array(pdb_res_iatom_ids, dtype=Npy.int32)
        self.pres_ie_ids_arr = _array(pdb_res_eatom_ids, dtype=Npy.int32)


    def setResCenter(self): # pass 
        """ set Monomer A SAX atoms residue centers, Monomer A geometry center as the Cartesian origin point(0, 0, 0) """
        res_centers = []
        for resid in self.resids:
            res_centers.append(self.residues[resid].getCenter())
        self.prct_arr = _array(res_centers)


if (__name__ == "__main__"):
    pdb_obj = PDB(fn_pdb='A-RNA.pdb')
    pdb_obj.start()
