#!/usr/bin/env python

"""
    Requirements: Python 3 or higher
                  Antechamber and related AmberTools
                  OpenBabel (strongly recommended for use with acpype)
				  acpype (latest version recommended with all its requirements)
				  Gromacs (Compulsory)
                  flask (Compulsory)
                  flaskwebgui (recommended)
                  pyfladesk (recommended)

    This code is released under GNU General Public License V3.

          <<<  NO WARRANTY AT ALL!!!  >>>

    It was inspired by:

    - CS50 online training for which this code serves as part of the final project

	- PLEASE Read the README.md file and also follow instructions on the GUI and/or Terminal

	<NAME>, B.Pharm. M.Sc. (Pharmacology) and Ph.D. (Science) Biochemistry
    Department of Pharmacology, Faculty of Pharmacy
    Obafemi Awolowo Universiy, Ile-Ife, Nigeria.
    >>http://www.oauife.edu.ng<<

    <EMAIL>; <EMAIL>
"""

import sys

if sys.version_info[0] < 3:
	raise Exception("Python 3 or a more recent version is required.")

import os
import subprocess
from pathlib import Path
import time
import shutil
import random
import string
import math
import glob
from colored import fore, back, style

from gmodsScripts.gmodsHelpers import indexoflines, printWarning, printNote

def OPLStop(LFtop, ff, name):
	topcwdir = Path.cwd()
	mergename = "opls" + name + ".top"
	oplstopopen = open(mergename, "+a")

	# Insert the header information into the opls[name].top file
	topheaders = ["; This file was generated by modifying tleap/acpype generated ligand topology", "; by replacing generated atom types with opls compatible atom types", "; as contained in the opls ffnonbonded.itp forcefield in gromacs.", "; Where opls compatible atom type(s) is/are not found, the tleap/acpype generated atom types will be retained"]

	for header in topheaders:
		oplstopopen.write(header)
		oplstopopen.write('\n')
	oplstopopen.write('\n')

	time.sleep(5)

	# We shall check for and remove duplicates in atomtypes between gmx standard and amber/tleap generated 
	if not ff[0:4] == 'opls':
		print(ff, "is not an opls forcefield. Generation can not continue")
		return LFtop

	else:
		print("Compatible forcefiled detected:", ff)

	# Get the absolute path to the forcefield directory and copy ffnonbonded.itp file
	gmxtopdir = " "
	try:
		gmxtopdir = os.path.join(os.environ.get('GMXDATA'), 'top', ff)
	except:
		print("We are unable to autodetect your forcefield directory")
		print("Please supply absolute path to your forcefield directory")
		print("Usually /path-to-gromacs/share/gromacs/top/your-selected-forcefield.ff")
		while True:
			gmxtopdir = input("Path to Forcefield: ")
			if not os.path.isdir(gmxtopdir):
				print("Directory you supplied does not exist. Please check and try again")
			else:
				break

	if not os.path.isdir(gmxtopdir):
		print(gmxtopdir, "that was autodetected, is not a valid forcefield directory")
		print("Please supply the correct absolute path to your forcefield directory")
		print("Usually /path-to-gromacs/share/gromacs/top/your-selected-forcefield.ff")
		while True:
			gmxtopdir = input("Path to Forcefield: ")
			if not os.path.isdir(gmxtopdir):
				print("Directory you supplied does not exist. Please check and try again")
			else:
				break

	print("Your topology directory is", gmxtopdir)
	time.sleep(5)

	lsgmxtopdir = os.listdir(gmxtopdir)
	for tp in lsgmxtopdir:
		if tp == "ffnonbonded.itp":
			shutil.copy(os.path.join(gmxtopdir, tp), './')

	# Get the ffnonbonded.itp opls atomtypes range to consider for use inside new file
	ffb = open("ffnonbonded.itp", "r")
	ffbreadlines = ffb.readlines()
	fn = 0
	for fb in ffbreadlines:
		if 'opls_128' in fb.split():
			break
		else:
			fn += 1
	ffb.seek(0)

	# Check the LFtop file to replace atomtypes and moleculetype atoms with opls types where possible
	lftindex = indexoflines(LFtop)
	atlp = int(lftindex['atomtypes'])
	mtlp = int(lftindex['moleculetype'])

	lftopopen = open(LFtop, "r")
	lftopreadlines = lftopopen.readlines()

	btlp = 0
	for bd in lftopreadlines:
		if 'bonds' in bd.split() and bd.split()[0] == '[' and bd.split()[1] == 'bonds':
			break
		else:
			btlp += 1

	lftopopen.seek(0)
	for aline in lftopreadlines[atlp:mtlp]:
		if not (aline.split() == [] or aline.split()[0][0] == ';' or aline.split()[0][0] == '['):
			atom = aline.split()[0]
			ffb.seek(0)
			afound = 0
			for fline in ffbreadlines[fn:]:
				if fline.split()[1] == atom or fline.split()[1].lower() == atom or fline.split()[1].upper() == atom.upper():
					afound += 1
			if not afound > 0:
				oplstopopen.write(aline)
		else:
			oplstopopen.write(aline)
	oplstopopen.write('\n')

	lftopopen.seek(0)
	for bline in lftopreadlines[mtlp:btlp]:
		if not (bline.split() == [] or bline.split()[0][0] == ';' or bline.split()[0][0] == '['):
			matom = bline.split()[1]
			ffb.seek(0)
			replace = " "
			bfound = 0
			for fline in ffbreadlines[fn:]:
				if fline.split()[1] == matom or fline.split()[1].lower() == matom or fline.split()[1] == matom.upper() and 'opls' in fline.split()[0].split('_'):
					replace = bline.replace(matom, fline.split()[0], 1)
					bfound += 1
					break
			if not bfound > 0:
				oplstopopen.write(bline)
			else:
				oplstopopen.write(replace)
		else:
			oplstopopen.write(bline)
	oplstopopen.write('\n')

	lftopopen.seek(0)
	for cline in lftopreadlines[btlp:]:
		oplstopopen.write(cline)

	ffb.close()
	lftopopen.close()
	oplstopopen.close()

	# Check to be sure oplstopol.top has been successfully generated
	checktlp = os.listdir()
	if not mergename in checktlp:
		printWarning("Something went wrong. Generating oplstopol.top file was not successful")
		return LFtop
	else:
		checkindex = indexoflines(mergename)
		a = int(checkindex['atomtypes'])
		b = int(checkindex['moleculetype'])
		c = int(checkindex['system'])
		checkopen = open(mergename, "r")
		checkreadlines = checkopen.readlines()
		if not ('atomtypes' in checkreadlines[a].split() and 'moleculetype' in checkreadlines[b].split() and 'system' in checkreadlines[c].split()):
			return LFtop
		else:
			printNote("opls compatible ligand topology has been generated successfully")
			return mergename


