"""
SCRIPT TO CONVERT WRITE CHARMM RTF AND PRM FILES 
FROM BOSS ZMATRIX
Created on Mon Feb 15 15:40:05 2016
@author: Leela S. Dodda leela.dodda@yale.edu
@author: William L. Jorgensen Lab 

Usage: python OPM_Routines.py -z phenol.z -r PHN
REQUIREMENTS:
BOSS (need to set BOSSdir in bashrc and cshrc)
Preferably Anaconda python with following modules
pandas 
argparse
numpy
"""

from LigParGen.BOSSReader import bossPdbAtom2Element,bossElement2Mass,ucomb
import pickle
import pandas as pd
import numpy as np


def retDihedImp(df):
    odihed = []
    if np.sum([df['V' + str(pot)] for pot in range(1, 5)]) != 0.0:
        for pot in range(1, 5):
            if (df['V' + str(pot)] != 0.0):
                odihed.append('IMPRoper %s %4.5f %d %4.5f \n' % (df['NAME'].replace(
                    "-", "   "), df['V' + str(pot)], pot, 180.00 * abs(pot % 2 - 1)))
    else:
        pot = 2
        odihed.append('IMPRoper %s %4.5f %d %4.5f \n' % (df['NAME'].replace(
            "-", "   "), df['V' + str(pot)], pot, 180.00 * abs(pot % 2 - 1)))
    return (odihed)


def retDihed(df):
    odihed = []
    for pot in range(1, 5):
        odihed.append('DIHEdral %s           %4.5f %d %4.5f \n' % (df['NAME'].replace(
            "-", "   "), df['V' + str(pot)], pot, 180.00 * abs(pot % 2 - 1)))
    return (odihed)


def Boss2CharmmRTF(num2typ2symb, Qs, resid, bnd_df, angs,props,imps):
    charges = [float(Qs[i][1]) for i in range(len(Qs))]
    rtf = open(resid + '.top', 'w+')
    rtf.write('Remarks LigParGen generated XPLOR-TOP file for Bonvin group (by Leela Dodda)\n')
    rtf.write('\nset echo=false end\n')
    rtf.write('\nautogenerate angles=True dihedrals=True end\n')
    rtf.write('{ atomType  mass }\n')
    Mass = ['MASS %s %3.4f \n' % (num2typ2symb[i][2], bossElement2Mass(bossPdbAtom2Element(num2typ2symb[i][0]))) for i in range(len(Qs))]
    for i in range(len(Mass)):
        rtf.write('%s' % Mass[i])
    rtf.write('\nRESIdue %5s\n' % (resid))
    rtf.write('\nGROUP\n')
    rtf.write('\n{ atomName  atomType  Charge } \n')
    for i in range(len(Qs)):
        rtf.write('ATOM %6s TYPE= %6s CHARGE= %8s END\n' % (num2typ2symb[i][0], bossPdbAtom2Element(
            num2typ2symb[i][0]) + num2typ2symb[i][1][-3:], Qs[i][1]))
    rtf.write('\n{ Bonds: atomName1  atomName2 } \n')
    for (x, y) in zip(bnd_df.cl1, bnd_df.cl2):
        rtf.write('BOND %s %s \n' % (num2typ2symb[x][0], num2typ2symb[y][0]))
    rtf.write('\n{ Angles: atomName1 atomName2 atomName3} \n')
    for a in angs: rtf.write('ANGLe %s \n'%a) 
    rtf.write('\n{ Proper Dihedrals: aName1 aName2 aName3 aName4 } \n')
    for i in props:
        rtf.write('DIHEdral %s \n' % (i.replace("-", " ")))
    rtf.write('\n{ Improper Dihedrals: aName1 aName2 aName3 aName4 }\n')
    for i in imps:
        rtf.write('IMPRoper %s \n' % (i.replace("-", " ")))
    rtf.write('\nEND {RESIdue UNK}\n')
    rtf.write('\nset echo=true end\n')
    rtf.close()
    return None


