import os
import numpy as np
import torch
import torch.nn as nn
from torch.optim import Adam
import torch.nn.functional as F
import dmc2gym
import utils

from iris_world_model import IRISWorldModel
from iris_tokenizer import Tokenizer

import argparse

import hydra
from omegaconf import OmegaConf

def make_env(env_name, seed, image_size=84, frame_stack=3):
    if env_name == 'ball_in_cup_catch':
        domain_name, task_name = 'ball_in_cup', 'catch'
    elif env_name == 'point_mass_easy':
        domain_name, task_name = 'point_mass', 'easy'
    else:
        domain_name = env_name.split('_')[0]
        task_name = '_'.join(env_name.split('_')[1:])
    camera_id = 2 if domain_name == 'quadruped' else 0
    env = dmc2gym.make(domain_name=domain_name,
                       task_name=task_name,
                       seed=seed,
                       visualize_reward=False,
                       from_pixels=True,
                       height=image_size,
                       width=image_size,
                       frame_skip=4,
                       camera_id=camera_id)
    env = utils.FrameStack(env, k=frame_stack)
    env.seed(seed)
    return env

class IRISBuffer:
    def __init__(self, obs_shape, action_shape, capacity, device):
        self.capacity = capacity
        self.device = torch.device(device)
        self.ptr = 0
        self.full = False

        self.obs = np.empty((capacity,) + obs_shape, dtype=np.uint8)  # 存uint8
        self.acts = np.empty((capacity, action_shape[0]), dtype=np.float32)
        self.rews = np.empty((capacity,), dtype=np.float32)
        self.dones = np.empty((capacity,), dtype=np.float32)


    def add(self, o, a, r, d):
        self.obs[self.ptr] = (o * 255.0).astype(np.uint8) if o.dtype != np.uint8 else o
        self.acts[self.ptr] = a
        self.rews[self.ptr] = r
        self.dones[self.ptr] = float(d)

        self.ptr += 1
        if self.ptr >= self.capacity:
            self.ptr = 0
            self.full = True

    def sample(self, B):
        hi = self.capacity if self.full else self.ptr
        idx = np.random.randint(0, hi - 1, size=B)

        o = torch.as_tensor(self.obs[idx], device=self.device).float() / 255.0
        a = torch.as_tensor(self.acts[idx], device=self.device).float()
        r = torch.as_tensor(self.rews[idx], device=self.device).float()
        d = torch.as_tensor(self.dones[idx], device=self.device).float()
        op = torch.as_tensor(self.obs[idx + 1], device=self.device).float() / 255.0

        return o, a, r, d, op

def load_victim_drq(cfg, env):
    cfg.agent.params.obs_shape = env.observation_space.shape
    cfg.agent.params.action_shape = env.action_space.shape
    cfg.agent.params.action_range = [
        float(env.action_space.low.min()),
        float(env.action_space.high.max())
    ]

    agent = hydra.utils.instantiate(cfg.agent)

    return agent

def train(env_name='cartpole_swingup', seed=1, device='cuda',
          buffer_steps=200000, batch=64,
          tok_epochs=5, trans_epochs=10,
          lr_tok=2e-4, lr_trans=2e-4,
          save_dir='./iris_ckpt', num_codes=1024, code_dim=128, d_model=512, nlayer=6, nhead=8):
        p_random = 0.1
        device = torch.device(device)
        os.makedirs(save_dir, exist_ok=True)
        env = make_env(env_name, seed)
        utils.set_seed_everywhere(seed)
        obs_shape = env.observation_space.shape
        action_dim = env.action_space.shape[0]
        buf = IRISBuffer(obs_shape, env.action_space.shape, buffer_steps+10, device)

        cfg = OmegaConf.load("config.yaml")
        victim_agent = load_victim_drq(cfg, env)
        
        model_path = "path_to_victim_model"
        victim_agent.load_model(model_path)
        victim_agent.train(False)
        
        o = env.reset()
        done = False
        while not buf.full and buf.ptr < buffer_steps:
            if np.random.rand() < p_random:
                a = env.action_space.sample()
            else:
                with torch.no_grad():
                    obs_tensor = torch.FloatTensor(o).to(device)
                    obs_tensor = obs_tensor.unsqueeze(0)
                    a = victim_agent.act(obs_tensor, sample=False)
            op, r, done, _ = env.step(a)
            buf.add(o, a, r, done)
            o = op if not done else env.reset()

        wm = IRISWorldModel(in_ch=obs_shape[0], action_dim=action_dim,
                            num_codes=num_codes, code_dim=code_dim,
                            d_model=d_model, nlayer=nlayer, nhead=nhead).to(device)

        opt_tok = Adam(wm.tokenizer.parameters(), lr=lr_tok)
        for ep in range(tok_epochs):
            for it in range(1000):
                o, _, _, _, _ = buf.sample(batch)
                x = o.to(device)
                x_hat, _, loss = wm.tokenizer(x)
                opt_tok.zero_grad()
                loss.backward()
                opt_tok.step()
            torch.save(wm.tokenizer.state_dict(), os.path.join(save_dir, f'tokenizer_{env_name}_ep{ep}.pt'))

        opt_trans = Adam(wm.trans.parameters(), lr=lr_trans)
        H,W = Tokenizer.tokens_hw(84,84,2)
        for ep in range(trans_epochs):
            for it in range(2000):
                o, a, r, d, op = buf.sample(batch)
                with torch.no_grad():
                    idx_t = wm.encode_tokens(o.to(device))
                    idx_p = wm.encode_tokens(op.to(device))
                loss = wm.trans_loss(idx_t.view(batch, H*W), a.to(device), idx_p, r.to(device), d.to(device))
                opt_trans.zero_grad()
                loss.backward()
                opt_trans.step()
            torch.save(wm.state_dict(), os.path.join(save_dir, f'wm_{env_name}_ep{ep}.pt'))

        return wm

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--env_name', type=str, default='cartpole_swingup')
    parser.add_argument('--seed', type=int, default=69)
    parser.add_argument('--device', type=str, default='cuda')
    parser.add_argument('--buffer_steps', type=int, default=200000)
    parser.add_argument('--batch', type=int, default=64)
    parser.add_argument('--tok_epochs', type=int, default=5)
    parser.add_argument('--trans_epochs', type=int, default=10)
    parser.add_argument('--save_dir', type=str, default='./iris_ckpt')
    args = parser.parse_args()

    train(env_name=args.env_name,
          seed=args.seed,
          device=args.device,
          buffer_steps=args.buffer_steps,
          batch=args.batch,
          tok_epochs=args.tok_epochs,
          trans_epochs=args.trans_epochs,
          save_dir=args.save_dir)
