"""Compact inference-only WorldDiT runtime for the released LIBERO checkpoints."""

from __future__ import annotations

import math
from functools import partial
from pathlib import Path

import clip
import torch
from einops import rearrange, repeat
from einops_exts import rearrange_many
from safetensors import safe_open
from timm.models.vision_transformer import Block, PatchEmbed
from torch import einsum, nn

SUITES = ("libero_10", "libero_spatial", "libero_goal", "libero_object")
CONTEXT_STEPS = 3
ACTION_HORIZON = 7
HIDDEN_DIM = 1024


class VisionEncoder(nn.Module):
    """The encoder half of the frozen MAE dependency."""

    def __init__(self):
        super().__init__()
        self.patch_embed = PatchEmbed(224, 16, 3, 768)
        self.cls_token = nn.Parameter(torch.zeros(1, 1, 768))
        self.pos_embed = nn.Parameter(torch.zeros(1, 197, 768), requires_grad=False)
        self.blocks = nn.ModuleList(
            [
                Block(
                    768,
                    12,
                    4,
                    qkv_bias=True,
                    norm_layer=partial(nn.LayerNorm, eps=1e-6),
                )
                for _ in range(12)
            ]
        )
        self.norm = nn.LayerNorm(768, eps=1e-6)

    def forward(self, images: torch.Tensor) -> torch.Tensor:
        patches = self.patch_embed(images) + self.pos_embed[:, 1:]
        # Retain the released encoder's mask_ratio=0 behavior, including RNG use.
        order = torch.rand(patches.shape[:2], device=patches.device).argsort(dim=1)
        patches = torch.gather(patches, 1, order.unsqueeze(-1).expand_as(patches))
        cls = (self.cls_token + self.pos_embed[:, :1]).expand(images.shape[0], -1, -1)
        tokens = torch.cat((cls, patches), dim=1)
        for block in self.blocks:
            tokens = block(tokens)
        return self.norm(tokens)


class PerceiverAttention(nn.Module):
    def __init__(self, dim: int = 768, dim_head: int = 64, heads: int = 8):
        super().__init__()
        self.scale = dim_head**-0.5
        self.heads = heads
        inner = dim_head * heads
        self.norm_media = nn.LayerNorm(dim)
        self.norm_latents = nn.LayerNorm(dim)
        self.to_q = nn.Linear(dim, inner, bias=False)
        self.to_kv = nn.Linear(dim, inner * 2, bias=False)
        self.to_out = nn.Linear(inner, dim, bias=False)

    def forward(self, media: torch.Tensor, latents: torch.Tensor) -> torch.Tensor:
        media, latents = self.norm_media(media), self.norm_latents(latents)
        query = self.to_q(latents)
        key, value = self.to_kv(torch.cat((media, latents), dim=-2)).chunk(2, dim=-1)
        query, key, value = rearrange_many(
            (query, key, value), "b t n (h d) -> b h t n d", h=self.heads
        )
        scores = einsum("... i d, ... j d -> ... i j", query * self.scale, key)
        weights = (scores - scores.amax(dim=-1, keepdim=True).detach()).softmax(-1)
        output = einsum("... i j, ... j d -> ... i d", weights, value)
        return self.to_out(rearrange(output, "b h t n d -> b t n (h d)"))


class PerceiverResampler(nn.Module):
    def __init__(self):
        super().__init__()
        self.latents = nn.Parameter(torch.randn(16, 768))
        self.layers = nn.ModuleList()
        for _ in range(3):
            feed_forward = nn.Sequential(
                nn.LayerNorm(768),
                nn.Linear(768, 3072, bias=False),
                nn.GELU(),
                nn.Linear(3072, 768, bias=False),
            )
            self.layers.append(nn.ModuleList((PerceiverAttention(), feed_forward)))
        self.norm = nn.LayerNorm(768)

    def forward(self, tokens: torch.Tensor) -> torch.Tensor:
        batch, steps = tokens.shape[:2]
        tokens = rearrange(tokens, "b t f v d -> b t (f v) d")
        latents = repeat(self.latents, "n d -> b t n d", b=batch, t=steps)
        for attention, feed_forward in self.layers:
            latents = attention(tokens, latents) + latents
            latents = feed_forward(latents) + latents
        return self.norm(latents)


def _time_embedding(timestep: torch.Tensor, dim: int) -> torch.Tensor:
    half = dim // 2
    frequency = torch.exp(
        -math.log(10000.0)
        * torch.arange(half, device=timestep.device, dtype=timestep.dtype)
        / max(half - 1, 1)
    )
    phase = timestep.unsqueeze(-1) * frequency.unsqueeze(0) * 1000.0
    embedding = torch.cat((torch.sin(phase), torch.cos(phase)), dim=-1)
    return embedding if dim % 2 == 0 else torch.nn.functional.pad(embedding, (0, 1))


