#!/usr/bin/env python3
"""Small local examples for local_eval.py."""

import csv
import json
import tempfile
from pathlib import Path

from local_eval import chrf_similarity, read_practice, read_submission, score_predictions


def write_csv(path: Path, fieldnames: list[str], rows: list[dict[str, str]]) -> None:
    with path.open("w", newline="", encoding="utf-8") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        writer.writeheader()
        writer.writerows(rows)


def main() -> None:
    assert chrf_similarity("ká", "ká") == 1.0
    assert chrf_similarity("", "ká") == 0.0
    assert 0.0 < chrf_similarity("kà", "ká") < 1.0

    with tempfile.TemporaryDirectory() as tmpdir:
        tmp = Path(tmpdir)
        practice_path = tmp / "practice.csv"
        submission_path = tmp / "submission.csv"

        write_csv(
            practice_path,
            ["id", "context", "query", "task_type", "eval_type", "gold"],
            [
                {
                    "id": "row1",
                    "context": "toy context",
                    "query": "1. x\n2. y",
                    "task_type": "translation",
                    "eval_type": "single",
                    "gold": json.dumps(["ká", "ŋa"], ensure_ascii=False),
                },
                {
                    "id": "row2",
                    "context": "toy context",
                    "query": "1. z",
                    "task_type": "translation",
                    "eval_type": "single",
                    "gold": json.dumps(["č"], ensure_ascii=False),
                },
            ],
        )
        write_csv(
            submission_path,
            ["id", "pred"],
            [
                {"id": "row1", "pred": json.dumps([" ká ", "ŋa"], ensure_ascii=False)},
                {"id": "row2", "pred": json.dumps(["c"], ensure_ascii=False)},
            ],
        )

        metrics, details = score_predictions(read_practice(practice_path), read_submission(submission_path))
        assert metrics["rows"] == 2.0
        assert metrics["answers"] == 3.0
        assert metrics["length_ok"] == 1.0
        assert metrics["answer_exact"] == 2 / 3
        assert metrics["row_exact"] == 0.5
        assert 0.0 < metrics["chrf"] < 1.0
        assert details[0]["pred"] == json.dumps([" ká ", "ŋa"], ensure_ascii=False)


if __name__ == "__main__":
    main()