#!/usr/bin/env python3
"""Compare task loss/score/metric semantics for two forward artifacts."""

from __future__ import annotations

import argparse
import json
from pathlib import Path
from types import SimpleNamespace

import torch
from sklearn.metrics import roc_auc_score

from validation.artifacts import load_artifact


def _evaluate(artifact, *, rewrite: bool):
    logits = torch.from_numpy(artifact.logits).to(torch.float32)
    labels = torch.from_numpy(artifact.labels).to(torch.long)
    weights = torch.from_numpy(artifact.weights).to(torch.float32)
    elementwise = torch.nn.functional.cross_entropy(logits, labels, reduction="none")
    loss = torch.zeros(())
    for label in torch.unique(labels):
        mask = labels == label
        loss = loss + (weights[mask] * elementwise[mask]).sum() / weights[mask].sum()
    loss = loss / len(torch.unique(labels))
    scores = torch.softmax(logits, dim=1)
    predictions = scores.argmax(dim=1)
    positive = weights > 0
    one_hot = torch.nn.functional.one_hot(labels, num_classes=scores.shape[1]).numpy()
    positive_numpy = positive.numpy()
    legacy_auc = float(
        roc_auc_score(
            one_hot[positive_numpy],
            scores.detach().numpy()[positive_numpy],
            multi_class="ovr",
            sample_weight=weights.numpy()[positive_numpy],
        )
    )
    if rewrite:
        from gnn4colliders.tasks import MulticlassClassificationTask

        batch = SimpleNamespace(
            labels=labels,
            metadata=SimpleNamespace(weight=weights),
        )
        auc = MulticlassClassificationTask().metrics(logits, batch)["roc_auc"]
    else:
        auc = legacy_auc
    return {
        "loss": float(loss),
        "accuracy": float((predictions == labels).float().mean()),
        "roc_auc": auc,
        "positive_weight_events": int(positive.sum()),
        "negative_weight_events": int((weights < 0).sum()),
        "zero_weight_events": int((weights == 0).sum()),
    }


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("legacy", type=Path)
    parser.add_argument("rewrite", type=Path)
    parser.add_argument("output", type=Path)
    args = parser.parse_args()
    left = _evaluate(load_artifact(args.legacy), rewrite=False)
    right = _evaluate(load_artifact(args.rewrite), rewrite=True)
    differences = {
        name: abs(left[name] - right[name]) for name in ("loss", "accuracy", "roc_auc")
    }
    report = {
        "schema_version": 1,
        "legacy": left,
        "rewrite": right,
        "absolute_differences": differences,
        "overall_status": "PASS"
        if all(value <= 1e-5 for value in differences.values())
        else "FAIL",
        "note": (
            "Legacy uses sklearn weighted OVR AUC; rewrite uses its task-owned "
            "weighted OVR implementation."
        ),
    }
    args.output.mkdir(parents=True, exist_ok=True)
    (args.output / "task_metrics.json").write_text(
        json.dumps(report, indent=2, sort_keys=True, allow_nan=True) + "\n"
    )
    print(json.dumps({"overall_status": report["overall_status"]}, indent=2))
    return 0 if report["overall_status"] == "PASS" else 1


if __name__ == "__main__":
    raise SystemExit(main())