import numpy as np
from multiprocessing import Pool
from typing import List, Callable

from RavenRL.Func import *
from RavenRL.Agent import CEM
from RavenRL.Env import setup_env
from RavenRL.Sampling import SamplingMethod
from RavenRL.Policy import setup_policy
from RavenRL.Utils import safety_test


class CEMSeldonian(CEM):
    """
    The Seldonian version of the Cross Entropy Method learning algorithm.
    """
    def __init__(self, epochs: int, pop_size: int, elite_ratio: float, n_sample: int,
                 ci_ub: Callable, ref_size: int, g_funcs: List[Callable], correction=1, delta=0.05,
                 gamma=0.9, extra_std=2.0, extra_decay_time=10, n_proc=8):
        """
        Parameters
        ----------
        epochs : int
            The total number of iterations that the algorithm will try to train and improve on its current policies.
        pop_size : int
            The total number of candidate policies (population) that will be generated at each epoch.
        elite_ratio : float
            The percentage of candidates that we will keep (elites) to improve on the next policy generation.
            The algorithm will keep the top `elite_ratio` percent of the population to update its mean and variance
                parameter.
        n_sample : int
            The total number of trials a candidate will be tested over the environment at each epoch.
            A big `n_sample` tends to estimate the policy performance more accurately, but also has a longer runtime
        ci_ub : Callable
            A function calculates the concentration upper bounds for safety test.
        ref_size : int
            The size of the dataset in the safety test.
            Used in the concentration bound calculation to prevent producing a bound that is overly conservative
        g_funcs : List[Callable]
            A :obj:`list` of user-defined constraint functions for safety test.
        correction : int
            A hyperparameter used to scale the concentration interval generated by `ci_ub`.
            For student t, use `2` in candidate selection; Otherwise use `1`.
        delta : float
            The significance level for the safety test in order to get a high confidence performance lower bound
            of a candidate policy
        gamma : float
            The discount factor. Mostly predefined by the MDP.
        extra_std : float
            A hyperparameter which adds extra variance to the Covariance matrix
        extra_decay_time : float
            A hyperparameter which controls the scale of the extra variance adds to the Covariance matrix
        n_proc : int
            The total number of processes that can be spawned for parallellization
       """
        super(CEMSeldonian, self).__init__(
            epochs, pop_size, elite_ratio, n_sample,
            gamma, extra_std, extra_decay_time, n_proc
        )
        self.sampler = None
        self.best_thetas = np.array([])
        self.ci_ub = ci_ub
        self.delta = delta
        self.ref_size = ref_size
        self.correction = correction
        self.g_funcs = g_funcs
        self.penalty = 1e5

    def load_sampler(self, sampler: SamplingMethod) -> None:
        """ Load the performance estimator for candidate evaluation

        Parameters
        ----------
        sampler : SamplingMethod
            An estimator used for candidate evaluation
        """
        self.sampler = sampler

    @staticmethod
    def get_best_idx(rewards: np.ndarray) -> np.ndarray:
        """ Compute any indices of the candidate policies with non-zero performance after safety test

        Parameters
        ----------
        rewards : np.ndarray
            A vector of estimated performance for each candidate policy
        """
        return np.where(rewards > 0)

    def candidate_eval(self, theta: np.ndarray, monitor=False, g_funcs=None) -> float:
        """ Run the given candidate policy over the environment and returns its average performance.

        Parameters
        ----------
        theta : np.ndarray
            A candidate policy generated by CEM
        monitor : bool
            A flag which enables the build-in animator of the environment
        g_funcs : List[Callable]
            A :obj:`list` of user-defined constraint functions for safety test.
        """
        assert isinstance(self.env_id, str)
        assert isinstance(self.sampler, SamplingMethod)
        env, _, _ = setup_env(self.env_id)
        policy = setup_policy(env, theta, env_id=self.env_id)
        self.sampler.load_eval_policy(policy)

        rewards = np.array([self.sampler.get_episodic_est() for _ in range(self.n_sample)])
        perf_est = np.average(rewards)

        for g in self.g_funcs:
            if self.ci_ub(g(rewards), ref_size=self.ref_size, correction=self.correction, delta=self.delta) > 0:
                perf_est -= self.penalty
                break
        return perf_est

    def train(self):
        """ Iterate over all candidate policies and update parameters using elite policies.
        """
        # Check if parameters are set
        assert all(isinstance(val, np.ndarray) for val in [self.means, self.stds])
        for epoch in range(self.epochs):

            extra_cov = max(1.0 - epoch / self.extra_decay_time, 0) * self.extra_std**2

            candidates = np.random.multivariate_normal(
                mean=self.means,
                cov=np.diag(np.array(self.stds**2) + extra_cov),
                size=self.pop_size
            )

            with Pool(self.n_proc) as p:
                g_candidates = p.map(self.candidate_eval, candidates)
            g_candidates = np.array(g_candidates).reshape(-1)

            elite_mask = self.get_elite_idx(g_candidates)
            best_mask = self.get_best_idx(g_candidates)
            self.elites = candidates[elite_mask]
            self.best_thetas = candidates[best_mask]

            self.update_params(self.elites)

    def get_best_candidates(self):
        """ Return the last generated elite policies by Seldonian CEM
        """
        if self.best_thetas.shape[0] == 0:
            return self.means
        else:
            return self.best_thetas

    def get_best_rewards(self, n: int):
        """ Return the performance estimation of top-n elite policies.

        Parameters
        ----------
        n : int
            The top-n policies that will be evaluated and monitored.
        """
        if self.best_thetas.shape[0] == 0:
            return self.means.reshape(1, -1)
        if len(self.best_thetas.shape) == 1:
            return self.best_thetas.reshape(1, -1)
        else:
            return [self.candidate_eval(theta, monitor=True) for theta in self.best_thetas[:n]]


# Add any constraint functions here
def g0(d: np.ndarray) -> np.ndarray:
    return 20 - d


if __name__ == "__main__":
    from RavenRL.Sampling import PDIS
    from RavenRL.Utils import generate_dataset
    from RavenRL.Policy import RandomPolicy

    env_id = "cartpole"
    n_proc = 8
    train_size = 5000
    test_size = 5000

    env, _, _ = setup_env(env_id)
    dataset_train = generate_dataset(env_id, train_size, n_proc=n_proc)
    dataset_test = generate_dataset(env_id, test_size, n_proc=n_proc)

    behv_policy = RandomPolicy(env)
    eval_policy = RandomPolicy(env)

    sampler_train = PDIS(dataset_train, behv_policy, gamma=1, n_proc=n_proc)
    sampler_test = PDIS(dataset_test, behv_policy, gamma=1, n_proc=n_proc)

    g_funcs = [g0]

    agent = CEMSeldonian(
        epochs=10, pop_size=100, elite_ratio=0.15, n_sample=50,
        ci_ub=studentT_ub, ref_size=5000, g_funcs=g_funcs, correction=2, delta=0.05,
        gamma=1, extra_std=2.0, extra_decay_time=10, n_proc=n_proc
    )

    agent.load_env(env_id)
    agent.load_sampler(sampler_train)
    agent.init_params()
    agent.train()
    thetas = agent.get_best_candidates()

    solutions = []
    for theta in thetas:
        sol = safety_test(
            env_id=env_id, theta=theta, sampler=sampler_test, ref_size=test_size,
            ci_ub=studentT_ub, g_funcs=g_funcs
        )
        solutions.append(sol)

    print(solutions)
