"""Load the immanuelpeter/C-RADIOv4-H-3.5-Super vision bundle.

No language model weights are requested.
"""

import json
from pathlib import Path

import torch
from safetensors.torch import load_file
from torch import nn
from transformers import AutoConfig
from transformers.dynamic_module_utils import get_class_from_dynamic_module


class RMSNorm(nn.Module):
    def __init__(self, width):
        super().__init__()
        self.weight = nn.Parameter(torch.ones(width))

    def forward(self, x):
        value = x.float()
        value = value * torch.rsqrt(value.square().mean(-1, keepdim=True) + 1e-5)
        return (self.weight.float() * value).to(x.dtype)


class SquaredReLU(nn.Module):
    def forward(self, x):
        return torch.relu(x).square()


def pixel_shuffle(x):
    n, h, w, c = x.shape
    x = x.reshape(n, h, w // 2, c * 2).permute(0, 2, 1, 3).contiguous()
    return x.reshape(n, w // 2, h // 2, c * 4).permute(0, 2, 1, 3).contiguous()


class Vision(nn.Module):
    def __init__(self, bundle, dtype=torch.bfloat16):
        super().__init__()
        bundle = Path(bundle)
        self.settings = json.loads((bundle / "config.json").read_text())
        config = AutoConfig.from_pretrained(
            self.settings["code_config_source"], revision=self.settings["code_config_revision"], trust_remote_code=True,
        )
        radio_cls = get_class_from_dynamic_module(
            "hf_model.RADIOModel", "nvidia/C-RADIOv4-H", revision=self.settings["radio_code_revision"],
        )
        self.tower = radio_cls(config.vision_config)
        generator = self.tower.radio_model.model.patch_generator
        generator.video_embedder = nn.Linear(2 * 3 * 16 * 16, 1280, bias=False)
        weights = load_file(bundle / "model.safetensors")
        self.tower.radio_model.summary_idxs = weights.pop("summary_idxs")
        weights["radio_model.summary_idxs"] = self.tower.radio_model.summary_idxs
        self.tower.load_state_dict(weights, strict=True)
        self.tower.to(dtype)
        conditioner = self.tower.radio_model.input_conditioner
        conditioner.float()
        conditioner.load_state_dict({
            name: weights[f"radio_model.input_conditioner.{name}"]
            for name in ("norm_mean", "norm_std")
        }, strict=True)
        conditioner.dtype = dtype
        self.final_norm = nn.LayerNorm(1280, eps=1e-6) if self.settings["vision_final_layernorm"] else nn.Identity()
        if self.settings["vision_final_layernorm"]:
            self.final_norm.load_state_dict(load_file(bundle / "final_norm.safetensors"), strict=True)
        self.final_norm.to(dtype)
        self.projector = nn.Sequential(
            RMSNorm(5120), nn.Linear(5120, 20480, bias=False),
            SquaredReLU(), nn.Linear(20480, 4096, bias=False),
        )
        self.projector.load_state_dict(load_file(bundle / "projector.safetensors"), strict=True)
        self.projector.to(dtype)
        self.eval()

    @torch.inference_mode()
    def forward(self, pixel_values):
        """Accept RGB float32 pixels in [0, 1] on a grid divisible by 32."""
        h, w = pixel_values.shape[-2:]
        if h % 32 or w % 32:
            raise ValueError("pixel dimensions must be divisible by 32")
        tower = self.tower(pixel_values.float()).features
        merged = pixel_shuffle(self.final_norm(tower).reshape(len(tower), h // 16, w // 16, 1280)).flatten(1, 2)
        return {"tower": tower, "merged": merged, "projected": self.projector(merged)}


def load_vision(bundle, dtype=torch.bfloat16):
    return Vision(bundle, dtype=dtype)