"""Preprocess GPT-OSS model: dequantize MXFP4 experts and unfuse into per-expert format.

This converts the GPT-OSS HF checkpoint from:
  - MXFP4 quantized fused expert weights (gate_up_proj_blocks/scales, down_proj_blocks/scales)
To:
  - BF16 per-expert weights (experts.{e}.gate_proj.weight, experts.{e}.up_proj.weight, etc.)

Usage:
    python tools/preprocess_gpt_oss.py \
        --input /path/to/gpt-oss-20b \
        --output /path/to/gpt-oss-20b-bf16
"""

import argparse
import json
import math
import os
import shutil
from collections import OrderedDict

import torch
from safetensors import safe_open
from safetensors.torch import save_file


def dequantize_mxfp4(
    blocks: torch.Tensor,
    scales: torch.Tensor,
    dtype: torch.dtype = torch.bfloat16,
) -> torch.Tensor:
    """Dequantize MXFP4 weights to BF16. Adapted from megatron.bridge."""
    FP4_VALUES = [
        +0.0,
        +0.5,
        +1.0,
        +1.5,
        +2.0,
        +3.0,
        +4.0,
        +6.0,
        -0.0,
        -0.5,
        -1.0,
        -1.5,
        -2.0,
        -3.0,
        -4.0,
        -6.0,
    ]
    scales = scales.to(torch.int32) - 127
    lut = torch.tensor(FP4_VALUES, dtype=dtype, device=blocks.device)

    *prefix_shape, G, B = blocks.shape
    rows_total = math.prod(prefix_shape) * G

    blocks_flat = blocks.reshape(rows_total, B)
    scales_flat = scales.reshape(rows_total, 1)

    out = torch.empty(rows_total, B * 2, dtype=dtype, device=blocks.device)

    rows_per_chunk = 32768 * 1024
    for r0 in range(0, rows_total, rows_per_chunk):
        r1 = min(r0 + rows_per_chunk, rows_total)
        blk = blocks_flat[r0:r1]
        exp = scales_flat[r0:r1]

        idx_lo = (blk & 0x0F).to(torch.long)
        idx_hi = (blk >> 4).to(torch.long)

        sub = out[r0:r1]
        sub[:, 0::2] = lut[idx_lo]
        sub[:, 1::2] = lut[idx_hi]
        torch.ldexp(sub, exp, out=sub)

    return out.reshape(*prefix_shape, G, B * 2).view(*prefix_shape, G * B * 2)


