import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import copy
import math

import utils
import hydra


class Encoder(nn.Module):
    """Convolutional encoder for image-based observations."""
    def __init__(self, obs_shape, feature_dim):
        super().__init__()

        assert len(obs_shape) == 3
        self.num_layers = 4
        self.num_filters = 32
        self.output_dim = 35
        self.output_logits = False
        self.feature_dim = feature_dim

        self.convs = nn.ModuleList([
            nn.Conv2d(obs_shape[0], self.num_filters, 3, stride=2),
            nn.Conv2d(self.num_filters, self.num_filters, 3, stride=1),
            nn.Conv2d(self.num_filters, self.num_filters, 3, stride=1),
            nn.Conv2d(self.num_filters, self.num_filters, 3, stride=1)
        ])

        self.head = nn.Sequential(
            nn.Linear(self.num_filters * 35 * 35, self.feature_dim),
            nn.LayerNorm(self.feature_dim))

        self.outputs = dict()

    def forward_conv(self, obs):
        obs = obs / 255.
        self.outputs['obs'] = obs

        conv = torch.relu(self.convs[0](obs))
        self.outputs['conv1'] = conv

        for i in range(1, self.num_layers):
            conv = torch.relu(self.convs[i](conv))
            self.outputs['conv%s' % (i + 1)] = conv

        h = conv.view(conv.size(0), -1)
        return h

    def forward(self, obs, detach=False):
        h = self.forward_conv(obs)

        if detach:
            h = h.detach()

        out = self.head(h)
        if not self.output_logits:
            out = torch.tanh(out)

        self.outputs['out'] = out

        return out

    def copy_conv_weights_from(self, source):
        """Tie convolutional layers"""
        for i in range(self.num_layers):
            utils.tie_weights(src=source.convs[i], trg=self.convs[i])

    def log(self, logger, step):
        for k, v in self.outputs.items():
            logger.log_histogram(f'train_encoder/{k}_hist', v, step)
            if len(v.shape) > 2:
                logger.log_image(f'train_encoder/{k}_img', v[0], step)

        for i in range(self.num_layers):
            logger.log_param(f'train_encoder/conv{i + 1}', self.convs[i], step)


class Actor(nn.Module):
    """torch.distributions implementation of an diagonal Gaussian policy."""
    def __init__(self, encoder_cfg, action_shape, hidden_dim, hidden_depth,
                 log_std_bounds):
        super().__init__()

        self.encoder = hydra.utils.instantiate(encoder_cfg)

        self.log_std_bounds = log_std_bounds
        self.trunk = utils.mlp(self.encoder.feature_dim, hidden_dim,
                               2 * action_shape[0], hidden_depth)

        self.outputs = dict()
        self.apply(utils.weight_init)

    def forward(self, obs, detach_encoder=False):
        obs = self.encoder(obs, detach=detach_encoder)

        mu, log_std = self.trunk(obs).chunk(2, dim=-1)

        # constrain log_std inside [log_std_min, log_std_max]
        log_std = torch.tanh(log_std)
        log_std_min, log_std_max = self.log_std_bounds
        log_std = log_std_min + 0.5 * (log_std_max - log_std_min) * (log_std +
                                                                     1)
        std = log_std.exp()

        self.outputs['mu'] = mu
        self.outputs['std'] = std

        dist = utils.SquashedNormal(mu, std)
        return dist

    def log(self, logger, step):
        for k, v in self.outputs.items():
            logger.log_histogram(f'train_actor/{k}_hist', v, step)

        for i, m in enumerate(self.trunk):
            if type(m) == nn.Linear:
                logger.log_param(f'train_actor/fc{i}', m, step)


