import wandb
import os
import numpy as np
import torch
from tensorboardX import SummaryWriter
from onpolicy.utils.shared_buffer import SharedReplayBuffer

def _t2n(x):
    """Convert torch tensor to a numpy array."""
    return x.detach().cpu().numpy()

class Runner(object):
    """
    Base class for training recurrent policies.
    :param config: (dict) Config dictionary containing parameters for training.
    """
    def __init__(self, config):

        self.all_args = config['all_args']
        self.args = self.all_args
        self.envs = config['envs']
        self.eval_envs = config['eval_envs']
        self.device = config['device']
        self.num_agents = config['num_agents']
        if config.__contains__("render_envs"):
            self.render_envs = config['render_envs']       

        # parameters
        self.env_name = self.all_args.env_name
        self.algorithm_name = self.all_args.algorithm_name
        self.experiment_name = self.all_args.experiment_name
        self.use_centralized_V = self.all_args.use_centralized_V
        self.use_obs_instead_of_state = self.all_args.use_obs_instead_of_state
        self.num_env_steps = self.all_args.num_env_steps
        self.episode_length = self.all_args.episode_length
        # print('self.episode_length = {}'.format(self.episode_length))
        self.n_rollout_threads = self.all_args.n_rollout_threads
        self.n_eval_rollout_threads = self.all_args.n_eval_rollout_threads
        self.n_render_rollout_threads = self.all_args.n_render_rollout_threads
        self.use_linear_lr_decay = self.all_args.use_linear_lr_decay
        self.hidden_size = self.all_args.hidden_size
        self.use_wandb = self.all_args.use_wandb
        self.use_render = self.all_args.use_render
        self.recurrent_N = self.all_args.recurrent_N

        # interval
        self.save_interval = self.all_args.save_interval
        self.use_eval = self.all_args.use_eval
        self.eval_interval = self.all_args.eval_interval
        self.log_interval = self.all_args.log_interval

        # dir
        self.model_dir = self.all_args.model_dir

        if self.use_wandb:
            self.save_dir = str(wandb.run.dir)
            self.run_dir = str(wandb.run.dir)
        else:
            self.run_dir = config["run_dir"]
            self.log_dir = str(self.run_dir / 'logs')
            print('run_dir = {}, log_dir = {}'.format(self.run_dir,self.log_dir))
            if not os.path.exists(self.log_dir):
                os.makedirs(self.log_dir)
            self.writter = SummaryWriter(self.log_dir)
            self.save_dir = str(self.run_dir / 'models')
            if not os.path.exists(self.save_dir):
                os.makedirs(self.save_dir)

        from onpolicy.algorithms.r_mappo.r_mappo import R_MAPPO as TrainAlgo
        from onpolicy.algorithms.r_mappo.algorithm.rMAPPOPolicy import R_MAPPOPolicy as Policy


        if self.env_name == 'MPE':
            if self.args.mpe_share_obs:
                share_observation_space = self.envs.share_observation_space[0]
            else:
                share_observation_space = self.envs.observation_space[0]
        else:
            share_observation_space = self.envs.share_observation_space[0] if self.use_centralized_V else self.envs.observation_space[0]

        # policy network
        # print('self.envs.action_space[0] = {}, self.envs.action_space = {}'.format(self.envs.action_space[0],self.envs.action_space))
        self.policy = Policy(self.all_args,
                            self.envs.observation_space[0],
                            share_observation_space,
                            self.envs.action_space[0],
                            device = self.device)

        if self.model_dir is not None:
            self.restore()

        # algorithm
        self.trainer = TrainAlgo(self.all_args, self.policy, device = self.device)
        
        # buffer
        self.buffer = SharedReplayBuffer(self.all_args,
                                        self.num_agents,
                                        self.envs.observation_space[0],
                                        share_observation_space,
                                        self.envs.action_space[0])

    def run(self):
        """Collect training data, perform training updates, and evaluate policy."""
        raise NotImplementedError

    def warmup(self):
        """Collect warmup pre-training data."""
        raise NotImplementedError

    def collect(self, step):
        """Collect rollouts for training."""
        raise NotImplementedError

    def insert(self, data):
        """
        Insert data into buffer.
        :param data: (Tuple) data to insert into training buffer.
        """
        raise NotImplementedError
    
    @torch.no_grad()
    def compute(self):
        """Calculate returns for the collected data."""
        self.trainer.prep_rollout()
        if self.all_args.use_q or self.all_args.sp_use_q:

            next_values_all, actions, _, _, _,next_target_values_all,_ = self.trainer.policy.get_actions(
                np.concatenate(self.buffer.share_obs[-1]),
                np.concatenate(self.buffer.obs[-1]),
                np.concatenate(self.buffer.rnn_states[-1]),
                np.concatenate(self.buffer.rnn_states_critic[-1]),
                np.concatenate(self.buffer.masks[-1]),target=True, target_rnn_states_critic=np.concatenate(self.buffer.target_rnn_states_critic[-1]) )
            act_idx = actions
            if self.args.sp_use_q:
                next_values = next_target_values_all
                curr_values = next_values_all
            else:
                next_values = torch.gather(next_target_values_all, index=act_idx, dim=-1)
                curr_values = torch.gather(next_values_all, index=act_idx, dim=-1)
        else:
            next_values = self.trainer.policy.get_values(np.concatenate(self.buffer.share_obs[-1]),
                                                         np.concatenate(self.buffer.rnn_states_critic[-1]),
                                                         np.concatenate(self.buffer.masks[-1]))
            # act_idx = torch.argmax(actions,dim = -1 ,keepdim= True)

        next_values = np.array(np.split(_t2n(next_values), self.n_rollout_threads))


        if self.all_args.use_q or self.all_args.sp_use_q:
            curr_values = np.array(np.split(_t2n(curr_values), self.n_rollout_threads))
            self.buffer.compute_returns(next_values, self.trainer.value_normalizer,curr_value = curr_values)
        else:
            self.buffer.compute_returns(next_values, self.trainer.value_normalizer)
    
    def train(self,time_steps=None):
        """Train policies with data in buffer. """
        self.trainer.prep_training()
        train_infos = self.trainer.train(self.buffer,time_steps = time_steps)
        self.buffer.after_update()
        return train_infos

    def save(self):
        """Save policy's actor and critic networks."""
        policy_actor = self.trainer.policy.actor
        policy_critic = self.trainer.policy.critic
        if self.all_args.idv_para:
            for i in range(self.all_args.n_agents):
                torch.save(policy_actor[i].state_dict(), str(self.save_dir) + "/actor_{}.pt".format(i))
                torch.save(policy_critic[i].state_dict(), str(self.save_dir) + "/critic_{}.pt".format(i))
        else:
            torch.save(policy_actor.state_dict(), str(self.save_dir) + "/actor.pt")
            torch.save(policy_critic.state_dict(), str(self.save_dir) + "/critic.pt")

    def restore(self):
        """Restore policy's networks from a saved model."""
        if self.all_args.idv_para:
            for i in range(self.all_args.n_agents):
                policy_actor_state_dict = torch.load(str(self.model_dir) + "/actor_{}.pt".format(i))
                self.policy.actor[i].load_state_dict(policy_actor_state_dict)
                if not self.all_args.use_render:
                    policy_critic_state_dict = torch.load(str(self.model_dir) + "/critic_{}.pt".format(i))
                    self.policy.critic[i].load_state_dict(policy_critic_state_dict)
        else:
            policy_actor_state_dict = torch.load(str(self.model_dir) + '/actor.pt')
            self.policy.actor.load_state_dict(policy_actor_state_dict)
            if not self.all_args.use_render:
                policy_critic_state_dict = torch.load(str(self.model_dir) + '/critic.pt')
                self.policy.critic.load_state_dict(policy_critic_state_dict)
 
    def log_train(self, train_infos, total_num_steps):
        """
        Log training info.
        :param train_infos: (dict) information about training update.
        :param total_num_steps: (int) total number of training env steps.
        """
        for k, v in train_infos.items():
            if self.use_wandb:
                wandb.log({k: v}, step=total_num_steps)
            else:
                self.writter.add_scalars(k, {k: v}, total_num_steps)

    def log_env(self, env_infos, total_num_steps):
        """
        Log env info.
        :param env_infos: (dict) information about env state.
        :param total_num_steps: (int) total number of training env steps.
        """
        for k, v in env_infos.items():
            if len(v)>0:
                if self.use_wandb:
                    wandb.log({k: np.mean(v)}, step=total_num_steps)
                else:
                    self.writter.add_scalars(k, {k: np.mean(v)}, total_num_steps)
