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

import numpy as np

import dmc2gym
import hydra
import torch
import torch.nn as nn
import torch.nn.functional as F
import utils

from replay_buffer import ReplayBuffer

from PIL import Image
import os

import GANmodels


torch.backends.cudnn.benchmark = True

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=cfg.action_repeat,
                       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()

        self.cfg = cfg

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

        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.agent = hydra.utils.instantiate(cfg.agent)

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

    def attack(self, adv_mode, if_save):
        obs = self.env.reset()
        done = False
        episode_reward = 0
        idx = 0
        while not done:
            with utils.eval_mode(self.agent):
                if adv_mode == "test_clean":
                    action, adv_obs_save = self.agent.test_clean(obs, if_save, sample=False)
                elif adv_mode == "adv_gen_based":
                    adv_fun = getattr(self.agent, adv_mode)
                    action, adv_obs_save = adv_fun(obs, self.Gen, if_save, sample=False)
                elif hasattr(self.agent, adv_mode):
                    adv_fun = getattr(self.agent, adv_mode)
                    action, adv_obs_save = adv_fun(obs, if_save, sample=False)
            idx += 1
            obs, reward, done, info = self.env.step(action)
            episode_reward += reward
        return episode_reward

    def run(self):
        model_path = "path_to_victim_model"
        self.agent.load_model(model_path)

        adv_dict = {
            # image based attacks
            0: "test_clean",
            1: "adv_pgd",
            2: "adv_cw",
            3: "simba_attack",
            4: "adv_square",

            # vector RL based attacks
            5: "adv_critic",
            6: "adv_mad",

            # Gen based: paad, optimal, seba
            7: "adv_gen_based",
        }
        for adv_mode_idx in range(0, 8):
            adv_mode = adv_dict[adv_mode_idx]
            if adv_mode == "adv_gen_based":
                image_nc = 9
                eps_ = 8/255
                generator_path = "path_to_seba_model"
                self.Gen = GANmodels.Generator(image_nc, image_nc).to(self.device)
                self.Gen.load_state_dict(torch.load(generator_path))
                self.Gen.eval()
            if_save = False
            seed_ = 1
            utils.set_seed_everywhere(seed_)
            self.env.seed(seed_)
            reward_ = self.attack(adv_mode, if_save)
            print(self.env_name, ": ", adv_mode, ": ", reward_)


@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()