def Boss2CharmmPRM(resid, num2typ2symb, Qs, bnd_df, ang_df, tor_df):
    #### COLLECTING NONBONDING PART #######
    prm = open(resid + '.param', 'w+')
    prm.write('Remarks LigParGen generated XPLOR-TOP file for Bonvin group (by Leela Dodda)\n')
    prm.write('\nset echo=false end \n')
    #    bnd_df = bnd_df.drop_duplicates(['TIJ'])
    prm.write('\n{ Bonds: atomType1 atomType2 kb r0 } \n')
    for i in bnd_df.index:
        prm.write('BOND %6s %6s %8.1f %8.4f \n' % (num2typ2symb[bnd_df.cl1[i]][
            2], num2typ2symb[bnd_df.cl2[i]][2], bnd_df.KIJ[i], bnd_df.RIJ[i]))
    #    ang_df = ang_df.drop_duplicates(['TY', 'R', 'K'])
    prm.write('\n{ Angles: aType1 aType2 aType3 kt t0 }\n')
    for i in ang_df.index:
        prm.write('ANGLe %5s %5s %5s %8.1f %8.2f \n' % (num2typ2symb[ang_df.cl1[i]][2], num2typ2symb[
            ang_df.cl2[i]][2], num2typ2symb[ang_df.cl3[i]][2], ang_df.K[i], ang_df.R[i]))
    prm.write('\n{ Proper Dihedrals: aType1 aType2 aType3 aType4 kt period phase } \n')
    if len(tor_df.index) > 0:
        tor_df = tor_df.drop_duplicates(['NAME', 'TY'])
    pro_df = tor_df[tor_df.TY == 'Proper']
    for i in list(pro_df.index):
        ndf = pro_df.ix[i]
        pro_out = retDihed(ndf.to_dict())
        for i in range(4):
            prm.write('%s' % pro_out[i])
    prm.write(
        'X    X    X    X    0.00000 1 0.000000 ! WILD CARD FOR MISSING TORSION PARAMETERS\n')
    prm.write('\n{ Improper Dihedrals: aType1 aType2 aType3 aType4 kt period phase }\n')
    imp_df = tor_df[tor_df.TY == 'Improper']
    for i in list(imp_df.index):
        ndf = tor_df.ix[i]
        imp_out = retDihedImp(ndf.to_dict())
        for i in range(len(imp_out)):
            prm.write('%s' % imp_out[i])
    prm.write(
        'IMPRoper X    X    X    X    0.00000 1 0.000000 ! WILD CARD FOR MISSING IMPROPER PARAMETERS \n')
    prm.write(
        '\n{ Nonbonded: Type Emin sigma; (1-4): Emin/2 sigma }\n')
    Qlines = ['NONBonded %5s %11.6f %11.6f %11.6f %11.6f \n' %
              (num2typ2symb[i][2], float(Qs[i][3]), float(Qs[i][2]), float(Qs[i][3]) * 0.50,float(Qs[i][2])) for i in range(len(Qs))]
    #    uniqQs = list(set(Qlines))
    for i in range(len(Qlines)):
        prm.write('%s' % Qlines[i])
    prm.write('\nset echo=true end \n')
    prm.close()
    return None


def Boss2CharmmTorsion(bnd_df, num2opls, st_no, molecule_data, num2typ2symb):
    #    print num2opls
    dhd = []
    for line in molecule_data.MolData['TORSIONS']:
        dt = [float(l) for l in line]
        dhd.append(dt)
    dhd = np.array(dhd)
    dhd = dhd  # kcal to kj conversion
    dhd = dhd / 2.0  # Komm = Vopls/2
    dhd_df = pd.DataFrame(dhd, columns=['V1', 'V2', 'V3', 'V4'])
    ats = []
    for line in molecule_data.MolData['ATOMS'][3:]:
        dt = [line.split()[0], line.split()[4],
              line.split()[6], line.split()[8]]
        dt = [int(d) for d in dt]
        ats.append(dt)
    for line in molecule_data.MolData['ADD_DIHED']:
        dt = [int(l) for l in line]
        ats.append(dt)
    assert len(ats) == len(
        dhd), 'Number of Dihedral angles in Zmatrix and Out file dont match'
    ats = np.array(ats) - st_no
    for i in range(len(ats)):
        for j in range(len(ats[0])):
            if ats[i][j] < 0:
                ats[i][j] = 0
    at_df = pd.DataFrame(ats, columns=['I', 'J', 'K', 'L'])
    final_df = pd.concat([dhd_df, at_df], axis=1, join_axes=[at_df.index])
    bndlist = list(bnd_df.UR) + (list(bnd_df.UR))
    final_df['TY'] = ['Proper' if ucomb(list([final_df.I[n], final_df.J[n], final_df.K[
        n], final_df.L[n]]), bndlist) == 3 else 'Improper' for n in range(len(final_df.I))]
#    final_df['SumV'] = np.abs(
#        final_df.V1) + np.abs(final_df.V2) + np.abs(final_df.V3) + np.abs(final_df.V4)
    #    final_df = final_df[final_df['SumV'] != 0.00]
    final_df['TI'] = [num2typ2symb[j][2] for j in final_df.I]
    final_df['TJ'] = [num2typ2symb[j][2] for j in final_df.J]
    final_df['TK'] = [num2typ2symb[j][2] for j in final_df.K]
    final_df['TL'] = [num2typ2symb[j][2] for j in final_df.L]
    final_df['SYMB'] = ['   '.join([num2typ2symb[final_df.I[i]][0], num2typ2symb[final_df.J[i]][
        0], num2typ2symb[final_df.K[i]][0], num2typ2symb[final_df.L[i]][0]]) for i in final_df.index]
    if len(final_df.index) > 0:
        final_df['NAME'] = final_df.TI + '-' + final_df.TJ + \
            '-' + final_df.TK + '-' + final_df.TL
    return final_df