class MLP(nn.Module):
    def __init__(self, input_dim: int, hidden_dim: int, output_dim: int):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.SiLU(),
            nn.Linear(hidden_dim, output_dim),
        )

    def forward(self, value: torch.Tensor) -> torch.Tensor:
        return self.net(value)


def _modulate(value: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor):
    return value * (1.0 + scale) + shift


class DFDiTBlock(nn.Module):
    def __init__(self):
        super().__init__()
        self.norm1 = nn.LayerNorm(HIDDEN_DIM, elementwise_affine=False)
        self.attn = nn.MultiheadAttention(HIDDEN_DIM, 16, batch_first=True)
        self.norm2 = nn.LayerNorm(HIDDEN_DIM, elementwise_affine=False)
        self.mlp = nn.Sequential(
            nn.Linear(HIDDEN_DIM, HIDDEN_DIM * 4),
            nn.GELU(approximate="tanh"),
            nn.Dropout(0.0),
            nn.Linear(HIDDEN_DIM * 4, HIDDEN_DIM),
        )
        self.adaLN_modulation = nn.Sequential(
            nn.SiLU(), nn.Linear(HIDDEN_DIM, 6 * HIDDEN_DIM)
        )

    def forward(self, tokens, modulation, mask):
        shift_a, scale_a, gate_a, shift_m, scale_m, gate_m = self.adaLN_modulation(
            modulation
        ).chunk(6, dim=-1)
        attention_input = _modulate(self.norm1(tokens), shift_a, scale_a)
        attention = self.attn(
            attention_input,
            attention_input,
            attention_input,
            attn_mask=mask,
            need_weights=False,
        )[0]
        tokens = tokens + gate_a * attention
        return tokens + gate_m * self.mlp(
            _modulate(self.norm2(tokens), shift_m, scale_m)
        )


