"""
Implement the MO-HCPI (H-OPT) here

This code is based on: https://github.com/RomainLaroche/SPIBB/blob/master/HCPI.py

"""
import numpy as np
import cvxpy as cp
from scipy import stats
from math import log,sqrt
import time

from gridworld.agents.tabular_base_agent import Agent, ReshapedAdvantagePIAgent
from gridworld.core.utils import estimate_model, direct_policy_evaluation, default_termination, bounded_successive_approximation

class ISEstimator:
    """
    The importance sampling estimator class.
    It estimates the return of a target policy pi_t
    given a set of trajectories generated by a baseline policy pi_b
    The available estimator are :
    - importance sampling;
    - per decision importance sampling;
    """
    def __init__(self, gamma: float,
                 pi_b: np.ndarray,
                 pi_t: np.ndarray,
                 estimator_type: str,
                 R = None,
                 C = None,
                 P_hat = None ,
                 list_format = True):
        """

        :param gamma: discount factor
        :param pi_b: the baseline policy
        :param pi_t: the target policy
        :param estimator_type: the kind of IS algorithm the estimator uses
        :param R: Reward function
        :param C: Constraint function
        :param list_format: whether to return in a list or return a scalar
        """
        self.estimator_dict = {"importance_sampling": self.importance_sampling,
                               "weighted_importance_sampling": self.weighted_importance_sampling,
                               "weighted_per_decision_IS": self.weighted_per_decision_importance_sampling,
                               "per_decision_IS": self.per_decision_importance_sampling,
                               "doubly_robust": self.doubly_robust
                              }
        self.gamma=gamma
        self.estimator_type = estimator_type
        self.pi_b = pi_b
        self.pi_t = pi_t
        self.list_format = list_format

        # we need these for DR estimator
        self.R = R
        self.C = C
        self.P_hat = P_hat

    def __call__(self, trajectories):
        """
        :param trajectories:
        :return:
        """
        return self.estimator_dict[self.estimator_type](trajectories)

    def importance_sampling(self, trajectories: list):
        """
        (\Prod_{j=0}^{T} \frac{pi_t(x_t, a_t)}{pi_b(x_t, a_t)}) \sum_t [\gamma^t R_t]

        :param trajectories:
        :return:
        """
        # to return a scalar or list
        if self.list_format:
            R_estimate = []
            C_estimate = []
        else:
            R_estimate = 0
            C_estimate = 0

        # estimate for each traj in the list
        for trajectory in trajectories:

            t = 0
            cum_rew = 0
            cum_cost = 0
            weight = 1

            for transition in trajectory:

                state, action, reward, cost, next_state = transition

                # update the cost and rewards
                cum_rew += self.gamma**t * reward
                cum_cost += self.gamma**t * cost

                # update the IS weight
                weight = weight*(self.pi_t[state,action]/self.pi_b[state,action])
                t += 1

            if self.list_format:
                R_estimate.append(weight*cum_rew)
                C_estimate.append(weight * cum_cost)
            else:
                R_estimate += weight*cum_rew
                C_estimate += weight * cum_cost

        # return avg if the return is a scalar
        if self.list_format:
            return R_estimate, C_estimate
        else:
            return R_estimate/len(trajectories), C_estimate/len(trajectories)

    def weighted_importance_sampling(self, trajectories: list):
        """
        Instead of average returns a normalized/weighted IS estimate

        :param trajectories:
        :return:
        """
        # to return a scalar or list
        if self.list_format:
            R_estimate = []
            C_estimate = []
        else:
            R_estimate = 0
            C_estimate = 0

        sum_weight = 0
        for trajectory in trajectories:
            t = 0
            cum_rew = 0
            cum_cost = 0
            weight = 1

            for transition in trajectory:
                state, action, reward, cost, next_state = transition
                # update the cost and rewards
                cum_rew += self.gamma ** t * reward
                cum_cost += self.gamma ** t * cost
                # update the weights
                weight = weight * (self.pi_t[state, action] / self.pi_b[state, action])
                t += 1

            if self.list_format:
                R_estimate.append(cum_rew)
                C_estimate.append(cum_cost)
            else:
                R_estimate += weight * cum_rew
                C_estimate += weight * cum_cost
                sum_weight += weight

        if self.list_format:
            return R_estimate, C_estimate
        else:
            return R_estimate/sum_weight, C_estimate/sum_weight

    def per_decision_importance_sampling(self, trajectories, max_steps = 200):
        """
        \sum_t[ (\Prod_{j=0}^{t} \frac{pi_t(x_t, a_t)}{pi_b(x_t, a_t)}) \gamma^t R_t]

        :param trajectories:
        :param max_steps: the maximum horizon of the CMDP
        :return:
        """
        # whether to return a scalar or not
        if self.list_format:
            R_estimate = []
            C_estimate = []
        else:
            R_estimate = 0
            C_estimate = 0

        # create an empty IS ratio book-keeping container
        rho = np.zeros((max_steps,len(trajectories)))

        n = 0
        for trajectory in trajectories:
            weight = 1
            t=0
            for transition in trajectory:
                state, action, reward, cost, next_state = transition
                current_weight = self.pi_t[state, action] / self.pi_b[state, action]
                weight = weight*current_weight
                rho[t,n] = weight
                t += 1
            n += 1

        # compute the returns here
        n = 0
        for trajectory in trajectories:
            t = 0
            cum_rew = 0
            cum_cost = 0
            for transition in trajectory:
                state, action, reward, cost, next_state = transition
                cum_rew += self.gamma ** t * reward * float(rho[t,n])
                cum_cost += self.gamma ** t * cost * float(rho[t, n])
                t += 1
            if self.list_format:
                R_estimate.append(cum_rew)
                C_estimate.append(cum_cost)
            else:
                R_estimate += cum_rew
                C_estimate += cum_cost
            n += 1

        if self.list_format:
            return R_estimate, C_estimate
        else:
            return float(R_estimate)/len(trajectories), float(C_estimate)/len(trajectories)

    def weighted_per_decision_importance_sampling(self, trajectories, max_steps = 200):
        """
        returns normalized

        :param trajectories:
        :param max_steps: the maximum horizon of the CMDP
        :return:
        """
        # whether to return a scalar or not
        if self.list_format:
            R_estimate = []
            C_estimate = []
        else:
            R_estimate = 0
            C_estimate = 0

        # create an empty IS ratio book-keeping container
        rho = np.zeros((max_steps,len(trajectories)))
        sum_per_decision = np.zeros((max_steps, 1))

        n = 0
        for trajectory in trajectories:
            weight = 1
            t=0
            for transition in trajectory:
                state, action, reward, cost, next_state = transition
                current_weight = self.pi_t[state, action] / self.pi_b[state, action]
                weight = weight*current_weight
                rho[t,n] = weight
                sum_per_decision[t] += weight * self.gamma ** t
                t += 1
            n += 1

        # add eps to prevent division from 0
        sum_per_decision += 1e-9

        # compute the returns here
        n = 0
        for trajectory in trajectories:
            t = 0
            cum_rew = 0
            cum_cost = 0
            for transition in trajectory:
                state, action, reward, cost, next_state = transition
                cum_rew += self.gamma ** t * reward * float(rho[t, n])/float(sum_per_decision[t])
                cum_cost += self.gamma ** t * cost * float(rho[t, n])/float(sum_per_decision[t])
                t += 1
            if self.list_format:
                R_estimate.append(cum_rew)
                C_estimate.append(cum_cost)
            else:
                R_estimate += cum_rew
                C_estimate += cum_cost
            n += 1

        if self.list_format:
            return R_estimate, C_estimate
        else:
            return float(R_estimate)/float(np.sum(sum_per_decision)), float(C_estimate)/float(np.sum(sum_per_decision))

    def doubly_robust(self, trajectories):
        """
        As implemented in Jiang and Li, 2015;
        Make use of a control variate build from an approximate
        model of the MDP
        :param trajectories: a set of trajectories
        :return: an estimate of the return
        """

        if self.list_format:
            R_estimate, C_estimate = self.compute_estimare_dr(trajectories)
            return R_estimate, C_estimate
        else:
            # We used the 2-fold DR as model fitting
            raise Exception("not-implemented")

    def compute_estimare_dr(self, trajectories_eval, is_list=True):
        if is_list:
            R_estimate = []
            C_estimate = []
        else:
            R_estimate = 0
            C_estimate = 0


        # compute the V_hat and Q_hat for the pi_t
        # calculate the estimates for the target policy
        # R
        rV_hat = direct_policy_evaluation(self.P_hat, self.R, self.gamma, self.pi_t)
        rQ_hat = self.R + self.gamma * np.einsum('sat,t -> sa', self.P_hat, rV_hat)
        # C
        cV_hat = direct_policy_evaluation(self.P_hat, self.C, self.gamma, self.pi_t)
        cQ_hat = self.C + self.gamma * np.einsum('sat,t -> sa', self.P_hat, cV_hat)

        for trajectory in trajectories_eval:
            R_estimate_trajectory = 0
            C_estimate_trajectory = 0
            for transition in trajectory[::-1]:
                state, action, reward, cost, next_state = transition

                # R
                R_estimate_trajectory = int(rV_hat[state]) + \
                        self.pi_t[state, action] / self.pi_b[state, action]* (reward +
                                              self.gamma * R_estimate_trajectory-int(rQ_hat[state, action]))

                # C
                C_estimate_trajectory = int(cV_hat[state]) + \
                                        self.pi_t[state, action] / self.pi_b[state, action] * ( cost +
                                                self.gamma * C_estimate_trajectory - int(cQ_hat[state, action]))

            if is_list:
                R_estimate.append(R_estimate_trajectory)
                C_estimate.append(C_estimate_trajectory)
            else:
                R_estimate += R_estimate_trajectory
                C_estimate += C_estimate_trajectory

        # do average if not in list format
        if not(is_list):
            R_estimate = R_estimate / len(trajectories_eval)
            C_estimate = C_estimate / len(trajectories_eval)

        return R_estimate, C_estimate


