#!/usr/bin/env python
############################################################################################
# Python Implementation of SHIFTY 
#
# AUTHOR: YIFENG LIU
# DATE: APRIL 2011
############################################################################################
import sys, os, re, datetime, time

###############################################################
# SETTINGS
###############################################################
# path for xalign and blast executibles and blast databases
XALIGN_EXE = "xalign_x" # executable for xalign
BLASTP_EXE = "./blast/bin/blastall -p blastp" # executable for blast

BLASTDB = "blastdb" # folder containing the blast database
SHIFTDB = "refdb" # folder containing the BMRB shift database

##############################################################
# BLAST OPTIONS
##############################################################
# BLASTP_OPTION = " -evalue 10E-3 -outfmt 5 -out " # evalue cut off
BLASTP_OPTION = " -m 7 -o " # linux
MAX_BLAST_HITS = -1 # maxinum number of blast hits to keep, set to -1 if all blast hits are used.
IDENTITY_CUTOFF = 40 # sequence identity cutoff from 1...100

# you may set different identity cut off for each atom types
IDENTITY_CUTOFF_H = IDENTITY_CUTOFF
IDENTITY_CUTOFF_C = IDENTITY_CUTOFF
IDENTITY_CUTOFF_N = IDENTITY_CUTOFF

##############################################################
# DERIVED PATHS
##############################################################

BLASTDB_H = BLASTDB + "/shiftyh"
BLASTDB_N = BLASTDB + "/shiftyn"
BLASTDB_C = BLASTDB + "/shiftyc"

DEBUG = False # turn on/off debug output
DEUTERATED = False # a protein is deuterated or not

##############################################################
# data structures
##############################################################

ATOM_TYPES = ['C', 'N', 'H']
resnum_list = []
path = "" # path to the folder containing all files
base = "" # basename of the query pdb file, normally query.pdb
MAXRECORD = 0
INIT_QUERY_RESNUM = None # initial residue number for the query sequence
INIT_DB_RESNUM = None # initial residue number for the database entry
resno_list = [] # residue number from the original PDB

TIMESTAMP = str(int(time.time())) # time stamp for the application
tempfolder = '/tmp/'

##############################################################
# Essential Utility Functions
##############################################################

# amino acid code dictionary
aa_code_dict = {
    # twenty standard amino acid abbreviations
    #'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',
    'HIE' : 'H', 'HIP' : 'H', 'HID' : 'H'
}

