#!/usr/bin/env python3
"""Score an out-of-fold predictions table against the canonical targets.

Standard library only.

    python3 evaluate.py --targets targets/targets_her.csv --column voltage_she_v \
        --predictions preds_her_loco.csv --scheme loco

The predictions file has one row per target composition with the columns
`sample_id`, `prediction` and `fold`. Metrics are computed against the
target values read from the targets file, never from a column in the
predictions file, so re-scoring a submission only needs the two files.

Checks:
  - every target composition has exactly one prediction (coverage);
  - under `loco`, each row's `fold` equals the `elements` label of its
    composition (leave-one-composition-out: every ratio of the same element
    set is held out together); under `loo`, every `fold` is unique.

Metrics (pooled over all out-of-fold predictions):
  - pearson_r: signed Pearson correlation; positive is required for model acceptance.
  - pearson_r2: squared Pearson correlation of prediction vs target. This is
    what the OCx24 paper reports as "R2" (an OLS trendline of predicted vs
    experimental), so it is the number to compare with the published
    baselines;
  - r2: the coefficient of determination 1 - SS_res/SS_tot;
  - mae: mean absolute error, in the target's unit;
  - spearman_rho: rank correlation, for ranking-oriented use.

Prints one JSON object. Exit status 1 when a check fails.
"""

from __future__ import annotations

import argparse
import csv
import json
import math
import sys


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


def pearson(xs: list[float], ys: list[float]) -> float:
    n = len(xs)
    mx = sum(xs) / n
    my = sum(ys) / n
    sxy = sum((x - mx) * (y - my) for x, y in zip(xs, ys))
    sxx = sum((x - mx) ** 2 for x in xs)
    syy = sum((y - my) ** 2 for y in ys)
    if sxx == 0 or syy == 0:
        return math.nan
    return sxy / math.sqrt(sxx * syy)


def ranks(values: list[float]) -> list[float]:
    order = sorted(range(len(values)), key=lambda i: values[i])
    out = [0.0] * len(values)
    i = 0
    while i < len(order):
        j = i
        while j + 1 < len(order) and values[order[j + 1]] == values[order[i]]:
            j += 1
        r = (i + j) / 2 + 1
        for k in range(i, j + 1):
            out[order[k]] = r
        i = j + 1
    return out


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--targets", required=True, help="targets_her.csv or targets_co2rr.csv from prepare_targets.py")
    ap.add_argument("--column", required=True, help="target column, e.g. voltage_she_v, h2_pr, co_pr, liquid_pr, fe_co")
    ap.add_argument("--predictions", required=True, help="CSV with sample_id, prediction, fold")
    ap.add_argument("--scheme", choices=["loco", "loo"], required=True)
    ap.add_argument("--log-target", action="store_true", help="score log10 of target and prediction (production rates span decades)")
    args = ap.parse_args()

    targets = {r["sample_id"]: r for r in read_csv(args.targets)}
    preds = read_csv(args.predictions)
    problems: list[str] = []

    seen: dict[str, dict[str, str]] = {}
    for p in preds:
        sid = p.get("sample_id", "")
        if sid in seen:
            problems.append(f"duplicate prediction for {sid}")
        seen[sid] = p
    missing = [sid for sid in targets if sid not in seen]
    extra = [sid for sid in seen if sid not in targets]
    if missing:
        problems.append(f"{len(missing)} target compositions have no prediction (first: {missing[:3]})")
    if extra:
        problems.append(f"{len(extra)} predictions do not name a target composition (first: {extra[:3]})")

    xs: list[float] = []
    ys: list[float] = []
    folds: list[str] = []
    for sid, t in targets.items():
        p = seen.get(sid)
        if p is None:
            continue
        tv = t.get(args.column, "")
        if tv == "":
            continue
        try:
            y = float(tv)
            x = float(p["prediction"])
        except (KeyError, ValueError):
            problems.append(f"non-numeric prediction for {sid}")
            continue
        if not math.isfinite(x) or not math.isfinite(y):
            problems.append(f"{sid}: non-finite target or prediction")
            continue
        fold = p.get("fold", "")
        if args.scheme == "loco" and fold != t["elements"]:
            problems.append(f"{sid}: fold {fold!r} is not its element system {t['elements']!r}")
        folds.append(fold)
        if args.log_target:
            if x <= 0 or y <= 0:
                problems.append(f"{sid}: log scoring needs positive values")
                continue
            x, y = math.log10(x), math.log10(y)
        xs.append(x)
        ys.append(y)
    if args.scheme == "loo" and len(set(folds)) != len(folds):
        problems.append("loo requires a distinct fold per row")

    result: dict = {
        "scheme": args.scheme,
        "target_column": args.column,
        "log_target": args.log_target,
        "n": len(xs),
        "n_folds": len(set(folds)),
        "checks_passed": not problems,
        "problems": problems,
    }
    if len(xs) >= 3:
        r = pearson(xs, ys)
        my = sum(ys) / len(ys)
        ss_res = sum((y - x) ** 2 for x, y in zip(xs, ys))
        ss_tot = sum((y - my) ** 2 for y in ys)
        result.update(
            pearson_r=round(r, 4) if math.isfinite(r) else None,
            pearson_r2=round(r * r, 4) if not math.isnan(r) else None,
            r2=round(1 - ss_res / ss_tot, 4) if ss_tot else None,
            mae=sum(abs(y - x) for x, y in zip(xs, ys)) / len(xs),
            spearman_rho=(lambda rho: None if math.isnan(rho) else round(rho, 4))(pearson(ranks(xs), ranks(ys))),
        )
    print(json.dumps(result, indent=2))
    return 0 if not problems else 1


if __name__ == "__main__":
    sys.exit(main())