#!/usr/bin/env python3
"""Compare regenerated numerical outputs with frozen package results."""

from __future__ import annotations

import argparse
from pathlib import Path
import sys

import numpy as np
import pandas as pd


PACKAGE = Path(__file__).resolve().parents[1]

FILE_PAIRS = [
    (
        "Study 1 interactions",
        "results/study1/study1_contrasts.csv",
        "study1_contrasts.csv",
    ),
    (
        "Study 1 gates",
        "results/study1/study1_gates.csv",
        "study1_gates.csv",
    ),
    (
        "Study 2 inference",
        "results/study2/study2_inference.csv",
        "study2_inference.csv",
    ),
    (
        "Study 2 gates",
        "results/study2/study2_gates.csv",
        "study2_gates.csv",
    ),
]


def compare_csv(label: str, frozen: Path, candidate: Path, *, rtol: float, atol: float) -> list[str]:
    issues: list[str] = []
    if not candidate.exists():
        return [f"{label}: missing candidate file {candidate}"]

    left = pd.read_csv(frozen)
    right = pd.read_csv(candidate)

    if list(left.columns) != list(right.columns):
        issues.append(f"{label}: column mismatch")
        return issues
    if len(left) != len(right):
        issues.append(f"{label}: row-count mismatch ({len(left)} != {len(right)})")
        return issues

    for column in left.columns:
        left_col = left[column]
        right_col = right[column]
        if pd.api.types.is_numeric_dtype(left_col) and pd.api.types.is_numeric_dtype(right_col):
            left_vals = left_col.to_numpy()
            right_vals = right_col.to_numpy()
            if not np.allclose(left_vals, right_vals, rtol=rtol, atol=atol, equal_nan=True):
                diff = np.nanmax(np.abs(left_vals - right_vals))
                issues.append(f"{label}: numeric mismatch in {column} (max abs diff {diff})")
        else:
            if not left_col.fillna("<NA>").astype(str).equals(right_col.fillna("<NA>").astype(str)):
                issues.append(f"{label}: text mismatch in {column}")
    return issues


def main() -> int:
    parser = argparse.ArgumentParser()
    parser.add_argument("--study1-dir", type=Path, required=True)
    parser.add_argument("--study2-dir", type=Path, required=True)
    parser.add_argument("--rtol", type=float, default=1e-8)
    parser.add_argument("--atol", type=float, default=1e-10)
    args = parser.parse_args()

    study_dirs = {
        "results/study1/study1_contrasts.csv": args.study1_dir / "study1_contrasts.csv",
        "results/study1/study1_gates.csv": args.study1_dir / "study1_gates.csv",
        "results/study2/study2_inference.csv": args.study2_dir / "study2_inference.csv",
        "results/study2/study2_gates.csv": args.study2_dir / "study2_gates.csv",
    }

    issues: list[str] = []
    for label, frozen_rel, _ in FILE_PAIRS:
        frozen = PACKAGE / frozen_rel
        candidate = study_dirs[frozen_rel]
        issues.extend(compare_csv(label, frozen, candidate, rtol=args.rtol, atol=args.atol))

    if issues:
        print("NUMERIC VERIFICATION: FAIL")
        for issue in issues:
            print(f"- {issue}")
        return 1

    print("NUMERIC VERIFICATION: PASS")
    print(f"Compared {len(FILE_PAIRS)} frozen result files with rtol={args.rtol} atol={args.atol}.")
    return 0


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