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

import argparse
import gc
import torch
import concurrent.futures
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")

    # Process 4 shards per GPU (8 total) for max parallelism
    shards_per_gpu = 4
    num_gpus = 2
    max_parallel = shards_per_gpu * num_gpus
    
    def process_shard(idx, shard_file, gpu_id):
        """Process a single shard (called per-GPU)."""
        shard_name = f"model-{idx:05d}-of-{len(shards):05d}.safetensors"
        if (output_path / shard_name).exists():
            return
        
        print(f"[{idx+1}/{len(shards)}] {Path(shard_file).name} (GPU {gpu_id})")
        
        # Load to assigned GPU
        state_dict = load_file(shard_file, device=f"cuda:{gpu_id}")
        
        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()
        }
        
        save_file(state_dict, output_path / shard_name)
        print(f"   ✓ Saved {shard_name}")
        
        del state_dict
        gc.collect()
        torch.cuda.empty_cache()
    
    # Process in batches (4 per GPU = 8 total)
    for i in range(0, len(shards), max_parallel):
        batch = shards[i:i+max_parallel]
        if all((output_path / f"model-{idx:05d}-of-{len(shards):05d}.safetensors").exists() for idx in range(i, i+len(batch))):
            print(f"Shards {i+1}-{i+len(batch)} already saved, skipping")
            continue
        
        with concurrent.futures.ThreadPoolExecutor(max_workers=max_parallel) as executor:
            futures = []
            for j, shard_file in enumerate(batch):
                idx = i + j
                gpu_id = (j // shards_per_gpu) % num_gpus  # Assign to GPU
                futures.append(executor.submit(process_shard, idx, shard_file, gpu_id))
            for future in concurrent.futures.as_completed(futures):
                future.result()

    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)