from .Scheduler import UniformScheduler, MemorylessScheduler
from ..base.BW import BW, NB_PROCESS
from ..base.Set import Set
from ..base.tools import resolveRandom
from multiprocessing import Pool
from random import random
from numpy import zeros, dot, array, argmin, inf
from sys import platform
from datetime import datetime

class ActiveLearningScheduler:
	"""
	Class for a scheduler used during the active learning sampling.
	"""
	def __init__(self,memoryless_scheduler: MemorylessScheduler,m) -> None:
		"""
		Create an ActiveLearningScheduler.

		Parameters
		----------
		memoryless_scheduler : MemorylessScheduler
			A memoryless scheduler generated by Active_BW_MDP._strategy
		m : MDP
			The current hypothesis.
		"""
		self.m = m
		self.nb_states = self.m.nb_states
		self.memoryless_scheduler = memoryless_scheduler
		self.reset()
		
	def reset(self) -> None:
		"""
		Resets the scheduler. Should be done before genereting a new trace.
		"""
		self.last_obs = None
		self.last_act = None
		self.alpha_matrix = array([self.m.pi(s) for s in range(self.nb_states)])

	def getAction(self) -> str:
		"""
		Returns the action chosen by the scheduler

		Returns
		-------
		str
			Action chosen by the scheduler.
		"""
		if self.last_obs != None:
			alpha_matrix = zeros(self.nb_states)
			for s in range(self.nb_states):
				p = array([self.m.tau(ss,self.last_act,s,self.last_obs) for ss in range(self.nb_states)])
				alpha_matrix[s] = dot(self.alpha_matrix,p)
			self.alpha_matrix = alpha_matrix

		tot = self.alpha_matrix.sum()
		if tot <= 0.0:
			t = [1/len(t) for i in self.alpha_matrix]
		else:
			t = [i/tot for i in self.alpha_matrix]
		s_i = resolveRandom(t)
		act = self.memoryless_scheduler.getAction(s_i)
		self.last_act = act
		return act

	def addObservation(self,obs: str) -> None:
		"""
		Updates the scheduler according to the given observation ``obs``.

		Parameters
		----------
		obs : str
			An observation.
		"""
		self.last_obs = obs