def preprocess_gpt_oss(input_dir: str, output_dir: str):
    os.makedirs(output_dir, exist_ok=True)

    # Load config
    with open(os.path.join(input_dir, "config.json")) as f:
        config = json.load(f)

    num_experts = config["num_local_experts"]
    intermediate_size = config["intermediate_size"]

    # Remove quantization config and ensure torch_dtype is bfloat16
    new_config = {k: v for k, v in config.items() if k != "quantization_config"}
    new_config["torch_dtype"] = "bfloat16"
    with open(os.path.join(output_dir, "config.json"), "w") as f:
        json.dump(new_config, f, indent=2)

    # Copy non-weight files
    for fname in os.listdir(input_dir):
        if fname in ("config.json", "model.safetensors.index.json"):
            continue
        if fname.endswith(".safetensors"):
            continue
        src = os.path.join(input_dir, fname)
        dst = os.path.join(output_dir, fname)
        if os.path.isfile(src) and not os.path.exists(dst):
            shutil.copy2(src, dst)

    # Process safetensors: collect all weight names
    index_path = os.path.join(input_dir, "model.safetensors.index.json")
    if os.path.exists(index_path):
        with open(index_path) as f:
            index = json.load(f)
        weight_map = index["weight_map"]
    else:
        # Single file
        weight_map = None

    # Group weights by safetensors file
    if weight_map:
        files = set(weight_map.values())
    else:
        files = [f for f in os.listdir(input_dir) if f.endswith(".safetensors")]

    all_output_tensors = OrderedDict()
    new_weight_map = {}

    for sf_file in sorted(files):
        sf_path = os.path.join(input_dir, sf_file)
        print(f"Processing {sf_file}...")
        with safe_open(sf_path, framework="pt", device="cpu") as f:
            keys = list(f.keys())
            for key in keys:
                tensor = f.get_tensor(key)

                # Check if this is a quantized expert weight
                if key.endswith("_blocks"):
                    base_name = key[: -len("_blocks")]
                    scales_key = base_name + "_scales"
                    # Load scales from same or different file
                    try:
                        scales = f.get_tensor(scales_key)
                    except Exception:
                        # Try loading from other files
                        scales = _load_tensor_from_files(input_dir, files, scales_key)

                    print(f"  Dequantizing {base_name}...")
                    dequantized = dequantize_mxfp4(tensor, scales)
                    _unfuse_experts(
                        base_name, dequantized, num_experts, intermediate_size, all_output_tensors, new_weight_map
                    )

                elif key.endswith("_scales"):
                    continue  # handled with _blocks

                elif ".mlp.experts.gate_up_proj_bias" in key:
                    # Unfuse bias: [E, 2*intermediate] -> per-expert gate_bias + up_bias
                    # GPT-OSS uses interleaved format: even=gate, odd=up
                    layer_prefix = key.rsplit(".mlp.experts.gate_up_proj_bias", 1)[0]
                    for e in range(num_experts):
                        gate_bias = tensor[e, 0::2]
                        up_bias = tensor[e, 1::2]
                        gname = f"{layer_prefix}.mlp.experts.{e}.gate_proj.bias"
                        uname = f"{layer_prefix}.mlp.experts.{e}.up_proj.bias"
                        all_output_tensors[gname] = gate_bias.contiguous()
                        all_output_tensors[uname] = up_bias.contiguous()
                        new_weight_map[gname] = "model.safetensors"
                        new_weight_map[uname] = "model.safetensors"

                elif ".mlp.experts.down_proj_bias" in key:
                    # Unfuse bias: [E, hidden] -> per-expert
                    layer_prefix = key.rsplit(".mlp.experts.down_proj_bias", 1)[0]
                    for e in range(num_experts):
                        dname = f"{layer_prefix}.mlp.experts.{e}.down_proj.bias"
                        all_output_tensors[dname] = tensor[e].contiguous()
                        new_weight_map[dname] = "model.safetensors"

                else:
                    all_output_tensors[key] = tensor
                    new_weight_map[key] = "model.safetensors"

    # Save output
    print(f"Saving {len(all_output_tensors)} tensors...")

    # Split into chunks of ~5GB each
    chunk_size = 5 * 1024 * 1024 * 1024  # 5GB
    chunks = []
    current_chunk = OrderedDict()
    current_size = 0

    for name, tensor in all_output_tensors.items():
        tensor_size = tensor.numel() * tensor.element_size()
        if current_size + tensor_size > chunk_size and current_chunk:
            chunks.append(current_chunk)
            current_chunk = OrderedDict()
            current_size = 0
        current_chunk[name] = tensor
        current_size += tensor_size

    if current_chunk:
        chunks.append(current_chunk)

    final_weight_map = {}
    if len(chunks) == 1:
        out_file = "model.safetensors"
        save_file(chunks[0], os.path.join(output_dir, out_file))
        for k in chunks[0]:
            final_weight_map[k] = out_file
    else:
        total = len(chunks)
        for i, chunk in enumerate(chunks):
            out_file = f"model-{i+1:05d}-of-{total:05d}.safetensors"
            save_file(chunk, os.path.join(output_dir, out_file))
            for k in chunk:
                final_weight_map[k] = out_file

        # Save index
        total_size = sum(t.numel() * t.element_size() for t in all_output_tensors.values())
        index_data = {
            "metadata": {"total_size": total_size},
            "weight_map": final_weight_map,
        }
        with open(os.path.join(output_dir, "model.safetensors.index.json"), "w") as f:
            json.dump(index_data, f, indent=2)

    print(f"Done! Output saved to {output_dir}")


def _unfuse_experts(base_name, dequantized, num_experts, intermediate_size, output_tensors, weight_map):
    """Unfuse 3D expert tensor into per-expert format."""
    layer_prefix = base_name.rsplit(".mlp.experts.", 1)[0]
    weight_type = base_name.rsplit(".mlp.experts.", 1)[1]

    if "gate_up_proj" in weight_type:
        # [E, 2*intermediate, hidden] -> per-expert gate + up
        for e in range(num_experts):
            expert_weight = dequantized[e]  # [2*intermediate, hidden]
            # GPT-OSS uses interleaved format: [g0,u0,g1,u1,...] (even=gate, odd=up)
            gate = expert_weight[0::2]  # [intermediate, hidden]
            up = expert_weight[1::2]  # [intermediate, hidden]
            gname = f"{layer_prefix}.mlp.experts.{e}.gate_proj.weight"
            uname = f"{layer_prefix}.mlp.experts.{e}.up_proj.weight"
            output_tensors[gname] = gate.contiguous()
            output_tensors[uname] = up.contiguous()
            weight_map[gname] = "model.safetensors"
            weight_map[uname] = "model.safetensors"
    elif "down_proj" in weight_type:
        # [E, hidden, intermediate] -> per-expert
        for e in range(num_experts):
            dname = f"{layer_prefix}.mlp.experts.{e}.down_proj.weight"
            output_tensors[dname] = dequantized[e].contiguous()
            weight_map[dname] = "model.safetensors"


def _load_tensor_from_files(input_dir, files, key):
    """Load a tensor by searching across safetensors files."""
    for sf_file in files:
        sf_path = os.path.join(input_dir, sf_file)
        with safe_open(sf_path, framework="pt", device="cpu") as f:
            if key in f.keys():
                return f.get_tensor(key)
    raise KeyError(f"Tensor {key} not found in any safetensors file")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Preprocess GPT-OSS model")
    parser.add_argument("--input", required=True, help="Input HF model directory")
    parser.add_argument("--output", required=True, help="Output directory for BF16 model")
    args = parser.parse_args()
    preprocess_gpt_oss(args.input, args.output)