import torch
import torch.nn as nn
import torch.nn.functional as F

import q_noisy_model
from GANtraining import RL_GAN_Attack

import copy
import math
import os
import pickle as pkl
import sys
import time

import numpy as np
import dmc2gym
import hydra
import utils

from replay_buffer import ReplayBuffer
from logger import Logger

from iris_world_model import IRISWorldModel

torch.backends.cudnn.benchmark = True

use_cuda=True
image_nc=9
BOX_MIN = 0
BOX_MAX = 1
loops = 20
H_wm = 4

LAST = 9
training_steps = 5000

def make_env(cfg):
    if cfg.env == 'ball_in_cup_catch':
        domain_name = 'ball_in_cup'
        task_name = 'catch'
    elif cfg.env == 'point_mass_easy':
        domain_name = 'point_mass'
        task_name = 'easy'
    else:
        domain_name = cfg.env.split('_')[0]
        task_name = '_'.join(cfg.env.split('_')[1:])
    camera_id = 2 if domain_name == 'quadruped' else 0

    env = dmc2gym.make(domain_name=domain_name,
                       task_name=task_name,
                       seed=cfg.seed,
                       visualize_reward=False,
                       from_pixels=True,
                       height=cfg.image_size,
                       width=cfg.image_size,
                       frame_skip=4,
                       camera_id=camera_id)

    env = utils.FrameStack(env, k=cfg.frame_stack)

    env.seed(cfg.seed)
    assert env.action_space.low.min() >= -1
    assert env.action_space.high.max() <= 1

    return env


class Workspace(object):
    def __init__(self, cfg):
        self.work_dir = os.getcwd()
        print(f'workspace: {self.work_dir}')

        self.cfg = cfg

        self.logger = Logger(self.work_dir,
            save_tb=cfg.log_save_tb,
            log_frequency=cfg.log_frequency_step,
            agent=cfg.agent.name,
            action_repeat=cfg.action_repeat)

        utils.set_seed_everywhere(cfg.seed)
        self.device = torch.device(cfg.device)
        self.env = make_env(cfg)
        self.env_name = cfg.env

        cfg.agent.params.obs_shape = self.env.observation_space.shape
        cfg.agent.params.action_shape = self.env.action_space.shape
        cfg.agent.params.action_range = [
            float(self.env.action_space.low.min()),
            float(self.env.action_space.high.max())
        ]
        self.victim_agent = hydra.utils.instantiate(cfg.agent)
        
        self.q_noisy = q_noisy_model.Q_Noisy(
                 victim_agent=self.victim_agent,
                 obs_shape=cfg.agent.params.obs_shape,
                 action_shape=cfg.agent.params.action_shape,
                 device=self.device,
                 encoder_cfg=cfg.agent.params.encoder_cfg,
                 critic_cfg=cfg.agent.params.critic_cfg,
                 discount=0.99,
                 init_temperature=0.1,
                 lr=1e-3,
                 critic_tau=0.01,
                 critic_target_update_frequency=2,
                 batch_size=128)
        
        self.gan = RL_GAN_Attack(self.device,
                                self.env,
                                self.env_name,
                                self.victim_agent,
                                self.q_noisy,
                                image_nc,
                                BOX_MIN,
                                BOX_MAX)

        self.replay_buffer = ReplayBuffer(self.env.observation_space.shape,
                    self.env.action_space.shape,
                    cfg.replay_buffer_capacity,
                    self.cfg.image_pad, self.device)

    def train_Q_noisy(self, loop_idx):
        self.step = 0
        episode_step, done = 0, True
        while self.step < training_steps:
            if done:
                obs = self.env.reset()
                done = False
                episode_step = 0

            with utils.eval_mode(self.victim_agent):
                action, adv_obs = self.victim_agent.adv_gan_based(
                    obs, self.gan.netG, if_save=True, sample=False)
                adv_obs = adv_obs.astype(np.uint8)

            if self.step > 1000 or loop_idx > 0:
                for _ in range(self.cfg.num_train_iters):
                    self.q_noisy.update(self.replay_buffer, self.logger, self.step)

            next_obs, reward, done, info = self.env.step(action)

            done = float(done)
            done_no_max = 0 if episode_step + 1 == self.env._max_episode_steps else done

            adv_next_obs = self.get_adv_obs(next_obs)
            self.replay_buffer.add(adv_obs, action, reward, adv_next_obs,
                                     done, done_no_max)

            obs = next_obs
            episode_step += 1
            self.step += 1
    
    def train_Q_noisy_wm(self, loop_idx, world_model=None):
        self.step = 0
        episode_step, done = 0, True
        while self.step < training_steps:
            if done:
                obs = self.env.reset()
                done = False
                episode_step = 0

            with utils.eval_mode(self.victim_agent):
                action, adv_obs = self.victim_agent.adv_gan_based(
                    obs, self.gan.netG, if_save=True, sample=False)
                adv_obs = adv_obs.astype(np.uint8)

            if self.step > 1000 or loop_idx > 0:
                for _ in range(self.cfg.num_train_iters):
                    self.q_noisy.update(self.replay_buffer, self.logger, self.step)

            next_obs_wm, reward_wm, done_wm, _ = world_model.step(obs, action)
            next_obs = next_obs_wm
            reward = reward_wm
            done = float(done_wm)
            done_no_max = done
            
            adv_next_obs = self.get_adv_obs(next_obs)

            self.replay_buffer.add(adv_obs, action, reward, adv_next_obs, done, done_no_max)

            if done:
                obs = self.env.reset()
            else:
                obs = next_obs

            episode_step += 1
            self.step += 1
    
    def get_adv_obs(self, obs):
        obs = torch.FloatTensor(obs).to(self.device)
        obs = obs.unsqueeze(0)   
        obs = obs / 255.
        eps_ = 8 / 255

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

        adv_obs_save = adv_obs.squeeze(0).detach().cpu().numpy()
        adv_obs_save = adv_obs_save.astype(np.uint8)

        return adv_obs_save 
    
    def run(self):
        model_path = "victim_model_path"
        self.victim_agent.load_model(model_path)
        
        wm_path = os.path.join('iris_ckpt', f'wm_{self.env_name}_ep{LAST}.pt')
        world_model = IRISWorldModel(in_ch=self.env.observation_space.shape[0], 
                                        action_dim=self.env.action_space.shape[0],
                                        num_codes=1024, code_dim=128,
                                        d_model=512, nlayer=6, nhead=8).to(self.device)
        world_model.load_state_dict(torch.load(wm_path, map_location=self.device))
        world_model.to(self.device)
        world_model.eval()

        attack_reward_ = 10000

        loop_idx = 0
        while loop_idx < loops:
            self.train_Q_noisy(loop_idx)
            attack_reward_ = self.gan.train(training_steps, attack_reward_, world_model=None, use_wm=False)
            for _ in range(H_wm):
                self.train_Q_noisy_wm(loop_idx, world_model=world_model)
                attack_reward_ = self.gan.train(training_steps, attack_reward_, world_model=world_model, use_wm=True)
            loop_idx += 1

@hydra.main(config_path='config.yaml', strict=True)
def main(cfg):
    current_file = os.path.basename(__file__).replace('.py', '')
    module = __import__(current_file)
    W = getattr(module, 'Workspace')
    workspace = W(cfg)

    workspace.run()


if __name__ == '__main__':
    main()
