#!/usr/bin/env python
"""Compare rerun outputs against the bundled formal DST-GNN release."""

from __future__ import annotations

import argparse
import json
import re
from pathlib import Path

import numpy as np
import pandas as pd


NUMERIC_CSV_FILES = {
    "observed_T1_sss.csv",
    "observed_T2_sss.csv",
    "observed_T3_sss.csv",
    "predicted_T2_sss.csv",
    "predicted_T3_sss.csv",
    "top_dynamic_nodes.csv",
    "top_edge_changes.csv",
    "training_history.csv",
    "explainer_edges_t2_to_t3.csv",
}
JSON_FILES = {"summary.json", "run_config.json"}
TEXT_FILES = {"checksums.sha256"}


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Verify DST-GNN outputs against the bundled release.")
    parser.add_argument("--expected-dir", type=Path, required=True, help="Directory containing the bundled release.")
    parser.add_argument("--actual-dir", type=Path, required=True, help="Directory containing rerun outputs.")
    parser.add_argument("--atol", type=float, default=1e-5, help="Absolute tolerance for numeric comparisons.")
    parser.add_argument("--rtol", type=float, default=1e-5, help="Relative tolerance for numeric comparisons.")
    return parser.parse_args()


def load_csv(path: Path) -> pd.DataFrame:
    """Load a CSV and normalize unnamed index columns."""

    frame = pd.read_csv(path)
    unnamed = [column for column in frame.columns if str(column).startswith("Unnamed:")]
    if unnamed:
        frame = frame.set_index(unnamed[0])
        frame.index.name = None
    return frame


def compare_numeric_csv(expected_path: Path, actual_path: Path, atol: float, rtol: float) -> list[str]:
    """Compare two CSV files with mixed numeric and non-numeric columns."""

    errors: list[str] = []
    expected = load_csv(expected_path)
    actual = load_csv(actual_path)

    if list(expected.columns) != list(actual.columns):
        return [f"{expected_path.name}: columns differ"]
    if list(expected.index) != list(actual.index):
        return [f"{expected_path.name}: row order/index differs"]

    for column in expected.columns:
        left = expected[column]
        right = actual[column]
        if pd.api.types.is_numeric_dtype(left) and pd.api.types.is_numeric_dtype(right):
            if not np.allclose(left.to_numpy(), right.to_numpy(), atol=atol, rtol=rtol, equal_nan=True):
                max_delta = float(np.nanmax(np.abs(left.to_numpy() - right.to_numpy())))
                errors.append(f"{expected_path.name}:{column} numeric mismatch (max_abs_delta={max_delta:.6g})")
        else:
            if not left.fillna("<NA>").equals(right.fillna("<NA>")):
                errors.append(f"{expected_path.name}:{column} text mismatch")
    return errors


def compare_summary(expected: dict[str, object], actual: dict[str, object]) -> list[str]:
    """Compare the stable fields in summary.json."""

    errors: list[str] = []
    for key in ("device", "seed", "node_count", "stage_order", "sample_counts"):
        if expected.get(key) != actual.get(key):
            errors.append(f"summary.json:{key} differs")
    return errors


def compare_run_config(expected: dict[str, object], actual: dict[str, object]) -> list[str]:
    """Compare the stable fields in run_config.json."""

    errors: list[str] = []

    for key in ("seed", "device"):
        if expected.get(key) != actual.get(key):
            errors.append(f"run_config.json:{key} differs")

    expected_model = expected.get("model", {})
    actual_model = actual.get("model", {})
    for key in ("hidden_channels", "dropout", "epochs", "lr", "weight_decay", "top_k", "run_explainer"):
        if expected_model.get(key) != actual_model.get(key):
            errors.append(f"run_config.json:model.{key} differs")

    expected_analysis = expected.get("analysis", {})
    actual_analysis = actual.get("analysis", {})
    for key in ("node_count", "stage_order", "sample_counts"):
        if expected_analysis.get(key) != actual_analysis.get(key):
            errors.append(f"run_config.json:analysis.{key} differs")

    expected_input = expected.get("input_csv", {})
    actual_input = actual.get("input_csv", {})
    if expected_input.get("sha256") != actual_input.get("sha256"):
        errors.append("run_config.json:input_csv.sha256 differs")

    expected_env = expected.get("environment", {})
    actual_env = actual.get("environment", {})
    for key in ("python", "torch", "torch_geometric", "pandas", "numpy"):
        if key not in actual_env:
            errors.append(f"run_config.json:environment.{key} missing in actual output")
        if key not in expected_env:
            errors.append(f"run_config.json:environment.{key} missing in expected output")

    return errors


def compare_json(expected_path: Path, actual_path: Path) -> list[str]:
    """Compare structured JSON outputs."""

    expected = json.loads(expected_path.read_text(encoding="utf-8"))
    actual = json.loads(actual_path.read_text(encoding="utf-8"))
    if expected_path.name == "summary.json":
        return compare_summary(expected, actual)
    if expected_path.name == "run_config.json":
        return compare_run_config(expected, actual)
    return [] if expected == actual else [f"{expected_path.name}: JSON differs"]


def parse_checksums(path: Path) -> dict[str, str]:
    """Parse a checksums.sha256 file."""

    entries: dict[str, str] = {}
    pattern = re.compile(r"^([0-9a-f]{64})\s{2}(.+)$")
    for line in path.read_text(encoding="utf-8").splitlines():
        match = pattern.match(line.strip())
        if not match:
            raise ValueError(f"{path} contains an invalid checksum line: {line!r}")
        entries[match.group(2)] = match.group(1)
    return entries


def compare_text(expected_path: Path, actual_path: Path) -> list[str]:
    """Compare non-CSV plain-text files."""

    if expected_path.name == "checksums.sha256":
        expected = parse_checksums(expected_path)
        actual = parse_checksums(actual_path)
        if set(expected) != set(actual):
            return ["checksums.sha256 labels differ"]
        return []
    return [] if expected_path.read_text(encoding="utf-8") == actual_path.read_text(encoding="utf-8") else [
        f"{expected_path.name}: text differs"
    ]


def main() -> None:
    args = parse_args()
    args.expected_dir = args.expected_dir.resolve()
    args.actual_dir = args.actual_dir.resolve()

    expected_files = sorted(path.name for path in args.expected_dir.iterdir() if path.is_file())
    missing_files = [name for name in expected_files if not (args.actual_dir / name).exists()]
    if missing_files:
        raise SystemExit("Missing files in actual output: " + ", ".join(missing_files))

    errors: list[str] = []
    for name in expected_files:
        expected_path = args.expected_dir / name
        actual_path = args.actual_dir / name
        if name in NUMERIC_CSV_FILES:
            errors.extend(compare_numeric_csv(expected_path, actual_path, args.atol, args.rtol))
        elif name in JSON_FILES:
            errors.extend(compare_json(expected_path, actual_path))
        elif name in TEXT_FILES:
            errors.extend(compare_text(expected_path, actual_path))
        else:
            if expected_path.read_bytes() != actual_path.read_bytes():
                errors.append(f"{name}: binary/text content differs")

    if errors:
        raise SystemExit("Verification failed:\n- " + "\n- ".join(errors))

    print(f"Verification passed for {args.actual_dir} against {args.expected_dir}")


if __name__ == "__main__":
    main()