#!/usr/bin/python
# Yifeng Liu
# August 10th, 2011
# refdb_resnum_checker.py
# 
# SHIFTY3 is reported to have degraded performance due to 
# inconsistent residue numbering in RefDB; residues in
# difference sections of the same BMRB file are being numbered
# differently. The purpose of this program is to ensure 
# that the first reside in the SEQ section is always numbered
# with the integer 1. The content section of the parsed 
# .cs file will be re-numbered accordingly.
#
# Algorithm:
#  1. a sequence is parsed from the content section
#  2. such sequence will be shorter than the complete sequence quoted 
#     in the SEQ section, then we compare the parsed sequence and SEQ
#     sequence to calculate an offset, for example, the first residue
#     in the parse sequence is labeled with old residue number 1,
#     but according to the SEQ sequence it is indeed sequence number 7
#     the offset will be 7 - 1 = 6.
#  3. The offset is added to the residue number of the content residues
#     and the entire file is reprinted to replace the original file
import os, sys, re

res_dict = {}
SEQ = ""

def matching_percentage(offset):
    ''' calculate the percentage of matching between the sequence 
    and the res_dict given the offset '''
    global res_dict, SEQ
    match = 0
    total = len(SEQ)
    for i in range(0, len(SEQ)):
        seq_res = SEQ[i]
        key = i + offset + 1
        if key in res_dict:
            dict_res = res_dict[key]
            if seq_res == dict_res:
                match += 1
        else:
            continue
    return 100.00*match/total        

def calc_offset():
    ''' check the residue numbers and calculate an offset '''
    global res_dict, SEQ
    best_offset = 0
    best_match = 0

    for i in range(-1000, 1000):
        curr_match = matching_percentage(i)
        if curr_match > best_match:
            best_match = curr_match
            best_offset = i
        # print i, curr_match
    return best_offset    

def build_resdict(lines):
    ''' fill res_dict with residue numbers and residue type '''
    global res_dict, SEQ
    for line in lines:
        if line.startswith('#'):
            if line.startswith('#SEQ'):
                SEQ = line.split(':')[1].strip()
            else:
                continue
        elif line.strip() == "":
            continue
        else:
            tmplist = line.split(',')
            resnum = int(tmplist[0])
            restype = tmplist[1]
            res_dict[resnum] = restype

def reprint(lines, offset, outfile):
    global res_dict, SEQ
    ''' reprint the lines with the given offset '''
    output = open(outfile, 'w')
    for line in lines:
        if line.startswith('#'):
            print >>output, line.strip()
        elif line.strip()=="":
            continue
        else:
            tmplist = line.split(',')
            tmplist[0] = str(int(tmplist[0]) - offset)
            outline = ','.join(tmplist)
            print >>output, outline.strip()
    output.close()

def check_and_fix_resnum(infile):
    ''' check and fix the residue number in the given file, if necesary '''
    global SEQ, res_dict

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

    build_resdict(lines)
    if SEQ.strip()=="":
        return None
    offset = calc_offset()
    if offset == 0:
        print infile, ": OK"
    else:
        print infile, ": Fixing it with offset", offset   
        reprint(lines, offset, infile)

def check_and_fix_resnum2(infile, outfile):
    ''' check and fix the residue number in the given file, if necesary '''
    global SEQ, res_dict

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

    build_resdict(lines)
    if SEQ.strip()=="":
        return None
    offset = calc_offset()
    if offset == 0:
        print infile, ": OK"
    else:
        print infile, ": Fixing it with offset", offset   
        reprint(lines, offset, outfile)

# main function
if __name__ == '__main__':
    #    print "Execute the following command within the RefDB folder full of .cs files\nUsage: %prog "
    if len(sys.argv)>=2:
        infile = sys.argv[1]
        outfile = sys.argv[2]
        check_and_fix_resnum2(infile, outfile)
    else:
        csfiles = os.listdir(os.curdir)
        for file in csfiles:
            if file.endswith(".cs"):
                check_and_fix_resnum(file)
