#!/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.nn import Linear4bit
from torch import nn
from transformers import AutoConfig


def quantize_tensor_to_4bit(tensor: torch.Tensor) -> torch.Tensor:
    """Quantize a single 2D weight tensor to 4-bit NF4."""
    in_features = tensor.size(1)
    out_features = tensor.size(0)

    # Create a temporary Linear4bit layer
    linear_4bit = Linear4bit(
        in_features,
        out_features,
        bias=False,
        compute_dtype=torch.bfloat16,
        quant_type="nf4",
    ).to(tensor.device)

    with torch.no_grad():
        linear_4bit.weight = nn.Parameter(tensor.clone())
        # Force quantization to happen
        _ = linear_4bit.weight.quant_state

    return linear_4bit.weight


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

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

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

    for idx, shard_path in enumerate(shards):
        print(f"[{idx+1}/{len(shards)}] Processing {Path(shard_path).name}...")

        # Load shard to GPU 0
        state_dict = load_file(shard_path, device="cuda:0")

        # Find all 2D weight tensors (Linear layers)
        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"   Found {len(weight_keys)} weight tensors to quantize")

        # Quantize
        for key in weight_keys:
            try:
                state_dict[key] = quantize_tensor_to_4bit(state_dict[key])
            except Exception as e:
                print(f"   Warning: Failed to quantize {key}: {e}")

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

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

        print(f"   Saved → {out_shard.name}")

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

    # Save config and index
    config.save_pretrained(output_path)

    # Create model.safetensors.index.json if needed (for multi-shard)
    print(f"\n✅ Streaming quantization complete!")
    print(f"   Output: {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)