#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
"""
A network which implements a cross entropy method-based planner

The planner plans the best next action based on simulation data generated by
an ensemble of world models.

The idea is inspired by: https://arxiv.org/abs/1805.12114
"""
import itertools
import logging
import random
from typing import List, Optional, Tuple

import numpy as np
import scipy.stats as stats
import torch
import torch.nn as nn
from reagent.core import types as rlt
from reagent.core.parameters import CONTINUOUS_TRAINING_ACTION_RANGE
from reagent.models.base import ModelBase
from reagent.models.world_model import MemoryNetwork
from reagent.training.utils import rescale_actions
from torch.distributions.bernoulli import Bernoulli
from torch.distributions.categorical import Categorical
from torch.distributions.normal import Normal


logger = logging.getLogger(__name__)


class CEMPlannerNetwork(nn.Module):
    def __init__(
        self,
        mem_net_list: List[MemoryNetwork],
        cem_num_iterations: int,
        cem_population_size: int,
        ensemble_population_size: int,
        num_elites: int,
        plan_horizon_length: int,
        state_dim: int,
        action_dim: int,
        discrete_action: bool,
        terminal_effective: bool,
        gamma: float,
        alpha: float = 0.25,
        epsilon: float = 0.001,
        action_upper_bounds: Optional[np.ndarray] = None,
        action_lower_bounds: Optional[np.ndarray] = None,
    ):
        """
        :param mem_net_list: A list of world models used to simulate trajectories
        :param cem_num_iterations: The maximum number of iterations for
            searching the best action
        :param cem_population_size: The number of candidate solutions to
            evaluate in each CEM iteration
        :param ensemble_population_size: The number of trajectories to be
            sampled to evaluate a CEM solution
        :param num_elites: The number of elites kept to refine solutions
            in each iteration
        :param plan_horizon_length: The number of steps to plan ahead
        :param state_dim: state dimension
        :param action_dim: action dimension
        :param discrete_action: If actions are discrete or continuous
        :param terminal_effective: If False, planning will stop after a
            predicted terminal signal
        :param gamma: The reward discount factor
        :param alpha: The CEM solution update rate
        :param epsilon: The planning will stop early when the solution
            variance drops below epsilon
        :param action_upper_bounds: Upper bound of each action dimension.
            Only effective when discrete_action=False.
        :param action_lower_bounds: Lower bound of each action dimension.
            Only effective when discrete_action=False.
        """
        super().__init__()
        self.mem_net_list = nn.ModuleList(mem_net_list)
        self.cem_num_iterations = cem_num_iterations
        self.cem_pop_size = cem_population_size
        self.ensemble_pop_size = ensemble_population_size
        self.num_elites = num_elites
        self.plan_horizon_length = plan_horizon_length
        self.state_dim = state_dim
        self.action_dim = action_dim
        self.terminal_effective = terminal_effective
        self.gamma = gamma
        self.alpha = alpha
        self.epsilon = epsilon
        self.discrete_action = discrete_action
        if not discrete_action:
            assert (
                (action_upper_bounds is not None)
                and (action_lower_bounds is not None)
                and (
                    action_upper_bounds.shape
                    == action_lower_bounds.shape
                    == (action_dim,)
                )
            )
            assert np.all(action_upper_bounds >= action_lower_bounds)
            self.action_upper_bounds = np.tile(
                action_upper_bounds, self.plan_horizon_length
            )
            self.action_lower_bounds = np.tile(
                action_lower_bounds, self.plan_horizon_length
            )
            self.orig_action_upper = torch.tensor(action_upper_bounds)
            self.orig_action_lower = torch.tensor(action_lower_bounds)

    # pyre-fixme[56]: Decorator `torch.no_grad(...)` could not be called, because
    #  its type `no_grad` is not callable.
    @torch.no_grad()
    def forward(self, state: rlt.FeatureData):
        assert state.float_features.shape == (1, self.state_dim)
        if self.discrete_action:
            return self.discrete_planning(state)
        return self.continuous_planning(state)

    # pyre-fixme[56]: Decorator `torch.no_grad(...)` could not be called, because
    #  its type `no_grad` is not callable.
    @torch.no_grad()
    def acc_rewards_of_one_solution(
        self, init_state: torch.Tensor, solution: torch.Tensor, solution_idx: int
    ):
        """
        ensemble_pop_size trajectories will be sampled to evaluate a
        CEM solution. Each trajectory is generated by one world model

        :param init_state: its shape is (state_dim, )
        :param solution: its shape is (plan_horizon_length, action_dim)
        :param solution_idx: the index of the solution
        :return reward: Reward of each of ensemble_pop_size trajectories
        """
        reward_matrix = np.zeros((self.ensemble_pop_size, self.plan_horizon_length))

        for i in range(self.ensemble_pop_size):
            state = init_state
            mem_net_idx = np.random.randint(0, len(self.mem_net_list))
            for j in range(self.plan_horizon_length):
                # state shape:
                # (1, 1, state_dim)
                # action shape:
                # (1, 1, action_dim)
                (
                    reward,
                    next_state,
                    not_terminal,
                    not_terminal_prob,
                ) = self.sample_reward_next_state_terminal(
                    state=rlt.FeatureData(state.reshape((1, 1, self.state_dim))),
                    action=rlt.FeatureData(
                        solution[j, :].reshape((1, 1, self.action_dim))
                    ),
                    mem_net=self.mem_net_list[mem_net_idx],
                )
                reward_matrix[i, j] = reward * (self.gamma ** j)

                if not not_terminal:
                    logger.debug(
                        f"Solution {solution_idx}: predict terminal at step {j}"
                        f" with prob. {1.0 - not_terminal_prob}"
                    )

                if not not_terminal:
                    break

                state = next_state

        return np.sum(reward_matrix, axis=1)

    # pyre-fixme[56]: Decorator `torch.no_grad(...)` could not be called, because
    #  its type `no_grad` is not callable.
    @torch.no_grad()
    def acc_rewards_of_all_solutions(
        self, state: rlt.FeatureData, solutions: torch.Tensor
    ) -> float:
        """
        Calculate accumulated rewards of solutions.

        :param state: the input which contains the starting state
        :param solutions: its shape is (cem_pop_size, plan_horizon_length, action_dim)
        :returns: a vector of size cem_pop_size, which is the reward of each solution
        """
        acc_reward_vec = np.zeros(self.cem_pop_size)
        init_state = state.float_features
        for i in range(self.cem_pop_size):
            if i % (self.cem_pop_size // 10) == 0:
                logger.debug(f"Simulating the {i}-th solution...")
            acc_reward_vec[i] = self.acc_rewards_of_one_solution(
                init_state, solutions[i], i
            )
        return acc_reward_vec

    # pyre-fixme[56]: Decorator `torch.no_grad(...)` could not be called, because
    #  its type `no_grad` is not callable.
    @torch.no_grad()
    def sample_reward_next_state_terminal(
        self, state: rlt.FeatureData, action: rlt.FeatureData, mem_net: MemoryNetwork
    ):
        """Sample one-step dynamics based on the provided world model"""
        wm_output = mem_net(state, action)
        num_mixtures = wm_output.logpi.shape[2]
        mixture_idx = (
            Categorical(torch.exp(wm_output.logpi.view(num_mixtures)))
            .sample()
            .long()
            .item()
        )
        next_state = Normal(
            wm_output.mus[0, 0, mixture_idx], wm_output.sigmas[0, 0, mixture_idx]
        ).sample()
        reward = wm_output.reward[0, 0]
        if self.terminal_effective:
            not_terminal_prob = torch.sigmoid(wm_output.not_terminal[0, 0])
            not_terminal = Bernoulli(not_terminal_prob).sample().long().item()
        else:
            not_terminal_prob = 1.0
            not_terminal = 1
        return reward, next_state, not_terminal, not_terminal_prob

    def constrained_variance(self, mean, var):
        lb_dist, ub_dist = (
            mean - self.action_lower_bounds,
            self.action_upper_bounds - mean,
        )
        return np.minimum(np.minimum((lb_dist / 2) ** 2, (ub_dist / 2) ** 2), var)

    # pyre-fixme[56]: Decorator `torch.no_grad(...)` could not be called, because
    #  its type `no_grad` is not callable.
    @torch.no_grad()
    def continuous_planning(self, state: rlt.FeatureData) -> torch.Tensor:
        # TODO: Warmstarts means and vars using previous solutions (T48841404)
        mean = (self.action_upper_bounds + self.action_lower_bounds) / 2
        var = (self.action_upper_bounds - self.action_lower_bounds) ** 2 / 16
        normal_sampler = stats.truncnorm(
            -2, 2, loc=np.zeros_like(mean), scale=np.ones_like(mean)
        )

        for i in range(self.cem_num_iterations):
            logger.debug(f"{i}-th cem iteration.")
            const_var = self.constrained_variance(mean, var)
            solutions = (
                normal_sampler.rvs(
                    size=[self.cem_pop_size, self.action_dim * self.plan_horizon_length]
                )
                * np.sqrt(const_var)
                + mean
            )
            action_solutions = torch.from_numpy(
                solutions.reshape(
                    (self.cem_pop_size, self.plan_horizon_length, self.action_dim)
                )
            ).float()
            acc_rewards = self.acc_rewards_of_all_solutions(state, action_solutions)
            elites = solutions[np.argsort(acc_rewards)][-self.num_elites :]
            new_mean = np.mean(elites, axis=0)
            new_var = np.var(elites, axis=0)
            mean = self.alpha * mean + (1 - self.alpha) * new_mean
            var = self.alpha * var + (1 - self.alpha) * new_var

            if np.max(var) <= self.epsilon:
                break

        # Pick the first action of the optimal solution
        solution = mean[: self.action_dim]
        raw_action = solution.reshape(-1)
        low = torch.tensor(CONTINUOUS_TRAINING_ACTION_RANGE[0])
        high = torch.tensor(CONTINUOUS_TRAINING_ACTION_RANGE[1])
        # rescale to range (-1, 1) as per canonical output range of continuous agents
        return rescale_actions(
            torch.tensor(raw_action),
            new_min=low,
            new_max=high,
            prev_min=self.orig_action_lower,
            prev_max=self.orig_action_upper,
        )

    # pyre-fixme[56]: Decorator `torch.no_grad(...)` could not be called, because
    #  its type `no_grad` is not callable.
    @torch.no_grad()
    def discrete_planning(self, state: rlt.FeatureData) -> Tuple[int, np.ndarray]:
        # For discrete actions, we use random shoots to get the best next action
        random_action_seqs = list(
            itertools.product(range(self.action_dim), repeat=self.plan_horizon_length)
        )
        random_action_seqs = random.choices(random_action_seqs, k=self.cem_pop_size)
        action_solutions = torch.zeros(
            self.cem_pop_size, self.plan_horizon_length, self.action_dim
        )
        for i, action_seq in enumerate(random_action_seqs):
            for j, act_idx in enumerate(action_seq):
                action_solutions[i, j, act_idx] = 1
        acc_rewards = self.acc_rewards_of_all_solutions(state, action_solutions)

        first_action_tally = np.zeros(self.action_dim)
        reward_tally = np.zeros(self.action_dim)
        for action_seq, acc_reward in zip(random_action_seqs, acc_rewards):
            first_action = action_seq[0]
            first_action_tally[first_action] += 1
            reward_tally[first_action] += acc_reward

        best_next_action_idx = np.nanargmax(reward_tally / first_action_tally)
        best_next_action_one_hot = torch.zeros(self.action_dim).float()
        best_next_action_one_hot[best_next_action_idx] = 1

        logger.debug(
            f"Choose action {best_next_action_idx}."
            f"Stats: {reward_tally} / {first_action_tally}"
            f" = {reward_tally/first_action_tally} "
        )
        return best_next_action_idx, best_next_action_one_hot
