import torch
from onpolicy.algorithms.r_mappo.algorithm.r_actor_critic import R_Actor, R_Critic, R_Attack
from onpolicy.algorithms.utils.act import ACTLayer
from onpolicy.utils.util import update_linear_schedule
from gym.spaces import Discrete
import numpy as np

class R_MAPPOPolicy:
    """
    MAPPO Policy  class. Wraps actor and critic networks to compute actions and value function predictions.

    :param args: (argparse.Namespace) arguments containing relevant model and policy information.
    :param obs_space: (gym.Space) observation space.
    :param cent_obs_space: (gym.Space) value function input space (centralized input for MAPPO, decentralized for IPPO).
    :param action_space: (gym.Space) action space.
    :param device: (torch.device) specifies the device to run on (cpu/gpu).
    """

    def __init__(self, args, obs_space, cent_obs_space, act_space, device=torch.device("cpu")):
        self.device = device
        self.lr = args.lr
        self.critic_lr = args.critic_lr
        self.opti_eps = args.opti_eps
        self.weight_decay = args.weight_decay

        self.obs_space = obs_space
        self.share_obs_space = cent_obs_space
        self.act_space = act_space
        self.attack_space = Discrete(9)

        self.act = ACTLayer(act_space, args.hidden_size, args.use_orthogonal, args.gain, obs_space[2][0], device)
        self.actor = R_Actor(args, self.obs_space, self.act_space, self.device)


        self.attack = R_Attack(args, self.obs_space, self.attack_space, self.device)
        
        self.critic = R_Critic(args, self.share_obs_space, self.device)

        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(),
                                                lr=self.lr, eps=self.opti_eps,
                                                weight_decay=self.weight_decay)
        self.act_optimizer = torch.optim.Adam(self.act.parameters(),
                                                lr=self.lr, eps=self.opti_eps,
                                                weight_decay=self.weight_decay)

        '''
        self.attack_optimizer = torch.optim.Adam(self.attack.parameters(),
                                                lr=self.lr, eps=self.opti_eps,
                                                weight_decay=self.weight_decay)
        '''
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(),
                                                 lr=self.critic_lr,
                                                 eps=self.opti_eps,
                                                 weight_decay=self.weight_decay)

        self.act.to(device)
    def lr_decay(self, episode, episodes):
        """
        Decay the actor and critic learning rates.
        :param episode: (int) current training episode.
        :param episodes: (int) total number of training episodes.
        """
        update_linear_schedule(self.actor_optimizer, episode, episodes, self.lr)
        update_linear_schedule(self.critic_optimizer, episode, episodes, self.critic_lr)

    def get_actions(self, cent_obs, obs, rnn_states_actor, rnn_states_critic, masks, available_actions=None,
                    deterministic=False):
        """
        Compute actions and value function predictions for the given inputs.
        :param cent_obs (np.ndarray): centralized input to the critic.
        :param obs (np.ndarray): local agent inputs to the actor.
        :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor.
        :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic.
        :param masks: (np.ndarray) denotes points at which RNN states should be reset.
        :param available_actions: (np.ndarray) denotes which actions are available to agent
                                  (if None, all actions available)
        :param deterministic: (bool) whether the action should be mode of distribution or should be sampled.

        :return values: (torch.Tensor) value function predictions.
        :return actions: (torch.Tensor) actions to take.
        :return action_log_probs: (torch.Tensor) log probabilities of chosen actions.
        :return rnn_states_actor: (torch.Tensor) updated actor network RNN states.
        :return rnn_states_critic: (torch.Tensor) updated critic network RNN states.
        """
        #split0 = available_actions[:,0:7]
        #split1 = np.zeros((21,3))
        #split1 = np.hstack((split0,split1))
        #split1 = np.zeros((21,9))
        #split1 = np.zeros((21,3))
        #split1 = np.hstack((split0,split1))
        #split0 = available_actions
        #split1 = split0
        '''
        actions, action_log_probs, rnn_states_actor = self.actor(obs,
                                                                 rnn_states_actor,
                                                                 masks,
                                                                 available_actions,
                                                                 deterministic)
        '''
        actor_features, ava_actions, deterministic, rnn_states_actor  = self.actor(obs,  
                                                                 rnn_states_actor,
                                                                 masks,
                                                                 available_actions,
                                                                 deterministic)
        actions, action_log_probs = self.act(actor_features, ava_actions, deterministic)
        '''
        attacks, attack_log_probs, rnn_states_attack = self.attack(obs,
                                                                 rnn_states_actor,
                                                                 masks,
                                                                 split1,
                                                                 deterministic)
        count = 0
        for act in actions:
            if act >= 6:
                 actions[count] = attacks[count] + actions[count]
                 #actions[count] = attacks[count]
                 #action_log_probs[count] = action_log_probs_1[count]
            count += 1
        count = 0
        for act in actions:
           actions[count] = attacks[count]
           count += 1
        '''
        attacks = actions
        attack_log_probs = action_log_probs
        values, rnn_states_critic = self.critic(cent_obs, rnn_states_critic, masks)
        return values, actions, action_log_probs, rnn_states_actor, rnn_states_critic, attacks, attack_log_probs

    def get_values(self, cent_obs, rnn_states_critic, masks):
        """
        Get value function predictions.
        :param cent_obs (np.ndarray): centralized input to the critic.
        :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic.
        :param masks: (np.ndarray) denotes points at which RNN states should be reset.

        :return values: (torch.Tensor) value function predictions.
        """
        values, _ = self.critic(cent_obs, rnn_states_critic, masks)
        return values

    def evaluate_actions(self, cent_obs, obs, rnn_states_actor, rnn_states_critic, action, masks,
                         available_actions=None, active_masks=None):
        """
        Get action logprobs / entropy and value function predictions for actor update.
        :param cent_obs (np.ndarray): centralized input to the critic.
        :param obs (np.ndarray): local agent inputs to the actor.
        :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor.
        :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic.
        :param action: (np.ndarray) actions whose log probabilites and entropy to compute.
        :param masks: (np.ndarray) denotes points at which RNN states should be reset.
        :param available_actions: (np.ndarray) denotes which actions are available to agent
                                  (if None, all actions available)
        :param active_masks: (torch.Tensor) denotes whether an agent is active or dead.

        :return values: (torch.Tensor) value function predictions.
        :return action_log_probs: (torch.Tensor) log probabilities of the input actions.
        :return dist_entropy: (torch.Tensor) action distribution entropy for the given inputs.
        """
        #action_log_probs, dist_entropy = self.actor.evaluate_actions(obs,
        actor_features,action,ava_actions,active_masks = self.actor.evaluate_actions(obs,
                                                                     rnn_states_actor,
                                                                     action,
                                                                     masks,
                                                                     available_actions,
                                                                     active_masks)
        action_log_probs, dist_entropy = self.act.evaluate_actions(actor_features,
                                                                   action, ava_actions,
                                                                   active_masks)
        values, _ = self.critic(cent_obs, rnn_states_critic, masks)
        return values, action_log_probs, dist_entropy

    def evaluate_attacks(self, cent_obs, obs, rnn_states_actor, rnn_states_critic, action, masks,
                         available_actions=None, active_masks=None):
        """
        Get action logprobs / entropy and value function predictions for actor update.
        :param cent_obs (np.ndarray): centralized input to the critic.
        :param obs (np.ndarray): local agent inputs to the actor.
        :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor.
        :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic.
        :param action: (np.ndarray) actions whose log probabilites and entropy to compute.
        :param masks: (np.ndarray) denotes points at which RNN states should be reset.
        :param available_actions: (np.ndarray) denotes which actions are available to agent
                                  (if None, all actions available)
        :param active_masks: (torch.Tensor) denotes whether an agent is active or dead.

        :return values: (torch.Tensor) value function predictions.
        :return action_log_probs: (torch.Tensor) log probabilities of the input actions.
        :return dist_entropy: (torch.Tensor) action distribution entropy for the given inputs.
        """
        #available_actions=None
        action_log_probs, dist_entropy = self.attack.evaluate_attacks(obs,
                                                                     rnn_states_actor,
                                                                     action,
                                                                     masks,
                                                                     available_actions,
                                                                     active_masks)
        
        values, _ = self.critic(cent_obs, rnn_states_critic, masks)
        #values = None
        return values, action_log_probs, dist_entropy



    def act(self, obs, rnn_states_actor, masks, available_actions=None, deterministic=False):
        """
        Compute actions using the given inputs.
        :param obs (np.ndarray): local agent inputs to the actor.
        :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor.
        :param masks: (np.ndarray) denotes points at which RNN states should be reset.
        :param available_actions: (np.ndarray) denotes which actions are available to agent
                                  (if None, all actions available)
        :param deterministic: (bool) whether the action should be mode of distribution or should be sampled.
        """
        actions, _, rnn_states_actor = self.actor(obs, rnn_states_actor, masks, available_actions, deterministic)
        return actions, rnn_states_actor