class ActionSampler(nn.Module):
    """Action sampler used by the released checkpoints."""

    def __init__(self):
        super().__init__()
        self.context_norm = nn.LayerNorm(HIDDEN_DIM)
        self.context_proj = nn.Linear(HIDDEN_DIM, HIDDEN_DIM)
        self.action_tokenizer = MLP(7, HIDDEN_DIM * 4, HIDDEN_DIM)
        self.action_decoder = MLP(HIDDEN_DIM, HIDDEN_DIM * 4, 7)
        self.action_pos = nn.Parameter(
            torch.randn(1, ACTION_HORIZON, HIDDEN_DIM) * 0.02
        )
        self.step_pos = nn.Parameter(
            torch.randn(1, CONTEXT_STEPS, 1, HIDDEN_DIM) * 0.02
        )
        self.context_type = nn.Parameter(torch.randn(1, 1, HIDDEN_DIM) * 0.02)
        self.action_type = nn.Parameter(torch.randn(1, 1, HIDDEN_DIM) * 0.02)
        self.register_type = nn.Parameter(torch.randn(1, 1, HIDDEN_DIM) * 0.02)
        self.register_tokens = nn.Parameter(torch.randn(1, 4, HIDDEN_DIM) * 0.02)
        self.time_proj = nn.Sequential(
            nn.Linear(HIDDEN_DIM, HIDDEN_DIM * 4),
            nn.SiLU(),
            nn.Linear(HIDDEN_DIM * 4, HIDDEN_DIM),
        )
        self.blocks = nn.ModuleList([DFDiTBlock() for _ in range(4)])
        self.final_norm = nn.LayerNorm(HIDDEN_DIM)

    def _times(self, timestep: torch.Tensor, count: int, dtype: torch.dtype):
        embedded = self.time_proj(_time_embedding(timestep, HIDDEN_DIM).to(dtype))
        return embedded.unsqueeze(1).expand(-1, count, -1)

    @staticmethod
    def _mask(steps: int, context_count: int, dtype, device):
        targets, registers = ACTION_HORIZON, 4
        block = context_count + targets + registers
        mask = torch.full(
            (steps * block, steps * block), -torch.inf, dtype=dtype, device=device
        )
        for step in range(steps):
            start = step * block
            context = slice(start, start + context_count)
            action = slice(start + context_count, start + context_count + targets)
            register = slice(start + context_count + targets, start + block)
            for source_step in range(step + 1):
                source = source_step * block
                visible_context = slice(source, source + context_count)
                mask[context, visible_context] = 0
                mask[action, visible_context] = 0
                mask[register, visible_context] = 0
            mask[action, action] = 0
            mask[action, register] = 0
            mask[register, action] = 0
            mask[register, register] = 0
        return mask

    def _velocity(self, context: torch.Tensor, noisy_action: torch.Tensor, timestep):
        batch, steps, context_count, _ = context.shape
        flat_batch = batch * steps
        dtype, device = context.dtype, context.device
        encoded_context = self.context_proj(self.context_norm(context))
        encoded_context = encoded_context + self.context_type.to(dtype)
        encoded_context = encoded_context + self.step_pos[:, :steps].to(dtype).expand(
            -1, -1, context_count, -1
        )
        zero_time = torch.zeros(flat_batch, dtype=dtype, device=device)
        context_mod = self._times(zero_time, context_count, dtype).view(
            batch, steps, context_count, HIDDEN_DIM
        ) + self.context_type.to(dtype)

        flat_action = noisy_action.reshape(flat_batch, ACTION_HORIZON, 7)
        action_time = self._times(timestep, ACTION_HORIZON, dtype)
        action = self.action_tokenizer(flat_action.to(dtype))
        action = (
            action
            + self.action_pos.to(dtype)
            + self.action_type.to(dtype)
            + action_time
        )
        action = action.view(batch, steps, ACTION_HORIZON, HIDDEN_DIM)
        action = action + self.step_pos[:, :steps].to(dtype).expand(
            -1, -1, ACTION_HORIZON, -1
        )
        action_mod = (action_time + self.action_type.to(dtype)).view(
            batch, steps, ACTION_HORIZON, HIDDEN_DIM
        )

        registers = self.register_tokens.to(dtype).expand(batch, steps, -1, -1)
        registers = registers + self.register_type.to(dtype)
        registers = registers + self.step_pos[:, :steps].to(dtype).expand(-1, -1, 4, -1)
        register_mod = self._times(zero_time, 4, dtype).view(
            batch, steps, 4, HIDDEN_DIM
        ) + self.register_type.to(dtype)

        tokens = torch.cat((encoded_context, action, registers), dim=2).flatten(1, 2)
        modulation = torch.cat((context_mod, action_mod, register_mod), dim=2).flatten(
            1, 2
        )
        mask = self._mask(steps, context_count, tokens.dtype, device)
        for block in self.blocks:
            tokens = block(tokens, modulation, mask)
        tokens = self.final_norm(tokens).view(batch, steps, -1, HIDDEN_DIM)
        return self.action_decoder(
            tokens[:, :, context_count : context_count + ACTION_HORIZON]
        )

    def forward(self, context: torch.Tensor, sampling_steps: int = 20):
        batch, steps = context.shape[:2]
        dtype, device = self.context_norm.weight.dtype, context.device
        context = context.to(dtype)
        action = torch.randn(
            batch, steps, ACTION_HORIZON, 7, dtype=dtype, device=device
        )
        for time in torch.linspace(0.0, 1.0, sampling_steps + 1, device=device)[:-1]:
            timestep = torch.full(
                (batch * steps,), float(time.item()), dtype=dtype, device=device
            )
            action = action + self._velocity(context, action, timestep) / sampling_steps
        return action