# 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 rcs_lookup(atom_type, residue_type):
    ''' Given atom type (C,H,N) and residue type (A-Z), lookup the corresponding 
    random coil shift. If no matching atom types or residues are found return an 
    empty tuple.'''

    # WUTRICH/BUNDI PARAMETERS FOR AMINO ACIDS IN TETRAPEPTIDES (C)
    # HA, CA, CB, CO
    rcs_c = {
        'A' :   (4.32,    52.50,    19.10,    177.80),
        'B' :   (0.00,     0.00,     0.00,      0.00),
        'C' :   (4.71,    55.40,    41.10,    174.60),
        'D' :   (4.64,    54.20,    41.10,    176.30),
        'E' :   (4.35,    56.60,    29.90,    176.60),
        'F' :   (4.62,    57.70,    39.60,    175.80),
        'G' :   (3.96,    45.10,     0.00,    174.90),
        'H' :   (4.73,    55.0,     29.00,    174.10),
        'I' :   (4.17,    61.10,    38.80,    176.40),
        'J' :   (0.00,     0.00,     0.00,      0.00),
        'K' :   (4.32,    56.20,    33.10,    176.60),
        'L' :   (4.34,    55.10,    42.40,    177.60),
        'M' :   (4.48,    55.40,    32.90,    176.30),
        'N' :   (4.74,    53.10,    38.90,    175.20),
        'O' :   (0.00,     0.00,     0.00,      0.00),
        'P' :   (4.42,    63.30,    32.10,    177.30),
        'Q' :   (4.34,    55.70,    29.40,    176.00),
        'R' :   (4.34,    56.00,    30.90,    176.30),
        'S' :   (4.47,    58.30,    63.80,    174.60),
        'T' :   (4.35,    61.80,    69.80,    174.70),
        'U' :   (0.00,     0.00,     0.00,      0.00),
        'V' :   (4.12,    62.20,    32.90,    176.30),
        'W' :   (4.66,    57.50,    29.60,    176.10),
        'X' :   (0.00,     0.00,     0.00,      0.00),
        'Y' :   (4.55,    57.90,    38.80,    175.90),
        'Z' :   (0.00,     0.00,     0.00,      0.00)
     }

    # WISHART PARAMETERS FOR AMINO ACIDS IN HEXAPEPTIDES
    #            H       HA      HB2      HB3     HG2	HG3	HD2	HD3     HE2     HE3
    rcs_h = {
        'A' :   (8.24,	4.32,	1.39,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00),
        'B' :   (0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00),
        'C' :   (8.43,	4.71,	3.25,	2.99,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00),
        'D' :   (8.34,	4.64,	2.72,	2.65,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00),
	'E' :   (8.42,	4.35,	2.06,	1.96,	2.31,	2.31,	0.00,	0.00,	0.00,	0.00),
	'F' :   (8.30,	4.62,	3.14,	3.04,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00),
	'G' :   (8.33,	3.96,	3.96,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00),
	'H' :   (8.42,	4.73,	3.29,	3.16,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00),
	'I' :   (8.00,	4.17,	1.87,	1.45,	1.16,	0.91,	0.86,	0.00,	0.00,	0.00),
	'J' :   (0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00),
	'K' :   (8.29,	4.32,	1.84,	1.75,	1.44,	1.44,	1.68,	1.68,	3.00,	3.00),
	'L' :   (8.16,	4.34,	1.62,	1.62,	1.59,	0.92,	0.87,	0.00,	0.00,	0.00),
	'M' :   (8.28,	4.48,	2.11,	2.01,	2.60,	2.54,	2.10,	0.00,	0.00,	0.00),
	'N' :   (8.40,	4.74,	2.83,	2.75,	7.59,	6.91,	0.00,	0.00,	0.00,	0.00),
	'O' :   (0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00),
	'P' :   (-9,    4.42,	2.29,	1.94,	2.02,	2.02,	3.63,	3.63,	0.00,	0.00),
	'Q' :   (8.32,	4.34,	2.12,	1.99,	2.36,	2.36,	7.52,	6.85,	0.00,	0.00),
	'R' :   (8.23,	4.34,	1.86,	1.76,	1.63,	1.63,	3.20,	3.20,	8.07,	0.00),
	'S' :   (8.31,	4.47,	3.89,	3.87,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00),
	'T' :   (8.15,	4.35,	4.24,	1.21,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00),
	'U' :   (0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00),
	'V' :   (8.03,	4.12,	2.08,	0.94,	0.93,	0.00,	0.00,	0.00,	0.00,	0.00),
	'W' :   (8.25,	4.66,	3.29,	3.27,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00),
	'X' :   (0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00),
	'Y' :   (8.12,	4.55,	3.03,	2.98,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00),
	'Z' :   (0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00,	0.00)
    }

    # WUTRICH/BUNDI PARAMETERS FOR AMINO ACIDS IN TETRAPEPTIDES (N)
    # N   NH
    rcs_n = {
        'A' :     (123.80,   8.24),
        'B' :     (0.00,     0.00),
        'C' :     (118.60,   8.43),
        'D' :     (120.40,   8.34),
        'E' :     (120.20,   8.42),
        'F' :     (120.30,   8.30),
        'G' :     (108.80,   8.33),
        'H' :     (118.20,   8.42),
        'I' :     (119.90,   8.30),
        'J' :     (0.00,     0.00),
        'K' :     (120.40,   8.29),
        'L' :     (121.80,   8.16),
        'M' :     (119.60,   8.28),
        'N' :     (118.70,   8.40),
        'O' :     (0.00,     0.00),
        'P' :     (0.00,     0.00),
        'Q' :     (119.80,   8.32),
        'R' :     (120.50,   8.23),
        'S' :     (115.70,   8.31),
        'T' :     (113.60,   8.15),
        'U' :     (0.00,     0.00),
        'V' :     (119.20,   8.03),
        'W' :     (121.30,   8.25),
        'X' :     (0.00,     0.00),
        'Y' :     (120.30,   8.12),
        'Z' :     (0.00,     0.00)
    }

    value = None
    if(atom_type == 'CA'):
    #    HA, CA, CB, CO        
        value = rcs_c[residue_type][1]
    elif (atom_type == 'CB'):
        value = rcs_c[residue_type][2]
    elif (atom_type == 'CO'):
        value = rcs_c[residue_type][3]
    # N   NH
    elif(atom_type == 'N'):
        value = rcs_n[residue_type][0]
    elif (atom_type == 'NH'):
        value = rcs_n[residue_type][1]
    #            H       HA      HB2      HB3     HG2	HG3	HD2	HD3     HE2     HE3
    elif(atom_type == 'H'):
        value = rcs_h[residue_type][0]
    elif(atom_type == 'HA'):
        value = rcs_h[residue_type][1]
    elif(atom_type == 'HB2'):
        value = rcs_h[residue_type][2]
    elif(atom_type == 'HB3'):
        value = rcs_h[residue_type][3]
    elif(atom_type == 'HG2'):
        value = rcs_h[residue_type][4]
    elif(atom_type == 'HG3'):
        value = rcs_h[residue_type][5]
    elif(atom_type == 'HD2'):
        value = rcs_h[residue_type][6]
    elif(atom_type == 'HD3'):
        value = rcs_h[residue_type][7]
    elif(atom_type == 'HE2'):
        value = rcs_h[residue_type][8]
    elif(atom_type == 'HE3'):
        value = rcs_h[residue_type][9]
    else:
        value = None
    return value

#########################################################################################
# HELPER FUNCTIONS
#########################################################################################

def count_match(match_seq):
    ''' calculate the number of matching residues in match_seq identified by = '''
    match = 0
    for c in match_seq:
        if c=='=':
            match += 1
    return match

# ''' append space characters to the end of list to maxlength '''
def append_space(list, maxlength):
    ''' append space characters to the end of list to maxlength '''
    if(len(list) < maxlength):
        for i in range(maxlength - len(list)):
            list.append(' ')
    return list

