# $CVSHeader: pal8e/src/CompareDSSP_SOV.py,v 1.3 2005/02/15 01:10:24 indraneel Rel $


import os
import sys
import string
from DSSP import DSSP
from PSSDParser import PSSD

replacements = {
    '1a6o':'1lr4',
    '1dlx':'1ly7',
    '1hu4':'1n5d',
    '2fxb':'1iqz',
    '1alo':'1vlb',
    '1qj2':'1n5w',
    '1b91':'1n72'
    }

class FileName:
    def getDict(self, alndir):
        fnames = os.listdir(alndir)
        fnames.sort()
        filedict = {}
        for value in fnames:
            buf = value.split('.')
            buf = buf[0].split('_')
            if (filedict.has_key(buf[1]) and \
                (buf[0] in filedict[buf[1]])):
                continue
            if not filedict.has_key(buf[0]):
                filedict[buf[0]] = []
            filedict[buf[0]].append(buf[1])
        return filedict

class Common:
    def __sort(self, a, b):
        sdict = self.sdict
        chain = self.chain
        return cmp(sdict[chain][a]['#'], sdict[chain][b]['#'])
    
    def seqExtract(self, sdict, chain, type):
        ks = sdict[chain].keys()
        self.sdict = sdict
        self.chain = chain
        ks.sort(self.__sort)
        seq = []
        for k in ks:
            if (type == 'SS'):
                if (sdict[chain][k][type]):
                    seq.append(sdict[chain][k][type][0])
                else: seq.append(' ')
            else:
                seq.append(sdict[chain][k][type])
        return string.join(seq, '')

    def getFSSPseq(self, fsspfname):
        fh = file(fsspfname)
        lines = fh.readlines()
        fh.close()
        if not lines: return None
        #print lines[0]
        #print lines[1]
        k = lines[0].split()[2]
        v = lines[1].split()[2]
        return [k, v]

    def runAlign(self, prog, fsspseq, aaseq):
        #(no insertions of deletions in fssp
        #file observed for this dataset)
        seqs = [[fsspseq[0], aaseq[0]], [fsspseq[1], aaseq[1]]]
        alns = []
        for seq in seqs:
            for i in range(2):
                fh = file('/var/tmp/im'+str(i), 'w')
                fh.writelines(['>'+str(i)+'\n', seq[i]+'\n'])
                fh.close()
            # run align and get the alignment
            command = prog + ' /var/tmp/im0 /var/tmp/im1 2>/dev/null'
            pipe = os.popen(command)
            lines = pipe.readlines()
            pipe.close()
            #for line in lines:
                #print line
            # remove temporary files
            os.remove('/var/tmp/im0')
            os.remove('/var/tmp/im1')
            # read the alignment lines and convert to 2 strings
            faln = ''
            aaln = ''
            for i in range(len(lines)):
                if lines[i-2][:6] == '/var/t' \
                       and lines[i][:6] == '1     ':
                    faln += lines[i-2][7:-1]
                    aaln += lines[i][7:-1]
            aln = [faln, aaln]
            alns.append(aln)
        return alns
        