import torch
from onpolicy.algorithms.r_mappo.algorithm.r_actor_critic import R_Actor, R_Critic, R_Attack
from onpolicy.algorithms.utils.act import ACTLayer
from onpolicy.utils.util import update_linear_schedule
from gym.spaces import Discrete
import numpy as np

class R_MAPPOPolicy:
    """
    MAPPO Policy  class. Wraps actor and critic networks to compute actions and value function predictions.

    :param args: (argparse.Namespace) arguments containing relevant model and policy information.
    :param obs_space: (gym.Space) observation space.
    :param cent_obs_space: (gym.Space) value function input space (centralized input for MAPPO, decentralized for IPPO).
    :param action_space: (gym.Space) action space.
    :param device: (torch.device) specifies the device to run on (cpu/gpu).
    """

    def __init__(self, args, obs_space, cent_obs_space, act_space, device=torch.device("cpu")):
        self.device = device
        self.lr = args.lr
        self.critic_lr = args.critic_lr
        self.opti_eps = args.opti_eps
        self.weight_decay = args.weight_decay

        self.obs_space = obs_space
        self.share_obs_space = cent_obs_space
        self.act_space = act_space
        self.attack_space = Discrete(9)

        self.act = ACTLayer(act_space, args.hidden_size, args.use_orthogonal, args.gain, obs_space[2][0], device)
        self.actor = R_Actor(args, self.obs_space, self.act_space, self.device)


        self.attack = R_Attack(args, self.obs_space, self.attack_space, self.device)
        
        self.critic = R_Critic(args, self.share_obs_space, self.device)

        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(),
                                                lr=self.lr, eps=self.opti_eps,
                                                weight_decay=self.weight_decay)
        self.act_optimizer = torch.optim.Adam(self.act.parameters(),
                                                lr=self.lr, eps=self.opti_eps,
                                                weight_decay=self.weight_decay)

        '''
        self.attack_optimizer = torch.optim.Adam(self.attack.parameters(),
                                                lr=self.lr, eps=self.opti_eps,
                                                weight_decay=self.weight_decay)
        '''
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(),
                                                 lr=self.critic_lr,
                                                 eps=self.opti_eps,
                                                 weight_decay=self.weight_decay)

        self.act.to(device)
    def lr_decay(self, episode, episodes):
        """
        Decay the actor and critic learning rates.
        :param episode: (int) current training episode.
        :param episodes: (int) total number of training episodes.
        """
        update_linear_schedule(self.actor_optimizer, episode, episodes, self.lr)
        update_linear_schedule(self.critic_optimizer, episode, episodes, self.critic_lr)

    def get_actions(self, cent_obs, obs, rnn_states_actor, rnn_states_critic, masks, available_actions=None,
                    deterministic=False):
        """
        Compute actions and value function predictions for the given inputs.
        :param cent_obs (np.ndarray): centralized input to the critic.
        :param obs (np.ndarray): local agent inputs to the actor.
        :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor.
        :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic.
        :param masks: (np.ndarray) denotes points at which RNN states should be reset.
        :param available_actions: (np.ndarray) denotes which actions are available to agent
                                  (if None, all actions available)
        :param deterministic: (bool) whether the action should be mode of distribution or should be sampled.

        :return values: (torch.Tensor) value function predictions.
        :return actions: (torch.Tensor) actions to take.
        :return action_log_probs: (torch.Tensor) log probabilities of chosen actions.
        :return rnn_states_actor: (torch.Tensor) updated actor network RNN states.
        :return rnn_states_critic: (torch.Tensor) updated critic network RNN states.
        """
        #split0 = available_actions[:,0:7]
        #split1 = np.zeros((21,3))
        #split1 = np.hstack((split0,split1))
        #split1 = np.zeros((21,9))
        #split1 = np.zeros((21,3))
        #split1 = np.hstack((split0,split1))
        #split0 = available_actions
        #split1 = split0
        '''
        actions, action_log_probs, rnn_states_actor = self.actor(obs,
                                                                 rnn_states_actor,
                                                                 masks,
                                                                 available_actions,
                                                                 deterministic)
        '''
        actor_features, ava_actions, deterministic, rnn_states_actor  = self.actor(obs,  
                                                                 rnn_states_actor,
                                                                 masks,
                                                                 available_actions,
                                                                 deterministic)
        actions, action_log_probs = self.act(actor_features, ava_actions, deterministic)
        '''
        attacks, attack_log_probs, rnn_states_attack = self.attack(obs,
                                                                 rnn_states_actor,
                                                                 masks,
                                                                 split1,
                                                                 deterministic)
        count = 0
        for act in actions:
            if act >= 6:
                 actions[count] = attacks[count] + actions[count]
                 #actions[count] = attacks[count]
                 #action_log_probs[count] = action_log_probs_1[count]
            count += 1
        count = 0
        for act in actions:
           actions[count] = attacks[count]
           count += 1
        '''
        attacks = actions
        attack_log_probs = action_log_probs
        values, rnn_states_critic = self.critic(cent_obs, rnn_states_critic, masks)
        return values, actions, action_log_probs, rnn_states_actor, rnn_states_critic, attacks, attack_log_probs

    def get_values(self, cent_obs, rnn_states_critic, masks):
        """
        Get value function predictions.
        :param cent_obs (np.ndarray): centralized input to the critic.
        :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic.
        :param masks: (np.ndarray) denotes points at which RNN states should be reset.

        :return values: (torch.Tensor) value function predictions.
        """
        values, _ = self.critic(cent_obs, rnn_states_critic, masks)
        return values

    def evaluate_actions(self, cent_obs, obs, rnn_states_actor, rnn_states_critic, action, masks,
                         available_actions=None, active_masks=None):
        """
        Get action logprobs / entropy and value function predictions for actor update.
        :param cent_obs (np.ndarray): centralized input to the critic.
        :param obs (np.ndarray): local agent inputs to the actor.
        :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor.
        :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic.
        :param action: (np.ndarray) actions whose log probabilites and entropy to compute.
        :param masks: (np.ndarray) denotes points at which RNN states should be reset.
        :param available_actions: (np.ndarray) denotes which actions are available to agent
                                  (if None, all actions available)
        :param active_masks: (torch.Tensor) denotes whether an agent is active or dead.

        :return values: (torch.Tensor) value function predictions.
        :return action_log_probs: (torch.Tensor) log probabilities of the input actions.
        :return dist_entropy: (torch.Tensor) action distribution entropy for the given inputs.
        """
        #action_log_probs, dist_entropy = self.actor.evaluate_actions(obs,
        actor_features,action,ava_actions,active_masks = self.actor.evaluate_actions(obs,
                                                                     rnn_states_actor,
                                                                     action,
                                                                     masks,
                                                                     available_actions,
                                                                     active_masks)
        action_log_probs, dist_entropy = self.act.evaluate_actions(actor_features,
                                                                   action, ava_actions,
                                                                   active_masks)
        values, _ = self.critic(cent_obs, rnn_states_critic, masks)
        return values, action_log_probs, dist_entropy

    def evaluate_attacks(self, cent_obs, obs, rnn_states_actor, rnn_states_critic, action, masks,
                         available_actions=None, active_masks=None):
        """
        Get action logprobs / entropy and value function predictions for actor update.
        :param cent_obs (np.ndarray): centralized input to the critic.
        :param obs (np.ndarray): local agent inputs to the actor.
        :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor.
        :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic.
        :param action: (np.ndarray) actions whose log probabilites and entropy to compute.
        :param masks: (np.ndarray) denotes points at which RNN states should be reset.
        :param available_actions: (np.ndarray) denotes which actions are available to agent
                                  (if None, all actions available)
        :param active_masks: (torch.Tensor) denotes whether an agent is active or dead.

        :return values: (torch.Tensor) value function predictions.
        :return action_log_probs: (torch.Tensor) log probabilities of the input actions.
        :return dist_entropy: (torch.Tensor) action distribution entropy for the given inputs.
        """
        #available_actions=None
        action_log_probs, dist_entropy = self.attack.evaluate_attacks(obs,
                                                                     rnn_states_actor,
                                                                     action,
                                                                     masks,
                                                                     available_actions,
                                                                     active_masks)
        
        values, _ = self.critic(cent_obs, rnn_states_critic, masks)
        #values = None
        return values, action_log_probs, dist_entropy



    def act(self, obs, rnn_states_actor, masks, available_actions=None, deterministic=False):
        """
        Compute actions using the given inputs.
        :param obs (np.ndarray): local agent inputs to the actor.
        :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor.
        :param masks: (np.ndarray) denotes points at which RNN states should be reset.
        :param available_actions: (np.ndarray) denotes which actions are available to agent
                                  (if None, all actions available)
        :param deterministic: (bool) whether the action should be mode of distribution or should be sampled.
        """
        actions, _, rnn_states_actor = self.actor(obs, rnn_states_actor, masks, available_actions, deterministic)
        return actions, rnn_states_actor

