# Copyright (C) 2006, Thomas Hamelryck (thamelry@binf.ku.dk)
# This code is part of the Biopython distribution and governed by its
# license.  Please see the LICENSE file that should have been included
# as part of this package.


import os
from string import upper
from Bio.PDB import *


def run_psea(fname):
    """
    Run PSEA and return output filename.
    """
    os.system("psea "+fname)
    last=fname.split("/")[-1]
    base=last.split(".")[0]
    return base+".sea"

def psea(pname):
    """
    Parse PSEA output file.
    """
    fname=run_psea(pname)
    start=0
    ss=""
    fp=open(fname, 'r')
    for l in fp.readlines():
        if l[0:6]==">p-sea":
            start=1
            continue
        if not start:
            continue
        if l[0]=="\n":
            break
        ss=ss+l[0:-1]
    fp.close()
    return ss

def psea2HEC(pseq):
    """
    Translate PSEA secondary structure string into HEC.
    """
    seq=[]
    for ss in pseq:
        if ss=="a":
            n="H"
        elif ss=="b":
            n="E"
        elif ss=="c":
            n="C"
        seq.append(n)
    return seq

def annotate(m, ss_seq): 
    c=m.get_list()[0]
    all=c.get_list()
    residues=[]
    # Now remove HOH etc.
    for res in all:
        if is_aa(res):
            residues.append(res)
    L=len(residues)
    if not (L==len(ss_seq)):
        print "Length mismatch", L, len(ss_seq)
        raise Exception
    for i in range(0, L):
        residues[i].xtra["SS_PSEA"]=ss_seq[i]
    #os.system("rm "+fname)

class PSEA:
    def __init__(self, model, filename):
        ss_seq=psea(filename)
        ss_seq=psea2HEC(ss_seq)
        annotate(model, ss_seq)
        self.ss_seq=ss_seq

    def get_seq(self):
        """
        Return secondary structure string.
        """
        return self.ss_seq
        

if __name__=="__main__":

    import sys
    from Bio.PDB import *

    # Parse PDB file
    p=PDBParser()
    s=p.get_structure('X', sys.argv[1])

    # Annotate structure with PSEA sceondary structure info
    PSEA(s[0], sys.argv[1])