"""
File containing functions for encoding a string of amino acids into a numeric vector.

.............................................................................
parrot was developed by the Holehouse lab
     Original release ---- 2020

Question/comments/concerns? Raise an issue on github:
https://github.com/idptools/parrot

Licensed under the MIT license. 
"""

import sys
import os

import numpy as np
import torch

ONE_HOT = {'A':0, 'C':1, 'D':2, 'E':3, 'F':4, 'G':5, 'H':6, 'I':7, 'K':8, 'L':9,
		   'M':10,'N':11,'P':12,'Q':13,'R':14,'S':15,'T':16,'V':17,'W':18,'Y':19}

def one_hot(seq):
	"""Convert an amino acid sequence to a PyTorch tensor of one-hot vectors

	Each amino acid is represented by a length 20 vector with a single 1 and
	19 0's Inputing a sequence with a nono-canonical amino acid letter will
	cause the program to exit.

	E.g. Glutamic acid (E) is encoded: [0 0 0 1 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0 0]

	Parameters
	----------
	seq : str
		An uppercase sequence of amino acids (single letter code)

	Returns
	-------
	torch.IntTensor
		a PyTorch tensor representing the encoded sequence
	"""

	l = len(seq)
	m = np.zeros((l, 20))
	try:
		for i in range(l):
			m[i, ONE_HOT[seq[i]]] = 1
	except:
		error_str = 'Invalid amino acid detected: ' + seq[i]
		raise ValueError(error_str)
	return torch.from_numpy(m)

def rev_one_hot(seq_vectors):
	"""Decode a list of one-hot sequence vectors into amino acid sequences

	Parameters
	----------
	seq_vectors : list of numpy arrays
		A list containing sequence vectors

	Returns
	-------
	list
		Strings of amino acid sequences
	"""

	REV_ONE_HOT = 'ACDEFGHIKLMNPQRSTVWY'
	sequences = []

	for seq_vector in seq_vectors:
		seq = []
		for residue in seq_vector:
			seq.append(REV_ONE_HOT[np.argmax(residue)])
		sequences.append("".join(seq))

	return sequences


############  Biophysical scale encoding  #############

# Map each amino acid to a vector of biophysical properties
# 0: Hydrophobicity
# 1: Charge
# 2: pI
# 3: Molecular weight (g/mol)
# 4: Aromatic amino acid
# 5: Capable of hydrogen bonding
# 6: Side chain SASA (measured from ACE-XX-NME dipeptide)
# 7: Backbone SASA (measured from ACE-XX-NME dipeptide)
# 8: Free energy of solvation
BIOPHYSICS = {	'A':[ 1.8,  0,  6.0,  89.1, 0, 0,  75.8,  76.1,    1.9], 
				'C':[ 2.5,  0,  5.1, 121.2, 0, 0, 115.4,  67.9,   -1.2], 
				'D':[-3.5, -1,  2.8, 133.1, 0, 1, 130.3,  71.8, -107.3], 
				'E':[-3.5, -1,  3.2, 147.1, 0, 1, 161.8,  68.1, -107.3], 
				'F':[ 2.8,  0,  5.5, 165.2, 1, 0, 209.4,  66.0,   -0.8], 
				'G':[-0.4,  0,  6.0,  75.1, 0, 0,   0.0, 115.0,    0.0], 
				'H':[-3.2,  1,  7.6, 155.2, 0, 1, 180.8,  67.5,  -52.7], # Avg of HIP and HIE 
				'I':[ 4.5,  0,  6.0, 131.2, 0, 0, 172.7,  60.3,    2.2], 
				'K':[-3.9,  1,  9.7, 146.2, 0, 1, 205.9,  68.7, -100.9], 
				'L':[ 3.8,  0,  6.0, 131.2, 0, 0, 172.0,  64.5,    2.3],
		   		'M':[ 1.9,  0,  5.7, 149.2, 0, 0, 184.8,  67.8,   -1.4],
		   		'N':[-3.5,  0,  5.4, 132.1, 0, 1, 142.7,  66.8,   -9.7],
		   		'P':[-1.6,  0,  6.3, 115.1, 0, 0, 134.3,  55.8,    2.0],
		   		'Q':[-3.5,  0,  5.7, 146.2, 0, 1, 173.3,  66.6,   -9.4],
		   		'R':[-4.5,  1, 10.8, 174.2, 0, 1, 236.5,  66.7, -100.9],
		   		'S':[-0.8,  0,  5.7, 105.1, 0, 1,  95.9,  72.9,   -5.1],
		   		'T':[-0.7,  0,  5.6, 119.1, 0, 1, 130.9,  64.1,   -5.0],
		   		'V':[ 4.2,  0,  6.0, 117.1, 0, 0, 143.1,  61.7,    2.0],
		   		'W':[-0.9,  0,  5.9, 204.2, 1, 1, 254.6,  64.3,   -5.9],
		   		'Y':[-1.3,  0,  5.7, 181.2, 1, 1, 222.5,  71.9,   -6.1]
		   		}

