"""Parity evaluation: run the *deployed* inference service (ONNX int8 + its own
chunking + aggregation) on the canonical EN document test set, and compare with
the torch numbers from the report.

This validates the full production path end to end — we deploy what we
evaluated, or we find out here.

Usage (from the repo root):
    uv run --group train python training/eval_service.py \
        --artifact models/cefr/en_chunked_weighted/onnx-int8 \
        --config training/configs/en_only.toml
"""

import argparse
import json
from pathlib import Path

from datasets import load_dataset
from evaluation import DOC_FORMATS
from train_cefr import build_parts, load_config

from tutor.ml.cefr.inference import CEFRClassifier
from tutor.ml.cefr.metrics import classification_report
from tutor.ml.cefr.preprocessing import CANONICAL_LEVELS, normalize_level

LEVEL_TO_RANK = {level: rank for rank, level in enumerate(CANONICAL_LEVELS)}


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--artifact", type=Path, required=True, help="onnx-int8 artifact dir")
    parser.add_argument("--config", type=Path, required=True)
    parser.add_argument("--lang", default="en")
    args = parser.parse_args()

    config = load_config(args.config)
    parts = build_parts(config)
    test_doc_ids = {
        p.doc_id for p in parts["test"] if p.source_format in DOC_FORMATS and p.lang == args.lang
    }
    print(f"Canonical test: {len(test_doc_ids)} {args.lang} documents")

    classifier = CEFRClassifier.from_dir(args.artifact)

    y_true: list[int] = []
    y_pred: list[int] = []
    for subset in config["data"]["subsets"]:
        corpus = subset.split("/")[-1]
        if not any(doc_id.startswith(f"{corpus}:") for doc_id in test_doc_ids):
            continue
        dataset = load_dataset(subset, split="train")
        for index, row in enumerate(dataset):
            doc_id = f"{corpus}:{index}"
            if doc_id not in test_doc_ids:
                continue
            level = normalize_level(row.get("cefr_level"))
            text = str(row.get("text") or "")
            if level is None or not text.strip():
                continue
            prediction = classifier.classify_text(text)
            y_true.append(LEVEL_TO_RANK[level])
            y_pred.append(LEVEL_TO_RANK[prediction.level])

    report = classification_report(y_true, y_pred)
    print(json.dumps({"service_int8_document": report}, indent=2))

    import mlflow

    tracking = config.get("tracking", {})
    mlflow.set_tracking_uri(tracking.get("uri", "sqlite:///mlflow.db"))
    mlflow.set_experiment(tracking.get("experiment", "cefr-classifier"))
    with mlflow.start_run(run_name=f"eval_service__{args.artifact.parent.name}_int8"):
        mlflow.log_param("artifact", str(args.artifact))
        mlflow.log_metrics({f"service_doc_{k}": v for k, v in report.items()})


if __name__ == "__main__":
    main()