#!/usr/bin/env python3
"""Export a q8 ONNX model and validate the browser-demo embeddings."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import onnx
from onnx import numpy_helper
from optimum.onnxruntime.configuration import AutoQuantizationConfig
from sentence_transformers import SentenceTransformer
from sentence_transformers.backend import export_dynamic_quantized_onnx_model

from goktugtr.text import harrier_query


EXAMPLE_QUERIES = [
    "Bir modeli küçük GPU'da nasıl verimli eğitebilirim?",
    "Arama sistemlerinde kalite nasıl ölçülür?",
    "LiDAR verisinden üç boyutlu harita üretmek",
    "Yapay zekâ sonuçları neden yeniden üretilebilir olmalı?",
]


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", type=Path, default=Path("outputs/goktugtr-270m/final"))
    parser.add_argument("--space-dir", type=Path, default=Path("space"))
    parser.add_argument(
        "--report", type=Path, default=Path("artifacts/onnx-browser-validation.json")
    )
    return parser.parse_args()


def encode(model: SentenceTransformer, texts: list[str]) -> np.ndarray:
    return model.encode(
        texts,
        batch_size=8,
        normalize_embeddings=True,
        convert_to_numpy=True,
        show_progress_bar=True,
    ).astype("float32")


def safe_quantization_config() -> object:
    """Quantize the large embedding table while preserving transformer fidelity."""
    return AutoQuantizationConfig.avx2(
        is_static=False,
        operators_to_quantize=["Gather"],
    )


def requantize_embedding_per_dimension(
    onnx_path: Path, pytorch_model: SentenceTransformer
) -> None:
    """Replace the lossy single-scale embedding quantizer with 640 scales."""
    embedding = (
        pytorch_model[0]
        .auto_model.get_input_embeddings()
        .weight.detach()
        .float()
        .cpu()
        .numpy()
    )
    scales = np.max(np.abs(embedding), axis=0).astype("float32") / 127.0
    scales = np.maximum(scales, np.finfo("float32").tiny)
    quantized = np.clip(np.rint(embedding / scales[None, :]), -127, 127).astype("int8")
    zero_points = np.zeros(scales.shape, dtype="int8")

    model = onnx.load(onnx_path)
    replacements = {
        "embed_tokens.weight_quantized": numpy_helper.from_array(
            quantized, name="embed_tokens.weight_quantized"
        ),
        "embed_tokens.weight_scale": numpy_helper.from_array(
            scales, name="embed_tokens.weight_scale"
        ),
        "embed_tokens.weight_zero_point": numpy_helper.from_array(
            zero_points, name="embed_tokens.weight_zero_point"
        ),
    }
    replaced: set[str] = set()
    for index, initializer in enumerate(model.graph.initializer):
        if initializer.name in replacements:
            model.graph.initializer[index].CopyFrom(replacements[initializer.name])
            replaced.add(initializer.name)
    if replaced != set(replacements):
        raise RuntimeError(f"Embedding quantizer tensors not found: {set(replacements) - replaced}")

    dequantizers = [
        node
        for node in model.graph.node
        if node.op_type == "DequantizeLinear"
        and node.name.endswith("Gather_output_0_DequantizeLinear")
    ]
    if len(dequantizers) != 1:
        raise RuntimeError(f"Expected one embedding dequantizer, found {len(dequantizers)}")
    del dequantizers[0].attribute[:]
    dequantizers[0].attribute.append(onnx.helper.make_attribute("axis", 2))
    onnx.checker.check_model(model)
    temporary = onnx_path.with_suffix(".tmp.onnx")
    onnx.save(model, temporary)
    temporary.replace(onnx_path)


def main() -> None:
    args = parse_args()
    corpus = json.loads((args.space_dir / "corpus.json").read_text(encoding="utf-8"))
    documents = [item["text"] for item in corpus]
    queries = [harrier_query(query) for query in EXAMPLE_QUERIES]

    pytorch_model = SentenceTransformer(str(args.model), device="cpu")
    pytorch_model.max_seq_length = 256
    reference = encode(pytorch_model, documents + queries)

    export_model = SentenceTransformer(str(args.model), backend="onnx", device="cpu")
    export_dynamic_quantized_onnx_model(
        export_model,
        quantization_config=safe_quantization_config(),
        model_name_or_path=str(args.model),
        file_suffix="quantized",
    )
    onnx_path = args.model / "onnx/model_quantized.onnx"
    requantize_embedding_per_dimension(onnx_path, pytorch_model)
    q8_model = SentenceTransformer(
        str(args.model),
        backend="onnx",
        device="cpu",
        model_kwargs={"file_name": "onnx/model_quantized.onnx"},
    )
    q8_model.max_seq_length = 256
    quantized = encode(q8_model, documents + queries)

    document_count = len(documents)
    reference_documents = reference[:document_count]
    reference_queries = reference[document_count:]
    quantized_documents = quantized[:document_count]
    quantized_queries = quantized[document_count:]
    embedding_cosines = np.sum(reference * quantized, axis=1)
    reference_top = np.argmax(reference_queries @ reference_documents.T, axis=1)
    quantized_top = np.argmax(quantized_queries @ quantized_documents.T, axis=1)

    report = {
        "format": "ONNX mixed int8 embedding / float transformer blocks",
        "quantization": "per-dimension int8 embedding; float transformer blocks",
        "documents": document_count,
        "queries": len(queries),
        "embedding_dimension": int(quantized.shape[1]),
        "mean_pytorch_to_q8_cosine": float(np.mean(embedding_cosines)),
        "minimum_pytorch_to_q8_cosine": float(np.min(embedding_cosines)),
        "maximum_absolute_component_difference": float(np.max(np.abs(reference - quantized))),
        "example_query_top1_agreement": float(np.mean(reference_top == quantized_top)),
        "reference_top1_indices": reference_top.tolist(),
        "quantized_top1_indices": quantized_top.tolist(),
    }
    args.report.parent.mkdir(parents=True, exist_ok=True)
    args.report.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
    (args.space_dir / "embeddings.json").write_text(
        json.dumps(np.round(quantized_documents, 7).tolist(), separators=(",", ":")) + "\n",
        encoding="utf-8",
    )
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()