MAPPO Policy  class. Wraps actor and critic networks to compute actions and value function predictions.      :param args: (argparse.Namespace) arguments containing relevant model and policy information.     :param obs_space: (gym.Space) observation space.     :param cent_obs_space: (gym.Space) value function input space (centralized input for MAPPO, decentralized for IPPO).     :param action_space: (gym.Space) action space.     :param device: (torch.device) specifies the device to run on (cpu/gpu). self.attack_optimizer = torch.optim.Adam(self.attack.parameters(),                                                 lr=self.lr, eps=self.opti_eps,                                                 weight_decay=self.weight_decay) Decay the actor and critic learning rates.         :param episode: (int) current training episode.         :param episodes: (int) total number of training episodes. Compute actions and value function predictions for the given inputs.         :param cent_obs (np.ndarray): centralized input to the critic.         :param obs (np.ndarray): local agent inputs to the actor.         :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor.         :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic.         :param masks: (np.ndarray) denotes points at which RNN states should be reset.         :param available_actions: (np.ndarray) denotes which actions are available to agent                                   (if None, all actions available)         :param deterministic: (bool) whether the action should be mode of distribution or should be sampled.          :return values: (torch.Tensor) value function predictions.         :return actions: (torch.Tensor) actions to take.         :return action_log_probs: (torch.Tensor) log probabilities of chosen actions.         :return rnn_states_actor: (torch.Tensor) updated actor network RNN states.         :return rnn_states_critic: (torch.Tensor) updated critic network RNN states. #split0 = available_actions[:,0:7] #split1 = np.zeros((21,3)) #split1 = np.hstack((split0,split1)) #split1 = np.zeros((21,9)) #split1 = np.zeros((21,3)) #split1 = np.hstack((split0,split1)) #split0 = available_actions #split1 = split0 actions, action_log_probs, rnn_states_actor = self.actor(obs,                                                                  rnn_states_actor,                                                                  masks,                                                                  available_actions,                                                                  deterministic) attacks, attack_log_probs, rnn_states_attack = self.attack(obs,                                                                  rnn_states_actor,                                                                  masks,                                                                  split1,                                                                  deterministic)         count = 0         for act in actions:             if act >= 6:                  actions[count] = attacks[count] + actions[count]                  #actions[count] = attacks[count]                  #action_log_probs[count] = action_log_probs_1[count]             count += 1         count = 0         for act in actions:            actions[count] = attacks[count]            count += 1 Get value function predictions.         :param cent_obs (np.ndarray): centralized input to the critic.         :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic.         :param masks: (np.ndarray) denotes points at which RNN states should be reset.          :return values: (torch.Tensor) value function predictions. Get action logprobs / entropy and value function predictions for actor update.         :param cent_obs (np.ndarray): centralized input to the critic.         :param obs (np.ndarray): local agent inputs to the actor.         :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor.         :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic.         :param action: (np.ndarray) actions whose log probabilites and entropy to compute.         :param masks: (np.ndarray) denotes points at which RNN states should be reset.         :param available_actions: (np.ndarray) denotes which actions are available to agent                                   (if None, all actions available)         :param active_masks: (torch.Tensor) denotes whether an agent is active or dead.          :return values: (torch.Tensor) value function predictions.         :return action_log_probs: (torch.Tensor) log probabilities of the input actions.         :return dist_entropy: (torch.Tensor) action distribution entropy for the given inputs. #action_log_probs, dist_entropy = self.actor.evaluate_actions(obs, Get action logprobs / entropy and value function predictions for actor update.         :param cent_obs (np.ndarray): centralized input to the critic.         :param obs (np.ndarray): local agent inputs to the actor.         :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor.         :param rnn_states_critic: (np.ndarray) if critic is RNN, RNN states for critic.         :param action: (np.ndarray) actions whose log probabilites and entropy to compute.         :param masks: (np.ndarray) denotes points at which RNN states should be reset.         :param available_actions: (np.ndarray) denotes which actions are available to agent                                   (if None, all actions available)         :param active_masks: (torch.Tensor) denotes whether an agent is active or dead.          :return values: (torch.Tensor) value function predictions.         :return action_log_probs: (torch.Tensor) log probabilities of the input actions.         :return dist_entropy: (torch.Tensor) action distribution entropy for the given inputs. #available_actions=None #values = None Compute actions using the given inputs.         :param obs (np.ndarray): local agent inputs to the actor.         :param rnn_states_actor: (np.ndarray) if actor is RNN, RNN states for actor.         :param masks: (np.ndarray) denotes points at which RNN states should be reset.         :param available_actions: (np.ndarray) denotes which actions are available to agent                                   (if None, all actions available)         :param deterministic: (bool) whether the action should be mode of distribution or should be sampled.