class Active_BW_MDP(BW):
	"""
	Class for the Active Baum-Welch algorithm on MDP.
	This algorithm is described here:
	https://arxiv.org/pdf/2110.03014.pdf
	"""
	def __init__(self):
		super().__init__()

	def fit(self,traces: Set, sul, lr, nb_iterations: int, nb_sequences: int,
			sequence_length: int = None, epsilon_greedy: float = 0.9,
			initial_model=None, nb_states: int=None,
			random_initial_state: bool=True, output_file: str=None,
			epsilon: float=0.01, max_it: int=inf,pp: str='',
			verbose: bool = True, return_data: bool= False, stormpy_output: bool = True):
		"""
		Fits the model according to ``traces``.

		Parameters
		----------
		traces : Set
			Training set.
		sul : MDP
			The system under learning.
			Used to generate new sequences.
		lr : float, or ``0``, or ``"dynamic"``
			Learning rate. If ``lr=0`` the current hypothesis will be updated
			on ``traces`` plus all the active samples. If `lr` is a float, the
			current hypothesis will be merged with an updated version of it.
			The updated version is a version of the current hypothesis fitted
			on the last active sample. If `lr="dynamic"` the learning rate
			decrease at each iteration.
		nb_iterations: int
			Number of active learning iterations.
		nb_sequences: int
			Size of one active sample. I.e: number of sequences generated at
			each active learning iteration.
		sequence_length: int, optional.
			Size of one sequence generated during an active learning iteration.
			If not set it will be equal to the size of the first sequence in
			``traces``.
		epsilon_greedy: float, optional.
			It will use the active learning scheduler with probability
			``epsilon_greedy`` while genereting the active learning sample, and
			with probability `1-epsilon_greedy` it will use an uniform scheduler.
			By default 0.9
		initial_model : MDP, optional.
			First hypothesis. If not set it will create a random MDP with
			``nb_states`` states. Should be set if ``nb_states`` is not set.
		nb_states: int
			If ``initial_model`` is not set it will create a random MDP with
			``nb_states`` states. Should be set if ``initial_model`` is not set.
			Default is None.
		random_initial_state: bool, optional.
			If ``initial_model`` is not set it will create a random MDP with
			random initial state according to this sequence of probabilities.
			Should be set if ``initial_model`` is not set.
			Default is False.
		output_file : str, optional
			If set path file of the output model. Otherwise the output model
			will not be saved into a text file.
		epsilon : float, optional
			The learning process stops when the difference between the
			loglikelihood of the training set under the two last hypothesis is
			lower than ``epsilon``. The lower this value the better the output,
			but the longer the running time. By default 0.01.
		max_it: int
			Maximal number of iterations for the passive learning round.
			The algorithm will stop after `max_it` iterations.
			Default is infinity.
		pp : str, optional
			Will be printed at each iteration. By default ''
		verbose: bool, optional
			Print or not a small recap at the end of the learning.
			Default is True.
		return_data: bool, optional
			If set to True, a dictionary containing following values will be
			returned alongside the hypothesis once the learning is done.
			'learning_rounds', 'learning_time', 'training_set_loglikelihood',
			Default is False.
		stormpy_output: bool, optional
			If set to True the output model will be a Stormpy sparse model.
			Default is True.

		Returns
		-------
		MDP
			fitted MDP.
		"""
		if stormpy_output:
			try:
				from ..with_stormpy import jajapyModeltoStormpy
			except ModuleNotFoundError:
				print("WARNING: stormpy not found. The output model will not be a stormpy sparse model")
				stormpy_output = False

		counter = 0
		self.sul = sul
		start_time = datetime.now()
		_, info = super().fit(traces, initial_model,nb_states,
					random_initial_state, output_file,epsilon,
					max_it,"Passive iteration:"+pp,
					return_data=True, stormpy_output=False)
		counter += info['learning_rounds']
	
		total_traces = traces
		
		if sequence_length == None:
			sequence_length = int(len(traces.sequences[0])/2)

		c = 1
		while c <= nb_iterations :
			traces = self._addTraces(sequence_length,nb_sequences,total_traces,epsilon_greedy)
			total_traces.addSet(traces)
			
			if lr == 0:
				_, info = super().fit(total_traces, initial_model=self.h,
							output_file=output_file, epsilon=epsilon,
							pp="Active iteration "+str(c)+"/"+str(nb_iterations)+": "+pp,
							return_data=True, stormpy_output=False, verbose=False)
				counter += info['learning_rounds']
			else:
				if lr == "dynamic":
					lr_it = sum(traces.times)/(sum(total_traces.times)-sum(traces.times))
				else:
					lr_it = lr
				old_h = self.h
				_, info = super().fit(traces, initial_model=self.h,
							output_file=output_file, epsilon=epsilon,
							pp="Active iteration "+str(c)+"/"+str(nb_iterations)+": "+pp,
							return_data=True, stormpy_output=False, verbose=False)
				counter += info['learning_rounds']
				self._mergeModels(old_h,lr_it)

			c += 1
		
		running_time = datetime.now()-start_time
		running_time = running_time.total_seconds()

		if verbose:
			self._endPrint(counter, running_time)

		if stormpy_output:
			self.h = jajapyModeltoStormpy(self.h)
			
		if return_data:
			info = {"learning_rounds":counter,
					"learning_time":running_time,
					"training_set_loglikelihood":self.h.logLikelihood(total_traces)}
			return self.h, info

		return self.h

	def _mergeModels(self,old_h,lr: float) -> None:
		"""
		Merges the model ``self.h`` with model ``old_h`` using the learning
		rate ``lr``.
		WARNING: we assume that self.h.alphabet == old_h.alphabet and
		self.h.actions == old_h.actions.

		Parameters
		----------
		old_h : MDP
			A MDP to merge with ``self.h``. 
		lr : float
			Learning rate. It is the weight of ``self.h`` in the output model.
		"""
		self.h.matrix = lr*self.h.matrix + (1-lr)*old_h

	def _addTraces(self,sequence_length: int,nb_sequence: int,
				  traces: Set,epsilon_greedy: float) -> list:
		"""
		Generates an active learning sample.

		Parameters
		----------
		sequence_length : int
			Length of each sequence in the sample.
		nb_sequence : int
			Number of sequences in the sample.
		traces : Set
			Training set until now. Used to compute the active learning
			sampling strategy.
		epsilon_greedy : It will use the active learning scheduler with probability
			``epsilon_greedy`` and use an uniform scheduler with probability
			`1-epsilon_greedy`.
			
		Returns
		-------
		list
			The active learning sample
		"""
		memoryless_scheduler = self._strategy(traces)
		scheduler_exploit = ActiveLearningScheduler(memoryless_scheduler,self.h)
		scheduler_explore = UniformScheduler(self.h.getActions())

		traces = Set([],t=1)
		for n in range(nb_sequence):			
			if random() > epsilon_greedy:
				seq = self.sul.run(sequence_length,scheduler_explore)
			else:
				seq = self.sul.run(sequence_length,scheduler_exploit)

			if not seq in traces.sequences:
				traces.sequences.append(seq)
				traces.times.append(0)
			traces.times[traces.sequences.index(seq)] += 1
		return traces

	def _strategy(self,traces:Set) -> MemorylessScheduler:
		if platform != "win32":
			p = Pool(processes = NB_PROCESS)
			tasks = []
			for seq,times in zip(traces.sequences,traces.times):
				tasks.append(p.apply_async(self._computeProbas, [seq, times,]))
			temp = array([res.get() for res in tasks])
		else:
			temp = array([self._computeProbas(seq, times) for seq,times in zip(traces.sequences,traces.times)])
		temp = temp.sum(axis=0)
		scheduler = [self.h.getActions()[argmin(temp[s])] for s in range(self.nb_states)]
		return MemorylessScheduler(scheduler)

	def _computeProbas(self,seq:list,time:int) -> array:
		sequence_actions = [seq[i+1] for i in range(0,len(seq)-1,2)]
		sequence_obs = [seq[i] for i in range(0,len(seq)-1,2)] + [seq[-1]]
		alpha_matrix = self._computeAlphas(sequence_obs,sequence_actions).T[:-1].T
		res = zeros(shape=(len(self.actions),self.nb_states))
		fact = alpha_matrix.sum(axis=0) # fact[t] proba to generate sequence_obs[:t]

		for ia,a in enumerate(self.actions):
			arr_dirak = array([1.0 if t == a else 0.0 for t in sequence_actions])
			arr_dirak = arr_dirak*time/fact
			res[ia] = dot(alpha_matrix,arr_dirak)
		return res.T
