#!/usr/bin/env python3
"""
Dequantize a bitsandbytes NF4 model to standard bf16/fp16 weights.

Usage:
    python dequantize_to_bf16.py                    # uses current directory
    python dequantize_to_bf16.py --input ./final-9b --output ./final-9b-bf16 --dtype bf16

Requires: torch, transformers, bitsandbytes, safetensors
"""

import argparse
import gc
import os
import torch
import safetensors.torch
from transformers import AutoModelForCausalLM, AutoConfig
from bitsandbytes.nn import Linear4bit


def dequantize_linear4bit(module: Linear4bit) -> torch.nn.Linear:
    """Convert a Linear4bit module to a standard nn.Linear with dequantized weights."""
    weight = module.weight
    # Dequantize the 4-bit weight
    if weight.quant_state is not None:
        dequantized = module.dequantize(weight)
    else:
        # Fallback: extract via bnb.functional
        import bitsandbytes.functional as bnb
        dequantized = bnb.functional.dequantize_4bit(
            weight.data,
            weight.quant_state
        )

    out_features, in_features = dequantized.shape
    new_linear = torch.nn.Linear(in_features, out_features, bias=module.bias is not None)
    new_linear.weight.data = dequantized.contiguous()
    if module.bias is not None:
        new_linear.bias.data = module.bias.contiguous()
    return new_linear


def dequantize_model(
    input_dir: str,
    output_dir: str,
    dtype: str = "bf16",
):
    device = "cpu"
    torch_dtype = torch.bfloat16 if dtype == "bf16" else torch.float16

    print(f"[*] Loading NF4 model from {input_dir} on CPU...")
    model = AutoModelForCausalLM.from_pretrained(
        input_dir,
        device_map=device,
        torch_dtype=torch_dtype,
    )
    print(f"    Loaded {sum(p.numel() for p in model.parameters()):,} total params")

    print("[*] Dequantizing Linear4bit modules...")
    for name, module in model.named_modules():
        if isinstance(module, Linear4bit):
            parent_name = ".".join(name.split(".")[:-1])
            child_name = name.split(".")[-1]
            if parent_name:
                parent = model.get_submodule(parent_name)
            else:
                parent = model
            new_linear = dequantize_linear4bit(module)
            setattr(parent, child_name, new_linear)
            print(f"    ✓ {name}")

    print(f"[*] Collecting state dict ({dtype})...")
    state_dict = model.state_dict()
    state_dict = {k: v.to(dtype=torch_dtype, device="cpu") for k, v in state_dict.items()}

    os.makedirs(output_dir, exist_ok=True)

    print(f"[*] Saving to {output_dir}...")
    safetensors.torch.save_file(state_dict, os.path.join(output_dir, "model.safetensors"))

    # Copy config files
    import shutil
    for f in ["config.json", "tokenizer.json", "tokenizer_config.json",
              "generation_config.json", "chat_template.jinja"]:
        src = os.path.join(input_dir, f)
        if os.path.exists(src):
            shutil.copy2(src, os.path.join(output_dir, f))

    print(f"[✓] Done! Model saved to {output_dir}")
    print(f"    Size: {sum(os.path.getsize(os.path.join(output_dir, f)) for f in os.listdir(output_dir) if os.path.isfile(os.path.join(output_dir, f))) / 1e9:.1f} GB")
    print(f"    Dtype: {dtype}")
    print(f"    Usage: AutoModelForCausalLM.from_pretrained('{output_dir}')")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Dequantize NF4 model to bf16/fp16")
    parser.add_argument("--input", default=".", help="Input NF4 model directory")
    parser.add_argument("--output", default="./dequantized", help="Output directory")
    parser.add_argument("--dtype", choices=["bf16", "fp16"], default="bf16",
                        help="Target dtype (default: bf16)")
    args = parser.parse_args()
    dequantize_model(args.input, args.output, args.dtype)