#!/usr/bin/env python
##################################################################
# pred_corr.py calculates correlation coefficient between
# two given files.
# 
# author: Yifeng Liu
# date: April 18th, 2011
# 
# note: required numpy
##################################################################
import sys, os, re
from numpy import *

def read_file(filename):
    ''' read the input file (in shiftx2 csv format) and convert it to a 
    dictionary with key "resnum-restype+atomname" and value "cs" '''

    tmphash = {}
    input = open(filename, 'r')
    for line in input:
        if line.strip() == "" or line.startswith("NUM"):
            continue
        tmplist = line.split(',')
        # if atomtype != "" and atomtype != tmplist[2].strip():
        
        if tmplist[1].strip() == "B":
            tmplist[1] = "C" # convert B to C
        key = tmplist[0] + '-' + tmplist[1] + '-' + tmplist[2]
        value = tmplist[3]
        #print key, value
        tmphash[key] = value
    input.close()
    return tmphash

def calc_corr(dict1, dict2, atomtype):
    ''' given two dictionaries of type: 'resnum-restype-atomname' : 'cs' 
    this function calculates the pearson correlation coefficients between the command set of predictions '''
    keyset = set(dict1.keys()) & set(dict2.keys()) # sets of keys for common predictions
    coverage1 = "%.4f" % (len(keyset) * 1.00 / len(dict1.keys()))
    coverage2 = "%.4f" % (len(keyset) * 1.00 / len(dict2.keys()))
    keylist = list(keyset)
    keylist.sort()

    rms_square = 0
    N = 0

    # convert the common predictions into two ordered list
    list1 = []
    list2 = []
    for key in keylist:
        tmplist = key.split('-')
        # print tmplist
        if atomtype=="ALL" or tmplist[2] == atomtype:
            list1.append(dict1[key])
            list2.append(dict2[key])
            diff = float(dict1[key]) - float(dict2[key])
            rms_square = rms_square + diff**2
            N = N + 1
        else:
            continue

    rms_square = rms_square / N
    rmsd = math.sqrt(rms_square)

    array1 = array(list1)
    array2 = array(list2)
    corr = corrcoef(array1, array2)[0, 1]    

    # print corr, coverage1, coverage2
    return (corr, rmsd)

# main function
if __name__ == "__main__":
    atomtypes = ['H', 'HA','CA','CB','C','N','ALL']
    atomtype = ""

    if len(sys.argv)<3:
        print "usage: python %prog [prediction1.cs] [prediction2.cs]"
        sys.exit(0)

    pred1 = read_file(sys.argv[1])
    pred2 = read_file(sys.argv[2])
    
    # obs = read_file("1ubq.bmrb.cs")
    
    corr = []
    rmsd = []
    for atype in atomtypes:
        (cc, rm) = calc_corr(pred1, pred2, atype)
        corr.append("%.4f" % cc)
        rmsd.append("%.4f" % rm)
        print atype, cc
    
    print '","'.join(corr)
