"""Load the mixed packed Inkling 2 checkpoint with the native Transformers model."""

from pathlib import Path

import torch
from huggingface_hub import hf_hub_download
from safetensors.torch import load_file
from transformers import Inkling2Config, Inkling2ForConditionalGeneration
from transformers.conversion_mapping import get_model_conversion_mapping
from transformers.core_model_loading import WeightConverter, WeightRenaming, rename_source_key


def dequantize(weight, scale, amax=None, denominator=None):
    if weight.dtype == torch.uint8:
        if amax is None or denominator is None:
            raise ValueError("NVFP4 weights require per-expert amax and an outer-scale denominator")
        codes = torch.stack((weight & 15, weight >> 4), dim=-1).flatten(-2)
        levels = weight.new_tensor([0, 0.5, 1, 1.5, 2, 3, 4, 6], dtype=torch.float32)
        values = levels[(codes & 7).long()] * torch.where(codes < 8, 1.0, -1.0)
        return values * scale.float().repeat_interleave(16, -1) * (amax / denominator).reshape(-1, 1, 1)
    if weight.dtype == torch.float8_e4m3fn:
        scales = torch.pow(2.0, scale.float() - 127).repeat_interleave(32, -1)
        return weight.float() * scales
    raise ValueError(f"Unsupported packed expert dtype: {weight.dtype}")


class Inkling2TinyForConditionalGeneration(Inkling2ForConditionalGeneration):
    @classmethod
    def from_pretrained(
        cls, pretrained_model_name_or_path, *model_args, config=None, dtype=torch.bfloat16,
        revision=None, cache_dir=None, local_files_only=False, token=None,
        attn_implementation="eager", **kwargs,
    ):
        if model_args:
            raise ValueError("Inkling 2 does not take positional construction arguments")
        options = dict(revision=revision, cache_dir=cache_dir, local_files_only=local_files_only, token=token)
        root = Path(pretrained_model_name_or_path)

        def resolve(filename):
            if root.is_dir():
                return root / filename
            return Path(hf_hub_download(str(pretrained_model_name_or_path), filename, **options))

        import json

        raw = json.loads(resolve("config.json").read_text())
        # The native config otherwise assumes the all-MXFP8 sibling. This
        # adapter dequantizes each tensor according to its actual stored dtype.
        raw.pop("quantization_config", None)
        config = Inkling2Config(**raw)
        config._attn_implementation = attn_implementation
        dtype = kwargs.pop("torch_dtype", dtype)
        if dtype in (None, "auto"):
            dtype = torch.bfloat16
        elif isinstance(dtype, str):
            dtype = getattr(torch, dtype)
        if dtype not in (torch.float32, torch.bfloat16, torch.float16):
            raise ValueError("The Transformers adapter requires a floating compute dtype")
        device_map = kwargs.pop("device_map", None)
        if device_map == "auto":
            device_map = "cuda" if torch.cuda.is_available() else "cpu"
        if isinstance(device_map, dict):
            if set(device_map) != {""}:
                raise ValueError("Use a single device for this tiny Transformers adapter")
            device_map = device_map[""]
        allowed = {"_from_auto", "_from_pipeline", "trust_remote_code", "weights_only", "adapter_kwargs", "_commit_hash"}
        unexpected = set(kwargs) - allowed
        if unexpected:
            raise TypeError(f"Unsupported loading arguments: {sorted(unexpected)}")
        model = Inkling2ForConditionalGeneration(config)
        transforms = get_model_conversion_mapping(model)
        renamings = [t for t in transforms if isinstance(t, WeightRenaming)]
        converters = [t for t in transforms if isinstance(t, WeightConverter)]
        source = load_file(str(resolve("model.safetensors")))
        state = {}
        for name, weight in source.items():
            if name.endswith(("_scale", "_amax", "_outer_scale_denominator")):
                continue
            if name + "_scale" in source:
                weight = dequantize(
                    weight, source[name + "_scale"], source.get(name + "_amax"),
                    source.get(name + "_outer_scale_denominator"),
                )
            target, _ = rename_source_key(name, renamings, converters)
            if name.endswith("w13_weight"):
                # The source interleaves gate and up; Transformers chunks them.
                e, rows, columns = weight.shape
                weight = weight.reshape(e, rows // 2, 2, columns).transpose(1, 2).reshape(e, rows, columns)
            if target in state:
                raise ValueError(f"Checkpoint keys map to the same parameter: {target}")
            keep_float32 = source[name].dtype == torch.float32 or any(
                part in target for part in model._keep_in_fp32_modules_strict
            )
            target_dtype = torch.float32 if keep_float32 else dtype
            state[target] = weight.to(target_dtype).contiguous()
        model.load_state_dict(state, strict=True, assign=True)
        model.eval()
        if device_map is not None:
            model.to(device_map)
        return model