"""Differentiable loader for MiniMax-M3 vision tower from staged safetensors."""
from __future__ import annotations

import json
from pathlib import Path

import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors import safe_open
from transformers.models.minimax_m3_vl.configuration_minimax_m3_vl import MiniMaxM3VLVisionConfig
from transformers.models.minimax_m3_vl.modeling_minimax_m3_vl import MiniMaxM3VLVisionModel

MMX_DIR = Path.home() / "workspace/mmx_m3"
WEIGHTS = MMX_DIR / "vision_tower.safetensors"
CONFIG = MMX_DIR / "config.json"
IMAGE_SIZE = 672
PATCH_SIZE = 14
TEMPORAL_PATCH_SIZE = 2
SPATIAL_MERGE_SIZE = 2
MEAN = (0.48145466, 0.4578275, 0.40821073)
STD = (0.26862954, 0.26130258, 0.27577711)


def _build_config() -> MiniMaxM3VLVisionConfig:
    vc = json.loads(CONFIG.read_text())["vision_config"]
    itc = vc.get("img_token_compression_config", {})
    return MiniMaxM3VLVisionConfig(
        hidden_size=vc["hidden_size"],
        num_attention_heads=vc["num_attention_heads"],
        num_hidden_layers=vc["num_hidden_layers"],
        intermediate_size=vc["intermediate_size"],
        patch_size=vc["patch_size"],
        image_size=vc["image_size"],
        num_channels=vc["num_channels"],
        hidden_act=vc["hidden_act"],
        layer_norm_eps=vc["layer_norm_eps"],
        attention_dropout=vc.get("attention_dropout", 0.0),
        temporal_patch_size=itc.get("temporal_patch_size", TEMPORAL_PATCH_SIZE),
        spatial_merge_size=itc.get("spatial_merge_size", SPATIAL_MERGE_SIZE),
        rope_parameters={"rope_theta": vc.get("rope_theta", 10000.0)},
    )


def _remap_state_dict() -> dict[str, torch.Tensor]:
    sd: dict[str, torch.Tensor] = {}
    with safe_open(str(WEIGHTS), framework="pt") as f:
        for key in f.keys():
            nk = key.removeprefix("vision_tower.vision_model.")
            nk = nk.replace("embeddings.patch_embedding", "embeddings.proj")
            nk = nk.replace("encoder.layers", "layers")
            sd[nk] = f.get_tensor(key)
    return sd


def load_minimax(dtype: torch.dtype = torch.bfloat16) -> nn.Module:
    model = MiniMaxM3VLVisionModel(_build_config())
    model.load_state_dict(_remap_state_dict(), strict=True)
    model.eval()
    model.requires_grad_(False)
    return model.to(device="cuda", dtype=dtype)


def _patchify(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
    """Match MiniMaxM3VLImageProcessor layout for a single still image."""
    b, c, h, w = x.shape
    tps = TEMPORAL_PATCH_SIZE
    merge = SPATIAL_MERGE_SIZE
    ps = PATCH_SIZE
    patches = x.unsqueeze(1)  # B, T=1, C, H, W
    if patches.shape[1] % tps != 0:
        pad = tps - (patches.shape[1] % tps)
        patches = torch.cat([patches, patches[:, -1:].expand(-1, pad, -1, -1, -1)], dim=1)
    _, t, _, _, _ = patches.shape
    grid_t = t // tps
    grid_h, grid_w = h // ps, w // ps
    patches = patches.view(
        b,
        grid_t,
        tps,
        c,
        grid_h // merge,
        merge,
        ps,
        grid_w // merge,
        merge,
        ps,
    )
    patches = patches.permute(0, 1, 4, 7, 5, 8, 3, 2, 6, 9)
    flat = patches.reshape(b, grid_t * grid_h * grid_w, c * tps * ps * ps)
    grid = torch.tensor([[grid_t, grid_h, grid_w]], device=x.device, dtype=torch.int32)
    return flat[0], grid


def image_feat_minimax(model: nn.Module, x: torch.Tensor) -> torch.Tensor:
    """x: (1,3,H,W) in [0,1] on cuda, requires_grad=True. Returns (1, D) mean-pooled feature."""
    assert x.shape[0] == 1 and x.dim() == 4
    if x.shape[-1] != IMAGE_SIZE or x.shape[-2] != IMAGE_SIZE:
        x = F.interpolate(x, size=(IMAGE_SIZE, IMAGE_SIZE), mode="bicubic", align_corners=False)
    mean = torch.tensor(MEAN, device=x.device, dtype=x.dtype).view(1, 3, 1, 1)
    std = torch.tensor(STD, device=x.device, dtype=x.dtype).view(1, 3, 1, 1)
    px = (x - mean) / std
    dtype = next(model.parameters()).dtype
    pixel_values, grid_thw = _patchify(px.to(dtype=dtype))
    out = model(pixel_values=pixel_values, grid_thw=grid_thw)
    hidden = out.last_hidden_state.squeeze(0).float()
    return hidden.mean(dim=0, keepdim=True)