class HCPIAgent(Agent):
    """
    The HCPI agent method
    """
    def __init__(self,
                 lower_bound_strategy: str,
                 estimator_type : str,
                 coeff_list=[],
                 training_size = 0.7,
                 nb_steps=5000,
                 ):
        """

        :param gamma: discount
        :param lower_bound_strategy: ??
        :param confidence: delta in student t-test
        :param estimator_type: the kind of IS estimator to use
        :param rho_min: for normalizing thes return from IS
        :param rho_max: for normalizing the return from IS
        :param nb_steps: number of steps for the q-learning updates
        """

        # super().__init__()
        self.strategy_dict = {"CI": self.confidence_interval_based,
                              "student_t_test": self.student_t_test,
                              # "BCa": self.bootstrap_method
                             }

        self.lower_bound_strategy = lower_bound_strategy
        self.estimator_type = estimator_type
        self.training_size = training_size
        self.nb_steps = nb_steps
        self.coeff_list = coeff_list

        # name takes into account the IS sampling and LB schemes
        self._name = "H_OPT" + "-" + estimator_type + "-"

    def __call__(self, trajectories):
        return self.strategy_dict[self.lower_bound_strategy](trajectories)

    def confidence_interval_based(self, trajectories, confidence, estimator,
                                  R_min, R_max,
                                  C_min, C_max,
                                  ):
        """
        :param trajectories: a batch of trajectories to evaluate
        :param confidence: delta parameter
        :param estimator: the off policy estimator object
        :param R_min: the min value of R
        :param R_max: the max value of R

        :return: a lower bound on the estimate return
        """
        # Sec 2.3 of HCPI
        c = 0.5

        # get the estimates, normalize and clip them
        R_list_estimates, C_list_estimates = estimator(trajectories)

        # w.r.t. R
        R_list_estimates = self.normalize_return(R_list_estimates, x_min=R_min, x_max=R_max)
        R_list_estimates_cut = [min(x, c) for x in R_list_estimates]
        # subtract.outer => Apply the subtract op to all pairs (a, b) with a in A and b in B.
        R_cross_substract = np.subtract.outer(np.square(R_list_estimates_cut), np.square(R_list_estimates_cut))
        R_cross_substract_squared = np.square(R_cross_substract)

        n = len(R_list_estimates)
        R_lower_bound = (1./n) * np.sum(R_list_estimates_cut) - \
               (7*c*log(2./confidence,2))/(3.*(n-1)) - \
                     n/c*sqrt((log(2/confidence))*np.sum(np.sum(R_cross_substract_squared)))

        # map back to normal range
        R_lower_bound = R_lower_bound * (R_max - R_min) + R_min

        # w.r.t. C
        C_list_estimates = self.normalize_return(C_list_estimates, x_min=C_min, x_max=C_max)
        C_list_estimates_cut = [min(x, c) for x in C_list_estimates]

        C_cross_substract = np.subtract.outer(np.square(C_list_estimates_cut), np.square(C_list_estimates_cut))
        C_cross_substract_squared = np.square(C_cross_substract)

        n = len(C_list_estimates)
        C_upper_bound = (1. / n) * np.sum(C_list_estimates_cut) + n / c * sqrt(
            (log(2 / confidence)) * np.sum(np.sum(C_cross_substract_squared))) - \
                        (7 * c * log(2. / confidence, 2)) / (3. * (n - 1))

        C_upper_bound = C_upper_bound * (C_max - C_min) + C_min

        return R_lower_bound, C_upper_bound

    def student_t_test(self, trajectories, confidence, estimator,
                       R_min, R_max,
                       C_min, C_max,
                       ):
        """
        Warning !! This method relies on the assumption that the return is normally distributed

        :param trajectories: a batch of trajectories
        :param confidence: the delta parameter
        :param estimator: the off policy evaluation estimator objec

        :return: a lower bound on the estimate return
        """

        # estimate the returns
        R_list_estimates, C_list_estimates = estimator(trajectories)

        # normalize
        R_list_estimates = self.normalize_return(R_list_estimates, x_min=R_min, x_max=R_max, )
        C_list_estimates = self.normalize_return(C_list_estimates, x_min=C_min, x_max=C_max, )

        # t-test
        # Sec 2.4 HCPI (Thomas 2015)
        n = len(R_list_estimates)
        # w.r.t. R
        R_estimated_return = np.mean(R_list_estimates)
        R_sigma = np.sqrt(1./(n-1) * np.sum(np.square(np.array(R_list_estimates) - R_estimated_return)))
        R_lower_bound = R_estimated_return - R_sigma/sqrt((n-1)) * stats.t.ppf(1 - confidence, n-1)
        # undo the normalization here
        R_lower_bound = R_lower_bound * (R_max - R_min) + R_min

        # w.r.t. C
        C_estimated_return = np.mean(C_list_estimates)
        C_sigma = np.sqrt(1. / (n - 1) * np.sum(np.square(np.array(C_list_estimates) - C_estimated_return)))
        C_upper_bound = C_estimated_return + C_sigma / sqrt((n - 1)) * stats.t.ppf(1 - confidence, n - 1)
        # undo the normalization here
        C_upper_bound = C_upper_bound * (C_max - C_min) + C_min

        return R_lower_bound, C_upper_bound

    def normalize_return(self, list_estimate, x_min, x_max):
        """
        normalizes the return in [0,1]
        :param list_estimate:
        :return:
        """
        return [(x - x_min) / (x_max - x_min) for x in list_estimate]

    def compute_policy(self, trajectories, pi_b, confidence, coeffs,
                       R, C, discount,
                       pib_R_est_performance, pib_C_est_performance,
                       R_min, R_max,
                       C_min, C_max,
                       q_learning_iterative_mode=False):
        """
        do the policy improvement here
        :return:
        """
        # do a train-test split
        training_index = int(self.training_size * len(trajectories))

        training_trajectories = trajectories[:training_index]
        testing_trajectories = trajectories[training_index:]

        # create the batch from training trajectories here
        batch = []
        for trajectory in training_trajectories:
            for [state, action, reward, cost, next_state] in trajectory:
                batch.append([state, action, reward, cost, next_state])

        nstates = R.shape[0]
        nactions = R.shape[1]
        P_hat = estimate_model(batch, nstates, nactions)

        # We can't solve the Eqn H-OPT exactly, two options:
        #  1) we solve the MLE advantage function to get the policy (ReshapedAdvantagePIAgent)
        #  2) get pi_hat via q_learning
        #  Once we have pi_hat, we regularize it
        # We found 1) works better, and we use that for our experiments

        if q_learning_iterative_mode:
            # compute the pi_hat using just the modified rewards
            _, pi_hat = self._q_lambda_learning(batch, coeffs, discount,
                                                nstates, nactions,
                                                nb_steps=self.nb_steps)
        else:
            # create a PI operator
            base_pi_agent = ReshapedAdvantagePIAgent(termination_condition=default_termination,
                                                     coeff_list=[(None,None)])

            approx_pi_operator = base_pi_agent.make_policy_iteration_operator(P=P_hat, R=R, C=C,
                                                                              discount=discount,
                                                                              baseline=pi_b,
                                                                              coeffs=coeffs)
            # find the optimal pi
            pi_hat = bounded_successive_approximation(pi_b,
                                                      operator=approx_pi_operator,
                                                      termination_condition=default_termination,
                                                      max_limit=10,)

        # Regularization on the optimal policy
        regularization_list = [0, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9]

        # Evaluate all the policy and return the one with the higher lower bound
        best_R_lower_bound = - np.inf
        best_C_upper_bound = + np.inf

        current_best_policy = pi_hat
        current_regularization = 0

        # try different linear combinations
        for regularization_parameter in regularization_list:
            # new candidate policy
            current_pi = (1 - regularization_parameter) * pi_hat + \
                             regularization_parameter * pi_b

            # Create the estimator here
            estimator = ISEstimator(gamma=discount, pi_b=pi_b,
                                    pi_t=current_pi,
                                    estimator_type=self.estimator_type,
                                    R=R, C=C, P_hat=P_hat,
                                    list_format=True)

            # compute the bounds based on the test batch
            R_lb, C_ub = self.strategy_dict[self.lower_bound_strategy](testing_trajectories,
                                                                       confidence, estimator,
                                                                       R_min, R_max,
                                                                       C_min, C_max)

            # if the bounds are finite, only then use them, else skip
            if (not np.isfinite(R_lb)) or (not np.isfinite(C_ub)):
                continue

            # accept the candidate if it has better performance w.r.t R
            if R_lb >= best_R_lower_bound and C_ub <= best_C_upper_bound:
                current_best_policy = current_pi
                best_R_lower_bound = R_lb
                best_C_upper_bound = C_ub
                current_regularization = regularization_parameter


        # check if the best candidate passes the safety test
        # The first instance, uses the estimated returns using P_hat, and the mean performance of the baseline
        pi_sol, reg_sol = self.safety_test(current_best_policy, best_R_lower_bound, best_C_upper_bound,
                                           current_regularization, pi_b,
                                           pib_R_est_performance, pib_C_est_performance)

        return pi_sol, reg_sol

    def safety_test(self, pi_t, R_lb_target, C_ub_target,
                    current_regularization, pi_b,
                    R_perf_baseline, C_perf_baseline):
        """
        The No Solution Found (NSF) solution

        Essentially the safety test
        :return:
        """
        if R_lb_target >= R_perf_baseline and C_ub_target <= C_perf_baseline:
            # passes the safety test
            return pi_t, current_regularization
        else:
            # return the original policy
            return pi_b, 1.0

    def _values_to_argmax(self, q):
        """
        utility for extracting greedy policy w.r.t. Q-function

        :param q: Q-function
        :return: pi
        """
        def f_axis(q):
            z = np.zeros_like(q)
            z[np.argmax(q)] = 1
            return z

        return np.apply_along_axis(f_axis, 1, q)

    def _q_lambda_learning(self, batch, coeffs,
                    gamma: float,
                    nb_states: int,
                    nb_actions: int,
                    nb_steps: int):
        """
        Uses the modified reward

        :param batch:
        :param gamma:
        :param nb_states:
        :param nb_actions:
        :param nb_steps:
        :return:
        """
        Q = np.zeros((nb_states, nb_actions))
        current_alpha = 0.1

        for step in range(int(nb_steps)):
            id_choice = np.random.randint(len(batch))
            [state, action, reward, cost, next_state] = batch[id_choice]
            lambda_reward = coeffs[0] * reward - coeffs[1] * cost
            Q[state, action] = Q[state, action] + current_alpha * (
                    lambda_reward + gamma * np.max(Q[next_state, :]) - Q[state, action])

        return Q, self._values_to_argmax(Q)