import re
import os
import subprocess
import numpy as np
import pandas as pd
import pickle

class mumu(object):
	def __init__(self, phaistos_bin_dir):
		self.phaistos_dir = phaistos_bin_dir



	def observable(self):

		projectRoot = os.getcwd()
		owd = projectRoot+"/mutation_files"
		x_folders = os.listdir(owd)

		# Find folders that have the key-word '_mutated'
		x_folders = filter(lambda k: '_mutated' in k, x_folders)

		# Here we open the pickle file that is generated by sorting.py.
		proteinPickle = pickle.load( open( projectRoot+"/save.p", "rb" ) )
		proteinsPickleKeys = proteinPickle.keys()

		# new dictionary
		proteinDict = {}
		# Now let's search in these folders and get logP values of all proteins

		for folder in x_folders:
			# Get protein name
			ppKey = folder[0:4]
			#print folder

			# Here we point back to the original dataframe. The index is reset.
			proteinDict[ppKey] = proteinPickle[ppKey]
			proteinDict[ppKey]['logP'] = "Non"
			proteinDict[ppKey] = proteinDict[ppKey].reset_index(drop=True)

			os.chdir(self.phaistos_dir)

			try:
				###################
				# Get wildtype mumu
				wtPath = owd+"/"+folder+"/wildtype_mut.pdb"
				cmd = "./evaluate_observable --pdb-file "+wtPath+" --observable-mumu"
				wildtypeLogP = subprocess.check_output(cmd, shell=True)


				# 'we will clean it up a bit and put it into a list'
				wtlog_val = []

				wtf_split_sep = wildtypeLogP.split(':')

				# 'new list containing only log values'
				for i in range(len(wtf_split_sep)):
					# 'converting -logp to logp'
					wtlog_val.append(wtf_split_sep[i].rstrip('\t\n').split(',')[:1][0])

				# 'remove header'
				del wtlog_val[0]

				wtlog_array = np.asarray(wtlog_val).astype(np.float)

				###################

				for files in os.listdir(owd+"/"+folder):
					ppFile = owd+"/"+folder+"/"+files


					try:

						# Now we get the first value of their saved names to get their positions in the original dataframe.
						pointBackPosition = re.findall("(?<!\[)#[0-9]+|#[0-9]+(?!\])", files)[0].split('#')[1]

						cmd = "./evaluate_observable --pdb-file "+ppFile+" --observable-mumu"
						direct_output = subprocess.check_output(cmd, shell=True)

						# 'we will clean it up a bit and put it into a list'
						log_val = []

						f_split_sep = direct_output.split(':')

						# 'new list containing only log values'
						for i in range(len(f_split_sep)):
							# 'converting -logp to logp'
							log_val.append(f_split_sep[i].rstrip('\t\n').split(',')[:1][0])

						# 'remove header'
						del log_val[0]

						log_array = np.asarray(log_val).astype(np.float)

						# The log of the probabilities make the addable.
						# The sum gives the joint log probability of the given amino acid.
						proteinDict[ppKey]['logP'].loc[int(pointBackPosition)] = sum(wtlog_array)-sum(log_array)


					except Exception:
						print "Skipped an error"

			except Exception:
				print "wildtype error"

		print proteinDict

		pickle.dump(proteinDict, open("mumuDataframe2.p", "wb"))
phaistos_bin_dir = "/Users/thorn/phaistos/build/bin/"
mumu(phaistos_bin_dir).observable()
