#!/usr/bin/env python3
"""Quantize Qwen2.5-0.5B-Instruct with UltraChat AutoRound plus a GPTQ lm_head."""

import argparse
import importlib.metadata
import platform
import shutil
from pathlib import Path

import torch
from compressed_tensors.utils import match_named_modules
from datasets import Dataset, load_dataset
from llmcompressor import oneshot
from llmcompressor.modifiers.autoround import AutoRoundModifier
from llmcompressor.modifiers.quantization import GPTQModifier
from transformers import AutoModelForCausalLM, AutoTokenizer


MODEL_LABEL = "Qwen2.5-0.5B-Instruct"
BASE_MODEL = "Qwen/Qwen2.5-0.5B-Instruct"
SOURCE_REVISION = "7ae557604adf67be50417f59c2c2f167def9a775"
CALIBRATION_DATASET = "HuggingFaceH4/ultrachat_200k"
CALIBRATION_SPLIT = "train_sft[:512]"
NUM_CALIBRATION_SAMPLES = 512
MAX_SEQ_LENGTH = 1024
AUTOROUND_BITS = 4
AUTOROUND_GROUP_SIZE = 128
AUTOROUND_ITERS = 200
AUTOROUND_BATCH_SIZE = 1
LM_HEAD_BITS = 4
LM_HEAD_GROUP_SIZE = 128
GPTQ_BLOCK_SIZE = 128
GPTQ_DAMPENING_FRAC = 0.01


def package_version(*names):
    for name in names:
        try:
            return importlib.metadata.version(name)
        except importlib.metadata.PackageNotFoundError:
            continue
    return "not-installed"


def weight_config(bits, group_size):
    return {
        "num_bits": bits,
        "type": "int",
        "symmetric": True,
        "group_size": group_size,
        "strategy": "group",
        "dynamic": False,
    }


def calibration_dataset(tokenizer):
    if not tokenizer.chat_template:
        raise ValueError(
            "Tokenizer has no chat template; refusing untemplated calibration."
        )
    source = load_dataset(CALIBRATION_DATASET, split=CALIBRATION_SPLIT)
    if "messages" not in source.column_names:
        raise ValueError(
            f"{CALIBRATION_DATASET} {CALIBRATION_SPLIT} lacks messages."
        )

    packed_ids, buffer = [], []
    for example in source:
        text = tokenizer.apply_chat_template(
            example["messages"], tokenize=False, add_generation_prompt=False
        )
        buffer.extend(tokenizer(text, add_special_tokens=False)["input_ids"])
        while (
            len(buffer) >= MAX_SEQ_LENGTH
            and len(packed_ids) < NUM_CALIBRATION_SAMPLES
        ):
            packed_ids.append(buffer[:MAX_SEQ_LENGTH])
            buffer = buffer[MAX_SEQ_LENGTH:]
        if len(packed_ids) == NUM_CALIBRATION_SAMPLES:
            break
    if len(packed_ids) != NUM_CALIBRATION_SAMPLES:
        raise ValueError(
            f"UltraChat did not provide {NUM_CALIBRATION_SAMPLES} complete "
            f"{MAX_SEQ_LENGTH}-token spans."
        )

    dataset = Dataset.from_dict(
        {
            "input_ids": packed_ids,
            "attention_mask": [[1] * MAX_SEQ_LENGTH for _ in packed_ids],
        }
    )
    dataset.set_format(type="torch")
    return dataset


def validate_targets(model):
    decoder = list(
        match_named_modules(model, targets=["Linear"], ignore=["lm_head"])
    )
    head = list(match_named_modules(model, targets=["lm_head"]))
    if not decoder or [name for name, _ in head] != ["lm_head"]:
        raise ValueError(
            f"Invalid targets: decoder={len(decoder)}, "
            f"head={[name for name, _ in head]}"
        )
    return len(decoder)


def recipe():
    return [
        AutoRoundModifier(
            targets=["Linear"],
            ignore=["lm_head"],
            iters=AUTOROUND_ITERS,
            batch_size=AUTOROUND_BATCH_SIZE,
            device_ids="0",
            enable_torch_compile=False,
            config_groups={
                "decoder_autoround": {
                    "targets": ["Linear"],
                    "weights": weight_config(
                        AUTOROUND_BITS, AUTOROUND_GROUP_SIZE
                    ),
                }
            },
        ),
        GPTQModifier(
            targets=["lm_head"],
            block_size=GPTQ_BLOCK_SIZE,
            dampening_frac=GPTQ_DAMPENING_FRAC,
            actorder="static",
            offload_hessians=False,
            config_groups={
                "lm_head_gptq": {
                    "targets": ["lm_head"],
                    "weights": weight_config(
                        LM_HEAD_BITS, LM_HEAD_GROUP_SIZE
                    ),
                }
            },
        ),
    ]