def OPLSacpype(LFtop, lig, name):
	tlpcwdir = Path.cwd()

	lftopindex = indexoflines(LFtop)
	mtt = lftopindex['moleculetype'] + 2
	mtopen = open(LFtop, "r")
	mtread = mtopen.readlines()
	mtname = mtread[mtt].split()[0]
	mtopen.close()

	try:
		subprocess.run('acpype -i ' + lig + ' -a amber2 -b ' + mtname, shell=True)
	except:
		try:
			os.system('acpype.py -i ' + lig + ' -a amber2 -b ' + mtname)
		except:
			return LFtop

	acpypefolder = mtname + '.acpype'
	os.chdir(acpypefolder)
	for acpfile in os.listdir():
		if not (Path(acpfile).suffix == ".itp" or Path(acpfile).suffix == ".top"):
			os.remove(acpfile)
	
	try:
		subprocess.run('mv *_GMX_OPLS.top oplstop.top', shell=True)
		subprocess.run('mv *_GMX_OPLS.itp oplstop.itp', shell=True)
	except:
		try:
			os.system('mv *_GMX.top oplstop.top')
			os.system('mv *_GMX.top oplstop.itp')
		except:
			return LFtop

	mergename = "opls" + name + ".top"
	mergeopls = open(mergename, "+a")
	oplstopopen = open("oplstop.top", "r")
	oplstopread = oplstopopen.readlines()
	oplsitpopen = open("oplstop.itp", "r")
	oplsitpread = oplsitpopen.readlines()
	
	q = 0
	for oline in oplstopread:
		if not oline == [] and "system" in oline.split():
			break
		else:
			q += 1

	for itpline in oplsitpread:
		mergeopls.write(itpline)
	mergeopls.write('\n')

	oplstopopen.seek(0)
	while q < len(oplstopread):
		mergeopls.write(oplstopread[q])
		q += 1

	oplstopopen.close()
	oplsitpopen.close()
	mergeopls.close()

	shutil.copy(mergename, '../')
	os.chdir('../')

	# Check to be sure opls compatible ligand topology has been successfully generated
	checktlp = os.listdir()
	mergename = "opls" + name + ".top"
	if not mergename in checktlp:
		printWarning("Something went wrong. Generating oplstopol.top file was not successful")
		return LFtop
	else:
		checkindex = indexoflines(mergename)
		a = int(checkindex['atomtypes'])
		b = int(checkindex['moleculetype'])
		c = int(checkindex['system'])
		checkopen = open(mergename, "r")
		checkreadlines = checkopen.readlines()
		if not ('atomtypes' in checkreadlines[a].split() and 'moleculetype' in checkreadlines[b].split() and 'system' in checkreadlines[c].split()):
			return LFtop
		else:
			printNote("opls compatible ligand topology has been generated successfully")
			return mergename