def merge_dicts(dict1, dict2):
    ''' add all entries in dict2 to dict1 and return dict1 '''
    keyset = dict2.keys()
    for key in keyset:
        dict1[key] = dict2[key]
    return dict1

############################################################
# functions for performing natural sort
############################################################
def try_int(s):
    '''convert to integer if possible.'''
    try: return int(s)
    except: return s

def natsort_key(s):
    '''Used internally to get a tuple by which s is sorted.'''
    return map(try_int, re.findall(r'(\d+|\D+)', s))

def natcmp(a, b):
    '''Natural string comparison, case sensitive.'''
    return cmp(natsort_key(a), natsort_key(b))

def natcasecmp(a, b):
    '''Natural string comparison, ignores case."'''
    return natcmp(a.lower(), b.lower())

def natsort(seq, cmp=natcmp):
    '''In-place natural string sort.'''
    seq.sort(cmp)
    
def natsorted(seq, cmp=natcmp):
    '''Returns a copy of seq, sorted by natural string sort.'''
    import copy
    temp = copy.copy(seq)
    natsort(temp, cmp)
    return temp

#########################################################################
# FUNCTIONS RELATED TO CONVERTING PDB FILE TO FASTA
#########################################################################

def join_aa_list(aa_list):
    ''' join the given amino acid list into a single sequence'''
    aa_seq = []
    global resno_list
    resno_list = []
    for (aa, id) in aa_list:
		if aa in aa_code_dict:
			aa_seq.append(aa_code_dict[aa])
		else:
			aa_seq.append('X') # default is an unknown residue
		resno_list.append(id)
    return ''.join(aa_seq)

def pdb_to_fasta(pdb_file, fasta_file):
    ''' print the first sequence (marked by ATOM tag) of a given pdb file into a fasta file '''
    global INIT_QUERY_RESNUM
    global resno_list

    pdb_file = pdb_file.rstrip()
    # fasta_file = pdb_file[:-3]+"fasta"

    input = open(pdb_file, 'r')
    output = open(fasta_file, 'w')
    current_aa = ()
    aa_list = []

    for line in input:
        if line.startswith('ATOM'):
            # print line.rstrip()
            # ATOM      1  N   LYS A   1      20.944   8.963  27.829  1.00 27.17           N
            aa = line[17:21].strip()
            id = int(line[22:26])
            if(current_aa != (aa, id)):
                aa_list.append((aa, id))
                resno_list.append((aa, id))
                current_aa = (aa, id)
        elif line.startswith('TER'):
            break
    # print join_aa_list(aa_list)
    (first_aa, first_id) = aa_list[0]
    INIT_QUERY_RESNUM = first_id

    comment = ">QUERY\t|" + pdb_file + "|\t" + str(first_id)
    print >>output, comment
    query_seq = join_aa_list(aa_list)
    print >>output, query_seq # join_aa_list(aa_list)
    output.close()
    return query_seq

#########################################################################
# FUNCTIONS FOR RUNNING BLAST AND PARSE OUTPUT
#########################################################################

def get_query_seq(query_file):
    ''' return the query protein sequence from the query_file '''
    query = open(query_file, 'r')
    query.readline() # discard comments
    lines = query.readlines()
    seq = ""
    for line in lines:
        seq = seq + line.strip()
    query.close()
    return seq    

# ''' get the blast hit list for all three atom types '''
def get_blasthits_list(query_file):
    ''' get the blast hit list for all three atom types '''
    hitlist = []
    evallist = []
    idenlist = []

    (hit_list, eval_list, iden_list) = get_blasthits('C', query_file)
    hitlist.append(hit_list)
    evallist.append(eval_list)
    idenlist.append(iden_list)

    (hit_list, eval_list, iden_list) = get_blasthits('H', query_file)
    hitlist.append(hit_list)
    evallist.append(eval_list)
    idenlist.append(iden_list)

    (hit_list, eval_list, iden_list) = get_blasthits('N', query_file)
    hitlist.append(hit_list)
    evallist.append(eval_list)
    idenlist.append(iden_list)
    
    #print hitlist
    #print evallist
    #print idenlist
    return (hitlist, evallist, idenlist)