class Critic(nn.Module):
    """Critic network, employes double Q-learning."""
    def __init__(self, encoder_cfg, action_shape, hidden_dim, hidden_depth):
        super().__init__()

        self.encoder = hydra.utils.instantiate(encoder_cfg)

        self.Q1 = utils.mlp(self.encoder.feature_dim + action_shape[0],
                            hidden_dim, 1, hidden_depth)
        self.Q2 = utils.mlp(self.encoder.feature_dim + action_shape[0],
                            hidden_dim, 1, hidden_depth)

        self.outputs = dict()
        self.apply(utils.weight_init)

    def forward(self, obs, action, detach_encoder=False):
        assert obs.size(0) == action.size(0)
        obs = self.encoder(obs, detach=detach_encoder)

        obs_action = torch.cat([obs, action], dim=-1)
        q1 = self.Q1(obs_action)
        q2 = self.Q2(obs_action)

        self.outputs['q1'] = q1
        self.outputs['q2'] = q2

        return q1, q2

    def log(self, logger, step):
        self.encoder.log(logger, step)

        for k, v in self.outputs.items():
            logger.log_histogram(f'train_critic/{k}_hist', v, step)

        assert len(self.Q1) == len(self.Q2)
        for i, (m1, m2) in enumerate(zip(self.Q1, self.Q2)):
            assert type(m1) == type(m2)
            if type(m1) is nn.Linear:
                logger.log_param(f'train_critic/q1_fc{i}', m1, step)
                logger.log_param(f'train_critic/q2_fc{i}', m2, step)


