#!/usr/bin/env python3
"""True streaming 4-bit NF4 quantization - one shard at a time."""

import argparse
import gc
import torch
from pathlib import Path
from safetensors.torch import load_file, save_file
from bitsandbytes.functional import quantize_nf4
from transformers import AutoConfig


def quantize_weight_nf4(weight: torch.Tensor):
    """Quantize a single weight tensor to NF4 using bitsandbytes functional API."""
    if weight.dim() != 2:
        return weight, None

    # quantize_nf4 returns (quantized_tensor, quant_state)
    qweight, quant_state = quantize_nf4(
        weight,
        blocksize=64,
        compress_statistics=True,
    )
    return qweight, quant_state


def streaming_quantize(model_path: str, output_path: str):
    print(f"Streaming NF4 quantization: {model_path}")
    output_path = Path(output_path)
    output_path.mkdir(parents=True, exist_ok=True)

    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)

    import glob
    shards = sorted(glob.glob(f"{model_path}/*.safetensors"))
    print(f"Found {len(shards)} shards\n")

    for idx, shard_file in enumerate(shards):
        # Skip already processed shards
        shard_name = f"model-{idx:05d}-of-{len(shards):05d}.safetensors"
        if (output_path / shard_name).exists():
            print(f"[{idx+1}/{len(shards)}] {Path(shard_file).name} (already saved, skipping)")
            continue
        
        print(f"[{idx+1}/{len(shards)}] {Path(shard_file).name}")

        state_dict = load_file(shard_file, device="cuda:0")

        weight_keys = [
            k for k, v in state_dict.items()
            if "weight" in k and isinstance(v, torch.Tensor) and v.dim() == 2
        ]

        print(f"   Quantizing {len(weight_keys)} tensors...")

        for key in weight_keys:
            try:
                weight = state_dict[key]
                qweight, qstate = quantize_weight_nf4(weight)
                state_dict[key] = qweight
                del weight, qweight, qstate
                gc.collect()
                torch.cuda.empty_cache()
            except Exception as e:
                print(f"   Warning: Failed on {key}: {e}")
                gc.collect()
                torch.cuda.empty_cache()

        # Move to CPU and save
        state_dict = {
            k: v.cpu() if isinstance(v, torch.Tensor) else v
            for k, v in state_dict.items()
        }

        shard_name = f"model-{idx:05d}-of-{len(shards):05d}.safetensors"
        save_file(state_dict, output_path / shard_name)

        print(f"   Saved {shard_name}")

        del state_dict
        gc.collect()
        torch.cuda.empty_cache()

    config.save_pretrained(output_path)
    print(f"\n✅ Done → {output_path}")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--model-path", type=str, default="/data/models/Ornith-1.0-35B")
    parser.add_argument("--output-path", type=str, default="/data/models/Ornith-1.0-35B-nf4")
    args = parser.parse_args()

    streaming_quantize(args.model_path, args.output_path)