def get_blasthits(atomtype, query_file):
    ''' Blast query_file against a blast database based on the given atom type. 
    Return the sequence ID for all blast hits'''
    global DEBUG

    # detect what computer platform we are using
    # platform = sys.platform
    
    # get the query sequence
    seq = get_query_seq(query_file)

    # blast query file against database
    blastout_file = query_file + ".blastout"

    idcutoff = IDENTITY_CUTOFF_C    
    blastdb = BLASTDB_C    
    if(atomtype == 'N'):
        blastdb = BLASTDB_N
        idcutoff = IDENTITY_CUTOFF_N
    elif(atomtype == 'H'):
        blastdb = BLASTDB_H
        idcutoff = IDENTITY_CUTOFF_H

    blastcmd = BLASTP_EXE + " -i " + query_file + " -d " + blastdb + BLASTP_OPTION + blastout_file # LINUX
    if not DEBUG:
        blastcmd =  blastcmd + " > /dev/null 2>&1" # suppress standard output and standard error
    #else:
    #    print blastcmd

    os.system(blastcmd)

    # parse the blast output file for hit IDs
    blastout = open(blastout_file, 'r')
    hit_list = []
    eval_list = []
    iden_list = []
    curr_hit_id = "" # current hit ID.

    for line in blastout:
        m = re.search('<Hit_def>(.*)<\/Hit_def>', line)
        eval = re.search('<Hsp_evalue>(.*)<\/Hsp_evalue>', line)
        iden = re.search('<Hsp_identity>(.*)<\/Hsp_identity>', line)
        if(m):
            hit_title = m.group(1)
            hit_id = hit_title.split('#|')[0]
            if(not (curr_hit_id == hit_id)):
                curr_hit_id = hit_id
        if(eval):
            eval_list.append(eval.group(1))
        if(iden):
            hit_list.append(curr_hit_id)
            iden_list.append(iden.group(1))

    blastout.close()
    if DEBUG:
        print hit_list, eval_list, iden_list
        print idcutoff
    
    #os.unlink(os.path.abspath(blastout_file)) # remove blast files
    
    # only keep MAX_BLAST_HITS number blast hits
    hitlist = []
    evallist = []
    idenlist = []
    if(MAX_BLAST_HITS>0 and len(hit_list)>MAX_BLAST_HITS):
        hitlist = hit_list[0:MAX_BLAST_HITS]
        evallist = eval_list[0:MAX_BLAST_HITS]
        idenlist = iden_list[0:MAX_BLAST_HITS]
    elif(MAX_BLAST_HITS<0): # percent identity cutoff
        for i in range(0, len(iden_list)):
            identity = int(iden_list[i]) * 100.0 / len(seq)
            if (identity >= idcutoff):
                # print "%.3f" % (identity),
                # print hit_list, i
                hitlist.append(hit_list[i])
                evallist.append(eval_list[i])
                idenlist.append(int(iden_list[i]) * 100.0 / len(seq))
    # print (hitlist, evallist, idenlist)
    return (hitlist, evallist, idenlist)

#########################################################################
# FUNCTIONS FOR RUNNING AND PARSING XALIGN OUTPUTS
#########################################################################

# '''Align the sequence in query_file and hit_file using xalign and put the output in xalign_outfile'''
def run_xalign(query_file, db_seq, xalign_outfile):
    '''Align the sequence in query_file and hit_file using xalign and put the output in xalign_outfile'''
    # temporary files
    global XALIGN_EXE, tempfolder
    
    xalign_option = query_file + ".option"    
    xalign_infile = query_file + ".xalignin"
    
    # get the primary sequence from the hit_file.
    # seq = get_sequence_from_shiftdb_file(hit_file)

    # print out xalign option file
    xalign_opt = open(xalign_option, 'w')
    print >>xalign_opt, "1\n" + xalign_infile + "\n1\n" + xalign_outfile
    xalign_opt.close()
    
    # print out xalign input file
    query = open(query_file, 'r')
    xalign_in = open(xalign_infile, 'w')
    for line in query:
        print >>xalign_in, line.rstrip()
    #hit_id = hit_file.split('/')[-1]
    #print >>xalign_in, ">" + hit_id + "\t|\n" + db_seq + "\n"
    print >>xalign_in, ">" + "DBENTRY" + "\t|\n" + db_seq + "\n"

    query.close()
    xalign_in.close()

    # call xalign 
    xalign_cmd = os.path.abspath(XALIGN_EXE) + " < " + xalign_option + " > /dev/null ";
    # print xalign_cmd
    # xalign_cmd = os.path.abspath(XALIGN_EXE) + " < " + xalign_option
    homedir = os.path.abspath('.')
    try:
        os.chdir(os.path.abspath(tempfolder))        
        os.system(xalign_cmd);
        os.chdir(homedir)
    except:
        print "[SHIFTY3] Error in xalign_cmd: " + xalign_cmd
        sys.exit(-1)
    
    #os.unlink(os.path.abspath(xalign_option))
    #os.unlink(os.path.abspath(xalign_infile))

def parse_xalignout(xalignout_file):
    ''' parse a given xalign output file. The xalign output file contains the global pairwise alignment
    between two peptide sequences. One of them is the query sequence (QUERY) and the other is a ShiftDB
    entry sequence. This method returns the percent identity, alignment score and the detailed pairwise 
    alignment section parsed from the xalign output file.'''

    xalignout = open(xalignout_file,'r')
    alignment_start = False
    percent_homology = 0
    alignment_score = 0
    sequence_title = ""
    query_seq = ""
    match_seq = ""
    db_seq = ""
    pairwise_alignment_section = ""

    for line in xalignout:
        # line = line.strip()
        if line == "":
            continue
        elif 'Pairwise Alignments' in line:
            alignment_start = True
            continue
        elif alignment_start and '------' in line:
            alignment_start = False
    
        if '*****' in line:
            continue

        if alignment_start:
        
            if 'Percent Sequence Homology' in line: # percent homology line
                match = re.search(r'(\d+.\d\d)', line) 
                if match:
                    percent_homology = float(match.group(1))
            elif 'Alignment Score' in line: # alignment score line
                match = re.search(r'(\d+.\d\d)', line)
                if match:
                    alignment_score = float(match.group(1))
            elif ('Sequence' not in line):
                pairwise_alignment_section = pairwise_alignment_section + line
                if 'Query' in line or 'QUERY' in line: # query line
                    query_seq = query_seq + line[11:61]
                elif 'Match' in line: # match line
                    match_seq = match_seq + line[11:61]
                else: # database sequence line
                    db_seq = db_seq + line[11:61]

    xalignout.close()
    return (percent_homology, alignment_score, query_seq.rstrip(), match_seq.rstrip(), db_seq.rstrip(), pairwise_alignment_section)

