#!/usr/bin/env python3
"""Streaming quantization: process one shard at a time to avoid OOM."""

import argparse
import gc
import torch
from pathlib import Path
from transformers import AutoConfig, AutoModelForCausalLM
from bitsandbytes.nn import Linear4bit
from torch import nn


def streaming_quantize(model_path, output_path):
    """Quantize model by processing one shard at a time, using both GPUs."""
    
    print(f"Loading config from: {model_path}")
    config = AutoConfig.from_pretrained(model_path, trust_remote_code=True)
    
    # Find all safetensors files
    import glob
    safetensors_files = sorted(glob.glob(f"{model_path}/*.safetensors"))
    print(f"Found {len(safetensors_files)} safetensors files")
    
    # Create output directory
    output_path = Path(output_path)
    output_path.mkdir(parents=True, exist_ok=True)
    
    # Process each shard
    for shard_idx, shard_file in enumerate(safetensors_files):
        print(f"\n{'='*60}")
        print(f"Processing shard {shard_idx + 1}/{len(safetensors_files)}: {shard_file}")
        print(f"{'='*60}")
        
        # Load shard to CPU
        print("   Loading shard to CPU...")
        shard_state_dict = torch.load(shard_file, map_location="cpu", weights_only=False)
        
        # Quantize Linear layers in this shard using both GPUs
        print("   Quantizing Linear layers (both GPUs)...")
        quantized_keys = 0
        
        # Get all weight tensors
        weight_keys = [k for k, v in shard_state_dict.items() if 'weight' in k and v.dim() == 2]
        
        # Distribute between GPUs
        gpu0_keys = weight_keys[::2]  # Even indices
        gpu1_keys = weight_keys[1::2]  # Odd indices
        
        # Quantize on GPU 0
        print("      GPU 0: Quantizing...", end=" ")
        for key in gpu0_keys:
            tensor = shard_state_dict[key]
            in_features = tensor.size(1)
            out_features = tensor.size(0)
            
            dummy_linear = Linear4bit(
                in_features,
                out_features,
                bias=False,
                compute_dtype=torch.float16,
                quant_type='nf4',
            )
            
            with torch.no_grad():
                dummy_linear.weight = nn.Parameter(tensor.clone().to("cuda:0"))
                _ = dummy_linear.weight.quant_state
            
            shard_state_dict[key] = dummy_linear.weight.to("cpu")
            quantized_keys += 1
        print(f"✓ {len(gpu0_keys)} layers")
        
        # Quantize on GPU 1
        print("      GPU 1: Quantizing...", end=" ")
        for key in gpu1_keys:
            tensor = shard_state_dict[key]
            in_features = tensor.size(1)
            out_features = tensor.size(0)
            
            dummy_linear = Linear4bit(
                in_features,
                out_features,
                bias=False,
                compute_dtype=torch.float16,
                quant_type='nf4',
            )
            
            with torch.no_grad():
                dummy_linear.weight = nn.Parameter(tensor.clone().to("cuda:1"))
                _ = dummy_linear.weight.quant_state
            
            shard_state_dict[key] = dummy_linear.weight.to("cpu")
            quantized_keys += 1
        print(f"✓ {len(gpu1_keys)} layers")
        
        print(f"   ✓ Total: {quantized_keys} layers quantized")
        
        # Save quantized shard
        shard_name = f"model_shard_{shard_idx:05d}.safetensors"
        shard_path = output_path / shard_name
        print(f"   Saving to: {shard_path}")
        torch.save(shard_state_dict, shard_path)
        
        # Free memory
        del shard_state_dict
        gc.collect()
        torch.cuda.empty_cache()
    
    # Save config
    print(f"\n{'='*60}")
    print("Saving model config...")
    config.save_pretrained(output_path)
    print(f"✓ Model saved to: {output_path}")
    
    # Free memory
    gc.collect()
    torch.cuda.empty_cache()
    
    print("\n✓ Streaming quantization complete!")
    print(f"   Used both GPUs in parallel for faster quantization")


def main():
    parser = argparse.ArgumentParser(description="Streaming quantization")
    parser.add_argument("--model-path", type=str, 
                       default="/data/models/Ornith-1.0-35B",
                       help="Path to bf16 model")
    parser.add_argument("--output-path", type=str,
                       default="/data/models/Ornith-1.0-35B-streaming-4bit",
                       help="Output path for quantized model")
    args = parser.parse_args()
    
    streaming_quantize(args.model_path, args.output_path)


if __name__ == "__main__":
    main()