class DRQAgent(object):
    """Data regularized Q: actor-critic method for learning from pixels."""
    def __init__(self, obs_shape, action_shape, action_range, device,
                 encoder_cfg, critic_cfg, actor_cfg, discount,
                 init_temperature, lr, actor_update_frequency, critic_tau,
                 critic_target_update_frequency, batch_size):
        self.action_range = action_range
        self.device = device
        self.discount = discount
        self.critic_tau = critic_tau
        self.actor_update_frequency = actor_update_frequency
        self.critic_target_update_frequency = critic_target_update_frequency
        self.batch_size = batch_size

        self.actor = hydra.utils.instantiate(actor_cfg).to(self.device)

        self.critic = hydra.utils.instantiate(critic_cfg).to(self.device)
        self.critic_target = hydra.utils.instantiate(critic_cfg).to(
            self.device)
        self.critic_target.load_state_dict(self.critic.state_dict())

        # tie conv layers between actor and critic
        self.actor.encoder.copy_conv_weights_from(self.critic.encoder)

        self.log_alpha = torch.tensor(np.log(init_temperature)).to(device)
        self.log_alpha.requires_grad = True
        # set target entropy to -|A|
        self.target_entropy = -action_shape[0]

        # optimizers
        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=lr)
        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(),
                                                 lr=lr)
        self.log_alpha_optimizer = torch.optim.Adam([self.log_alpha], lr=lr)

        self.train()
        self.critic_target.train()

    def train(self, training=True):
        self.training = training
        self.actor.train(training)
        self.critic.train(training)

    @property
    def alpha(self):
        return self.log_alpha.exp()

    def act(self, obs, sample=False):
        dist = self.actor(obs)
        action = dist.sample() if sample else dist.mean
        action = action.clamp(*self.action_range)
        assert action.ndim == 2 and action.shape[0] == 1
        return utils.to_np(action[0])
    
    def test_clean(self, obs, if_save=False, sample=False):
        obs = torch.FloatTensor(obs).to(self.device)
        obs = obs.unsqueeze(0)
        dist = self.actor(obs)
        action = dist.sample() if sample else dist.mean
        action = action.clamp(*self.action_range)
        assert action.ndim == 2 and action.shape[0] == 1
        obs_save = None
        if if_save:
            obs_save = obs.squeeze(0).cpu().numpy()
        return utils.to_np(action[0]), obs_save
    
    def adv_pgd(self, obs, if_save=False, sample=False, epsilon=8, alpha=2, num_steps=10):
        obs = torch.FloatTensor(obs).to(self.device)
        obs = obs.unsqueeze(0)
        ori_obs = obs.clone().detach()
        
        delta = torch.zeros_like(obs).to(self.device)
        delta.uniform_(-epsilon, epsilon)
        delta = torch.clamp(obs + delta, 0, 255) - obs
        delta.requires_grad = True

        for _ in range(num_steps):
            adv_obs = obs + delta
            dist = self.actor(adv_obs, detach_encoder=False)
            action = dist.rsample()
            log_prob = dist.log_prob(action).sum(-1, keepdim=True)

            q1, q2 = self.critic(adv_obs, action, detach_encoder=False)
            q = torch.min(q1, q2)
            loss = (self.alpha.detach() * log_prob - q).mean()

            grad = torch.autograd.grad(loss, delta, retain_graph=False, create_graph=False)[0]
            delta.data = delta.data + alpha * grad.sign()
            delta.data = torch.clamp(delta.data, -epsilon, epsilon)
            delta.data = torch.clamp(obs + delta.data, 0, 255) - obs

        adv_obs = (obs + delta).detach()

        dist = self.actor(adv_obs)
        action = dist.sample() if sample else dist.mean
        action = action.clamp(*self.action_range)

        adv_obs_save = None
        if if_save:
            adv_obs_save = adv_obs.squeeze(0).cpu().numpy()

        return utils.to_np(action[0]), adv_obs_save


    def simba_attack(self, obs, if_save=False, sample=False, epsilon=8, num_steps=100):
        obs = torch.FloatTensor(obs).to(self.device)
        obs = obs.unsqueeze(0)  
        obs = obs.clone().detach().to(self.device)
        obs.requires_grad = True
        
        eps_ = epsilon

        for _ in range(num_steps):
            grad_sign = torch.zeros_like(obs).to(self.device)
            for dim in range(obs.shape[1]): # 9
                perturbed_obs_pos = obs.clone()
                perturbed_obs_pos[:, dim] += eps_
                perturbed_obs_neg = obs.clone()
                perturbed_obs_neg[:, dim] -= eps_

                action_pos = self.actor(perturbed_obs_pos).rsample()
                action_neg = self.actor(perturbed_obs_neg).rsample()

                q1_pos, q2_pos = self.critic(perturbed_obs_pos, action_pos)
                q1_neg, q2_neg = self.critic(perturbed_obs_neg, action_neg)


                q_min_pos = torch.min(q1_pos, q2_pos)
                q_min_neg = torch.min(q1_neg, q2_neg)

                if q_min_pos < q_min_neg:
                    grad_sign[:, dim] = 1  
                else:
                    grad_sign[:, dim] = -1 

            obs = obs + grad_sign * eps_
            obs = torch.clamp(obs, 0, 255) 

        dist = self.actor(obs)
        action = dist.sample() if sample else dist.mean
        action = action.clamp(*self.action_range)

        adv_obs_save = None
        if if_save:
            adv_obs_save = obs.squeeze(0).cpu().numpy()

        return utils.to_np(action[0]), adv_obs_save


    def adv_cw(self, obs, if_save=False, sample=False, epsilon=8, alpha=1, num_steps=10, c=1e-4):
        obs = torch.FloatTensor(obs).to(self.device)
        obs = obs.unsqueeze(0)
        ori_obs = obs.clone().detach()

        delta = torch.zeros_like(obs).to(self.device)
        delta.uniform_(-epsilon, epsilon)
        delta = torch.clamp(obs + delta, 0, 255) - obs
        delta.requires_grad = True

        for _ in range(num_steps):
            adv_obs = obs + delta
            dist = self.actor(adv_obs, detach_encoder=False)
            action = dist.rsample()

            q1, q2 = self.critic(adv_obs, action, detach_encoder=False)
            q = torch.min(q1, q2)

            l2_loss = F.mse_loss(delta, torch.zeros_like(delta))
            loss = (c * l2_loss - q.mean())

            grad = torch.autograd.grad(loss, delta, retain_graph=False, create_graph=False)[0]
            delta.data = delta.data + alpha * grad.sign()
            delta.data = torch.clamp(delta.data, -epsilon, epsilon)
            delta.data = torch.clamp(obs + delta.data, 0, 255) - obs

        adv_obs = (obs + delta).detach()
        dist = self.actor(adv_obs)
        action = dist.sample() if sample else dist.mean
        action = action.clamp(*self.action_range)

        adv_obs_save = None
        if if_save:
            adv_obs_save = adv_obs.squeeze(0).cpu().numpy()

        return utils.to_np(action[0]), adv_obs_save

    def adv_square(self, obs, if_save=False, sample=False, epsilon=8, num_steps=100, block_size=12):
        obs = torch.FloatTensor(obs).to(self.device).unsqueeze(0)
        adv_obs = obs.clone().detach()
        best_obs = adv_obs.clone().detach()

        with torch.no_grad():
            dist = self.actor(obs)
            action = dist.mean
            q1, q2 = self.critic(obs, action)
            best_score = torch.min(q1, q2).item()

        c, h, w = adv_obs.shape[1:]
        for _ in range(num_steps):
            x = torch.randint(0, h - block_size + 1, (1,)).item()
            y = torch.randint(0, w - block_size + 1, (1,)).item()
            sign = torch.randint(0, 2, (1,)).item() * 2 - 1
            candidate = adv_obs.clone()
            candidate[:, :, x:x+block_size, y:y+block_size] += sign * epsilon
            candidate = torch.clamp(candidate, 0, 255)

            with torch.no_grad():
                dist = self.actor(candidate)
                action = dist.mean
                q1, q2 = self.critic(candidate, action)
                score = torch.min(q1, q2).item()

            if score < best_score:
                best_score = score
                adv_obs = candidate.clone()
                best_obs = candidate.clone()

        dist = self.actor(best_obs)
        action = dist.sample() if sample else dist.mean
        action = action.clamp(*self.action_range)

        adv_obs_save = None
        if if_save:
            adv_obs_save = best_obs.squeeze(0).cpu().numpy()

        return utils.to_np(action[0]), adv_obs_save


    def adv_gan_based(self, obs, advGen, if_save=False, sample=False):
        obs = torch.FloatTensor(obs).to(self.device)
        obs = obs.unsqueeze(0)
        obs = obs / 255.
        eps_ = 8 / 255

        perturbation = advGen(obs)
        adv_obs = torch.clamp(perturbation, -eps_, eps_) + obs
        adv_obs = torch.clamp(adv_obs, 0, 1)
        adv_obs = adv_obs * 255.

        dist = self.actor(adv_obs)
        action = dist.sample() if sample else dist.mean
        action = action.clamp(*self.action_range)
        assert action.ndim == 2 and action.shape[0] == 1
        adv_obs_save = None
        if if_save:
            adv_obs_save = adv_obs.squeeze(0).detach().cpu().numpy()
        return utils.to_np(action[0]), adv_obs_save
    
    def get_adv_obs(self, obs):
        eps_ = 8 / 255 * 255
        alpha_ = 2 / 255 * 255
        steps_ = 10
        random_start_ = True

        obs = torch.FloatTensor(obs).to(self.device)
        obs = obs.unsqueeze(0)
        obs = obs.clone().detach().to(self.device)
        adv_obs = obs.clone().detach()

        if random_start_:
            adv_obs = adv_obs + torch.empty_like(adv_obs).uniform_(
                -eps_, eps_
            )
            adv_obs = torch.clamp(adv_obs, min=0, max=255).detach()

        for _ in range(steps_):
            adv_obs.requires_grad = True

            # Calculate gradient
            dist = self.actor(adv_obs, detach_encoder=False)
            action = dist.rsample()
            log_prob = dist.log_prob(action).sum(-1, keepdim=True)

            actor_Q1, actor_Q2 = self.critic(adv_obs, action, detach_encoder=False)
            actor_Q = torch.min(actor_Q1, actor_Q2)

            actor_loss = (self.alpha.detach() * log_prob - actor_Q).mean()
            grad = torch.autograd.grad(
                actor_loss, adv_obs, retain_graph=False, create_graph=False, allow_unused=True
            )[0]

            adv_obs = adv_obs.detach() + alpha_ * grad.sign()
            delta = torch.clamp(adv_obs - obs, min=-eps_, max=eps_)
            adv_obs = torch.clamp(obs + delta, min=0, max=255).detach()

        adv_obs = adv_obs.squeeze(0).cpu().numpy()
        adv_obs = adv_obs.astype(np.uint8)
        return adv_obs
    
    def update_critic(self, obs, obs_aug, action, reward, next_obs,
                      next_obs_aug, not_done, logger, step):
        with torch.no_grad():
            dist = self.actor(next_obs)
            next_action = dist.rsample()
            log_prob = dist.log_prob(next_action).sum(-1, keepdim=True)
            target_Q1, target_Q2 = self.critic_target(next_obs, next_action)
            target_V = torch.min(target_Q1,
                                 target_Q2) - self.alpha.detach() * log_prob
            target_Q = reward + (not_done * self.discount * target_V)

            dist_aug = self.actor(next_obs_aug)
            next_action_aug = dist_aug.rsample()
            log_prob_aug = dist_aug.log_prob(next_action_aug).sum(-1,
                                                                  keepdim=True)
            target_Q1, target_Q2 = self.critic_target(next_obs_aug,
                                                      next_action_aug)
            target_V = torch.min(
                target_Q1, target_Q2) - self.alpha.detach() * log_prob_aug
            target_Q_aug = reward + (not_done * self.discount * target_V)

            target_Q = (target_Q + target_Q_aug) / 2

        # get current Q estimates
        current_Q1, current_Q2 = self.critic(obs, action)
        critic_loss = F.mse_loss(current_Q1, target_Q) + F.mse_loss(
            current_Q2, target_Q)

        Q1_aug, Q2_aug = self.critic(obs_aug, action)

        critic_loss += F.mse_loss(Q1_aug, target_Q) + F.mse_loss(
            Q2_aug, target_Q)

        logger.log('train_critic/loss', critic_loss, step)

        # Optimize the critic
        self.critic_optimizer.zero_grad()
        critic_loss.backward()
        self.critic_optimizer.step()

        self.critic.log(logger, step)

    def update_actor_and_alpha(self, obs, logger, step):
        # detach conv filters, so we don't update them with the actor loss
        dist = self.actor(obs, detach_encoder=True)
        action = dist.rsample()
        log_prob = dist.log_prob(action).sum(-1, keepdim=True)
        # detach conv filters, so we don't update them with the actor loss
        actor_Q1, actor_Q2 = self.critic(obs, action, detach_encoder=True)

        actor_Q = torch.min(actor_Q1, actor_Q2)

        actor_loss = (self.alpha.detach() * log_prob - actor_Q).mean()

        logger.log('train_actor/loss', actor_loss, step)
        logger.log('train_actor/target_entropy', self.target_entropy, step)
        logger.log('train_actor/entropy', -log_prob.mean(), step)

        # optimize the actor
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()

        self.actor.log(logger, step)

        self.log_alpha_optimizer.zero_grad()
        alpha_loss = (self.alpha *
                      (-log_prob - self.target_entropy).detach()).mean()
        logger.log('train_alpha/loss', alpha_loss, step)
        logger.log('train_alpha/value', self.alpha, step)
        alpha_loss.backward()
        self.log_alpha_optimizer.step()

    def update(self, replay_buffer, logger, step):
        obs, action, reward, next_obs, not_done, obs_aug, next_obs_aug = replay_buffer.sample(
            self.batch_size)

        logger.log('train/batch_reward', reward.mean(), step)

        self.update_critic(obs, obs_aug, action, reward, next_obs,
                           next_obs_aug, not_done, logger, step)

        if step % self.actor_update_frequency == 0:
            self.update_actor_and_alpha(obs, logger, step)

        if step % self.critic_target_update_frequency == 0:
            utils.soft_update_params(self.critic, self.critic_target,
                                     self.critic_tau)

    def update_Q(self, replay_buffer, logger, step):
        obs, action, reward, next_obs, not_done, obs_aug, next_obs_aug = replay_buffer.sample(
            self.batch_size)

        logger.log('train/batch_reward', reward.mean(), step)

        self.update_critic(obs, obs_aug, action, reward, next_obs,
                           next_obs_aug, not_done, logger, step)

        if step % self.critic_target_update_frequency == 0:
            utils.soft_update_params(self.critic, self.critic_target,
                                     self.critic_tau)
    
    def save_model(self, save_path, task, step):
        state = {'actor':self.actor.state_dict(),
         'critic':self.critic.state_dict(),
         'alpha':self.log_alpha,
         }
        filename = save_path + task + "_step_" + str(step) + ".pth"
        torch.save(state, filename)

    def load_model(self, model_path):
        checkpoint = torch.load(model_path)
        self.actor.load_state_dict(checkpoint['actor'])
        self.critic.load_state_dict(checkpoint['critic'])
        self.log_alpha = checkpoint['alpha']

    def load_actor_dict(self, model_path):
        checkpoint = torch.load(model_path)
        self.actor.load_state_dict(checkpoint['actor'])