#############################################################
# FUNCTIONS FOR PARSING SHIFTDB FILES
#############################################################
def parse_dbfile(hit_file):
    ''' return the primary sequence from the shiftdb hit file '''
    global INIT_DB_RESNUM

    dbentry = {} # entry for database predictions "resnum, restype" : (atomtype, cs)

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

    ## print "## 592 SHIFTY3 hit_file:" + hit_file  ## print out hit the file in refdb

    resnum = None
    new_resnum = 0
    reslist = []
    seq = ""
    inti_res = 1

    for line in lines:
        if line.startswith('#'):
            if 'SEQ' in line:
                seq = line.split(':')[1].strip()
            if 'INIT_RES' in line:
                init_res = int(line.split(':')[1].strip())
            else:
                continue
        else:
            # 1,K,H,8.81
            tmplist = line.rstrip().split(',')
            curr_resnum = tmplist[0]
            res_type = tmplist[1]
            atom_type = tmplist[2]
            cs = tmplist[3]

            # set the initial database entry residue number
            #if resnum==None:
            #    INIT_DB_RESNUM = curr_resnum

            if curr_resnum != resnum:
                resnum = curr_resnum
                reslist.append(res_type)
                new_resnum += 1

            # add dbentry to the dictionary
            # we reassign the residue a number to ensure the residue numbers
            # form a consecutive sequence
            key = str(curr_resnum) + ',' + res_type 
            value = (atom_type, cs)
            if key not in dbentry:
                dbentry[key] = []
                # print new_resnum, res_type
                # print key
            dbentry[key].append(value)
            
    #if seq=="":
    #    seq = ''.join(reslist)
    for i in range(0, len(seq)):
        key = str(i+1) + ',' + seq[i]
        if key not in dbentry:
            dbentry[key] = []
            #print key
            
    return (seq, dbentry)

#############################################################
# MAIN FUNCTIONS
#############################################################

# ''' predict chemical shifts based on the shifty algorithm and format output string'''
def predict_shifts(shift_type, query_seq, match_seq, db_seq, dbentry):
    ''' predict chemical shifts based on the shifty algorithm and format output string
    shifty_type = C | N | H
    query_seq : query sequence
    match_seq : the sequence of match symbols
    db_seq : database entry sequence
    db_entry : a dictionary containing the database sequence
    '''
    global INIT_QUERY_RESNUM, INIT_DB_RESNUM
    global resno_list

    # print dbentry
    # dictionary for predictions
    prediction_dict = {}
    
    maxlength = max([len(query_seq), len(match_seq), len(db_seq)])
    querylist = append_space(list(query_seq), maxlength)
    matchlist = append_space(list(match_seq), maxlength)
    dblist = append_space(list(db_seq), maxlength)

    #print ''.join(querylist)
    #print ''.join(matchlist)
    #print ''.join(dblist)
    #print INIT_QUERY_RESNUM
    #print INIT_DB_RESNUM
    # sys.exit(0)

    # format output for each residue.
    dbindex = 0
    query_index = 0 # int(INIT_QUERY_RESNUM)
    dbentry_index = 1 #int(INIT_DB_RESNUM)
    # print dbentry

    for index in range(maxlength):
        queryaa = querylist[index]
        dbaa = dblist[index]
        cshifts = []
        ss = " "
        if(queryaa == ' ') or (dbaa == ' '):
            if (dbaa != ' '): # either sequence or db entry is blank. no prediction
                dbentry_index += 1
            if (queryaa != ' '):
                query_index += 1
            continue
        else:
            db_key = str(dbentry_index) + "," + dbaa
            if db_key in dbentry:
                cs_list = dbentry[db_key]
                if queryaa == dbaa: # match residue, transfer all shifts without modification
                    #print "matching residue", queryaa, dbaa
                    #print "transfering all shifts"
                    for (atomtype, cs) in cs_list:
                        if atomtype.startswith(shift_type):
                            key = str(resno_list[query_index]) + "," + queryaa + "," + atomtype
                            prediction_dict[key] = cs
                            #print key, cs
                else: # mismatch, must transfer shifts with RCS modifications.
                    #print "mismatch residue", queryaa, dbaa
                    #print "transfer modified shifts"
                    
                    for (atomtype, cs) in cs_list:                   
                        mod_cs = calculate_diffshifts(atomtype, queryaa, dbaa, cs)
                        if mod_cs and atomtype.startswith(shift_type):
                            key = str(resno_list[query_index]) + "," + queryaa + "," + atomtype
                            prediction_dict[key] = mod_cs
                            #print key, mod_cs
            else: # a missing residue
                #print "Missing", db_key
                continue
            dbentry_index += 1
            query_index += 1

    return prediction_dict

