#!/usr/bin/env python3
"""Quantize bf16 model to BnB 4-bit by replacing Linear layers."""

import argparse
import gc
import torch
from pathlib import Path
from transformers import AutoModelForCausalLM, AutoConfig


def quantize_model(model_path, output_path):
    """Load bf16 model, quantize to BnB 4-bit, save."""
    
    print(f"Loading model from: {model_path}")
    model = AutoModelForCausalLM.from_pretrained(
        model_path,
        device_map="cpu",
        torch_dtype=torch.bfloat16,
        trust_remote_code=True,
        low_cpu_mem_usage=True,
    )
    print(f"✓ Model loaded to CPU (~70GB bf16)")
    
    # Count parameters
    total_params = sum(p.numel() for p in model.parameters())
    print(f"   Total parameters: {total_params / 1e9:.2f}B")
    
    # Replace Linear layers with Linear4bit
    print("\nReplacing Linear layers with Linear4bit...")
    from bitsandbytes.nn import Linear4bit
    from torch import nn
    
    quantized_count = 0
    for name, module in list(model.named_modules()):
        if isinstance(module, nn.Linear) and 'lm_head' not in name:
            # Create 4-bit version
            new_module = Linear4bit(
                module.in_features,
                module.out_features,
                bias=module.bias is not None,
                compute_dtype=torch.float16,
            )
            
            # Copy weights
            with torch.no_grad():
                new_module.weight = nn.Parameter(
                    module.weight.data.clone()
                )
                if module.bias is not None:
                    new_module.bias = nn.Parameter(
                        module.bias.data.clone()
                    )
            
            # Replace in model
            layers = name.split('.')
            parent = model
            for layer in layers[:-1]:
                parent = getattr(parent, layer)
            setattr(parent, layers[-1], new_module)
            quantized_count += 1
    
    print(f"✓ Quantized {quantized_count} linear layers to 4-bit")
    
    # Count 4-bit parameters
    bnb_params = sum(
        p.numel() for p in model.parameters() 
        if hasattr(p, 'quant_state')
    )
    print(f"   4-bit parameters: {bnb_params / 1e9:.2f}B")
    
    # Save model
    print(f"\nSaving to: {output_path}")
    model.save_pretrained(output_path)
    print("✓ Model saved")
    
    # Free memory
    del model
    gc.collect()
    torch.cuda.empty_cache()
    
    print("\nDone! Model is ready for QLoRA training.")


def main():
    parser = argparse.ArgumentParser(description="Quantize model to BnB 4-bit")
    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-bnb-4bit",
                       help="Output path for quantized model")
    args = parser.parse_args()
    
    Path(args.output_path).mkdir(parents=True, exist_ok=True)
    quantize_model(args.model_path, args.output_path)


if __name__ == "__main__":
    main()