import numpy as np
import torch
import torch.nn as nn
import logging

import gym.spaces as spaces
from rlberry.agents import IncrementalAgent
from rlberry.agents.utils.memories import CEMMemory
from rlberry.agents.utils.torch_training import optimizer_factory
from rlberry.agents.utils.torch_models import default_policy_net_fn
from rlberry.utils.torch import choose_device
from rlberry.utils.writers import PeriodicWriter

logger = logging.getLogger(__name__)


class CEMAgent(IncrementalAgent):
    """
    Parameters
    ----------
    env : Model
        Environment for training.
    n_episodes : int
        Number of training episodes.
    horizon : int
        Maximum length of a trajectory.
    gamma : double
        Discount factor in [0, 1].
    entr_coef : double
        Entropy coefficient.
    batch_size : int
        Number of trajectories to sample at each iteration.
    percentile : int
        Percentile used to remove trajectories with low rewards.
    learning_rate : double
        Optimizer learning rate
    optimizer_type: str
        Type of optimizer. 'ADAM' by defaut.
    on_policy : bool
        If True, updates are done only with on-policy data.
    policy_net_fn : function(env, **kwargs)
        Function that returns an instance of a policy network (pytorch).
        If None, a default net is used.
    policy_net_kwargs : dict
        kwargs for policy_net_fn
    device : str
        Device to put the tensors on
    """

    name = "CrossEntropyAgent"
    fit_info = ("n_episodes", "episode_rewards")

    def __init__(self,
                 env,
                 n_episodes=1000,
                 horizon=100,
                 gamma=0.99,
                 entr_coef=0.1,
                 batch_size=16,
                 percentile=70,
                 learning_rate=0.01,
                 optimizer_type='ADAM',
                 on_policy=False,
                 policy_net_fn=None,
                 policy_net_kwargs=None,
                 device="cuda:best",
                 **kwargs):
        IncrementalAgent.__init__(self, env, **kwargs)

        # check environment
        assert isinstance(self.env.observation_space, spaces.Box)
        assert isinstance(self.env.action_space, spaces.Discrete)

        # parameters
        self.gamma = gamma
        self.entr_coef = entr_coef
        self.batch_size = batch_size
        self.n_episodes = n_episodes
        self.percentile = percentile
        self.learning_rate = learning_rate
        self.horizon = horizon
        self.on_policy = on_policy
        self.policy_net_kwargs = policy_net_kwargs or {}
        self.policy_net_fn = policy_net_fn or default_policy_net_fn
        self.optimizer_kwargs = {'optimizer_type': optimizer_type,
                                 'lr': learning_rate}
        self.device = choose_device(device)
        self.reset()

    def reset(self, **kwargs):
        # policy net
        self.policy_net = self.policy_net_fn(
                            self.env,
                            **self.policy_net_kwargs
                            ).to(self.device)

        # loss function and optimizer
        self.loss_fn = nn.CrossEntropyLoss()
        self.optimizer = optimizer_factory(
                                    self.policy_net.parameters(),
                                    **self.optimizer_kwargs)

        # memory
        self.memory = CEMMemory(self.batch_size)

        # default writer
        self.writer = PeriodicWriter(self.name,
                                     log_every=5*logger.getEffectiveLevel())

        #
        self.episode = 0
        self._rewards = np.zeros(self.n_episodes)
        self._cumul_rewards = np.zeros(self.n_episodes)

    def partial_fit(self, fraction: float, **kwargs):
        assert 0.0 < fraction <= 1.0
        n_episodes_to_run = int(np.ceil(fraction * self.n_episodes))
        count = 0
        while count < n_episodes_to_run and self.episode < self.n_episodes:
            self._run_episode()
            count += 1

        info = {"n_episodes": self.episode,
                "episode_rewards": self._rewards[:self.episode]}
        return info

    def policy(self, observation, **kwargs):
        state = torch.from_numpy(observation).float().to(self.device)
        action_dist = self.policy_net(state)
        action = action_dist.sample().item()
        return action

    def _process_batch(self):
        rewards = np.array(self.memory.rewards)

        reward_bound = np.percentile(rewards, self.percentile)
        reward_mean = float(np.mean(rewards))

        train_states = []
        train_actions = []

        for ii in range(self.memory.size):
            if rewards[ii] < reward_bound:
                continue
            train_states.extend(self.memory.states[ii])
            train_actions.extend(self.memory.actions[ii])

        train_states_tensor = torch.FloatTensor(train_states).to(self.device)
        train_actions_tensor = torch.LongTensor(train_actions).to(self.device)

        # states in last trajectory
        last_states = self.memory.states[-1]
        last_states_tensor = torch.FloatTensor(last_states).to(self.device)

        return train_states_tensor, train_actions_tensor, \
            reward_bound, reward_mean, last_states_tensor

    def _run_episode(self):
        # interact for H steps
        episode_rewards = 0
        episode_states = []
        episode_actions = []
        state = self.env.reset()
        for _ in range(self.horizon):
            # take action according to policy_net
            action = self.policy(state)
            next_state, reward, done, _ = self.env.step(action)

            # save
            episode_states.append(state)
            episode_actions.append(action)

            # increment rewards
            episode_rewards += reward

            if done:
                break
            state = next_state

        self.memory.append(episode_states, episode_actions, episode_rewards)

        # update
        if (self.episode % self.batch_size == 0) or (not self.on_policy):
            self._update()

        #
        ep = self.episode
        self._rewards[ep] = episode_rewards
        self._cumul_rewards[ep] = episode_rewards + \
            self._cumul_rewards[max(0, ep - 1)]

        # increment ep and write
        self.episode += 1

        if self.writer is not None:
            self.writer.add_scalar("episode", self.episode, None)
            self.writer.add_scalar("ep reward", episode_rewards)

        return episode_rewards

    def _update(self):
        train_states_tensor, train_actions_tensor,\
            reward_bound, reward_mean, \
            last_states_tensor = self._process_batch()
        self.optimizer.zero_grad()
        action_scores = self.policy_net.action_scores(train_states_tensor)
        loss = self.loss_fn(action_scores, train_actions_tensor)

        # entropy in last trajectory
        action_dist_last_traj = self.policy_net(last_states_tensor)
        entropy = action_dist_last_traj.entropy().mean()

        loss = loss - self.entr_coef*entropy
        loss.backward()
        self.optimizer.step()

        return loss.item(), reward_mean, reward_bound

    #
    # For hyperparameter optimization
    #
    @classmethod
    def sample_parameters(cls, trial):
        batch_size = trial.suggest_categorical('batch_size',
                                               [10, 20, 50, 100, 200])
        gamma = trial.suggest_categorical('gamma',
                                          [0.9, 0.95, 0.99])
        learning_rate = trial.suggest_loguniform('learning_rate', 1e-5, 1)

        entr_coef = trial.suggest_loguniform('entr_coef', 1e-8, 0.1)

        on_policy = trial.suggest_categorical('on_policy',
                                              [False, True])
        return {
                'batch_size': batch_size,
                'gamma': gamma,
                'learning_rate': learning_rate,
                'entr_coef': entr_coef,
                'on_policy': on_policy
                }
