#!/usr/bin/python 
############################################################
# bmrb_parser.py
# parses all RefDB corrected data files in the current dir
# and store the chemcial shift assignment into CSV files
# 
# resno, restype, atom, atomtype, cs
# 
# For example, the BMRB file bmr979.str.corr will be parsed
# and converted to CSV file bmr979.cs
# 
# the primary sequences will be stored in three fasta files
# for C, H, and N atoms.
# 
# auhor: Yifeng Liu
# date: Jan 2011
# 
# Details for Parsing:
# 
# #Corrected using PDB structure: 1SJQA
#   _BMRB_flat_file_name     bmr6178.str
#   _Name_common                                "polypyrimidine tract binding protein 1 domain 1"
#
#    _Residue_count                               148
#   _Mol_residue_sequence                       
# ;
# ADQLTEEQIAEFKEAFSLFD
# KDGDGTITTKELGTVMRSLG
# QNPTEAELQDMINEVDADGD
# GTIDFPEFLTMMARKMKDTD
# SEEEIREAFRVFDKDGNGYI
# SAAELRHVMTNLGEKLTDEE
# VDEMIREADIDGDGQVNYEE
# FVQMMTAK
# ;
# 
# Side chain carbon shifts are corrected using the CA 
# re-reference value, if available.
# 
############################################################
import sys, os, re

aa_code_dict = {
    #'Ala' : 'A', 'Cys' : 'C', 'Asp' : 'D', 'Glu' : 'E', 'Phe' : 'F',
    #'Gly' : 'G', 'His' : 'H', 'ILE' : 'I', 'Lys' : 'K', 'Leu' : 'L',
    #'Met' : 'M', 'Asn' : 'N', 'Pro' : 'P', 'Gln' : 'Q', 'Arg' : 'R',
    #'Ser' : 'S', 'Thr' : 'T', 'Val' : 'V', 'Trp' : 'W', 'Tyr' : 'Y',

    'ALA' : 'A', 'VAL' : 'V', 'PHE' : 'F', 'PRO' : 'P', 'MET' : 'M',
    'ILE' : 'I', 'LEU' : 'L', 'ASP' : 'D', 'GLU' : 'E', 'LYS' : 'K',
    'ARG' : 'R', 'SER' : 'S', 'THR' : 'T', 'TYR' : 'Y', 'HIS' : 'H',
    'CYS' : 'C', 'ASN' : 'N', 'GLN' : 'Q', 'TRP' : 'W', 'GLY' : 'G',
    '2AS' : 'D', '3AH' : 'H', '5HP' : 'E', 'ACL' : 'R', 'AIB' : 'A',
    'ALM' : 'A', 'ALO' : 'T', 'ALY' : 'K', 'ARM' : 'R', 'ASA' : 'D',
    'ASB' : 'D', 'ASK' : 'D', 'ASL' : 'D', 'ASQ' : 'D', 'AYA' : 'A',
    'BCS' : 'C', 'BHD' : 'D', 'BMT' : 'T', 'BNN' : 'A', 'BUC' : 'C',
    'BUG' : 'L', 'C5C' : 'C', 'C6C' : 'C', 'CCS' : 'C', 'CEA' : 'C',
    'CHG' : 'A', 'CLE' : 'L', 'CME' : 'C', 'CSD' : 'A', 'CSO' : 'C',
    'CSP' : 'C', 'CSS' : 'C', 'CSW' : 'C', 'CXM' : 'M', 'CY1' : 'C',
    'CY3' : 'C', 'CYG' : 'C', 'CYM' : 'C', 'CYQ' : 'C', 'DAH' : 'F',
    'DAL' : 'A', 'DAR' : 'R', 'DAS' : 'D', 'DCY' : 'C', 'DGL' : 'E',
    'DGN' : 'Q', 'DHA' : 'A', 'DHI' : 'H', 'DIL' : 'I', 'DIV' : 'V',
    'DLE' : 'L', 'DLY' : 'K', 'DNP' : 'A', 'DPN' : 'F', 'DPR' : 'P',
    'DSN' : 'S', 'DSP' : 'D', 'DTH' : 'T', 'DTR' : 'W', 'DTY' : 'Y',
    'DVA' : 'V', 'EFC' : 'C', 'FLA' : 'A', 'FME' : 'M', 'GGL' : 'E',
    'GLZ' : 'G', 'GMA' : 'E', 'GSC' : 'G', 'HAC' : 'A', 'HAR' : 'R',
    'HIC' : 'H', 'HIP' : 'H', 'HMR' : 'R', 'HPQ' : 'F', 'HTR' : 'W',
    'HYP' : 'P', 'IIL' : 'I', 'IYR' : 'Y', 'KCX' : 'K', 'LLP' : 'K',
    'LLY' : 'K', 'LTR' : 'W', 'LYM' : 'K', 'LYZ' : 'K', 'MAA' : 'A',
    'MEN' : 'N', 'MHS' : 'H', 'MIS' : 'S', 'MLE' : 'L', 'MPQ' : 'G',
    'MSA' : 'G', 'MSE' : 'M', 'MVA' : 'V', 'NEM' : 'H', 'NEP' : 'H',
    'NLE' : 'L', 'NLN' : 'L', 'NLP' : 'L', 'NMC' : 'G', 'OAS' : 'S',
    'OCS' : 'C', 'OMT' : 'M', 'PAQ' : 'Y', 'PCA' : 'E', 'PEC' : 'C',
    'PHI' : 'F', 'PHL' : 'F', 'PR3' : 'C', 'PRR' : 'A', 'PTR' : 'Y',
    'SAC' : 'S', 'SAR' : 'G', 'SCH' : 'C', 'SCS' : 'C', 'SCY' : 'C',
    'SEL' : 'S', 'SEP' : 'S', 'SET' : 'S', 'SHC' : 'C', 'SHR' : 'K',
    'SOC' : 'C', 'STY' : 'Y', 'SVA' : 'S', 'TIH' : 'A', 'TPL' : 'W',
    'TPO' : 'T', 'TPQ' : 'A', 'TRG' : 'K', 'TRO' : 'W', 'TYB' : 'Y',
    'TYQ' : 'Y', 'TYS' : 'Y', 'TYY' : 'Y', 'AGM' : 'R', 'GL3' : 'G',
    'SMC' : 'C', 'ASX' : 'B', 'CGU' : 'E', 'CSX' : 'C', 'GLX' : 'Z',
    'UNK' : 'X', 'SEB' : 'S', 'FTR' : 'W', 'BFD' : 'D', 'CAS' : 'C',
    'LYR' : 'K', 'GYC' : 'F', 'ARG_NH2' : 'R', 'TYR_NH2' : 'Y', 
    'CYS-NH2' : 'C', 'LEU_NH2' : 'L'
}

