from __future__ import annotations

import argparse
import json
import math
from pathlib import Path
from statistics import mean
from typing import Any


def load_jsonl(path: Path) -> list[dict[str, Any]]:
    return [json.loads(line) for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]


def parse_prediction(row: dict[str, Any]) -> dict[str, Any] | None:
    text = row.get("predictionText", "")
    if isinstance(text, dict):
        return text
    start = str(text).find("{")
    end = str(text).rfind("}")
    if start < 0 or end <= start:
        return None
    try:
        parsed = json.loads(str(text)[start : end + 1])
    except json.JSONDecodeError:
        return None
    return parsed if isinstance(parsed, dict) else None


def invite(value: dict[str, Any]) -> bool:
    return bool(value.get("invite_to_speed_date") or value.get("passed"))


def main() -> None:
    parser = argparse.ArgumentParser(description="Evaluate matchmaker predictions against teacher or feedback labels.")
    parser.add_argument("--dataset", type=Path, required=True, help="Test/eval JSONL produced by prepare_dataset.py.")
    parser.add_argument("--predictions", type=Path, help="Prediction JSONL produced by predict.py.")
    args = parser.parse_args()

    dataset = {row["pairHash"]: row for row in load_jsonl(args.dataset)}
    if not args.predictions:
        scores = [float(row["target"]["compatibility"]) for row in dataset.values()]
        labels = {row["target"]["label"] for row in dataset.values()}
        print(
            json.dumps(
                {
                    "pairs": len(dataset),
                    "meanTeacherCompatibility": round(mean(scores), 3) if scores else 0,
                    "positiveRate": round(sum(invite(row["target"]) for row in dataset.values()) / max(1, len(dataset)), 4),
                    "labels": sorted(labels),
                },
                indent=2,
                sort_keys=True,
            )
        )
        return

    pred_rows = load_jsonl(args.predictions)
    mae_values: list[float] = []
    exact_invite = 0
    valid_json = 0
    evaluated = 0
    for pred_row in pred_rows:
        target_row = dataset.get(pred_row["pairHash"])
        if not target_row:
            continue
        evaluated += 1
        parsed = parse_prediction(pred_row)
        if not parsed:
            continue
        valid_json += 1
        target_score = float(target_row["target"]["compatibility"])
        pred_score = float(parsed.get("compatibility", math.nan))
        if not math.isnan(pred_score):
            mae_values.append(abs(pred_score - target_score))
        exact_invite += int(invite(parsed) == invite(target_row["target"]))

    print(
        json.dumps(
            {
                "evaluated": evaluated,
                "validJsonRate": round(valid_json / max(1, evaluated), 4),
                "compatibilityMae": round(mean(mae_values), 3) if mae_values else None,
                "inviteAccuracy": round(exact_invite / max(1, evaluated), 4),
            },
            indent=2,
            sort_keys=True,
        )
    )


if __name__ == "__main__":
    main()