class WorldDiTPolicy(nn.Module):
    """Released policy with only modules used by inference."""

    def __init__(self, mae_path: Path, clip_path: Path):
        super().__init__()
        self.text_projector = nn.Linear(512, HIDDEN_DIM)
        self.arm_state_encoder = nn.Linear(6, HIDDEN_DIM)
        self.gripper_state_encoder = nn.Linear(2, HIDDEN_DIM)
        self.state_projector = nn.Linear(HIDDEN_DIM * 2, HIDDEN_DIM)
        self.vision_encoder = VisionEncoder()
        self.perceiver_resampler = PerceiverResampler()
        self.image_primary_projector = nn.Linear(768, HIDDEN_DIM)
        self.cls_token_primary_projector = nn.Linear(768, HIDDEN_DIM)
        self.image_wrist_projector = nn.Linear(768, HIDDEN_DIM)
        self.cls_token_wrist_projector = nn.Linear(768, HIDDEN_DIM)
        self.embedding_layer_norm = nn.LayerNorm(HIDDEN_DIM)
        self.transformer_backbone_position_embedding = nn.Parameter(
            torch.zeros(1, CONTEXT_STEPS, 1, HIDDEN_DIM)
        )
        self.unified_action_world_head = ActionSampler()

        mae = torch.load(mae_path, map_location="cpu", weights_only=False)
        self.vision_encoder.load_state_dict(mae["model"], strict=False)
        self.clip_model, self.image_processor = clip.load(str(clip_path), device="cpu")
        self.vision_encoder.requires_grad_(False)
        self.clip_model.requires_grad_(False)

    def _images(self, primary: torch.Tensor, wrist: torch.Tensor):
        batch, steps = primary.shape[:2]
        vision_dtype = next(self.vision_encoder.parameters()).dtype
        with torch.no_grad():
            primary_tokens = self.vision_encoder(primary.flatten(0, 1).to(vision_dtype))
            wrist_tokens = self.vision_encoder(wrist.flatten(0, 1).to(vision_dtype))
        cls_primary = primary_tokens[:, :1]
        cls_wrist = wrist_tokens[:, :1]
        resampler_dtype = next(self.perceiver_resampler.parameters()).dtype
        cls_primary = cls_primary.to(resampler_dtype)
        cls_wrist = cls_wrist.to(resampler_dtype)
        primary_tokens = primary_tokens[:, 1:].to(resampler_dtype)
        wrist_tokens = wrist_tokens[:, 1:].to(resampler_dtype)
        primary_latents = self.perceiver_resampler(
            primary_tokens.reshape(batch * steps, 196, 768).unsqueeze(1).unsqueeze(1)
        )
        wrist_latents = self.perceiver_resampler(
            wrist_tokens.reshape(batch * steps, 196, 768).unsqueeze(1).unsqueeze(1)
        )
        images = torch.cat(
            (
                self.image_primary_projector(primary_latents.flatten(0, 2)).view(
                    batch, steps, 16, HIDDEN_DIM
                ),
                self.image_wrist_projector(wrist_latents.flatten(0, 2)).view(
                    batch, steps, 16, HIDDEN_DIM
                ),
            ),
            dim=2,
        )
        cls = torch.cat(
            (
                self.cls_token_primary_projector(cls_primary).view(
                    batch, steps, 1, HIDDEN_DIM
                ),
                self.cls_token_wrist_projector(cls_wrist).view(
                    batch, steps, 1, HIDDEN_DIM
                ),
            ),
            dim=2,
        )
        return images, cls

    def forward(self, image_primary, image_wrist, state, text_token):
        batch, steps = state.shape[:2]
        if steps != CONTEXT_STEPS:
            raise ValueError(f"expected {CONTEXT_STEPS} context frames, got {steps}")
        with torch.no_grad():
            text = self.clip_model.encode_text(text_token.flatten(0, 1)).type_as(state)
        text = self.text_projector(text).view(batch, steps, 1, HIDDEN_DIM)
        flat_state = state.flatten(0, 1)
        state_token = self.state_projector(
            torch.cat(
                (
                    self.arm_state_encoder(flat_state[:, :6]),
                    self.gripper_state_encoder(flat_state[:, 6:]),
                ),
                dim=1,
            )
        ).view(batch, steps, 1, HIDDEN_DIM)
        images, cls = self._images(image_primary, image_wrist)
        context = torch.cat((text, state_token, images, cls), dim=2)
        context = context + self.transformer_backbone_position_embedding[:, :steps].to(
            context.dtype
        )
        context = self.embedding_layer_norm(
            context.to(self.embedding_layer_norm.weight.dtype)
        )
        return self.unified_action_world_head(context, sampling_steps=20)


def load_model(
    model_root: str | Path,
    suite: str,
    device: str | torch.device = "cuda",
    *,
    vision_bfloat16: bool = True,
) -> WorldDiTPolicy:
    """Load one suite checkpoint from a downloaded model-repository directory."""
    if suite not in SUITES:
        raise ValueError(f"unknown suite {suite!r}; choose one of {SUITES}")
    root = Path(model_root).expanduser().resolve()
    paths = {
        "checkpoint": root / "checkpoints" / suite / "model.safetensors",
        "mae": root / "dependencies" / "mae_pretrain_vit_base.pth",
        "clip": root / "dependencies" / "ViT-B-32.pt",
    }
    missing = [str(path) for path in paths.values() if not path.is_file()]
    if missing:
        raise FileNotFoundError(f"missing model files: {missing}")

    model = WorldDiTPolicy(paths["mae"], paths["clip"])
    with safe_open(paths["checkpoint"], framework="pt", device="cpu") as source:
        released = {
            name.removeprefix("module."): source.get_tensor(name)
            for name in source.keys()
        }
    current = model.state_dict()
    retained = {name: value for name, value in released.items() if name in current}
    expected = {
        name
        for name in current
        if not name.startswith(("vision_encoder.", "clip_model."))
    }
    missing_parameters = sorted(expected - retained.keys())
    if missing_parameters:
        raise RuntimeError(
            f"checkpoint is missing inference parameters: {missing_parameters}"
        )
    model.load_state_dict(retained, strict=False)
    model.float()
    if vision_bfloat16:
        model.vision_encoder.bfloat16()
    model.to(torch.device(device)).eval()
    return model