import sys, os
from pmx import *
from pmx.ffparser import RTPParser, NBParser
from pmx.rotamer import _aa_chi
from pmx.parser import kickOutComments, readSection, parseList
import tempfile

standard_pair_list = [
    ('N','N'),
    ('H','H'),
    ('CA','CA'),
    ('C','C'),
    ('O','O'),
    ('HA','HA'),
    ('CB','CB'),
    ('1HB','1HB'),
    ('2HB','2HB'),
    ('CG','CG')
    ]

standard_pair_list_charmm = [
    ('N','N'),
    ('HN','HN'),
    ('CA','CA'),
    ('C','C'),
    ('O','O'),
    ('HA','HA'),
    ('CB','CB'),
    ('1HB','1HB'),
    ('2HB','2HB'),
    ('CG','CG')
    ]

standard_pair_listB = [
    ('N','N'),
    ('H','H'),
    ('CA','CA'),
    ('C','C'),
    ('O','O'),
    ('HA','HA'),
    ('CB','CB'),
    ('1HB','1HB'),
    ('2HB','2HB')
    ]

standard_pair_list_charmmB = [
    ('N','N'),
    ('HN','HN'),
    ('CA','CA'),
    ('C','C'),
    ('O','O'),
    ('HA','HA'),
    ('CB','CB'),
    ('1HB','1HB'),
    ('2HB','2HB')
    ]

standard_pair_listC = [
    ('N','N'),
    ('H','H'),
    ('CA','CA'),
    ('C','C'),
    ('O','O'),
    ('HA','HA'),
    ('CB','CB')
    ]

standard_pair_list_charmmC = [
    ('N','N'),
    ('HN','HN'),
    ('CA','CA'),
    ('C','C'),
    ('O','O'),
    ('HA','HA'),
    ('CB','CB')
    ]

standard_pair_listD = [
    ('N','N'),
    ('H','H'),
    ('CA','CA'),
    ('C','C'),
    ('O','O')
    ]

standard_pair_list_charmmD = [
    ('N','N'),
    ('HN','HN'),
    ('CA','CA'),
    ('C','C'),
    ('O','O')
    ]
    
use_standard_pair_list = {
    'PHE': [ 'TRP','HIP','HID','HIE','HSP','HSD','HSE','HIS1','HISH','HISE'],
    'TYR': [ 'TRP','HIP','HID','HIE','HSP','HSD','HSE','HIS1','HISH','HISE'],
    'TRP': [ 'PHE','TYR','HIP','HID','HSE','HSP','HSD','HSE','HIS1','HISH','HISE','HIE'],
    'HID': [ 'PHE','TYR','TRP'], #[ 'PHE','TYR','HIP','TRP','HIE'],
    'HIE': [ 'PHE','TYR','TRP'], #[ 'PHE','TYR','HIP','HID','TRP'],
    'HIP': [ 'PHE','TYR','TRP'], #,'HID','HIE'],
    'HSD': [ 'PHE','TYR','TRP'], #[ 'PHE','TYR','HIP','TRP','HIE'],
    'HSE': [ 'PHE','TYR','TRP'], #[ 'PHE','TYR','HIP','HID','TRP'],
    'HSP': [ 'PHE','TYR','TRP'], #,'HID','HIE'],
    'HIS1': [ 'TRP','PHE','TYR'],
    'HISE': [ 'TRP','PHE','TYR'],
    'HISH': [ 'TRP','PHE','TYR']
    }

res_with_rings = [ 'HIS','HID','HIE','HIP','HISE','HISH','HIS1','HSE','HSD','HSP',
		   'PHE','TYR','TRP','PRO' ]

res_diff_Cb = [ 'THR', 'ALA', 'VAL', 'ILE' ]

res_gly_pro = [ 'GLY', 'PRO' ]

merge_by_name_list = {
    'PHE':['TYR'],
    'TYR':['PHE'],
    'HID':['HIP','HIE'],
    'HIE':['HIP','HID'],
    'HIP':['HID','HIE'],
    'HSD':['HSP','HSE'],
    'HSE':['HSP','HSD'],
    'HSP':['HSD','HSE'],
    'HIS1':['HISE','HISH'],
    'HISE':['HIS1','HISH'],
    'HISH':['HIS1','HISE']
}
    

mol_branch = {
    'ILE':2,
    'VAL':2,
    'LEU':3,
    'GLN':3,
    'GLU':3,
    'GLUH':3,
    'ASP':2,
    'ASN':2,
    'PHE':2,
    'TYR':2,
    'TRP':2,
    'HIS':2,
    'HIS1':2,
    'HISE':2,
    'HISH':2,
    'THR':2,
    'ALA':5,
    'SER':5,
    'GLY':5,
    'CYS':5,
    'MET':5,
    'ARG':5,
    'LYS':5,
    'LYN':5,
    }

def max_rotation(dihedrals):
    m = 0
    for d in dihedrals:
        if d[-2] == 1 and d[-1]>0: m = d[-1]
        if d[-2] == 0 and d[-1]!=-1: return d[-1]
    return m+1

def get_dihedrals(resname):
    return library._aa_dihedrals[resname]

def set_dihedral(atoms,mol,phi):
    print atoms[0].name,atoms[1].name,atoms[2].name
    a1 = atoms[0]
    a2 = atoms[1]
    a3 = atoms[2]
    a4 = atoms[3]
