from iotbx import pdb
from phaser import mmt
import math, os, sys
import subprocess
from iotbx.bioinformatics import clustal_alignment_parse


class Vector():
    def __init__(self, x,y,z):
        self.x = x
        self.y = y
        self.z = z

    def L2Norm(self):
        return math.sqrt(self.x*self.x +self.y*self.y +self.z*self.z)


    def unitvector(self):
        norm = self.L2Norm()
        u = Vector(self.x/norm, self.y/norm, self.z/norm)
        return u

    def __mul__(self, a):
        r = Vector(self.x*a, self.y*a, self.z*a)
        return r



def listCalphas(pdbname, chainid = ""):
    mypdb = pdb.input(file_name=pdbname)
    root = mypdb.construct_hierarchy()
    calphas = []

    for chain in root.models()[0].chains():
        if len(chainid) > 0 and chainid.upper() != chain.id.upper():
            continue # only list a specific sequence if asked to do so
# otherwise list all of them
        for residue in chain.residue_groups():
            for atm in residue.atoms():
                if atm.name==" CA ":
                    calphas.append((atm, residue.unique_resnames()[0], residue.resseq_as_int()) )

    return calphas



def MakeCaplhaDistanceMatrix(Calphas):
    nr = len(Calphas)
    matrix= []
    for i,a in enumerate(Calphas):
        matrix.append([])
        for j in range(0, nr):
            matrix[i].insert(j,None)

    for i,a in enumerate(Calphas):
        for j in range(i, nr):
            b = Calphas[j]
            D = Vector(a[0].xyz[0] - b[0].xyz[0],
                       a[0].xyz[1] - b[0].xyz[1],
                       a[0].xyz[2] - b[0].xyz[2])

            matrix[i][j] = (a[1], b[1], D, D.L2Norm())
            matrix[j][i] = (b[1], a[1], D*(-1), D.L2Norm())

    return matrix



def GetContactVector(matrix):
    dim = len(matrix)
    gap = 2
    cv = []
    for i in range(0,dim):
        cv.append(0)

    for i in range(gap,dim):
        for j in range(i+gap,dim):
            if matrix[j][j-i][3] < 9.0:
                cv[i] += 1

    return cv



def GetContactVectorFromPDB(pdbname, chainid=""):
    calphas = listCalphas(pdbname, chainid)
    Dm = MakeCaplhaDistanceMatrix(calphas)
    cv = GetContactVector(Dm)

    return cv



def PrintCalphasFromPDB(pdbname, chainid=""):
    calphas = listCalphas(pdbname, chainid)
    for ca in calphas:
        print "%f \t %f \t %f \t %s" %(ca[0].xyz[0],ca[0].xyz[1],ca[0].xyz[2], ca[1])



def PrintContactVectorFromPDB(pdbname, chainid=""):
    cv = GetContactVectorFromPDB(pdbname, chainid)
    for i in range(0, len(cv)):
       print "%d \t %d" %(i, cv[i])



def PrintContactMatrix(pdbname, chainid=""):
    calphas = listCalphas(pdbname, chainid)
    Dm = MakeCaplhaDistanceMatrix(calphas)

    nameprefix = os.path.splitext(os.path.basename(pdbname))[0]

    outfile = open(nameprefix + "listcontactmatrix.txt","w")

    dim = len(Dm)
    s = ""
    for i in range(0,dim):
        s += "j= %d \t" %i

    s += "\n"

    for i in range(0,dim):
        s += "i= %d\t" %i
        for j in range(0,dim):
            length = Dm[i][j][2].L2Norm()
            s += " %f\t" %length
            outfile.write("%d\t %d\t %f\n" %(i, j, -length))

        s += "\n"

    outfile.close()

    mfile = open(nameprefix + "contactmatrix.txt","w")
    mfile.write(s)
    mfile.close()

    return s






def GetContactVectorDistance(pdb1name, chain1id, pdb2name, chain2id):
    cv1 = GetContactVectorFromPDB(pdb1name, chain1id)
    if len(cv1) == 0:
        print "invalid chain %s in %s" %(chain1id, pdb1name)
        return None

    cv2 = GetContactVectorFromPDB(pdb2name, chain2id)
    if len(cv2) == 0:
        print "invalid chain %s in %s" %(chain2id, pdb2name)
        return None

# extend contact vectors with zeros to same dimensions
    d = len(cv1) - len(cv2)
    if d > 0:
        for i in range(0,d):
            cv2.append(0)
    else:
        for i in range(0,d):
            cv1.append(0)

    s = 0.0
    for i in range(0, len(cv1)):
        s += math.pow(cv1[i] - cv2[i], 2)

    return math.sqrt(s)




def GetSequenceFromPDB(pdbfname, chainid=""):
    ms = mmt.PROTEIN
    trans = ms.get_translation()
    known = set( ms.three_letter_codes ) # all 20 amino acids

    mypdb = pdb.input(file_name = pdbfname)
    root = mypdb.construct_hierarchy()

    seqstr = ""

    for chain in root.models()[0].chains():
        if len(chainid) > 0 and chainid.upper() != chain.id.upper():
            continue # only print a specific sequence if asked to do so
