#!/usr/bin/env python3
"""Test quantization on single shard."""

import gc
import torch
from pathlib import Path
from safetensors.torch import load_file
from bitsandbytes.functional import quantize_nf4


def test_one_shard(model_path, output_dir):
    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    
    import glob
    shards = sorted(glob.glob(f"{model_path}/*.safetensors"))
    
    if len(shards) == 0:
        print("No shards found!")
        return
    
    shard_file = shards[0]  # First shard only
    print(f"Testing shard: {shard_file}")
    
    state_dict = load_file(shard_file, device="cuda:0")
    
    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\n")
    
    quant_states = {}
    quantized = 0
    failed = 0
    
    for key in weight_keys:
        try:
            weight = state_dict[key]
            qweight, qstate = quantize_nf4(weight, blocksize=64, compress_statistics=True)
            state_dict[key] = qweight
            if qstate is not None:
                quant_states[f"{key}.quant_state"] = qstate
            quantized += 1
            print(f"✓ {key} -> {tuple(qweight.shape)}")
            del weight, qweight, qstate
            gc.collect()
            torch.cuda.empty_cache()
        except Exception as e:
            failed += 1
            print(f"✗ {key}: {e}")
            gc.collect()
            torch.cuda.empty_cache()
    
    print(f"\nResults: {quantized} quantized, {failed} failed, {len(quant_states)} quant states")
    
    # Save test output as model.safetensors
    state_dict_cpu = {k: v.cpu() if isinstance(v, torch.Tensor) else v for k, v in state_dict.items()}
    test_output = output_dir / "model.safetensors"
    torch.save({**state_dict_cpu, **quant_states}, test_output)
    print(f"Saved to: {test_output}")


if __name__ == "__main__":
    test_one_shard("/data/models/Ornith-1.0-35B", "/data/models/test_quantize")