# $CVSHeader: pal8e/ext/add_ss_to_aln.py,v 1.2 2005/02/15 01:10:24 indraneel Rel $

#!/usr/bin/env python

'''
extract manual alignments from a list file and run gdt on the common region.
GDT runs on the corresponding PDB files with only those residues (CA) that
are in the common region of alignment. Also one of the PDB files is
altered to have the same residue names as the other.
'''

import sys, os, os.path, shutil, string
from Bio.PDB.PDBParser import PDBParser
from Bio.PDB.PDBIO import PDBIO

sys.path.append('/home/indraneel/wrk/ssd/src')
from SSD import NonStopSSD

tld = '../../tmp/gdt-20021210/allinputfiles/'
outdir = '../../tmp/gdt-20021210/alloutputfiles/'

def dump_aln(aln):
    for attr in dir(aln):
        if attr[:2] != '__':
            print attr, getattr(aln, attr)
    print

def read_aln_file(param):
    ifh = file(param.alnfname)
    lines = ifh.readlines()
    ifh.close()
    alignments = []
    class tmp:
        pass
    aln = tmp()
    counter = 0
    for i in range(len(lines)):
        line = lines[i]
        if line[:2] == '##' and lines[i+1][:2] == '##':
            continue
        if line[:2] == '##':
            if line.find('.dali') != -1:
                aln.skip = 1
                continue
            if line.find('circular') != -1:
                aln.circular = 1
            parts = line[2:].split('#')
            aln.tgtfname = parts[0].split('-')[0]
            aln.tgtch = parts[0].split('-')[1]
            aln.pdbfname = parts[1].split('-')[1]
            aln.pdbch = parts[1].split('-')[2]
            continue
        if line[:2] == '//':
            if 'skip' in dir(aln):
                aln = tmp()
                continue
            if len(aln.tgt) != len(aln.pdb):
                print 'error in listfile: line: ', i
            else:
                counter += 1
                aln.counter = counter
                ## store the length
                aln.length = len(aln.pdb)
                alignments.append(aln)
            aln = tmp()
            continue
        if lines[i+1][:2] == '//':
            continue
        ## sequence on this and next line
        parts = line.split()
        aln.tgtst = int(parts[1])
        aln.tgt = parts[2]
        aln.tgt = string.join(aln.tgt.split('.'),'-')
        buf = ''
        for r in aln.tgt:
            if r.isalpha() and r.islower():
                buf += '-'
            else: buf += r
        aln.tgt = buf
        aln.tgtlen = len(string.join(aln.tgt.split('-'),''))
        ## get the pdb name
        line = lines[i+1]
        parts = line.split()
        aln.pdbst = int(parts[1])
        aln.pdb = parts[2]
        aln.pdb = string.join(aln.pdb.split('.'),'-')
        buf = ''
        for r in aln.pdb:
            if r.isalpha() and r.islower():
                buf += '-'
            else: buf += r
        aln.pdb = buf
        aln.pdblen = len(string.join(aln.pdb.split('-'),''))
    #for aln in alignments:
    #    dump_aln(aln)
    return alignments

def run_ssd(aln):
    nonstopssd = NonStopSSD()
    nonstopssd.setnew(tld+aln.tgtfname+'.pdb', aln.tgtfname)
    aln.tgtstr = nonstopssd.run()
    nonstopssd.setnew(tld+aln.pdbfname+'.pdb', aln.pdbfname)
    aln.pdbstr = nonstopssd.run()

def tag_secstr(aln):
    reslst = []
    peptides = aln.tgtstr.get_CA_peptides()
    for peptide in peptides:
        for res in peptide:
            reslst.append(res)
    for i in range(len(reslst)):
        if reslst[i].id[1] == aln.tgtst:
            break
    reslst = reslst[i:i+aln.tgtlen]
    aln.tgtsec = ''
    ct = 0
    for r in aln.tgt:
        if r == '-':
            aln.tgtsec += '-'
            continue
        res = reslst[ct]
        if res.H and res.E: secstr = 'O'
        elif res.H: secstr = 'H'
        elif res.E: secstr = 'E'
        else: secstr = ' '
        aln.tgtsec += secstr
        ct += 1
    
    reslst = []
    peptides = aln.pdbstr.get_CA_peptides()
    for peptide in peptides:
        for res in peptide:
            reslst.append(res)
    for i in range(len(reslst)):
        if reslst[i].id[1] == aln.pdbst:
            break
    reslst = reslst[i:i+aln.pdblen]
    aln.pdbsec = ''
    ct = 0
    for r in aln.pdb:
        if r == '-':
            aln.pdbsec += '-'
            continue
        res = reslst[ct]
        if res.H and res.E: secstr = 'O'
        elif res.H: secstr = 'H'
        elif res.E: secstr = 'E'
        else: secstr = ' '
        aln.pdbsec += secstr
        ct += 1

def run_one(alignment):
    run_ssd(alignment)
    tag_secstr(alignment)

def mkoutput(alignments):
    outbuf = []
    for aln in alignments:
        keys = aln.tgtfname+' '+aln.tgtch+' '+str(aln.tgtst)+' '
        keys += aln.pdbfname+' '+aln.pdbch+' '+str(aln.pdbst)
        outbuf += [keys+'\n']
        outbuf += [aln.tgtsec+'\n']
        outbuf += [aln.tgt+'\n']
        outbuf += [aln.pdb+'\n']
        outbuf += [aln.pdbsec+'\n']
        outbuf += ['\n']
    ofh = file(outdir+'out', 'w')
    ofh.writelines(outbuf)
    ofh.close()

def run_all():
    class parm:
        pass
    param = parm()
    param.alnfname = sys.argv[1]
    alignments = read_aln_file(param)
    for alignment in alignments:
        ## temporary
        #if alignment.counter not in [1]:
        #    continue
        ## temporary
        run_one(alignment)
    mkoutput(alignments)

if __name__ == '__main__':
    run_all()