# otherwise print all of them
        substr = ""
        for res in chain.residue_groups():
# only print amino acids codes, not water, ligands, etc.
            if res.unique_resnames()[0] in known:
                substr += trans.get( res.unique_resnames()[0], ms.unknown_one_letter )

        if len(substr) > 0:
            seqstr += ">" + os.path.basename(pdbfname)[:4].upper() \
             + "_Chain" + chain.id + "\n" + substr + "\n\n"

    return seqstr




def GetSequenceContactVector(pdbfname, chainid=""):
    return (GetSequenceFromPDB(pdbfname, chainid), GetContactVectorFromPDB(pdbfname, chainid))




def GetSequenceDistance(pdb1name, chain1id, pdb2name, chain2id):
    clustalw = r"E:\Programs\clustalw2\clustalw2.exe"

    seq1 = GetSequenceFromPDB(pdb1name, chain1id)
    seq2 = GetSequenceFromPDB(pdb2name, chain2id)
    if seq2 == seq1:
        return 0.0

    fileout = open("clustinput.txt","w")
    fileout.write(seq1 + seq2)
    fileout.close()

    command = (clustalw, '-INFILE=clustinput.txt' ,'-ALIGN',\
     '-OUTFILE=clust.aln')

    #print "Executing command:" + subprocess.list2cmdline(command)
    output = subprocess.Popen(command,shell=False, stdout=subprocess.PIPE).communicate()[0]
    #print output

    alnfile = open("clust.aln","r")
    ali = clustal_alignment_parse(alnfile.read())
    alnfile.close()

    if ali == (None, ''):
        return None
# 1 is perfect match and 0 is no match. So reverse this
    percentdiff = 1.0 - float(ali[0].identity_count())/len(ali[0].alignments[0])
# map values close to 1 asymptotically towards infinity
    D = math.tan(percentdiff*math.pi/2.0)

    return D




def GetCVSeqDistance(pdb1name, chain1id, pdb2name, chain2id):
    seqd = GetSequenceDistance(pdb1name, chain1id, pdb2name, chain2id)
    cvd = GetContactVectorDistance(pdb1name, chain1id, pdb2name, chain2id)

    return seqd+cvd




def FrobeniusNorm(matrix):
    dim = len(matrix)
    s =0.0

    for i in range(0,dim):
        for j in range(0,dim):
            d= matrix[i][j]
            s += d*d

    return math.sqrt(s)




def GetDiffMatrix(matrix1, matrix2,dim2, offset):
#subtract matrix2 from matrix1 given that dimension of matrix1 is bigger than dimension of matrix2
    diffmatrix = []
    for i in range(0, dim2):
        diffmatrix.append([])
        for j in range(0, dim2):
            diffmatrix[i].insert(j,None)
            diffmatrix[i][j] = matrix1[i+offset-1][j+offset-1][3] - matrix2[i][j][3]

    return diffmatrix




def GetDiffMatrixNormGraph(pdb1name, pdb2name, chain1id="", chain2id=""):
    calphas1 = listCalphas(pdb1name, chain1id)
    Dm1 = MakeCaplhaDistanceMatrix(calphas1)
    dim1 = len(Dm1)

    calphas2 = listCalphas(pdb2name, chain2id)
    Dm2 = MakeCaplhaDistanceMatrix(calphas2)
    dim2 = len(Dm2)

    print "dims= %d  %d" %(dim1,dim2)
    Dm3 = []

    nullelement = (None, None, Vector(0,0,0),0)

    diffdim = abs(dim1-dim2)
    if dim2 > dim1: # want dim2 to be smaller than dim1 so swap them if not
        calphas2 = listCalphas(pdb1name, chain1id)
        Dm2 = MakeCaplhaDistanceMatrix(calphas1)
        dim2 = len(Dm1)

        calphas1 = listCalphas(pdb2name, chain2id)
        Dm1 = MakeCaplhaDistanceMatrix(calphas2)
        dim1 = len(Dm2)

    dim3 = dim1 + 2*dim2
    for i in range(0, dim3):
        Dm3.append([])
        for j in range(0, dim2):
            Dm3[i].insert(j, nullelement)

        for j in range(0, dim1):
            dm1 = Dm1[i-dim2][j] if i > dim2 and i < (dim1+dim2) else nullelement
            Dm3[i].insert(j+dim2, dm1)

        for j in range(0, dim2):
            Dm3[i].insert(j+dim1+dim2, nullelement)

    for offset in range(0, dim1+dim2):
        diffmatrix = GetDiffMatrix(Dm3,Dm2,dim2, offset)
        norm = FrobeniusNorm(diffmatrix)/((dim2-offset)*(dim2-offset)  + dim2*dim2)
        print "%d \t%f" %(offset,norm)






if __name__=="__main__":
    while len(sys.argv) < 5:
        sys.argv.append("")

    GetDiffMatrixNormGraph(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4] )