#!/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."""
    
    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=True)
        
        # Quantize Linear layers in this shard
        print("   Quantizing Linear layers...")
        quantized_keys = 0
        for key, tensor in list(shard_state_dict.items()):
            if 'weight' in key and tensor.dim() == 2:
                # This is a Linear layer weight
                # Create a dummy Linear4bit to get the quantization format
                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',
                )
                
                # Quantize the weight
                with torch.no_grad():
                    dummy_linear.weight = nn.Parameter(tensor.clone())
                    # Force quantization by accessing quant_state
                    _ = dummy_linear.weight.quant_state
                
                # Replace with quantized version
                shard_state_dict[key] = dummy_linear.weight
                quantized_keys += 1
        
        print(f"   ✓ Quantized {quantized_keys} weights")
        
        # 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()
    
    # 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!")


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()