def biophysics(seq):
	"""Convert an amino acid sequence to a PyTorch tensor with biophysical encoding

	Each amino acid is represented by a length 9 vector with each value representing
	a biophysical property. The nine encoded biophysical scales are Kyte-Doolittle
	hydrophobicity, charge, isoelectric point, molecular weight, aromaticity, 
	h-bonding ability, side chain solvent accessible surface area, backbone SASA, and 
	free energy of solvation. Inputing a sequence with a nono-canonical amino acid 
	letter will cause the program to exit.

	E.g. Glutamic acid (E) is: [-3.5, -1,  3.2, 147.1, 0, 1, 161.8,  68.1, -107.3]

	Parameters
	----------
	seq : str
		An uppercase sequence of amino acids (single letter code)

	Returns
	-------
	torch.FloatTensor
		a PyTorch tensor representing the encoded sequence
	"""
	l = len(seq)
	m = np.zeros((l, len(BIOPHYSICS['A'])))
	try:
		for i in range(l):
			m[i] = BIOPHYSICS[seq[i]]
	except:
		error_str = 'Invalid amino acid detected: ' + seq[i]
		raise ValueError(error_str)
	return torch.from_numpy(m)

def rev_biophysics(seq_vectors):
	"""Decode a list of biophysically-encoded sequence vectors into amino acid sequences

	Parameters
	----------
	seq_vectors : list of numpy arrays
		A list containing sequence vectors

	Returns
	-------
	list
		Strings of amino acid sequences
	"""

	REV_BIOPHYSICS = {}
	for key, value in BIOPHYSICS.items():
		REV_BIOPHYSICS[str(value[6])] = key

	sequences = []
	for seq_vector in seq_vectors:
		seq = []
		for residue in seq_vector:
			seq.append(REV_BIOPHYSICS[str(residue[6])])
		sequences.append("".join(seq))

	return sequences


################## User-specified encoding ####################

def parse_encode_file(file):
	"""Helper function to convert an encoding file into key:value dictionary"""

	with open(file) as f:
		lines = [x.strip().split() for x in f]

	l = len(lines[0]) - 1
	d = {}
	for line in lines:
		d[line[0]] = line[1:]

		if len(line) - 1 != l:
			raise ValueError('Vectors in encoding file do not have same length.')

	return d, l

class UserEncoder():
	"""User-specified amino acid-to-vector encoding scheme object

	Attributes
	----------
	encode_file : str
		A path to a file that describes the encoding scheme
	encode_dict : dict
		A dictionary that maps each amino acid to a numeric vector
	input_size : int
		The length of the encoding vector used for each amino acid
	"""

	def __init__(self, encode_file):
		"""
		Parameters
		----------
		encode_file : str
			A path to a file that describes the encoding scheme
		"""

		self.encode_file = os.path.abspath(encode_file)
		if not os.path.isfile(self.encode_file):
			raise FileNotFoundError('Encoding file does not exist.')

		self.encode_dict, self.input_size = parse_encode_file(self.encode_file)

	def __len__(self):
		"""Get length of encoding scheme"""

		return self.input_size

	def encode(self, seq):
		"""Convert an amino acid sequence into this encoding scheme

		Parameters
		----------
		seq : str
			An uppercase sequence of amino acids (single letter code)

		Returns
		-------
		torch.FloatTensor
			a PyTorch tensor representing the encoded sequence
		"""

		l = len(seq)
		m = np.zeros((l, self.input_size))

		try:
			for i in range(l):
				m[i] = self.encode_dict[seq[i]]
		except:
			error_str = 'Invalid amino acid detected: ' + seq[i]
			raise ValueError(error_str)
		return torch.from_numpy(m)

	def decode(self, seq_vectors):
		"""Converts a list of sequence vectors back to a list of protein sequences

		Parameters
		----------
		seq_vectors : list of numpy arrays
			A list containing sequence vectors

		Returns
		-------
		list
			Strings of amino acid sequences
		"""

		# Create a reverse of the encode_dict using hashing
		rev_dict = {}
		for key, value in self.encode_dict.items():
			value = np.array(value, dtype=np.float32)
			rev_dict[hash(tuple(value))] = key

		sequences = []
		for seq_vector in seq_vectors:
			seq = []
			for residue in seq_vector:
				seq.append(rev_dict[hash(tuple(residue))])
			sequences.append("".join(seq))

		return sequences