class Interface:
    def run(self, alndir, dsspdir, pssddir, outdir, alignprog):
        fnames = FileName()
        filedict = fnames.getDict(alndir)
        outfiles = os.listdir(outdir)
        ctr = 1
        for (k, vals) in filedict.items():
            #print ctr, k
            ctr += 1
            for v in vals:
                #print ctr, v
                common = Common()
                fsspfname = k+'_'+v+'.aln'
                #if fsspfname[:-4]+'_dssp' in outfiles: continue
                #if fsspfname[:-4]+'_pssd' in outfiles: continue
                fsspseq = common.getFSSPseq(alndir+fsspfname)
                if not fsspseq: continue
                if (len(fsspseq[0]) != len(fsspseq[1])): continue
                kfname = k[:4]
                if kfname in replacements.keys():
                    kfname = replacements[kfname]
                vfname = v[:4]
                if vfname in replacements.keys():
                    vfname = replacements[vfname]                
                if len(k) == 4: kchn = ' '
                else: kchn = k[4]
                if len(v) == 4: vchn = ' '
                else: vchn = v[4]
                kdssp = DSSP()
                kdssp.store(dsspdir+kfname+'.dssp')
                vdssp = DSSP()
                vdssp.store(dsspdir+vfname+'.dssp')
                kpssd = PSSD()
                kpssd.store(pssddir+kfname+'.ssd')
                vpssd = PSSD()
                vpssd.store(pssddir+vfname+'.ssd')
                if not kdssp.keys(): continue
                if not vdssp.keys(): continue
                if kchn not in kdssp.keys():
                    kchn = 'A'
                if kchn not in kdssp.keys():
                    kchn = kdssp.keys()[0]
                if vchn not in vdssp.keys():
                    vchn = 'A'
                if vchn not in vdssp.keys():
                    vchn = vdssp.keys()[0]
                kdsspSS = common.seqExtract(kdssp, kchn, 'SS')
                vdsspSS = common.seqExtract(vdssp, vchn, 'SS')
                kdsspAA = common.seqExtract(kdssp, kchn, 'AA')
                vdsspAA = common.seqExtract(vdssp, vchn, 'AA')
                kpssdSS = common.seqExtract(kpssd, kchn, 'SS')
                vpssdSS = common.seqExtract(vpssd, vchn, 'SS')
                kpssdAA = common.seqExtract(kpssd, kchn, 'AA')
                vpssdAA = common.seqExtract(vpssd, vchn, 'AA')                
                #print len(kdsspAA), len(kdsspSS), len(kpssdSS),
                #print len(vdsspAA), len(vdsspSS), len(vpssdSS)
                AAalns = common.runAlign(alignprog, \
                                         fsspseq, [kpssdAA, vpssdAA])
                #print AAalns
                #mapping from fssp file (not from alignment)
                kpssdAA2 = ''
                vpssdAA2 = ''
                kdsspSS2 = ''
                vdsspSS2 = ''
                kpssdSS2 = ''
                vpssdSS2 = ''
                kdsspctr = -1
                vdsspctr = -1
                kpssdctr = -1
                vpssdctr = -1
                #print len(fsspseq[0])
                #print len(fsspseq[1])
                #print fsspseq[0]
                #print fsspseq[1]
                #print kpssdAA
                #print vpssdAA
                # check if Align is really required
                c=0
                for i in fsspseq[0]:
                    if i!='.': c+=1
                if (c != len(kpssdAA)):
                    print >> sys.stderr, "Seq Lengths Different - k"
                    continue
                c=0
                for i in fsspseq[1]:
                    if i!='.': c+=1
                if (c != len(vpssdAA)):
                    print >> sys.stderr, "Seq Lengths Different - v"
                    continue
                for i in range(len(fsspseq[0])):
                    a1 = fsspseq[0][i]
                    a2 = fsspseq[1][i]
                    if ((a1 == '.') and (a2 == '.')): continue
                    if (a1 != '.'):
                        kdsspctr += 1
                        kpssdctr += 1
                    if (a2 != '.'):
                        vdsspctr += 1
                        vpssdctr += 1
                    if ((a1 == '.') or \
                        (a2 == '.')):
                        continue
                    if ((a1.lower() == a1) or \
                        (a2.lower() == a2)):
                        continue
                    kpssdAA2 += kpssdAA[kpssdctr]
                    vpssdAA2 += vpssdAA[vpssdctr]
                    kdsspSS2 += kdsspSS[kdsspctr]
                    vdsspSS2 += vdsspSS[vdsspctr]
                    kpssdSS2 += kpssdSS[kpssdctr]
                    vpssdSS2 += vpssdSS[vpssdctr]
                #print kpssdAA2
                #print vpssdAA2
                #print kdsspSS2
                #print vdsspSS2
                #print kpssdSS2
                #print vpssdSS2
                ttable = string.maketrans(' BGITS', 'CCCCCC')
                kdsspSS2 = kdsspSS2.translate(ttable)
                vdsspSS2 = vdsspSS2.translate(ttable)
                kpssdSS2 = kpssdSS2.translate(ttable)
                vpssdSS2 = vpssdSS2.translate(ttable)
                dsspoutlines = ['AA  OSEC PSEC\n']
                pssdoutlines = ['AA  OSEC PSEC\n']
                for i in range(len(kpssdAA2)):
                     dsspoutlines.append(kpssdAA2[i]+'   '+ \
                                         kdsspSS2[i]+'    '+vdsspSS2[i]+'\n')
                     pssdoutlines.append(kpssdAA2[i]+'   '+ \
                                         kpssdSS2[i]+'    '+vpssdSS2[i]+'\n')
                ofh = file(outdir+fsspfname[:-4]+'_dssp', 'w')
                ofh.writelines(dsspoutlines)
                ofh.close()
                ofh = file(outdir+fsspfname[:-4]+'_pssd', 'w')
                ofh.writelines(pssdoutlines)
                ofh.close()
            #break

if __name__ == '__main__':
    alndir = '/mnt/e3/data/indraneel/fssp_490/'
    dsspdir = '/mnt/e13/data/dssp/'
    pssddir = '/mnt/e13/data/pssd/'
    alignprog = '~/align0'
    outdir = '/mnt/e3/data/indraneel/compaln/'
    sovoutdir = '/mnt/e3/data/indraneel/compsov/'
    interface = Interface()
    interface.run(alndir, dsspdir, pssddir, outdir, alignprog)
