"""
Random shooting algorithm that utilizes learned NN sampling policy
"""
from copy import deepcopy
import numpy as np
import torch
import torch.nn.functional as F
from torch.optim import Adam

from mjmpc.control.controller import Controller
from mjmpc.control.softqmpc.algs.sac import ReplayMemoryTraj, soft_update, hard_update
from mjmpc.utils.control_utils import cost_to_go, generate_noise, scale_ctrl
from mjmpc.utils import helpers

# class SACMPC(Controller):
#     def __init__(self, 
#                  d_state,
#                  d_obs,
#                  d_action,                
#                  action_lows,
#                  action_highs,
#                  horizon,
#                  gamma,
#                  step_size,
#                  filter_coeffs,
#                  n_iters,
#                  num_particles,
#                  init_cov,
#                  base_action,
#                  policy,
#                  critic,
#                  lr,
#                  alpha,
#                  tau,
#                  train_batch_size,
#                  updates_per_step,
#                  start_steps,
#                  target_update_interval,
#                  replay_size,
#                  automatic_entropy_tuning,
#                  set_sim_state_fn=None,
#                  rollout_fn=None,
#                  sample_mode='mean',
#                  batch_size=1,
#                  seed=0):
class SACMPC(Controller):
    def __init__(self, param_dict):
        """
        Parameters
        __________

        """

        super(SACMPC, self).__init__(param_dict['d_state'],
                                     param_dict['d_obs'],
                                     param_dict['d_action'],
                                     param_dict['action_lows'], 
                                     param_dict['action_highs'],
                                     param_dict['horizon'],
                                     param_dict['gamma'],  
                                     param_dict['n_iters'],
                                     param_dict['set_sim_state_fn'],
                                     param_dict['rollout_fn'],
                                     param_dict['sample_mode'],
                                     param_dict['batch_size'],
                                     param_dict['seed'])
        self.param_dict = param_dict
        torch.manual_seed(seed)
        self.init_cov = np.array([param_dict['init_cov']] * self.d_action)
        self.mean_action = np.zeros(shape=(self.horizon, self.d_action))
        self.num_particles = param_dict['num_particles']
        self.base_action = param_dict['base_action']
        self.cov_action = np.diag(self.init_cov)
        self.step_size = param_dict['step_size']
        self.filter_coeffs = param_dict['filter_coeffs']
        self.lr = param_dict['lr']
        self.alpha = param_dict['alpha']
        self.tau = param_dict['tau']
        self.automatic_entropy_tuning = param_dict['automatic_entropy_tuning']
        self.train_batch_size = param_dict['train_batch_size']
        self.updates_per_step = param_dict['updates_per_step']
        self.start_steps = param_dict['start_steps']
        self.target_update_interval = param_dict['target_update_interval']
        self.replay_size = param_dict['replay_size']
        self._get_sim_obs_fn = None

        #policy and qfunction objects
        self.policy = policy
        self.critic = critic
        self.policy_optim = Adam(self.policy.parameters(), lr=self.lr)
        self.critic_optim = Adam(self.critic.parameters(), lr=self.lr)
        self.critic_target = deepcopy(self.critic)

        self.memory = ReplayMemoryTraj(self.replay_size)

        # Target Entropy = −dim(A) (e.g. , -6 for HalfCheetah-v2) as given in the paper
        if self.automatic_entropy_tuning is True:
            self.target_entropy = -torch.prod(torch.Tensor(self.d_action)).item() #.to(self.device)
            self.log_alpha = torch.zeros(1, requires_grad=True) #, device=self.device
            self.alpha_optim = Adam([self.log_alpha], lr=self.lr)

    @property
    def get_sim_obs_fn(self):
        return self._get_sim_obs_fn
    
    @get_sim_obs_fn.setter
    def get_sim_obs_fn(self, fn):
        self._get_sim_obs_fn = fn

    def _get_next_action(self, mode='mean'):
        with torch.no_grad():
            obs_torch = torch.FloatTensor(self.start_obs).unsqueeze(0)
            next_action, _ = self.policy.get_action(obs_torch, mode)
            next_action = next_action.numpy().copy().reshape(self.d_action,) 

        # if mode == 'mean':
        #     next_action = self.mean_action[0].copy()
        # elif mode == 'sample':
        #     delta = generate_noise(self.cov_action, self.filter_coeffs,
        #                            shape=(1, 1), base_seed=self.seed + self.num_steps)
        #     next_action = self.mean_action[0].copy() + delta.reshape(self.d_action).copy()
        # else:
        #     raise ValueError('Unidentified sampling mode in get_next_action')
        return next_action
    
    def _sample_noise(self):
        delta = generate_noise(self.cov_action, self.filter_coeffs,
                               shape=(self.num_particles, self.horizon), 
                               base_seed = self.seed_val + self.num_steps)
        return delta.copy()
        
    def generate_rollouts(self, state):
        """
            Samples a batch of actions, rolls out trajectories for each particle
            and returns the resulting observations, costs,  
            actions

            Parameters
            ----------
            state : dict or np.ndarray
                Initial state to set the simulation env to
         """
        self._set_sim_state_fn(deepcopy(state)) #set state of simulation
        self.start_obs = (self.get_sim_obs_fn()[0]).copy()
        # delta = self._sample_noise() #sample noise sequence
        delta = None
        obs_seq, act_seq, logprob_seq, cost_seq, done_seq, next_obs_seq, info_seq = self._rollout_fn(mode='sample',
                                                                                             noise=delta)

        trajectories = dict(
            observations=obs_seq,
            next_observations=next_obs_seq,
            actions=act_seq,
            log_probs=logprob_seq,
            costs=cost_seq,
            dones=done_seq,
            infos=helpers.stack_tensor_dict_list(info_seq)
        )
        return trajectories

    def _update_distribution(self, trajectories):
        """
        Update current control distribution using 
        rollout trajectories
        
        Parameters
        -----------
        trajectories : dict
            Rollout trajectories. Contains the following fields
            observations : np.ndarray ()
                observations along rollouts
            actions : np.ndarray 
                actions sampled from control distribution along rollouts
            costs : np.ndarray 
                step costs along rollouts
            dones : np.ndarray
                bool signalling end of episode
        """
        # costs = trajectories["costs"].copy()
        # actions = trajectories["actions"].copy()
        # Q = cost_to_go(costs, self.gamma_seq)
        # best_id = np.argmin(Q, axis = 0)[0]
        # self.mean_action = (1.0 - self.step_size) * self.mean_action +\
        #                     self.step_size * actions[best_id]

        # Q = cost_to_go(trajectories['costs'], self.gamma_seq)
        # log_prob_to_go = cost_to_go(trajectories['log_probs'], self.gamma_seq)
        # trajectories["cost_to_go"] = Q
        # trajectories["log_prob_to_go"] = log_prob_to_go
        # for k in trajectories.keys():
        #     if k is not 'infos':
        #         trajectories[k] = np.concatenate(trajectories[k], axis=0)
        self.memory.push(trajectories)
        self.update_parameters(True)

    def update_parameters(self, suppress_stdout=True):
        state_batch, action_batch, reward_batch, next_state_batch, mask_batch = self.memory.sample(batch_size=self.train_batch_size)

        state_batch = torch.FloatTensor(state_batch)#.to(self.device)
        next_state_batch = torch.FloatTensor(next_state_batch)#.to(self.device)
        action_batch = torch.FloatTensor(action_batch)#.to(self.device)
        reward_batch = torch.FloatTensor(reward_batch).unsqueeze(1)#.to(self.device).unsqueeze(1)
        mask_batch = torch.FloatTensor(mask_batch).unsqueeze(1)#.to(self.device).unsqueeze(1)

        with torch.no_grad():
            next_state_action, next_state_log_pi, _ = self.policy.sample(next_state_batch)
            qf1_next_target, qf2_next_target = self.critic_target(next_state_batch, next_state_action)
            min_qf_next_target = torch.min(qf1_next_target, qf2_next_target) - self.alpha * next_state_log_pi
            next_q_value = reward_batch + mask_batch * self.gamma * (min_qf_next_target)
        qf1, qf2 = self.critic(state_batch, action_batch)  # Two Q-functions to mitigate positive bias in the policy improvement step
        qf1_loss = F.mse_loss(qf1, next_q_value)  # JQ = 𝔼(st,at)~D[0.5(Q1(st,at) - r(st,at) - γ(𝔼st+1~p[V(st+1)]))^2]
        qf2_loss = F.mse_loss(qf2, next_q_value)  # JQ = 𝔼(st,at)~D[0.5(Q1(st,at) - r(st,at) - γ(𝔼st+1~p[V(st+1)]))^2]
        qf_loss = qf1_loss + qf2_loss

        self.critic_optim.zero_grad()
        qf_loss.backward()
        self.critic_optim.step()

        pi, log_pi, _ = self.policy.sample(state_batch)

        qf1_pi, qf2_pi = self.critic(state_batch, pi)
        min_qf_pi = torch.min(qf1_pi, qf2_pi)

        policy_loss = ((self.alpha * log_pi) - min_qf_pi).mean() # Jπ = 𝔼st∼D,εt∼N[α * logπ(f(εt;st)|st) − Q(st,f(εt;st))]

        self.policy_optim.zero_grad()
        policy_loss.backward()
        self.policy_optim.step()

        if self.automatic_entropy_tuning:
            alpha_loss = -(self.log_alpha * (log_pi + self.target_entropy).detach()).mean()

            self.alpha_optim.zero_grad()
            alpha_loss.backward()
            self.alpha_optim.step()

            self.alpha = self.log_alpha.exp()
            alpha_tlogs = self.alpha.clone() # For TensorboardX logs
        else:
            alpha_loss = torch.tensor(0.).to(self.device)
            alpha_tlogs = torch.tensor(self.alpha) # For TensorboardX logs


        if self.num_steps % self.target_update_interval == 0:
            soft_update(self.critic_target, self.critic, self.tau)

        return qf1_loss.item(), qf2_loss.item(), policy_loss.item(), alpha_loss.item(), alpha_tlogs.item()

    def get_agent_infos(self):
        return None

    # def fit(self, traj_data, suppress_stdout=True):
    #     epoch_losses_qfunc = np.zeros(shape=(self.num_epochs,))
    #     epoch_losses_policy = np.zeros(shape=(self.num_epochs,))
    #     for i in range(self.num_epochs):
    #         batch_losses_qfunc = []
    #         batch_losses_policy = []
    #         for batch_data in self.train_batches(traj_data, self.train_batch_size, 1):
    #             obs_batch = torch.FloatTensor(batch_data['obs_batch'])
    #             next_obs_batch = torch.FloatTensor(batch_data['next_obs_batch'])
    #             action_batch = torch.FloatTensor(batch_data['action_batch'])
    #             reward_batch = -1.0 * torch.FloatTensor(batch_data['cost_batch'])
    #             done_batch = torch.FloatTensor(batch_data['done_batch'])
    #             #Update Q-function
    #             #Generate Q-targets
    #             with torch.no_grad():
    #                 next_action, next_log_pi, _ = self.policy.sample(next_obs_batch)
    #                 # qf1_next_target, qf2_next_target = self.qfunc(next_obs_batch, next_obs_action)
    #                 qf1_next_target, qf2_next_target = self.qfunc_target(next_obs_batch, next_action)
    #                 min_qf_next_target = torch.min(qf1_next_target, qf2_next_target) - self.lam * next_log_pi
    #                 next_q_value = reward_batch + self.gamma * (min_qf_next_target) 

    #             qf1, qf2 = self.qfunc(obs_batch, action_batch)  # Two Q-functions to mitigate positive bias in the policy improvement step
    #             qf1_loss = F.mse_loss(qf1, next_q_value)  # JQ = 𝔼(st,at)~D[0.5(Q1(st,at) - r(st,at) - γ(𝔼st+1~p[V(st+1)]))^2]
    #             qf2_loss = F.mse_loss(qf2, next_q_value)  # JQ = 𝔼(st,at)~D[0.5(Q1(st,at) - r(st,at) - γ(𝔼st+1~p[V(st+1)]))^2]
    #             qf_loss = qf1_loss + qf2_loss

    #             self.qfunc_optim.zero_grad()
    #             qf_loss.backward()
    #             self.qfunc_optim.step()

    #             pi, log_pi, _ = self.policy.sample(obs_batch)

    #             qf1_pi, qf2_pi = self.qfunc(obs_batch, pi)
    #             min_qf_pi = torch.min(qf1_pi, qf2_pi)

    #             policy_loss = ((self.lam * log_pi) - min_qf_pi).mean() # Jπ = 𝔼st∼D,εt∼N[α * logπ(f(εt;st)|st) − Q(st,f(εt;st))]

    #             self.policy_optim.zero_grad()
    #             policy_loss.backward()
    #             self.policy_optim.step()

    #             batch_losses_qfunc.append(qf_loss.item())
    #             batch_losses_policy.append(policy_loss.item())
            
    #         epoch_losses_qfunc[i] = np.average(batch_losses_qfunc)
    #         epoch_losses_policy[i] = np.average(batch_losses_policy)

    #         if not suppress_stdout:
    #             print('Training epoch = {0}, Avg. Policy Loss = {1}, Avg. Qfunc Loss = {2}'.format(i, 
    #                                                                                         epoch_losses_policy[i],
    #                                                                                         epoch_losses_qfunc[i]))
        
    #     #update target network
    #     if self.num_steps%self.update_target_steps == 0:
    #         self.qfunc_target.load_state_dict(self.qfunc.state_dict())
        
    #     return epoch_losses_policy, epoch_losses_qfunc

    
    # def train_batches(self, trajectories, batch_size, ensemble_size):
    #     """

    #     """
    #     #unpack data
    #     observations = trajectories["observations"]; actions = trajectories["actions"]
    #     costs = trajectories["costs"]; log_probs = trajectories["log_probs"]; 
    #     cost_to_go = trajectories["cost_to_go"]; log_prob_to_go = trajectories["log_prob_to_go"]
    #     dones = trajectories["dones"]; next_observations = trajectories["next_observations"]
    #     num, d_obs = observations.shape
    #     indices = [self.np_random.permutation(range(num)) for _ in range(ensemble_size)]
    #     indices = np.stack(indices).T
    #     for i in range(0, num, batch_size):
    #         j = min(num, i + batch_size)

    #         if (j - i) < batch_size and i != 0:
    #             # drop incomplete last batch
    #             return

    #         batch_size = j - i

    #         batch_indices = indices[i:j]
    #         batch_indices = batch_indices.flatten()

    #         obs_batch      = observations[batch_indices]
    #         action_batch   = actions[batch_indices]
    #         cost_batch     = costs[batch_indices]
    #         log_prob_batch = log_probs[batch_indices]
    #         cost_to_go_batch = cost_to_go[batch_indices]
    #         log_prob_to_go_batch = log_prob_to_go[batch_indices]
    #         done_batch     = dones[batch_indices] 
    #         next_obs_batch = next_observations[batch_indices]

    #         if ensemble_size>1:shape=[ensemble_size, batch_size]
    #         else: shape= [batch_size]          

    #         obs_batch      = obs_batch.reshape(*shape, d_obs)
    #         action_batch   = action_batch.reshape(*shape, 1)
    #         cost_batch     = cost_batch.reshape(*shape, 1)
    #         log_prob_batch = log_prob_batch.reshape(*shape, 1)
    #         cost_to_go_batch = cost_to_go_batch.reshape(*shape, 1)
    #         log_prob_to_go_batch = log_prob_to_go_batch.reshape(*shape, 1)
    #         done_batch     = done_batch.reshape(*shape, 1)
    #         next_obs_batch = next_obs_batch.reshape(*shape, d_obs)
            
    #         yield {'obs_batch': obs_batch, 
    #                 'action_batch': action_batch,
    #                 'next_obs_batch': next_obs_batch,
    #                 'cost_batch': cost_batch,
    #                 'cost_to_go_batch': cost_to_go_batch,
    #                 'log_prob_batch': log_prob_batch,
    #                 'log_prob_to_go_batch': log_prob_to_go_batch,
    #                 'done_batch': done_batch,
    #                 }

    def _shift(self):
        """
            Predict good parameters for the next time step by
            shifting the mean forward one step
        """
        # self.mean_action[:-1] = self.mean_action[1:]
        # if self.base_action == 'random':
        #     self.mean_action[-1] = np.random.normal(0, self.init_cov, self.d_action)
        # elif self.base_action == 'null':
        #     self.mean_action[-1] = np.zeros((self.d_action, ))
        # elif self.base_action == 'repeat':
        #     self.mean_action[-1] = self.mean_action[-2]
        # else:
        #     raise NotImplementedError("invalid option for base action during shift")
        pass

    def reset(self):
        """
        Reset the controller
        """
        self.num_steps = 0
        self.mean_action = np.zeros(shape=(self.horizon, self.d_action))
        self.cov_action = np.diag(self.init_cov)

    def _calc_val(self, cost_seq, act_seq):
        """
        Calculate value of state given 
        rollouts from a policy
        """
        pass

    def check_convergence(self):
        """
        Checks if controller has converged
        Returns False by default
        """
        return False
