import torch
import torch.nn as nn
import torch.nn.functional as F
from iris_tokenizer import Tokenizer
from iris_transformer import ActionCondTransformer

class IRISWorldModel(nn.Module):
    def __init__(self, in_ch=9, action_dim=1, num_codes=1024, code_dim=128, d_model=512, nlayer=6, nhead=8, beta=0.25):
        super().__init__()
        self.tokenizer = Tokenizer(in_ch, hidden=128, code_dim=code_dim, num_codes=num_codes)
        H, W = Tokenizer.tokens_hw(84,84,2)
        self.HW = (H, W)
        self.trans = ActionCondTransformer(num_codes, code_dim, H*W, action_dim, d_model, nlayer, nhead)

    @torch.no_grad()
    def encode_tokens(self, x):
        return self.tokenizer.encode_indices(x)

    @torch.no_grad()
    def decode_tokens(self, idx):
        return self.tokenizer.decode_indices(idx)

    def tokenizer_loss(self, x):
        _, _, loss = self.tokenizer(x)
        return loss

    def trans_loss(self, idx_t, act_t, idx_next, rew_next, done_next):
        B = idx_t.size(0)
        logits, rew_pred, done_pred = self.trans(idx_t.view(B, -1), act_t)
        ce = F.cross_entropy(logits.view(B*self.HW[0]*self.HW[1], -1),
                             idx_next.view(B*self.HW[0]*self.HW[1]), reduction='mean')
        rew_loss = F.mse_loss(rew_pred, rew_next)
        done_loss = F.binary_cross_entropy_with_logits(done_pred, done_next)
        return ce + rew_loss + done_loss

    @torch.no_grad()
    def rollout_decode(self, x0, actions, H):
        x = x0
        imgs = [x0]
        for t in range(H):
            idx = self.encode_tokens(x)
            logits, _, _ = self.trans(idx.view(idx.size(0), -1), actions[:,t,:])
            pred_idx = logits.argmax(-1).view(idx.shape)
            x = self.decode_tokens(pred_idx)
            imgs.append(x)
        return imgs

    
    @torch.no_grad()
    def step_tensor(self, obs_tensor, action_tensor, return_uint8=True, done_threshold=0.5):
        assert obs_tensor.dim() == 4
        assert action_tensor.dim() == 2
        device = next(self.parameters()).device
        x = obs_tensor.to(device)
        a = action_tensor.to(device)
        idx_t = self.encode_tokens(x)
        B = idx_t.size(0)
        logits, rew_pred, done_logit = self.trans(idx_t.view(B, -1), a)
        pred_idx = logits.argmax(dim=-1).view_as(idx_t)
        next_x = self.decode_tokens(pred_idx).clamp(0.0, 1.0)
        if rew_pred.dim() == 2 and rew_pred.size(-1) == 1:
            rew_pred = rew_pred.squeeze(-1)
        reward = rew_pred
        if done_logit.dim() == 2 and done_logit.size(-1) == 1:
            done_logit = done_logit.squeeze(-1)
        done_prob = torch.sigmoid(done_logit)
        done = (done_prob > done_threshold)
        if return_uint8:
            next_np = (next_x * 255.0).round().byte().cpu().numpy()
            return next_np, reward.detach().cpu(), done.detach().cpu()
        return next_x.detach(), reward.detach(), done.detach()

    @torch.no_grad()
    def step(self, obs_numpy, action_numpy, done_threshold=0.5):
        device = next(self.parameters()).device
        if obs_numpy.ndim == 3:
            obs_numpy = obs_numpy[None, ...]
        if action_numpy.ndim == 1:
            import numpy as np
            action_numpy = np.asarray(action_numpy, dtype=np.float32)[None, ...]
        obs_t = torch.as_tensor(obs_numpy, dtype=torch.float32, device=device)
        if obs_t.max() > 1.5:
            obs_t = obs_t / 255.0
        act_t = torch.as_tensor(action_numpy, dtype=torch.float32, device=device)
        next_obs_np, reward_t, done_t = self.step_tensor(obs_t, act_t, return_uint8=True, done_threshold=done_threshold)
        if next_obs_np.shape[0] == 1:
            next_obs_np = next_obs_np[0]
            reward = float(reward_t[0].item())
            done = bool(done_t[0].item())
        else:
            reward = reward_t.numpy()
            done = done_t.numpy().astype(bool)
        info = {}
        return next_obs_np, reward, done, info

