"""score.py — reference scorer for the Perfect Audit Crime Challenge.

Scores a participant's submission (a list of journal-entry `document_id`s flagged as fraudulent,
optionally with a ranking score) against the held-out per-JE ground truth, on the imbalanced fraud
class. Primary metric: PR-AUC (when scores are supplied); otherwise precision / recall / F1 on the
flagged set. Reports per-fraud-family and per-observability-class recall — the scientific payoff: the
leaderboard reveals *which* family each method catches, and that the mimetic perfect-crime family is
uncatchable in the ledger-only track.

Labels file (private; CSV or JSON) — one row per journal entry:
    document_id, is_fraud (0/1), family (str), observability (str)
Submission file (CSV or JSON):
    document_id [, score]    — rows are the entries the participant flags as fraudulent.

    python score.py --labels labels.csv --submission sub.csv
    python score.py --selftest
"""
from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
import pandas as pd


def _read(path: Path) -> pd.DataFrame:
    path = Path(path)
    if path.suffix.lower() == ".json":
        return pd.DataFrame(json.loads(path.read_text()))
    return pd.read_csv(path, low_memory=False)


def load_labels(path) -> pd.DataFrame:
    df = _read(path)
    df["document_id"] = df["document_id"].astype(str)
    df["is_fraud"] = df["is_fraud"].astype(int)
    for c in ("family", "observability"):
        if c not in df.columns:
            df[c] = "unknown"
    return df[["document_id", "is_fraud", "family", "observability"]]


def load_submission(path) -> pd.DataFrame:
    df = _read(path)
    if "document_id" not in df.columns:
        # tolerate a single-column id list
        df = df.rename(columns={df.columns[0]: "document_id"})
    df["document_id"] = df["document_id"].astype(str)
    df = df.drop_duplicates("document_id")
    if "score" not in df.columns:
        df["score"] = 1.0
    return df[["document_id", "score"]]


def score(submission: pd.DataFrame, labels: pd.DataFrame) -> dict:
    """Score a submission against the labels. Submissions reference only flagged entries; every
    unflagged ledger entry is treated as predicted-negative (score 0)."""
    truth = labels.set_index("document_id")
    y = truth["is_fraud"].to_numpy()
    n_fraud = int(y.sum())
    n = len(truth)

    # align submission scores onto the full ledger (unflagged → 0)
    s = pd.Series(0.0, index=truth.index)
    sub = submission[submission["document_id"].isin(truth.index)]
    s.loc[sub["document_id"]] = pd.to_numeric(sub["score"], errors="coerce").fillna(1.0).to_numpy()
    scores = s.to_numpy()

    flagged = set(sub["document_id"])
    pred = truth.index.isin(flagged).astype(int)
    tp = int(((pred == 1) & (y == 1)).sum())
    fp = int(((pred == 1) & (y == 0)).sum())
    fn = int(((pred == 0) & (y == 1)).sum())
    precision = tp / (tp + fp) if (tp + fp) else 0.0
    recall = tp / (tp + fn) if (tp + fn) else 0.0
    f1 = 2 * precision * recall / (precision + recall) if (precision + recall) else 0.0

    out = {
        "ledger_entries": n,
        "fraud_entries": n_fraud,
        "submitted": len(sub),
        "true_positives": tp,
        "false_positives": fp,
        "precision": round(precision, 4),
        "recall": round(recall, 4),
        "f1": round(f1, 4),
    }
    # PR-AUC when the submission carries a non-degenerate ranking
    if scores.max() > scores.min() and n_fraud > 0:
        try:
            from sklearn.metrics import average_precision_score
            out["pr_auc"] = round(float(average_precision_score(y, scores)), 4)
        except Exception:
            pass

    # per-family / per-observability recall (the frontier breakdown)
    def by(col):
        g = {}
        for key, grp in labels[labels["is_fraud"] == 1].groupby(col):
            ids = set(grp["document_id"])
            caught = len(ids & flagged)
            g[str(key)] = {"caught": caught, "total": len(ids),
                           "recall": round(caught / len(ids), 3) if len(ids) else 0.0}
        return g

    out["recall_by_family"] = by("family")
    out["recall_by_observability"] = by("observability")
    return out


# --------------------------------------------------------------------------- self-test
def _selftest():
    rng = np.random.default_rng(0)
    rows = []
    for i in range(1000):
        fam = "clean"
        if i < 30:
            fam = "structured"
        elif i < 50:
            fam = "campaign"
        elif i < 60:
            fam = "mimetic"
        rows.append({
            "document_id": f"D{i:05d}",
            "is_fraud": 1 if fam != "clean" else 0,
            "family": fam,
            "observability": {"structured": "per_je_density", "campaign": "memory_only",
                              "mimetic": "memory_only", "clean": "none"}[fam],
        })
    labels = pd.DataFrame(rows)

    # a detector that catches structured + campaign but misses mimetic, with some false positives
    flagged = [r["document_id"] for r in rows if r["family"] in ("structured", "campaign")]
    flagged += [f"D{i:05d}" for i in range(900, 920)]  # 20 false positives
    sub = pd.DataFrame({"document_id": flagged, "score": 0.9})
    res = score(load_submission_df(sub), labels)
    print(json.dumps(res, indent=2))
    assert res["fraud_entries"] == 60
    assert res["recall_by_family"]["structured"]["recall"] == 1.0
    assert res["recall_by_family"]["mimetic"]["recall"] == 0.0, "mimetic must be uncaught here"
    assert res["false_positives"] == 20
    print("SCORE_SELFTEST_OK")


def load_submission_df(df: pd.DataFrame) -> pd.DataFrame:
    df = df.copy()
    df["document_id"] = df["document_id"].astype(str)
    if "score" not in df.columns:
        df["score"] = 1.0
    return df[["document_id", "score"]]


def main(argv=None):
    ap = argparse.ArgumentParser(description=__doc__)
    ap.add_argument("--labels", type=Path)
    ap.add_argument("--submission", type=Path)
    ap.add_argument("--selftest", action="store_true")
    a = ap.parse_args(argv)
    if a.selftest:
        _selftest(); return
    assert a.labels and a.submission, "need --labels and --submission (or --selftest)"
    res = score(load_submission(a.submission), load_labels(a.labels))
    print(json.dumps(res, indent=2))


if __name__ == "__main__":
    main()