import os
import time
import torch
import numpy as np

from tqdm import tqdm
from algos.algo_base import AlgoBase
from agents.cem_agent import CEM
from agents.random_shooting_agent import RandomShooting
from utils.sampler import Sampler


class AlgoShootingMPC(AlgoBase):
    """
    Online planning algorithm that uses MPC with shooting methods to plan action sequences.
    """
    def __init__(self, env, *args, **kwargs):
        super().__init__(*args, **kwargs)

        if self.args.controller_type == 'random_shooting':
            self.agent = RandomShooting(self.device,
                                        self.args.action_dim,
                                        self.observation_encoder,
                                        self.reward_model,
                                        self.rssm,
                                        self.args.horizon,
                                        self.args.num_control_samples,
                                        self.args.exploration_noise_var)
        elif self.args.controller_type == 'cem':
            self.agent = CEM(self.device,
                             self.args.action_dim,
                             self.observation_encoder,
                             self.reward_model,
                             self.rssm,
                             self.args.horizon,
                             self.args.num_control_samples,
                             self.args.max_iterations,
                             self.args.num_elites,
                             self.args.exploration_noise_var)
        else:
            raise ValueError('unknown controller type')

        self.sampler = Sampler(env, self.replay_buffer, self.agent)

    def train(self):
        episode = 0

        # collect initial episodes
        _, _, _, all_episode_rewards = self.sampler.collect_episodes(self.args.init_episodes,
                                                                     random=True,
                                                                     render=self.args.render_training)
        episode += self.args.init_episodes
        mean_ep_reward = np.mean(all_episode_rewards)
        best_ep_reward = np.max(all_episode_rewards)
        self.logger.log('train/mean_ep_reward', mean_ep_reward, episode)
        self.logger.log('train/best_ep_reward', best_ep_reward, episode)

        # main training loop
        for it in tqdm(range(self.args.training_iterations), desc='Training progress'):
            itr_start_time = time.time()

            # collect experience on policy
            _, _, _, all_episode_rewards = self.sampler.collect_episodes(self.args.agent_episodes,
                                                                         exploration=True,
                                                                         render=self.args.render_training)
            episode += self.args.agent_episodes
            mean_ep_reward = np.mean(all_episode_rewards)
            best_ep_reward = np.max(all_episode_rewards)
            self.logger.log('train/mean_ep_reward', mean_ep_reward, episode)
            self.logger.log('train/best_ep_reward', best_ep_reward, episode)

            # model training loop
            for _ in range(self.args.model_iterations):
                self.optimize_model()

            # save model frequently
            if self.args.save_iter_model and episode % self.args.save_iter_model_freq == 0:
                self.save_model(episode, 'model_iter_' + str(episode))

            itr_time = time.time() - itr_start_time
            self.logger.log('train/itr_time', itr_time, episode)

            # evaluate policy
            if it % self.args.eval_freq == 0:
                eval_start_time = time.time()

                # collect experience on policy
                actions, obs, _, all_episode_rewards = self.sampler.collect_episodes(self.args.eval_episodes,
                                                                                     render=self.args.render_training)
                mean_ep_reward = np.mean(all_episode_rewards)
                best_ep_reward = np.max(all_episode_rewards)
                video = self.get_video(actions, obs)

                # log results
                eval_time = time.time() - eval_start_time
                self.logger.log('eval/mean_ep_reward', mean_ep_reward, episode)
                self.logger.log('eval/best_ep_reward', best_ep_reward, episode)
                self.logger.log('eval/eval_time', eval_time, episode)
                self.logger.log_video('eval/ep_video', video, episode)

    def save_model(self, model_path, model_name):
        torch.save({
            'observation_encoder_state_dict': self.observation_encoder.to('cpu').state_dict(),
            'observation_decoder_state_dict': self.observation_decoder.to('cpu').state_dict(),
            'reward_model_state_dict': self.reward_model.to('cpu').state_dict(),
            'rssm_state_dict': self.rssm.to('cpu').state_dict(),
        }, os.path.join(model_path, model_name))

    def load_model(self, model_path):
        if self.device.type == 'cuda':
            checkpoint = torch.load(model_path, map_location="cuda:0")
        else:
            checkpoint = torch.load(model_path)
        self.observation_encoder.load_state_dict(checkpoint['observation_encoder_state_dict'])
        self.observation_decoder.load_state_dict(checkpoint['observation_decoder_state_dict'])
        self.reward_model.load_state_dict(checkpoint['reward_model_state_dict'])
        self.rssm.load_state_dict(checkpoint['rssm_state_dict'])
