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

from optimal_adversary import build_optimal_adversary

torch.backends.cudnn.benchmark = True

def make_env(cfg):
    """Helper function to create dm_control environment"""
    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:])

    # per dreamer: https://github.com/danijar/dreamer/blob/02f0210f5991c7710826ca7881f19c64a012290c/wrappers.py#L26
    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.opt_adv = 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 run(self):        
        model_path = "victim_model_path"
        self.agent.load_model(model_path)

        in_channels = self.env.observation_space.shape[0]
        eps = 8/255
        self.opt_adv = build_optimal_adversary(
            in_channels=in_channels,
            eps=eps,
            victim_agent=self.agent,
            device=self.device,
            lr=1e-4
        )

        opt_niter = getattr(self.cfg, 'opt_niter', 5000) 
        opt_batch_steps = getattr(self.cfg, 'opt_batch_steps', 128)
        self.opt_adv.train(self.env, n_iter=opt_niter, B=opt_batch_steps, logger=None)

        save_dir = getattr(self.cfg, "save_dir", "./opt_adv_ckpt")
        os.makedirs(save_dir, exist_ok=True)
        torch.save(self.opt_adv.pi_adv.state_dict(), os.path.join(save_dir, f"{self.env_name}_opt_adv.pth"))
        print("saved optimal adversary to:", os.path.join(save_dir, f"{self.env_name}_opt_adv.pth"))


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