def write_release_files(source, output, decoder_count):
    for item in source.iterdir():
        if (
            item.is_file()
            and item.suffix in {".json", ".txt", ".model", ".jinja"}
            and not item.name.endswith(".index.json")
            and not (output / item.name).exists()
        ):
            shutil.copy2(item, output / item.name)
    shutil.copy2(Path(__file__).resolve(), output / "quantize.py")
    output.joinpath("recipe.yaml").write_text(
        f"# Effective {MODEL_LABEL} UltraChat AutoRound recipe\n"
        "decoder_autoround:\n"
        "  method: AutoRound\n"
        "  targets: Linear excluding lm_head\n"
        f"  iters: {AUTOROUND_ITERS}\n"
        f"  batch_size: {AUTOROUND_BATCH_SIZE}\n"
        f"  weights: {{num_bits: {AUTOROUND_BITS}, type: int, symmetric: true, "
        f"strategy: group, group_size: {AUTOROUND_GROUP_SIZE}}}\n"
        "  calibration:\n"
        f"    dataset: {CALIBRATION_DATASET}\n"
        f"    split: {CALIBRATION_SPLIT}\n"
        "    preprocessing: apply_chat_template; deterministic packed spans; "
        "add_special_tokens=false\n"
        f"    samples: {NUM_CALIBRATION_SAMPLES}\n"
        f"    sequence_length: {MAX_SEQ_LENGTH}\n"
        "    shuffle: false\n"
        "lm_head_gptq:\n"
        "  method: GPTQ\n"
        "  target: lm_head\n"
        f"  block_size: {GPTQ_BLOCK_SIZE}\n"
        f"  dampening_frac: {GPTQ_DAMPENING_FRAC}\n"
        "  actorder: static\n"
        f"  weights: {{num_bits: {LM_HEAD_BITS}, type: int, symmetric: true, "
        f"strategy: group, group_size: {LM_HEAD_GROUP_SIZE}}}\n"
        f"decoder_target_count: {decoder_count}\n"
        f"base_model: {BASE_MODEL}\n"
        f"source_revision: {SOURCE_REVISION}\n"
    )
    output.joinpath("versions.txt").write_text(
        "\n".join(
            (
                f"Python: {platform.python_version()}",
                f"torch: {torch.__version__}",
                f"CUDA: {torch.version.cuda or 'not-built'}",
                f"transformers: {package_version('transformers')}",
                f"llm-compressor: {package_version('llmcompressor', 'llm-compressor')}",
                f"auto-round: {package_version('auto-round', 'auto_round')}",
                f"compressed-tensors: {package_version('compressed-tensors')}",
            )
        )
        + "\n"
    )


def validate_scales(output):
    from safetensors.torch import safe_open

    count = 0
    for shard in output.glob("*.safetensors"):
        with safe_open(shard, framework="pt", device="cpu") as tensors:
            for name in tensors.keys():
                if "scale" not in name.lower():
                    continue
                value = tensors.get_tensor(name)
                count += 1
                if value.is_floating_point() and not torch.isfinite(
                    value.float()
                ).all():
                    raise ValueError(f"Non-finite scale: {shard.name}:{name}")
    if not count:
        raise ValueError("No saved scale tensors found.")


def smoke_test(output):
    tokenizer = AutoTokenizer.from_pretrained(output, trust_remote_code=False)
    model = AutoModelForCausalLM.from_pretrained(
        output,
        device_map="auto",
        torch_dtype="auto",
        trust_remote_code=False,
    )
    prompt = tokenizer.apply_chat_template(
        [{"role": "user", "content": "Reply with one word: ready."}],
        add_generation_prompt=True,
        tokenize=False,
    )
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    with torch.inference_mode():
        generated = model.generate(
            **inputs, max_new_tokens=4, do_sample=False
        )
    if generated.shape[-1] <= inputs["input_ids"].shape[-1]:
        raise ValueError("Generation smoke test produced no new tokens.")


def quantize(model_path, output_dir):
    source, output = Path(model_path), Path(output_dir)
    if output.exists() and any(output.iterdir()):
        raise FileExistsError(f"Refusing to overwrite: {output}")

    model = AutoModelForCausalLM.from_pretrained(
        source,
        device_map="cpu",
        torch_dtype="auto",
        trust_remote_code=False,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        source, trust_remote_code=False
    )
    decoder_count = validate_targets(model)
    oneshot(
        model=model,
        dataset=calibration_dataset(tokenizer),
        recipe=recipe(),
        max_seq_length=MAX_SEQ_LENGTH,
        num_calibration_samples=NUM_CALIBRATION_SAMPLES,
        shuffle_calibration_samples=False,
    )

    output.mkdir(parents=True, exist_ok=True)
    model.save_pretrained(output, save_compressed=True)
    tokenizer.save_pretrained(output)
    write_release_files(source, output, decoder_count)
    validate_scales(output)
    smoke_test(output)
    print(f"Quantized model written to: {output}")


def main():
    parser = argparse.ArgumentParser(
        description=f"{MODEL_LABEL} UltraChat AutoRound/GPTQ quantization"
    )
    parser.add_argument("--model-path", required=True)
    parser.add_argument("--output-dir", required=True)
    args = parser.parse_args()
    quantize(args.model_path, args.output_dir)


if __name__ == "__main__":
    main()