# ''' Calculate chemical shifts when query residue differ from database residue: 
# lookup the random coil shifts for both queryaa and dbaa, subtract dbaa RCS from 
# cshifts and add queryaa RCS'''
# Method: If query residue differs from a database entry residue: 
# Query-residue-cs <- DB-residue-cs - DB-residue-rcs + Query-residue-rcs
def calculate_diffshifts(atomtype, queryaa, dbaa, db_cs):
    ''' Calculate chemical shifts when query residue differ from database residue: 
    lookup the random coil shifts for both queryaa and dbaa, subtract database shfits
    from dbaa's RCS and add queryaa's RCS'''
    queryaa_rcs = rcs_lookup(atomtype, queryaa)
    dbaa_rcs = rcs_lookup(atomtype, dbaa)
    queryaa_cs = None
    # atomtype is one of the backbone and popular side chain type
    if queryaa_rcs!=None and dbaa_rcs!=None: 
        queryaa_cs = float(dbaa_rcs) - float(db_cs) + float(queryaa_rcs)
        if queryaa_cs < 0:
            queryaa_cs = None
        #print db_cs, dbaa_rcs, queryaa_rcs, queryaa_cs
    else:
        #print "omitting", atomtype
        # option 1
        # queryaa_cs = db_cs # in the case of mismatch residue, sidechains shifts are transfer exactly
        # option 2
        queryaa_cs = None # in the case of mismatch residue, sidechains shifts are discarded
        # What is the RCS for side chains????
    return queryaa_cs

# ''' Combine xalign output file and database entry into shifty output format '''
# [percent_homology, alignment_score, title, ref_info, num_match, rmsd, 
# pairwise_alignment_section, prediction, prediction_list]
# return a prediction ditionary 
def combine_xalign_dbentry(shift_type, xalignout_file, db_entry):
    ''' Combine xalign output file and database entry into shifty output format '''
    
    # parse the xalign output file
    (percent_homology, alignment_score, query_seq, match_seq, db_seq, pairwise_alignment_section) = \
        parse_xalignout(xalignout_file)

    # parse the database entry file
    # (title, ref_info, dbarray) = parse_dbentry(shiftdb_file)
    match_seq = match_seq.replace('|','=')

    # combine prediction results
    prediction_dict = predict_shifts(shift_type, query_seq, match_seq, db_seq, db_entry) # REWRITE!!
    
    # count number of matches and estimate rmsd
    num_match = count_match(match_seq) 
    
    return prediction_dict

def fine_tune_predictions(pred_dict):
    ''' fine tune the given dictionary of generated predictions,
    1. negative shifts should be removed
    2. invalid shifts should be removed
    3. if the deuteration option is selected, we make the corresponding changes '''
    # F should have HD2, HE2
    # W should have HE3
    # valine doesn't have HB3?
    # R should not have HE2
    # Y should have HD2 and HE2 but we need more verifications.

    global DEUTERATED 
    #new_dict = {}
    keys = pred_dict.keys()

    for key in keys:
        [resnum, restype, atomname] = key.split(',')
        # print "\t## 796" 
        cs = float(pred_dict[key])
        # print "\t## 798  cs: %f key: %s" % (cs, key)  

        ## excluded because when a cs is negative, it is tried to delete again; then makes an error   
        ## by Beomsoo Han 03 OCT 2011
        ## if cs < 0:
        ##    del pred_dict[key]

        if atomname == 'CA' and DEUTERATED:
            pred_dict[key] = cs - 0.43 # deuterated CA
        elif atomname == 'CB' and DEUTERATED:
            pred_dict[key] = cs - 0.82 # deuterated CB
        elif atomname == 'N' and DEUTERATED:
            pred_dict[key] = cs - 0.23 # deuterated N
        elif atomname == 'HB3' and (restype in ['A','G','V']):
            del pred_dict[key]
        elif atomname == 'HG2' and (restype not in ['V','E','N','I','L','M','P','Q','R','K']):
            del pred_dict[key]
        elif atomname == 'HG3' and (restype not in ['E','N','I','L','M','P','Q','R','K']):
            del pred_dict[key]
        elif atomname == 'HD2' and (restype not in ['F', 'I','L','M','P','Q','R','K']):
            del pred_dict[key]
        elif atomname == 'HD3' and (restype not in ['P','Q','R','K']):
            del pred_dict[key]
        elif atomname == 'HE2' and (restype not in ['F', 'K']):
            del pred_dict[key]
        elif atomname == 'HE3' and (restype not in ['W', 'K']):
            del pred_dict[key]
        elif restype == 'P' and (atomname in ['N','NH']):
            del pred_dict[key] # proline has no N or NH
        elif restype == 'HB2' and (restype in ['T', 'V', 'I']):
            restype = 'HB' # T and V and I have HB instead of HB2
            del pred_dict[key]
            new_key = resnum + ',' + restype + ',' + atomname
            pred_dict[new_key] = cs
        else:
            continue

    return pred_dict