#    a1,a2,a3,a4 = atoms
    d = a1.dihedral(a2,a3,a4)
    r = Rotation(a2.x,a3.x)
    rot = d-phi
#    print a2.name, a3.name
    for atom
 in mol.atoms:
        if atom.order > a3.order:
            if a3.long_name[3]==' ':
#                print 'rotating', atom.n
ame
                atom.x = r.apply(atom.x,-rot)
            else:
                if atom.long_name[3]==a3.long_name[3] \
                   or atom.long_name[3]==' ':
#                    print 'rotating', atom.name
                    atom.x = r.apply(atom.x,-rot)
#    print a1.dihedral(a2,a3,a4)

def is_number(s):
    try:
        float(s)
        return True
    except ValueError:
        return False

def do_fit(m1,dihed1,m2,dihed2):
    bonds = []
    diheds = []
    #if not diheds: # we take the chi's
    diheds = []
    for dih in dihed1:
        chi = dih[-1]
        if chi  > 0:
            for dih2 in dihed2:
                if dih2[-1] == chi and dih2[-2] == 1:
                    diheds.append((dih,dih2))
    for dih1, dih2 in diheds:
#	for d in range(0,4):
#	   foo = dih2[d]
#	   if is_number(foo[-1]) and is_number(foo[-2]):
#		dih2[d] = foo[-1] + foo[0] + foo[1] + foo[2]
#	   elif foo[0]=='H' and foo[1]=='B':
#		dih2[d] = foo[-1] + foo[0] + foo[1]
#           foo = dih1[d]
#           if is_number(foo[-1]) and is_number(foo[-2]):
#                dih1[d] = foo[-1] + foo[0] + foo[1] + foo[2]
#           elif foo[0]=='H' and foo[1]=='B':
#                dih1[d] = foo[-1] + foo[0] + foo[1]

        atoms1 = m1.fetchm(dih1)
	print "fetching"
        atoms2 = m2.fetchm(dih2)
	print dih2
	print atoms2
        a1,a2,a3,a4 = atoms1
        if (a2.name, a3.name) not in bonds:
            phi = a1.dihedral(a2,a3,a4)
#            print phi
#            print 'rot1', a1.name,a2.name,a3.name,a4.name
	    set_dihedral(atoms2,m2,phi)
#            a1,a2,a3,a4 = atoms1
#            print a1.dihedral(a2,a3,a4)
    op = open('check.pdb','w')
    for atom in m2.atoms:
        print >>op, atom


def tag(atom):
    s = '%s|%s|%s|%s' %( atom.resname, atom.name, atom.atomtype, atom.atype)
    return '%-20s' % s

def align_sidechains(r1, r2):
    #MS: if you add new force field, check pmx/molecule.py, make sure res is
    #MS: correctly set into get_real_resname 
    for i in range(r1.nchi()):
        phi = r1.get_chi(i+1, degree = True)
        r2.set_chi(i+1,phi)

def assign_rtp_entries( mol, rtp):
    entr = rtp[mol.resname]
    neigh=[]
    # atoms
    for atom_entry in entr['atoms']:
        atom_name = atom_entry[0]
        atom_type = atom_entry[1]
        atom_q    = atom_entry[2]
        atom_cgnr = atom_entry[3]
        atom = mol.fetch( atom_name )[0]
        atom.atomtype = atom_type
        atom.q = atom_q
        atom.cgnr = atom_cgnr
    # bonds
    for a1, a2 in entr['bonds']:
	#MS charmm uses next residue in bonds (+), amber pervious (-)
	#MS also write rtp is affected now, since normally -C N is added
	#MS charmm used C +N, problem is that you run into trouble on the
	#MS termini
	bmin=not '-' in a1 and not '-' in a2
	bplus=not '+' in a1 and not '+' in a2
        if bmin and bplus:
            atom1, atom2 = mol.fetchm( [a1, a2] )
            atom1.bonds.append( atom2 )
            atom2.bonds.append( atom1 )
	else :
	    neigh.append([a1,a2])
    return neigh
        

def assign_branch(mol):
    for atom in mol.atoms:
        if atom.long_name[-1] == ' ':
            atom.branch = 0
        elif atom.long_name[-1] == '1':
            atom.branch = 1
        else:
            atom.branch = 2
        
def get_atoms_by_order(mol,order):
    res = []
    for atom in mol.atoms:
        if atom.order == order:
            res.append(atom)
    return res

def get_atoms_by_order_and_branch( mol, order, branch, merged_atoms ):
    res = []
    for atom in mol.atoms:
        if atom.order == order and atom not in merged_atoms:
            if atom.branch in [0,branch] or atom.branch < mol_branch[mol.real_resname] + 1:
                res.append(atom)
    return res
    

def last_atom_is_morphed( atom, merged_list ):
    for at in atom.bonds:
        if at.order < atom.order and at in merged_list:
            return True
    return False

def cmp_mol2_types( type1, type2 ):
    if type1 == type2 : return True
    if type1 == 'H' or type2 == 'H': return True
    tp1_ext = type1.split('.')[1]
    tp2_ext = type2.split('.')[1]
    if tp1_ext in ['2','3'] and \
       tp2_ext in ['2','3']:
        return False
    elif tp1_ext in ['am','co2'] and \
         tp2_ext in ['am','co2']:
        return True
    elif type1 in ['O.2','O.co2'] and \
         type2 in ['O.2','O.co2']:
        return True
    elif (type1[0] == 'S' and type2[0] != 'S') or \
         (type1[0] != 'S' an