#!/usr/bin/env python
####################################################################
# shiftx2_util.py: utility functions supporting shiftx2_main.py
# 
# author: Yifeng Liu
# date: April 13th, 2011
####################################################################
import sys, os
from natsorted import *

# FUNCTIONS FOR SPLITING PDB FILE WITH MULTIPLE NMR MODELS

def print_model(model_lines, file_prefix, num):
    ''' model_lines contains all lines for a given protein model 
    this method will print the lines into a single file '''
    outfile = file_prefix + "-m" + str(num) + ".pdb"
    output = open(outfile, 'w')
    for line in model_lines:
        print >>output, line.rstrip()
    output.close()
    return outfile

def split_model(infile):
    ''' split the given PDB file with multiple NMR models into 
    a list of temporary files, for example, 1UBQ.pdb
    will be split as 1UBQ-m1.pdb, 1UBQ-m2.pdb etc...
    the return value is a list of model files
    '''
        
    file_prefix = infile[0:-4] # filename without .pdb is the prefix
    input = open(infile, 'r')
    lines = input.readlines()
    input.close()

    model_num = 0
    model_lines = []
    INMODEL = False
    model_files = []

    for line in lines:
        if line.startswith("ENDMDL"):
            model_file = print_model(model_lines, file_prefix, model_num)
            model_files.append(model_file)
            model_lines = []
            INMODEL = False
            continue
        if INMODEL:
            model_lines.append(line)
        elif line.startswith("MODEL"):
            model_lines.append(line)
            INMODEL = True
            model_num += 1
        else:
            continue
            
    #print "detected", model_num, "models"
    return model_files

# FUNCTIONS FOR AVERAGING MODEL FILES
# Given a list of model files, this program will calculate the average 
# of all predictions, a scenario in predicting average of multiple
# NMR-solved protein models

# 6,PRO,P,HA : [4.4532, ... ]
pred_dict = {} # a dictionary from key to a list of predicted values
ave_dict = {} # a dictionary from key to a single averaged predicted value

def add_predictions(infile):
    ''' add the prediction in the input file to the prediction list '''
    input = open(infile, 'r')
    lines = input.readlines()
    for i in range(1, len(lines)):
        line = lines[i].strip()
        tmplist = line.split(',')
        key = ','.join(tmplist[0:4])
        value = tmplist[4]
        if key not in pred_dict:
            pred_dict[key] = []
        pred_dict[key].append(value)
    input.close()

def average(alist):
    ''' return the average of the given list '''
    ave = 0.00
    if len(alist)==0:
        return 0.00
    for item in alist:
        ave += float(item)
    return ave/len(alist)
    
def average_prediction():
    ''' calculate the average prediction for each entry in pred_dict 
    and store the result in ave_dict '''
    entry_keys = pred_dict.keys()
    for key in entry_keys:
        value = average(pred_dict[key])
        ave_dict[key] = value
        
# main function
def model_average(infiles, outfile):
    ''' given a list of prediction files for each NMR model
    calculate the average predictions and print the result
    to the output file'''

    for infile in infiles:
        add_predictions(infile)
    average_prediction()
    output = open(outfile, 'w')
    entry_keys = ave_dict.keys()
    #entry_keys.sort()
    entry_keys = natsorted(entry_keys)
    for key in entry_keys:
        outstr = "%s,%.4f" % (key, ave_dict[key])
        print >>output, outstr
    output.close()

# support functions for multi-chain mode
def merge_chains(infile, outfile):
    ''' merge all PDB chains in the input file,  and return a dictionary of residue numbers and chains '''
    # renumbering of residues
    # the residue number for the first chain is unchanged
    # residue number for the second chain will be the original residue number plus the number of residue in previous chain
    # subsequent residue numbers are also original residue numbers plus the total number of residues in previous chain

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

    dict = {} # dictionary from newly assigned residue number to "CHAIN_ID, ORIGINAL_RESNUM" 
    curr_chain = ""
    res_count = 0 # residue count in the current chain.
    curr_res = "" # current residue number
    sum_resnum = 0 # previous residue number 

    for line in lines:
        if line.startswith("ATOM"):
            chainid = line[20:22].strip()
            resnum = line[22:26].strip()

            if chainid != curr_chain:
                curr_chain = chainid
                sum_resnum = res_count

            new_resnum = str(sum_resnum + int(resnum))
            dict[new_resnum] = chainid + "," + resnum

            if resnum != curr_res:
                # print chainid, resnum, new_resnum
                res_count += 1
                curr_res = resnum
                # print chainid, resnum, new_resnum
                
            # print line.rstrip()
            newline = line[0:20] + " A" + "%4s" % new_resnum + line[26:]
            outlines.append(newline.rstrip())
    outlines.append("END")

    output = open(outfile, 'w')
    print >>output, "\n".join(outlines)
    output.close()

    return dict

def restore_chain_resnum(infile, outfile, dict, format):
    ''' restore the original residue number with chain number with the information given in the dictionary 
    (RESNUM -> CHAINID,OLD_RESNUM), and prints out prediction file according to different formats '''
    input = open(infile, 'r')
    lines = input.readlines()
    input.close()


    outlines = []
    if format=="CSV":
        for line in lines:
            line = line.rstrip()
            if line.strip()=="":
                outlines.append("\t")
                continue
            tmplist = line.split(",")
            resnum = tmplist[0]
            if resnum in dict:
                tmplist[0] = dict[resnum]
            elif tmplist[0] == "NUM":    ## Add by B HAN to print out column header
                tmplist[0] = "CHAIN,NUM"

            #print tmplist
            outlines.append(",".join(tmplist))

    elif format=="TAB": # tabular format line starts with resnum or white space.
        # BACKBONE ATOMS
        # Num  RES       CA       CB       CO        N        H       HA
        # 1    S  58.5588  63.6231 174.9481 116.1368   8.3760   4.3310
        for line in lines:
            line = line.rstrip()
            if line.strip()=="":
                outlines.append("\t")
                continue
            if line.startswith("Num"):
                outlines.append("C " + line)
            elif not line[0].isalpha(): # starts with number of space
                resnum = line[0:3].strip()
                if resnum in dict:
                    chainid, oldres = dict[resnum].split(",")
                #print line
                newline = "%1s %3s" % (chainid, oldres) + line[3:]
                outlines.append(newline)
                #print newline   
            else:
                outlines.append(line)
                
    elif format=="BMRB": # NMR START format, there are no positions for chain ID in NMR STAR format, 
        # so the output is unchanged
        for line in lines:
            line = line.rstrip()
            outlines.append(line)
        
    output = open(outfile, 'w')
    print >>output, "\n".join(outlines)
    output.close()

    
            