# run shifty using one blast hit for the given shift_type 
def shifty_one_hit(shift_type, query_file, hit):
    ''' EXPERIMENT ONLY! RUN SHIFTY USING ONE BLAST HIT WITH THE GIVEN IDENTITY TYPE'''
    global SHIFTDB
    xalign_outfile = query_file + ".xalignout"
    hit_file = SHIFTDB + "/" + hit
    
    (db_seq, dbentry) = parse_dbfile(hit_file)
    
    # print dbentry
    run_xalign(query_file, db_seq, xalign_outfile)

    ## check the existance of xalign_outfile (add by Beomsoo Han, 27 Sep 2011)
    if os.path.exists(xalign_outfile) == False:
        print "[SHIFTY+] Error! There is no xaligned outfile"
        sys.exit(-1)
    
    # combine xalign result and db entry.
    prediction_dict = combine_xalign_dbentry(shift_type, xalign_outfile, dbentry)
    
    # filter negative, invalid predictions and add deuteration constants 
    prediction_dict = fine_tune_predictions(prediction_dict)

    return prediction_dict

# shifty with the top most blast hit above identity cutoff
def shifty_top_hit(query_file, output):
    ''' RUN SHIFTY USING THE BEST BLAST HIT WITH SEQUENCE IDENTITY ABOVE CUTOFF FOR ALL ATOM TYPES '''
    global DEUTERATED, IDENTITY_CUTOFF
    (hits,evals,idens) = get_blasthits_list(query_file) # GET THE BLAST HIT LITS

    # cshifts = {}
    # prediction_lists = []
    for i in range(0, 3):

        shift_type = 'C'
        if i==1:
            shift_type = 'H'
        elif i==2:
            shift_type = 'N'

        tmp_hit = hits[i]
        tmp_eval = evals[i]
        tmp_iden = idens[i]

        if DEBUG:
            print tmp_iden

        HAS_PREDICTION = False
        for j in range(0, len(tmp_iden)):
            iden = tmp_iden[j]
            if ((iden > IDENTITY_CUTOFF) and (HAS_PREDICTION == False)):
                hit = tmp_hit[j]
                eval = tmp_eval[j]
                iden = tmp_iden[j]
                if DEBUG:
                    print iden
                
                #print "USING HOMOLOG", hit
                pred_dict = shifty_one_hit(shift_type, query_file, hit)
                # cshifts = merge_dicts(cshifts, pred_dict)
                print_csvfile(pred_dict, iden, output)
                
                #print_iden_hits(shift_type, prediction_list, iden, output) # REWRITE!!
                #prediction_lists.append(prediction_list) # REWRITE!!

                # REWRITE:
                # 1. PARSE SEQUENCE AND SHIFTS FROM HIT FILE - done
                # 2. RUN XALIGN FOR ALIGNMENT - done
                # 3. PARSE XALIGN RESULTS FOR PARING SEQUENCE - done
                # 4. TRANSFER THE SHIFTS - done
                # 5. FILTER RESULTS - done
                # 6. DEUTERATION - done
                # 7. PRINT PREDICTIONS TO CSV WITH IDENTITY SCORE - working

                if len(pred_dict) > 0:
                    HAS_PREDICTION = True
            else:
                continue

def shifty_top_noiden_hit(query_file, output):
    ''' RUN SHIFTY USING THE BEST BLAST HIT WITH SEQUENCE IDENTITY ABOVE CUTOFF 
    BUT BELOW 100 SEQUENCE IDENTITY FOR ALL ATOM TYPES; THIS FUNCTION SHOULD 
    BE USED FOR ASSESSING PERFORMANCE ONLY AND NOT FOR PRODUCTION '''
    global DEUTERATED, IDENTITY_CUTOFF
    (hits,evals,idens) = get_blasthits_list(query_file) # GET THE BLAST HIT LITS
    # cshifts = {}
    # prediction_lists = []
    for i in range(0, 3):
        shift_type = 'C'
        if i==1:
            shift_type = 'H'
        elif i==2:
            shift_type = 'N'
        tmp_hit = hits[i]
        tmp_eval = evals[i]
        tmp_iden = idens[i]
        
        ### print "## 931: shift_type: %s" % shift_type
        if DEBUG:
            print tmp_iden
        HAS_PREDICTION = False
        for j in range(0, len(tmp_iden)):
            iden = tmp_iden[j]
            if ((iden > IDENTITY_CUTOFF) and (iden < 100) and (HAS_PREDICTION == False)):
                hit = tmp_hit[j]
                eval = tmp_eval[j]
                iden = tmp_iden[j]
                if DEBUG:
                    print iden

                pred_dict = shifty_one_hit(shift_type, query_file, hit)

                print_csvfile(pred_dict, iden, output)

                if len(pred_dict) > 0:
                    HAS_PREDICTION = True
                    break ## If shifty predicted, then go to next atom (added by Beomsoo Han)
            else:
                continue

def print_csvfile(pred_dict, identity, output):
    ''' print predictions in pred_dict with sequence identity to a file output in CSV format '''
    keyset = pred_dict.keys()
    keyset = natsorted(keyset) # sort the keyset in natural order
    outlines = []
    for key in keyset:
        [resnum, restype, atomname] = key.split(',')
        cs = pred_dict[key]
        outstr = str(resnum) + ',' + restype + ',' + atomname + ',' + str(cs) + ',' + "%.2f" % (identity)
        outlines.append(outstr)
    print >> output, "\n".join(outlines)
    return "\n".join(outlines)

