"""Stage-aware comparison of normalized legacy and rewrite artifacts."""

from __future__ import annotations

import json
from pathlib import Path
from typing import Any

import numpy as np

from .artifacts import ValidationArtifact, load_artifact

STATUS = ("PASS", "PASS_WITH_EXPECTED_DIFFERENCE", "FAIL", "NOT_APPLICABLE", "NOT_RUN")


def _numeric(
    left: np.ndarray, right: np.ndarray, atol: float, rtol: float
) -> dict[str, Any]:
    if left.shape != right.shape:
        return {
            "status": "FAIL",
            "reason": "shape mismatch",
            "left_shape": left.shape,
            "right_shape": right.shape,
        }
    delta = np.abs(left.astype(np.float64) - right.astype(np.float64))
    scale = np.maximum(
        np.abs(left.astype(np.float64)), np.abs(right.astype(np.float64))
    )
    relative = np.divide(delta, scale, out=np.zeros_like(delta), where=scale != 0)
    passed = np.isclose(left, right, atol=atol, rtol=rtol)
    location = (
        np.unravel_index(int(delta.argmax()), delta.shape) if delta.size else None
    )
    return {
        "status": "PASS" if bool(np.all(passed)) else "FAIL",
        "max_abs": float(delta.max()) if delta.size else 0.0,
        "max_rel": float(relative.max()) if relative.size else 0.0,
        "outside_tolerance": int((~passed).sum()),
        "elements": int(delta.size),
        "location": list(location) if location is not None else None,
        "atol": atol,
        "rtol": rtol,
    }


def _exact(left: np.ndarray, right: np.ndarray) -> dict[str, Any]:
    return {
        "status": "PASS" if np.array_equal(left, right) else "FAIL",
        "elements": int(left.size),
        "shape": left.shape,
        "reason": "exact mismatch" if not np.array_equal(left, right) else None,
    }


def _ids(left: ValidationArtifact, right: ValidationArtifact) -> dict[str, Any]:
    left_ids, right_ids = left.sample_id.tolist(), right.sample_id.tolist()
    left_set, right_set = set(left_ids), set(right_ids)
    duplicates = sorted(
        {item for item in left_ids if left_ids.count(item) > 1}
        | {item for item in right_ids if right_ids.count(item) > 1}
    )
    return {
        "status": "PASS" if left_set == right_set and not duplicates else "FAIL",
        "left_count": len(left_ids),
        "right_count": len(right_ids),
        "missing_in_rewrite": sorted(left_set - right_set),
        "missing_in_legacy": sorted(right_set - left_set),
        "duplicates": duplicates,
        "order_equal": left_ids == right_ids,
    }


def compare_artifacts(
    left: ValidationArtifact | str | Path,
    right: ValidationArtifact | str | Path,
    *,
    tolerances: dict[str, float] | None = None,
) -> dict[str, Any]:
    left = load_artifact(left) if not isinstance(left, ValidationArtifact) else left
    right = load_artifact(right) if not isinstance(right, ValidationArtifact) else right
    left.validate()
    right.validate()
    tol = {
        "feature_atol": 1e-6,
        "feature_rtol": 1e-6,
        "logit_atol": 1e-5,
        "logit_rtol": 1e-5,
        **(tolerances or {}),
    }
    report: dict[str, Any] = {
        "schema_version": 1,
        "dataset": _ids(left, right),
        "stages": {},
    }
    stages = report["stages"]
    stages["labels"] = _exact(left.labels, right.labels)
    stages["folds"] = _exact(left.folds, right.folds)
    stages["weights"] = _numeric(
        left.weights, right.weights, tol["feature_atol"], tol["feature_rtol"]
    )
    stages["globals"] = _numeric(
        left.globals, right.globals, tol["feature_atol"], tol["feature_rtol"]
    )
    stages["node_offsets"] = _exact(left.node_offsets, right.node_offsets)
    stages["node_features"] = _numeric(
        left.node_features_flat,
        right.node_features_flat,
        tol["feature_atol"],
        tol["feature_rtol"],
    )
    stages["edge_offsets"] = _exact(left.edge_offsets, right.edge_offsets)
    stages["edge_src"] = _exact(left.edge_src_flat, right.edge_src_flat)
    stages["edge_dst"] = _exact(left.edge_dst_flat, right.edge_dst_flat)
    stages["edge_features"] = _numeric(
        left.edge_features_flat,
        right.edge_features_flat,
        tol["feature_atol"],
        tol["feature_rtol"],
    )
    if left.logits is None or right.logits is None:
        stages["logits"] = {"status": "NOT_APPLICABLE"}
    else:
        stages["logits"] = _numeric(
            left.logits, right.logits, tol["logit_atol"], tol["logit_rtol"]
        )
    if left.scores is None or right.scores is None:
        stages["scores"] = {"status": "NOT_APPLICABLE"}
    else:
        stages["scores"] = _numeric(
            left.scores, right.scores, tol["logit_atol"], tol["logit_rtol"]
        )
    if left.predictions is None or right.predictions is None:
        stages["predictions"] = {"status": "NOT_APPLICABLE"}
    else:
        stages["predictions"] = _exact(left.predictions, right.predictions)
    report["overall_status"] = (
        "PASS"
        if all(
            item["status"]
            in ("PASS", "PASS_WITH_EXPECTED_DIFFERENCE", "NOT_APPLICABLE")
            for item in [report["dataset"], *stages.values()]
        )
        else "FAIL"
    )
    return report


def write_reports(report: dict[str, Any], directory: str | Path) -> tuple[Path, Path]:
    directory = Path(directory)
    directory.mkdir(parents=True, exist_ok=True)
    json_path = directory / "comparison.json"
    md_path = directory / "comparison.md"
    json_path.write_text(
        json.dumps(report, indent=2, sort_keys=True, default=_json_default) + "\n"
    )
    lines = [
        "# End-to-end parity report",
        "",
        f"Overall status: **{report['overall_status']}**",
        "",
        "| Stage | Status | Max abs | Outside tolerance |",
        "|---|---|---:|---:|",
    ]
    for name, result in report.get("stages", {}).items():
        lines.append(
            f"| {name} | {result['status']} | "
            f"{result.get('max_abs', '—')} | "
            f"{result.get('outside_tolerance', '—')} |"
        )
    md_path.write_text("\n".join(lines) + "\n")
    return json_path, md_path


def _json_default(value: Any) -> Any:
    """Serialize NumPy scalar/shape values present in diagnostic reports."""
    if isinstance(value, np.generic):
        return value.item()
    if isinstance(value, tuple):
        return list(value)
    raise TypeError(f"cannot serialize {type(value).__name__}")