# amino acid name dictionary
aa_name_dict = {
    'A' : 'ALA',     'C' : 'CYS',    'D' : 'ASP',    'E' : 'GLU',
    'F' : 'PHE',     'G' : 'GLY',    'H' : 'HIS',    'I' : 'ILE',
    'K' : 'LYS',     'L' : 'LEU',    'M' : 'MET',    'N' : 'ASN',
    'P' : 'PRO',     'Q' : 'GLN',    'R' : 'ARG',    'S' : 'SER',
    'T' : 'THR',     'V' : 'VAL',    'W' : 'TRP',    'Y' : 'TYR'
}

def convert_csv(infile):
    ''' convert a given infile from the bmrb*.str.corr format to 
    a comma separated format: resno, restype, atomname, cs '''

    input = open(infile, 'r')
    lines = input.readlines()
    input.close()

    outfile = infile[0:-9] + ".cs"
    output = open(outfile, 'w')

    # info fields to fill in
    bmrb_num = "" # bmrb accession number
    common_name = ""
    sequence = ""
    ph = None
    temp = None # temperature
    num_res = None # number of residues
    first_res = None # sequence id of the first residue
    shift_types = set()
    carbon_correction = 0.00 # correction for carbon alpha shifts and sidechain carbon shifts

    INSECTION = False
    INSEQUENCE = False

    semicolon_count = 0
    for i in range(0, len(lines)):
        line = lines[i]
        if line.strip() == "":
            continue
        elif '  _BMRB_flat_file_name' in line: # parse the BMRB file name
            bmrb_num = line.replace("_BMRB_flat_file_name","").strip()
            print >>output, "#%s: %s" % ("BMRB_ID", bmrb_num)
        elif '_Name_common' in line: # parse molecule common name
            common_name = line.replace("_Name_common","").strip()
            print >>output, "#%s: %s" % ("COMMON_NAME", common_name.upper())
        elif ' pH ' in line: # pH
            ph = re.split("\s+", line.strip())[1]
            print >>output, "#%s: %s" % ("PH", ph)
        elif ' temperature ' in line: # temperature
            temp = re.split("\s+", line.strip())[1]
            print >>output, "#%s: %s" % ("TEMP", temp)
        elif  '_Residue_count' in line:
            res_num = line.replace('_Residue_count','').strip()
            print >>output, "#%s: %s" % ("NUM_RES", res_num)        
        elif '_Mol_residue_sequence' in line: # parse out the primary sequence
            INSEQUENCE = True
        elif INSEQUENCE:
            if 'loop_' in line:
                INSEQUENCE = False
            if line.strip()==';':
                if semicolon_count==0:
                    semicolon_count = semicolon_count + 1
                else:
                    INSEQUENCE = False
                    print >>output, "#%s: %s" % ("SEQ", sequence)
            else: # in sequence section so read out the sequence
                sequence += line.strip()

        # extract the carbon alpha shifts correction value from line        
        #bmr1399.str.corr chemical shifts have been re-referenced with the following 
        #offsets (these values have been added to the original bmr1399.str file):
        #HA       CA       CB      CO       N      HN
        #N/A     N/A      N/A     N/A     N/A     N/A
        if line.startswith('#') and 're-referenced' in line:
            refline = lines[i+3]
            alist = re.split("\\s+", refline)
            if alist[1].strip()!="N/A":
                carbon_correction = float(alist[1])
                # print carbon_correction 

        # detect chemical shift sections
        if "_Chem_shift_ambiguity_code" in line:
            INSECTION = True
        elif "stop_" in line:
            INSECTION = False
        elif INSECTION:
            line = line.strip()
            #print line
            tmplist = re.split("\s+", line)
            # to structure resid, restype, atom, cs
            rescode = "NONE"
            if first_res == None:
                first_res = tmplist[1] # set the first residue number 
                print >>output, "#%s: %s" % ("INIT_RES", first_res)
            if tmplist[2] in aa_code_dict:
                # 1141   107   PHE       H     H      7.58    0.030     1
                rescode = aa_code_dict[tmplist[2]]
                atomtype = tmplist[3]
                if atomtype.startswith('C') and atomtype!="CA" and atomtype!="CB" and atomtype !="CO":
                    #print "corrected shift", line,
                    tmplist[5] = float(tmplist[5]) + carbon_correction 
                    #print tmplist[5]
                outstr = tmplist[1] + "," + rescode + "," + tmplist[3] + "," + str(tmplist[5])
                #print tmplist[4]
                shift_types.add(tmplist[4].strip())
                print >>output, outstr
        else:
            continue
    output.close()

    # print out fasta entry
    types = ",".join(list(shift_types))
    outstr = ">%s\t| %s | %s | %s\n%s" % (infile[0:-9]+".cs", common_name.upper(), first_res, types, sequence)
    # print outstr
    return outstr

if __name__ == "__main__":
    if len(sys.argv)<1:
        print "Usage: python bmrb_parser.py"
        sys.exit(0)

    cfasta = open("refdb.c.fasta", "w")
    hfasta = open("refdb.h.fasta", "w")
    nfasta = open("refdb.n.fasta", "w")

    files = os.listdir(os.curdir)
    files.sort()
    for file in files:
        if file.endswith(".corr"):
            print file
            outstr = convert_csv(file)
            # print outstr
            types = outstr.split("\n")[0].split('|')[3].split(',') # >bmr915.cs      | 'T4 LYSOZYME' | 1 | H,N
            # removing leading and trailing space from residue type
            for i in range(0, len(types)):
                types[i] = types[i].strip()

            # print types
            if 'C' in types:
                print >>cfasta, outstr
            if 'H' in types:
                print >>hfasta, outstr
            if 'N' in types:
                print >>nfasta, outstr

    cfasta.close()
    hfasta.close()
    nfasta.close()