def cleanup(path):
    ''' cleanup temporary files '''
    for f in ['.query.fasta', '.query.fasta.blastout', '.query.fasta.option', '.query.fasta.xalignin', \
        '.query.fasta.xalignout']:
        tmpfile = os.path.join(path, TIMESTAMP + f)
        # print tmpfile
        if os.path.exists(tmpfile):
            os.unlink(os.path.abspath(tmpfile))
    for f in ['xalign.parms', 'xalign.parms.tmp', 'wt.rbo', 'xalign_x']:
        tmpfile = os.path.join(path, f)
        if os.path.exists(tmpfile):
            os.unlink(os.path.abspath(tmpfile))
    # finally removing the temporary folder
    os.rmdir(os.path.abspath(path))
    return None

def copy_xalign_tmpfiles():
    ''' copy temp files needed by xalign_x to the temporary folder for easy clean-up
    this feature is added to ensure smooth paralllel processing '''
    global tempfolder, XALIGN_EXE
    cmd = 'cp xalign_x ' + tempfolder
    os.system(cmd)
    cmd = 'cp xalign.parms ' + tempfolder
    os.system(cmd)
    cmd = 'cp wt.rbo ' + tempfolder
    os.system(cmd)
    XALIGN_EXE = tempfolder + 'xalign_x'
    # print "in copy xalign tmpfiles", tempfolder, XALIGN_EXE
    return None
    
##########################################################################################
# ACTUAL MAIN PROGRAM
##########################################################################################
def main(query_file, output_file, cutoff):
    global path, base, tempfolder
    global IDENTITY_CUTOFF, IDENTITY_CUTOFF_H, IDENTITY_CUTOFF_C, IDENTITY_CUTOFF_N

    path = os.path.dirname(query_file)
    base = os.path.basename(query_file)

    #output_file = os.path.join(path, "query.shifty")
    # fasta_file = os.path.join(path, TIMESTAMP + ".query.fasta")
    fasta_file = os.path.join(tempfolder, TIMESTAMP + ".query.fasta")
    
    # take input from pdb and store in a tmp fasta file
    pdb_to_fasta(query_file, fasta_file) 
    query_file = fasta_file 
    
    # set identity cutoff to the customized identity
    IDENTITY_CUTOFF = IDENTITY_CUTOFF_H = IDENTITY_CUTOFF_C = IDENTITY_CUTOFF_N = cutoff
    output = open(output_file, 'w') 

    shifty_top_hit(query_file, output)       # rewrite this function. 
    ## shifty_top_noiden_hit(query_file, output)
    output.close()

    # cleanup temporary files
    cleanup(tempfolder)

########################################################################################
# main function: for command line interface
########################################################################################
import getopt
from optparse import OptionParser

if __name__ == '__main__': 
    # python shifty2.py -i query.pdb -o output.shifty -c id_cutoff -d 
    # global tempfolder
    
    usage = "usage %prog -i query.pdb -o output.shifty -c id_cutoff -d"
    parser = OptionParser(usage)
    parser.add_option("-i", "--infile", dest="infile", help="Input PDB File")
    parser.add_option("-o", "--outfile", dest="outfile", help="Output CSV File Name")
    parser.add_option("-c", "--idcutoff", dest="idcut", default=40, \
                          help="Sequence Similarity Cutoff, default 35")
    parser.add_option("-d", "--deuterated", dest="deuterated", action="store_true", \
                          help="Protein is deuterated or not, default False")
    parser.add_option("-v", "--debug", dest="debug", action="store_true", \
                          help="Turn on debug output")
    parser.add_option("-t", "--temp", dest="tempfolder", default="/tmp/", \
                        help="Specify which folder to store temporary files")
    (options, args) = parser.parse_args()

    if options.tempfolder:
        if not os.path.exists(options.tempfolder):
            parser.error("the specified temporary folder does not exist")
        elif os.path.abspath(options.tempfolder)==os.path.abspath('.'):
            parser.error("the specified temporary folder cannot be the same as the current folder")
        else:
            tempfolder = os.path.abspath(options.tempfolder)
    else:
        tempfolder = '/tmp/'
    tempfolder = tempfolder + '/' + TIMESTAMP + '/'
    os.mkdir(tempfolder) # creating the temporary folder for temp files
    # print "in main", tempfolder
    copy_xalign_tmpfiles()
    
    if not options.infile:
        parser.error("no input pdb file specified")

    query_file = options.infile

    output_file = ""
    if not options.outfile:
        output_file = query_file.split('.')[0] + ".shifty"
    else:
        output_file = options.outfile

    cutoff = int(options.idcut)
    # global DEUTERATED, DEBUG
    DEUTERATED = options.deuterated
    DEBUG = options.debug

    try:
        main(query_file, output_file, cutoff)
    except: # something is wrong in shifty+ 
        # cleanup(os.path.dirname(query_file)) # clean up temp file in exception.
        cleanup(tempfolder)
        print "[SHIFTY+] Error: execution halted due to unexpected errors. \
               \n  ->> SHIFTX2 is calculating with only SHIFTX+'s prediction results"
        sys.exit(-1)
    

