# -*- coding: utf-8 -*-

# Copyright (c) 2018: Elodie Laine
# This code is part of the gemme package and governed by its license.
# Please see the LICENSE.txt file included as part of this package.

import sys
import os
import argparse
import re
import math
import subprocess
import shutil


# Extract the query sequence from the input alignment
def extractQuerySeq(filename):

	fIN = open(filename,"r")
	lines = fIN.readlines()
	fIN.close()

	if lines[0][0]!=">":
		raise Exception('bad FASTA format')
	else:
		prot = re.compile("[^A-Z0-9a-z]").split(lines[0][1:])[0]
		fOUT = open(prot+".fasta","w")
		fOUT.write(">"+prot+"\n")
		seq=""
		i = 1
		while lines[i][0]!=">":
			seq = seq + lines[i].strip().strip(".").strip("-")
			fOUT.write(lines[i])
			i = i + 1
	fOUT.close()
	return prot,seq,i-1

# Get the number of sequences in a multi-fasta file
def getNbSeq(filename):

	if filename!='':
		proc=subprocess.Popen("grep -c '^>' "+filename,stdout=subprocess.PIPE,shell=True)
		return int(proc.stdout.read())
	else:
		return 0


# Create a PDB file with dummy CA atoms based on the query sequence
def createPDB(prot,seq):
	"""
		If there is not a real PDB file for a given sequence,
		create a fake PDB containing only CA atoms.  
	"""
	d = {'C': 'CYS', 'D': 'ASP', 'S': 'SER',  'Q': 'GLN', 'K': 'LYS',
	'I': 'ILE', 'P': 'PRO', 'T': 'THR', 'F': 'PHE', 'N': 'ASN', 
	'G': 'GLY',  'H': 'HIS', 'L': 'LEU', 'R': 'ARG', 'W': 'TRP', 
	'A': 'ALA', 'V': 'VAL', 'E': 'GLU', 'Y': 'TYR', 'M': 'MET'}

	fOUT = open(prot+'.pdb','w')
	i = 1
	for let in seq:
		fOUT.write('ATOM%7d  CA  %s A%4d      43.524  70.381  46.465  1.00   0.0\n'%(i,d[let],i))
		i += 1
	fOUT.close()

# Edit JET configuration file with correct number of Seqs & MSA
def editConfJET(N):
	reCode=subprocess.call("sed -i 's/results\t\t5000/results\t\t"+str(N)+"/' default.conf",shell=True)
	return(reCode)

