"""Evaluate shopper agent predictions against labeled scenarios."""

from __future__ import annotations

import csv
import io
import json
from collections import defaultdict
from dataclasses import dataclass, field
from pathlib import Path
from typing import Any


REQUIRED_PREDICTION_FIELDS = {"record_id", "purchase_probability"}


@dataclass
class EvaluationResult:
    n_records: int
    mae: float
    rmse: float
    brier: float
    persona_mae: dict[str, float] = field(default_factory=dict)
    category_mae: dict[str, float] = field(default_factory=dict)
    design_effect_mae: float = 0.0
    missing_record_ids: list[str] = field(default_factory=list)
    extra_record_ids: list[str] = field(default_factory=list)

    def to_dict(self) -> dict[str, Any]:
        return {
            "n_records": self.n_records,
            "mae": round(self.mae, 4),
            "rmse": round(self.rmse, 4),
            "brier": round(self.brier, 4),
            "persona_mae": {k: round(v, 4) for k, v in self.persona_mae.items()},
            "category_mae": {k: round(v, 4) for k, v in self.category_mae.items()},
            "design_effect_mae": round(self.design_effect_mae, 4),
            "missing_record_ids": self.missing_record_ids,
            "extra_record_ids": self.extra_record_ids,
        }


def load_labels_csv(path: Path) -> list[dict[str, Any]]:
    with path.open(newline="", encoding="utf-8") as handle:
        return list(csv.DictReader(handle))


def load_predictions(path: Path) -> list[dict[str, Any]]:
    if path.suffix.lower() == ".jsonl":
        rows: list[dict[str, Any]] = []
        for line in path.read_text(encoding="utf-8").splitlines():
            line = line.strip()
            if line:
                rows.append(json.loads(line))
        return rows
    if path.suffix.lower() == ".json":
        data = json.loads(path.read_text(encoding="utf-8"))
        if isinstance(data, list):
            return data
        if isinstance(data, dict) and "predictions" in data:
            return data["predictions"]
        raise ValueError("JSON file must be a list or {\"predictions\": [...]}")
    if path.suffix.lower() == ".csv":
        with path.open(newline="", encoding="utf-8") as handle:
            return list(csv.DictReader(handle))
    raise ValueError(f"Unsupported predictions format: {path}")


def _looks_like_csv(text: str) -> bool:
    first = next((line for line in text.splitlines() if line.strip()), "")
    if not first or first.lstrip().startswith(("{", "[")):
        return False
    return "record_id" in first.lower() and "," in first


def parse_predictions_text(text: str, *, filename: str | None = None) -> list[dict[str, Any]]:
    text = text.strip()
    if not text:
        raise ValueError("Empty predictions payload")

    suffix = Path(filename).suffix.lower() if filename else ""
    if suffix == ".csv" or _looks_like_csv(text):
        return list(csv.DictReader(io.StringIO(text)))

    if text.startswith("["):
        data = json.loads(text)
        if not isinstance(data, list):
            raise ValueError("JSON payload must be a list of predictions")
        return data
    rows: list[dict[str, Any]] = []
    for line in text.splitlines():
        line = line.strip()
        if line:
            rows.append(json.loads(line))
    if not rows:
        raise ValueError("No predictions found in JSONL payload")
    return rows


def _validate_predictions(predictions: list[dict[str, Any]]) -> None:
    for idx, row in enumerate(predictions):
        missing = REQUIRED_PREDICTION_FIELDS - set(row.keys())
        if missing:
            raise ValueError(f"Prediction row {idx} missing fields: {sorted(missing)}")
        prob = float(row["purchase_probability"])
        if prob < 0 or prob > 1:
            raise ValueError(
                f"Prediction for {row['record_id']}: probability {prob} out of range [0, 1]"
            )


def _mean(values: list[float]) -> float:
    return sum(values) / len(values) if values else 0.0


def _design_effect_errors(
    labels: list[dict[str, Any]],
    pred_by_id: dict[str, float],
) -> list[float]:
    """Per persona×category×pair: |true(B-A) - pred(B-A)| when both designs are predicted."""
    true_by_pair: dict[tuple[str, str, str], dict[str, float]] = defaultdict(dict)
    pred_by_pair: dict[tuple[str, str, str], dict[str, float]] = defaultdict(dict)

    for row in labels:
        key = (row["persona_id"], row["category"], row["scenario_pair_id"])
        design = row["design_variant"]
        true_by_pair[key][design] = float(row["purchase_probability"])
        record_id = row["record_id"]
        if record_id in pred_by_id:
            pred_by_pair[key][design] = pred_by_id[record_id]

    errors: list[float] = []
    for key, design_probs in true_by_pair.items():
        if "A" not in design_probs or "B" not in design_probs:
            continue
        true_effect = design_probs["B"] - design_probs["A"]
        pred_design = pred_by_pair.get(key, {})
        if "A" not in pred_design or "B" not in pred_design:
            continue
        pred_effect = pred_design["B"] - pred_design["A"]
        errors.append(abs(true_effect - pred_effect))
    return errors


def evaluate_predictions(
    predictions: list[dict[str, Any]],
    labels: list[dict[str, Any]],
    *,
    require_all: bool = False,
) -> EvaluationResult:
    _validate_predictions(predictions)

    label_by_id = {row["record_id"]: row for row in labels}
    expected_ids = set(label_by_id)
    pred_by_id: dict[str, float] = {}

    for row in predictions:
        record_id = str(row["record_id"])
        if record_id in pred_by_id:
            raise ValueError(f"Duplicate prediction for record_id {record_id}")
        pred_by_id[record_id] = float(row["purchase_probability"])

    missing = sorted(expected_ids - set(pred_by_id))
    extra = sorted(set(pred_by_id) - expected_ids)

    if require_all and missing:
        raise ValueError(f"Missing predictions for {len(missing)} records, e.g. {missing[:5]}")

    matched_ids = sorted(expected_ids & set(pred_by_id))
    if not matched_ids:
        raise ValueError("No overlapping record_ids between predictions and labels")

    abs_errors: list[float] = []
    sq_errors: list[float] = []
    persona_errors: dict[str, list[float]] = defaultdict(list)
    category_errors: dict[str, list[float]] = defaultdict(list)

    for record_id in matched_ids:
        truth = float(label_by_id[record_id]["purchase_probability"])
        pred = pred_by_id[record_id]
        err = pred - truth
        abs_errors.append(abs(err))
        sq_errors.append(err**2)
        persona_errors[label_by_id[record_id]["persona_id"]].append(abs(err))
        category_errors[label_by_id[record_id]["category"]].append(abs(err))

    design_errors = _design_effect_errors(labels, pred_by_id)

    return EvaluationResult(
        n_records=len(matched_ids),
        mae=_mean(abs_errors),
        rmse=(sum(sq_errors) / len(sq_errors)) ** 0.5,
        brier=_mean([(pred_by_id[r] - float(label_by_id[r]["purchase_probability"])) ** 2 for r in matched_ids]),
        persona_mae={k: _mean(v) for k, v in persona_errors.items()},
        category_mae={k: _mean(v) for k, v in category_errors.items()},
        design_effect_mae=_mean(design_errors),
        missing_record_ids=missing,
        extra_record_ids=extra,
    )


def evaluate_files(
    predictions_path: Path,
    labels_path: Path,
    *,
    require_all: bool = False,
) -> EvaluationResult:
    return evaluate_predictions(
        load_predictions(predictions_path),
        load_labels_csv(labels_path),
        require_all=require_all,
    )
