import torch.nn as nn
import torch
import numpy as np
import GANmodels
import torch.nn.functional as F
import torchvision
import os
import utils
from iris_world_model import IRISWorldModel

models_path = 'g_models/'
batch_size = 128

def weights_init(m):
    classname = m.__class__.__name__
    if classname.find('Conv') != -1:
        nn.init.normal_(m.weight.data, 0.0, 0.02)
    elif classname.find('BatchNorm') != -1:
        nn.init.normal_(m.weight.data, 1.0, 0.02)
        nn.init.constant_(m.bias.data, 0)

class RL_GAN_Attack:
    def __init__(self,
                 device,
                 env,
                 env_name,
                 victim_agent,
                 q_noisy,
                 image_nc,
                 box_min,
                 box_max):
        output_nc = image_nc
        self.device = device
        self.env = env
        self.env_name = env_name
        self.victim_agent = victim_agent
        self.q_noisy = q_noisy
        self.input_nc = image_nc
        self.output_nc = output_nc
        self.box_min = box_min
        self.box_max = box_max
        self.epsilon = 8.0 / 255.0

        self.gen_input_nc = image_nc
        self.netG = GANmodels.Generator(self.gen_input_nc, image_nc).to(device)
        self.netDisc = GANmodels.Discriminator(image_nc).to(device)

        # initialize all weights
        self.netG.apply(weights_init)
        self.netDisc.apply(weights_init)

        # initialize optimizers
        self.optimizer_G = torch.optim.Adam(self.netG.parameters(),
                                            lr=0.001)
        self.optimizer_D = torch.optim.Adam(self.netDisc.parameters(),
                                            lr=0.001)

        if not os.path.exists(models_path):
            os.makedirs(models_path)

    def train_step(self, obs):
        x = obs / 255.
        perturbation = self.netG(x)
        adv_images = torch.clamp(perturbation, -self.epsilon, self.epsilon) + x
        adv_images = torch.clamp(adv_images, self.box_min, self.box_max)

        # optimize D
        for i in range(1):
            self.optimizer_D.zero_grad()
            pred_real = self.netDisc(x)
            loss_D_real = F.mse_loss(pred_real, torch.ones_like(pred_real, device=self.device))
            loss_D_real.backward()

            pred_fake = self.netDisc(adv_images.detach())
            loss_D_fake = F.mse_loss(pred_fake, torch.zeros_like(pred_fake, device=self.device))
            loss_D_fake.backward()

            loss_D_GAN = loss_D_fake + loss_D_real
            self.optimizer_D.step()

        # optimize G
        self.optimizer_G.zero_grad()
        total_loss_G = 0
        total_loss_adv = 0

        for i in range(1):
            # cal G's loss in GAN
            g_lambda = 1
            pred_fake = self.netDisc(adv_images)
            loss_G = F.mse_loss(pred_fake, torch.ones_like(pred_fake, device=self.device))
            total_loss_G += g_lambda * loss_G

            # cal attack loss
            adv_lambda = 1
            adv_obs = adv_images * 255.
            with utils.eval_mode(self.victim_agent):
                dist = self.victim_agent.actor(adv_obs)
                action = dist.rsample()
                Q1_, Q2_ = self.q_noisy.critic(adv_obs, action)
                # Q1_, Q2_ = self.victim_agent.critic(adv_obs, action)
                Q_min_ = torch.min(Q1_, Q2_)
                loss_adv = Q_min_[0][0]
            total_loss_adv += adv_lambda * loss_adv
        total_loss_G.backward(retain_graph=True)
        total_loss_adv.backward()
        self.optimizer_G.step()

        return loss_D_GAN.item(), total_loss_G.item(), total_loss_adv.item()

    def train(self, train_epochs, attack_reward_, world_model=None, use_wm=False):
        done = True
        attack_reward = attack_reward_
        obs_buffer = []

        for epoch in range(train_epochs):
            if done:
                obs = self.env.reset()
                done = False

            obs_buffer.append(obs)
            if len(obs_buffer) >= batch_size or epoch == train_epochs - 1:
                obs_t = torch.FloatTensor(np.stack(obs_buffer)).to(self.device)  # [B,C,H,W]
                loss_D, loss_G, loss_adv = self.train_step(obs_t)
                obs_buffer = []

            with utils.eval_mode(self.victim_agent):
                action = self.victim_agent.act(torch.FloatTensor(obs).to(self.device).unsqueeze(0), sample=False)
                action_np = action

            if use_wm and (world_model is not None):
                obs, reward, done, info = world_model.step(obs, action_np)
            else:
                obs, reward, done, info = self.env.step(action_np)

            if epoch % 2000 == 0:
                reward_ = 0
                for _ in range(5):
                    reward_ += self.eval_gan()
                reward_ /= 5
                if reward_ < attack_reward:
                    attack_reward = reward_
                    print("min attack reward: ", attack_reward)
                    netG_file_name = models_path + self.env_name + '.pth'
                    self.save_model(netG_file_name)
                    print("successfully save gan model!")
        return attack_reward


    def eval_gan(self):
        obs = self.env.reset()
        done = False
        episode_reward = 0
        while not done:
            with utils.eval_mode(self.victim_agent):
                action, _ = self.victim_agent.adv_gan_based(obs, self.netG, if_save=False, sample=False)
            obs, reward, done, info = self.env.step(action)
            episode_reward += reward
        return episode_reward
    
    def save_model(self, model_path):
        torch.save(self.netG.state_dict(), model_path)


