import numpy as np
from typing import List, Callable

from lightRaven.sampling import SamplingBase
from lightRaven.policy import FAPolicy, TabularPolicy, PolicyBase


class CEMSeldonian:
    """
    The Seldonian version of the Cross Entropy Method learning algorithm.
    """
    def __init__(self, epochs: int, pop_size: int, elite_ratio: float,
                 ci_ub: Callable, ref_size: int, g_funcs: List[Callable], c=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.
        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.
        c : 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
       """
        self.epochs = epochs
        self.pop_size = pop_size
        self.elite_ratio = elite_ratio
        assert 0 < elite_ratio < 1
        self.elite_size = int(self.pop_size * self.elite_ratio)
        self.extra_std = extra_std
        self.extra_decay_time = extra_decay_time
        self.n_proc = n_proc
        self.gamma = gamma

        self.p_type = None
        self.sampler = None
        self.obs_shape = None
        self.act_shape = None

        self.theta_dim = None
        self.means = None
        self.stds = None
        self.elites = None

        self.sampler = None
        self.best_thetas = np.array([])
        self.ci_ub = ci_ub
        self.delta = delta
        self.ref_size = ref_size
        self.c = c
        self.g_funcs = g_funcs
        self.penalty = 1e5

    @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 load_setup(self, obs_shape: int, act_shape: int, p_type: str) -> None:
        """ Load a supported Gym environment into the algorithm

        Parameters
        ----------
        obs_shape : int
            Observation space shape
        act_shape : int
            Action space shape
        p_type : str
            Indicating policy type from 'fa' (functional approximation), 'tb' (tabular method)
        """
        self.obs_shape, self.act_shape = obs_shape, act_shape
        if p_type not in ['fa', 'tb']:
            ValueError(f'policy type {p_type} is not supported!')
        self.p_type = p_type

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

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

    def init_params(self, theta=None) -> None:
        """ Define the policy shape. Initialize the mean and std vector for policy generation of CEM.
        """
        assert all(isinstance(val, int) for val in [self.obs_shape, self.act_shape])
        self.theta_dim = self.obs_shape * self.act_shape
        if theta is not None:
            self.means = theta
        else:
            self.means = np.random.uniform(size=self.theta_dim)
        self.stds = np.ones(self.theta_dim)

    def update_params(self, elites: np.ndarray) -> None:
        """ Use a predefined policy matrix to initialize the mean and std vector of CEM.

        Parameters
        ----------
        elites : np.ndarray
            A predefined policy that is previously trained by CEM
        """
        self.means = np.mean(elites, axis=0)
        self.stds = np.std(elites, axis=0)

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

        Parameters
        ----------
        theta : np.ndarray
            A candidate policy generated by CEM
        """
        assert isinstance(self.sampler, SamplingBase)

        policy = None
        if self.p_type == 'fa':
            policy = FAPolicy(self.obs_shape, self.act_shape, theta)
        elif self.p_type == 'tb':
            policy = TabularPolicy(self.obs_shape, self.act_shape, theta)

        assert isinstance(policy, PolicyBase)
        self.sampler.load_eval_policy(policy)

        rewards = self.sampler.get_est()
        perf_est = rewards.mean()

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

    def get_elite_idx(self, rewards: np.ndarray) -> np.ndarray:
        """ Compute the indices of the candidate policies by their estimated performance (descending order)

        Parameters
        ----------
        rewards : np.ndarray
            A vector of estimated performance for each candidate policy
        """
        return np.argsort(rewards)[::-1][:self.elite_size]

    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
            )

            g_candidates = list(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) for theta in self.best_thetas[:n]]


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


def g1(d: np.ndarray) -> np.ndarray:
    return 25 - d


def g2(d: np.ndarray) -> np.ndarray:
    return 30 - d