def boss2CharmmBond(molecule_data, st_no):
    bdat = molecule_data.MolData['BONDS']
    bdat['cl1'] = [x - st_no if not x - st_no < 0 else 0 for x in bdat['cl1']]
    bdat['cl2'] = [x - st_no if not x - st_no < 0 else 0 for x in bdat['cl2']]
    bnd_df = pd.DataFrame(bdat)
    bnd_df['UF'] = ((bnd_df.cl1 + bnd_df.cl2) *
                    (bnd_df.cl1 + bnd_df.cl2 + 1) * 0.5) + bnd_df.cl2
    bnd_df['UR'] = ((bnd_df.cl1 + bnd_df.cl2) *
                    (bnd_df.cl1 + bnd_df.cl2 + 1) * 0.5) + bnd_df.cl1
#    bnd_df.to_csv('bos_bonds.csv', index=False)
    hb_df = bnd_df.drop(['cl1', 'cl2', 'UF', 'UR'], 1)
    hb_df = hb_df.drop_duplicates()
    return bnd_df


def boss2CharmmAngle(anglefile, num2opls, st_no,num2typ2symb):
    adat = anglefile
    adat['cl1'] = [x - st_no if not x - st_no < 0 else 0 for x in adat['cl1']]
    adat['cl2'] = [x - st_no if not x - st_no < 0 else 0 for x in adat['cl2']]
    adat['cl3'] = [x - st_no if not x - st_no < 0 else 0 for x in adat['cl3']]
    ang_df = pd.DataFrame(adat)
    ang_df = ang_df[ang_df.K > 0]
#    ang_df.to_csv('bos_angles.csv', index=False)
    ang_df['TY'] = np.array([num2opls[i] + '-' + num2opls[j] + '-' + num2opls[k]
                             for i, j, k in zip(ang_df.cl1, ang_df.cl2, ang_df.cl3)])
    ang_df['TI']=[num2typ2symb[ang_df.cl1[i]][2] for i in ang_df.index]
    ang_df['TJ']=[num2typ2symb[ang_df.cl2[i]][2] for i in ang_df.index]
    ang_df['TK']=[num2typ2symb[ang_df.cl3[i]][2] for i in ang_df.index]
    ang_df['TY'] = np.array([i + '  ' + j + '  ' + k
                             for i, j, k in zip(ang_df.TI, ang_df.TJ, ang_df.TK)])
    return ang_df


def bossData(molecule_data):
    ats_file = molecule_data.MolData['ATOMS']
    types = []
    for i in enumerate(ats_file):
        types.append([i[1].split()[1], 'opls_' + i[1].split()[2]])
    st_no = 3
    Qs = molecule_data.MolData['Q_LJ']
    assert len(Qs) == len(types), 'Please check the at_info and Q_LJ_dat files'
    num2opls = {}
    for i in range(0, len(types)):
        num2opls[i] = Qs[i][0]
    num2typ2symb = {i: types[i] for i in range(len(Qs))}
    for i in range(len(Qs)):
        num2typ2symb[i].append(bossPdbAtom2Element(
            num2typ2symb[i][0]) + num2typ2symb[i][1][-3:])
        num2typ2symb[i].append(bossPdbAtom2Element(num2typ2symb[i][0]))
        num2typ2symb[i].append(bossElement2Mass(num2typ2symb[i][3]))
        num2typ2symb[i].append(Qs[i][0])
    return (types, Qs, num2opls, st_no, num2typ2symb)


def Boss2Charmm(resid, molecule_data):
    types, Qs, num2opls, st_no, num2typ2symb = bossData(molecule_data)
    bnd_df = boss2CharmmBond(molecule_data, st_no)
    ang_df = boss2CharmmAngle(molecule_data.MolData['ANGLES'], num2opls, st_no,num2typ2symb)
    tor_df = Boss2CharmmTorsion(bnd_df, num2opls, st_no,
                                molecule_data, num2typ2symb)
    Boss2CharmmRTF(num2typ2symb, Qs, resid, bnd_df, list(ang_df['TY']), list(
        tor_df[tor_df.TY == 'Proper']['SYMB']),list(tor_df[tor_df.TY == 'Improper']['SYMB']))
    Boss2CharmmPRM(resid, num2typ2symb, Qs, bnd_df, ang_df, tor_df)
    return None


def mainBOSS2XPLOR(resid, clu=False):
    mol = pickle.load(open(resid + ".p", "rb"))
    Boss2Charmm(resid, mol)
    return None
