import numpy as np
import math
from modules.evaluation import Evaluation, Sov

class Gor(Sov):
	"""
	GOR is a second generation secondary structure predictor. It takes into consideration the context. Therefore,
	the conformation of the residue depends on the surrounding sequence. The training part calculates the probability
	each secondary structure depending on a window of positions that contain the probability of each residue on each
	position. The probabilities come from the profiles generated by the psiblast production pssm files. There are 4
	windows, coil helix, beta sheets and general window. According to the occurance, the windows are filled with
	probabilities and then normalized. The general window calculates the probability of residue on a position
	regardless the secondary structure
	"""
	def __init__(self, window_size: int):
		super(Gor, self).__init__()
		self.residues = ["A", "R", "N", "D", "C", "Q", "E", "G", "H", "I", "L", "K", "M", "F", "P", "S", "T", "W", \
						 "Y", "V"]
		self.window_size = window_size
		self.c_window = np.zeros((window_size, 20), float)
		self.h_window = np.zeros((window_size, 20), float)
		self.e_window = np.zeros((window_size, 20), float)
		self.common_window = np.zeros((window_size, 20), float)
		self.c_count = 0
		self.h_count = 0
		self.e_count = 0
		self.c_probability = 0
		self.e_probability = 0
		self.h_probability = 0
		self.total_count = 0

	def windows_calculator(self, profile, secondary_structure_seq):
		"""it reads the secondary structure of the residues given by the dssp files and check the conformation of the
		residues by checking the signs(H, - or E).
		it increments the count of the residues' structure and the residues generally by one.
		Those will be used later to calculate the general probability of the structure and for normalization.
		It will also consider the surrounding residues as those play a role in determining the conformation of the
		central residue. The context(segment) is involved in the calculation of probabilities of residues contained in
		a window with a central residue having a specific conformation(H, - or E)
		"""
		# all the probabilities of residues in the profiles(residues) on a specific position(context) in the
		# sequence window modify the probabilities inside the general window of the structure(c_window)
		# also the general probability of a residue occurring on a specific position is calculated(common_window)
		for ss_pos in range(len(secondary_structure_seq)):
			if secondary_structure_seq[ss_pos] == "C" or secondary_structure_seq[ss_pos] == "-":
				self.c_count += 1
				self.total_count += 1
				for context in range(ss_pos - self.window_size // 2, ss_pos + self.window_size // 2 + 1):
					for residues in range(len(profile[context])):
						self.c_window[context - ss_pos + self.window_size // 2][residues] += \
							profile[context + self.window_size // 2][residues]
						self.common_window[context - ss_pos + self.window_size // 2][residues] += \
							profile[context + self.window_size // 2][residues]
			elif secondary_structure_seq[ss_pos] == "H":
				self.h_count += 1
				self.total_count += 1
				for context in range(ss_pos - self.window_size // 2, ss_pos + self.window_size // 2 + 1):
					for residues in range(len(profile[context])):
						self.h_window[context - ss_pos + self.window_size // 2][residues] += \
							profile[context + self.window_size // 2][residues]
						self.common_window[context - ss_pos + self.window_size // 2][residues] += \
							profile[context + self.window_size // 2][residues]

			elif secondary_structure_seq[ss_pos] == "E":
				self.e_count += 1
				self.total_count += 1
				for context in range(ss_pos - self.window_size // 2, ss_pos + self.window_size // 2 + 1):
					for residues in range(len(profile[context])):
						self.e_window[context - ss_pos + self.window_size // 2][residues] += \
							profile[context + self.window_size // 2][residues]
						self.common_window[context - ss_pos + self.window_size // 2][residues] += \
							profile[context + self.window_size // 2][residues]

	def normalize(self):
		"""the normalization is for the probability calculation"""
		total = self.total_count
		for positions in range(len(self.common_window)):
			for probabilities in range(len(self.common_window[positions])):
				self.c_window[positions][probabilities] = self.c_window[positions][probabilities] / total
				self.e_window[positions][probabilities] = self.e_window[positions][probabilities] / total
				self.h_window[positions][probabilities] = self.h_window[positions][probabilities] / total
				self.common_window[positions][probabilities] = self.common_window[positions][probabilities] / total
		self.e_probability = self.e_count / total
		self.c_probability = self.c_count / total
		self.h_probability = self.h_count / total

	def information(self, residue, fasta_seq, profile):
		"""Information is the key for the prediction of the conformation for a specific residue. It reads the context
		by reading the profile that provides the probability of each residue. Second, according to information
		provided by the calculated windows(P(residue_profile)*log(P(residue_window_structure)/P(structure)*P(
		residue_window_common). The 3(window_H,window_C,window_E) information are compared and the highest is chosen to
		be the conformation of the central residue inside the window"""
		info_h = 0
		info_c = 0
		info_e = 0
		info = []
		# 7 positions with zeros were added in the beginning and the end of the profiles to make the calculations
		# easier so the number in the indexing was added. to keep the index within the size of the window the index of
		# the residue substracted the indices of the surrounding residues and half of the window size was added in the
		# indexing to keep the central residue index in the center of the window and the context in the correct positions
		for context in range(residue, residue + self.window_size):
			for probabilities in range(len(profile[context])):
				info_h += profile[context - residue][probabilities] * \
						  math.log10(self.h_window[context-residue][probabilities] / \
									 (self.h_probability * self.common_window[context-residue][probabilities]))
				info_c += profile[context][probabilities] * \
						  math.log10(self.c_window[context-residue][probabilities] / \
									 (self.c_probability * self.common_window[context-residue][probabilities]))
				info_e += profile[context][probabilities] * \
						  math.log10(self.e_window[context-residue][probabilities] / \
									 (self.e_probability * self.common_window[context-residue][probabilities]))
		info.append(info_c), info.append(info_h), info.append(info_e)
		return info

	def predict(self, fasta_seq, profile):
		"""simply take the information calculated in the previous method and provide the prediction accordingly"""
		secondary_structure_seq = ""
		for residue in range(len(fasta_seq)):
			info = self.information(residue, fasta_seq, profile)
			if sum(info) == 0 or max(info) == info[0]:
				secondary_structure_seq += "-"
			elif max(info) == info[1]:
				secondary_structure_seq += "H"
			elif max(info) == info[2]:
				secondary_structure_seq += "E"
		return secondary_structure_seq