# Run JET to compute TJET values
def launchJET(prot, retMet, bFile, fFile, pdbfile, n, N, nl):
	"""
		Call JET2 and produce prot+"_jet.res" file.

		prot+"_jet.res" will be used in the following steps (in launchPred)
		to calculate independent and epistatic models.
		Ideally, this call to JET2 should be from Dockers or Singularity 
		because installing all requirements of JET2 is a pain in the ass!

	Parameters
    ----------
    prot: string ???
        Name of the protein ???
    retMet: string
		Retreival method of multiple sequence alignments file
		It can be 'input', 'local' or 'server'. Default is local.
	bFile: string 
        A multiple sequence alignment file obtained with psiblast.
		It is used only if the retMet (explained above) is input.
	fFile: string 
        A multiple sequence alignment file obtained with psiblast.
		It is used only if the retMet (explained above) is input.
	pdbfile: string 
        a Protein Data Bank file obtained with rcsb.org or any 
		computational method like alphafold.
		If it is None, only JET and PC scores are calculated.
		Otherwise, CV and other structural-dynamical features also
		can be calculated.
	n: int
		Number of JET2 iterations.
	N: int
		Default 20000
	nl: int
		Number of lines after > character in the query sequences file.
		It is obtained in extractQuerySeq() function.

    Returns
    -------
	Nothing 
	"""

	subprocess.call("cp $GEMME_PATH/default.conf .",shell=True)
	if retMet=="input":
		if bFile!='':
			if(bFile == prot+"_A.psiblast"):
				shutil.copy2(bFile, bFile+".orig")
				shutil.copy2(bFile+".orig ", prot+"_A.psiblast")
			else:
				shutil.copy2(bFile+" ", prot+"_A.psiblast")
			if(pdbfile == None):
				jetcmd = "java -Xmx1000m -cp $JET2_PATH:$JET2_PATH/jet/extLibs/vecmath.jar jet.JET -c default.conf -i "+\
						prot+".pdb -o `pwd` -p J -r input -b "+prot+"_A.psiblast -d chain -n "+n+" > "+prot+".out"
			else:
				jetcmd = "java -Xmx1000m -cp $JET2_PATH:$JET2_PATH/jet/extLibs/vecmath.jar jet.JET -c default.conf -i "+\
						prot+".pdb -o `pwd` -p AVJ -r input -b "+prot+"_A.psiblast -d chain -n "+n+" > "+prot+".out"

		else:
			print(N)
			editConfJET(N)

			if(fFile == prot+"_A.fasta"):
				shutil.copy2(fFile, fFile+".orig")
				#I think this subprocess call causes overwriting of the fasta file. 
				#subprocess.call("cp "+fFile+" "+prot+"_A.fasta",shell=True)
				grpcmd="grep -m "+str(int(N)+1)+" -A "+str(nl)+" '^>' "+fFile+".orig > "+prot+"_A.fasta"
			else:
				#subprocess.call("cp "+fFile+" "+prot+"_A.fasta",shell=True)
				grpcmd="grep -m "+str(int(N)+1)+" -A "+str(nl)+" '^>' "+fFile+" > "+prot+"_A.fasta"
			
			print("\nRunning:\n"+grpcmd)
			subprocess.call(grpcmd,shell=True)
			if(pdbfile == None):
				jetcmd = "java -Xmx1000m -cp $JET2_PATH:$JET2_PATH/jet/extLibs/vecmath.jar jet.JET -c default.conf -i "+\
				     	prot+".pdb -o `pwd` -p J -r input -f "+prot+"_A.fasta -d chain -n "+n+" > "+prot+".out"
			else:
				jetcmd = "java -Xmx1000m -cp $JET2_PATH:$JET2_PATH/jet/extLibs/vecmath.jar jet.JET -c default.conf -i "+\
				     	prot+".pdb -o `pwd` -p AVJ -r input -f "+prot+"_A.fasta -d chain -n "+n+" > "+prot+".out"
	else:
		if(pdbfile == None):
			jetcmd = "java -Xmx1000m -cp $JET2_PATH:$JET2_PATH/jet/extLibs/vecmath.jar jet.JET -c default.conf -i "+\
			    	prot+".pdb -o `pwd` -p J -r "+retMet+" -d chain -n "+n+" > "+prot+".out"
		else:
			jetcmd = "java -Xmx1000m -cp $JET2_PATH:$JET2_PATH/jet/extLibs/vecmath.jar jet.JET -c default.conf -i "+\
			    	prot+".pdb -o `pwd` -p AVJ -r "+retMet+" -d chain -n "+n+" > "+prot+".out"
	print("\nRunning:\n"+jetcmd)
	reCode=subprocess.call(jetcmd,shell=True)
	if os.path.isfile(prot+"/"+prot+"_jet.res"):
		os.rename(prot+"/"+prot+"_jet.res",prot+"_jet.res")
	return(reCode)

# Run Rscript to compute predictions
def launchPred(prot,inAli,mutFile, normWeightMode):

	if mutFile!='':
		rcmd="Rscript --save $GEMME_PATH/computePred.R "+prot+" "+inAli+" FALSE "+mutFile+" "+normWeightMode
	else:
		rcmd="Rscript --save $GEMME_PATH/computePred.R "+prot+" "+inAli+" TRUE none "+normWeightMode

	print("\nRunning: \n"+rcmd)
	reCode=subprocess.call(rcmd,shell=True)

	#Add plots here with gemmemore
	return(reCode)

# Remove temporary files
def cleanTheMess(prot,bFile,fFile):

	if bFile!='':
		if bFile!=prot+"_A.psiblast":
			os.remove(prot+"_A.psiblast")
	else:
		if os.path.isfile(prot+"/"+prot+"_A.psiblast"):
			os.rename(prot+"/"+prot+"_A.psiblast",prot+"_A.psiblast")
	if fFile!='':
		if fFile!=prot+"_A.fasta":
			if os.path.isfile(prot+"_A.fasta"):
				os.remove(prot+"_A.fasta")		
#	if os.path.isfile(prot+"/"+prot+"_jet.res"):
#		os.rename(prot+"/"+prot+"_jet.res",prot+"_jet.res")
#	os.remove(prot+".pdb")
	dir_name = prot+"/"
	if os.path.isdir(dir_name):
		for f in os.listdir(dir_name): 
			f_path = os.path.join(dir_name, f)
			if os.path.isfile(f_path):
				os.remove(f_path)
		os.rmdir(dir_name)
