#!/usr/bin/env python3
"""
check_correction_data.py — Validate step-wise correction dataset.

Checks:
  1. Global statistics (elements, marks, overlap, IoU, correction ratios)
  2. Per-step trend analysis (correction signal vs AMR step)
  3. Spatial visualization of 5 random samples
  4. Field & shape consistency
  5. Final verdict: PASS / WARNING / FAIL

Usage:
  python outlook/check_correction_data.py \
      --data-dir outlook/data_correction \
      --output-dir outlook/data_correction_check
"""

import argparse
import sys
from pathlib import Path

import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np

# ──────────────────────────────────────────────────────────────────────────────
# Required fields and their expected properties
# ──────────────────────────────────────────────────────────────────────────────

REQUIRED_FIELDS = {
    "features":          {"ndim": 2, "dtype_kind": "f"},
    "edge_index":        {"ndim": 2, "dtype_kind": "i"},
    "physics_score":     {"ndim": 1, "dtype_kind": "f"},
    "teacher_eta":       {"ndim": 1, "dtype_kind": "f"},
    "teacher_mark":      {"ndim": 1, "dtype_kind": "i"},
    "physics_mark":      {"ndim": 1, "dtype_kind": "i"},
    "correction_label":  {"ndim": 1, "dtype_kind": "i"},
    "elements":          {"ndim": 0},
    "step":              {"ndim": 0},
    "aw_rel_before":     {"ndim": 0},
    "max_err_before":    {"ndim": 0},
    "k":                 {"ndim": 0},
    "eps_r":             {"ndim": 0},
    "cx":                {"ndim": 0},
    "cy":                {"ndim": 0},
    "radius":            {"ndim": 0},
}


# ──────────────────────────────────────────────────────────────────────────────
# Helpers
# ──────────────────────────────────────────────────────────────────────────────


def iou_binary(a: np.ndarray, b: np.ndarray) -> float:
    """IoU between two binary (0/1) arrays."""
    a = a.astype(bool)
    b = b.astype(bool)
    inter = np.sum(a & b)
    union = np.sum(a | b)
    return float(inter) / float(union) if union > 0 else 1.0


def load_all_samples(data_dir: Path):
    """Load all sampleXXXX_stepYYY.npz files, return list of dicts."""
    files = sorted(data_dir.glob("sample*_step*.npz"))
    samples = []
    for f in files:
        try:
            d = dict(np.load(f, allow_pickle=True))
            d["_path"] = str(f)
            d["_name"] = f.name
            samples.append(d)
        except Exception as e:
            print(f"  [WARN] Cannot load {f.name}: {e}")
    return samples


# ──────────────────────────────────────────────────────────────────────────────
# 1. Field & shape consistency
# ──────────────────────────────────────────────────────────────────────────────


def check_fields(samples, issues):
    """Check every file has the required fields with correct ndim."""
    print("\n" + "=" * 70)
    print("  [1/5]  Field & Shape Consistency")
    print("=" * 70)

    n_files = len(samples)
    n_issues_before = len(issues)

    for s in samples:
        name = s["_name"]
        for field, spec in REQUIRED_FIELDS.items():
            if field not in s:
                issues.append(f"MISSING field '{field}' in {name}")
                continue
            arr = s[field]
            if arr.ndim != spec["ndim"]:
                issues.append(
                    f"BAD ndim for '{field}' in {name}: "
                    f"expected {spec['ndim']}, got {arr.ndim}"
                )
            if "dtype_kind" in spec and arr.dtype.kind != spec["dtype_kind"]:
                issues.append(
                    f"BAD dtype for '{field}' in {name}: "
                    f"expected kind={spec['dtype_kind']}, got {arr.dtype}"
                )

        # Cross-check: features rows == len(physics_score) == elements
        if "features" in s and "physics_score" in s:
            nf = s["features"].shape[0]
            ns = len(s["physics_score"])
            ne = int(s["elements"])
            if not (nf == ns == ne):
                issues.append(
                    f"SIZE MISMATCH in {name}: features={nf}, "
                    f"physics_score={ns}, elements={ne}"
                )

    n_new = len(issues) - n_issues_before
    if n_new == 0:
        print(f"  All {n_files} files have consistent fields & shapes.")
    else:
        print(f"  Found {n_new} field/shape issues.")


# ──────────────────────────────────────────────────────────────────────────────
# 2. Global statistics
# ──────────────────────────────────────────────────────────────────────────────


def compute_global_stats(samples):
    """Compute and print global statistics across all files."""
    print("\n" + "=" * 70)
    print("  [2/5]  Global Statistics")
    print("=" * 70)

    elems = []
    teacher_ratios = []
    physics_ratios = []
    overlaps = []
    overlap_teacher = []
    overlap_physics = []
    ious = []
    pos_ratios = []
    neg_ratios = []

    for s in samples:
        n = int(s["elements"])
        tm = s["teacher_mark"]
        pm = s["physics_mark"]
        cl = s["correction_label"]

        elems.append(n)
        n_tm = int(np.sum(tm))
        n_pm = int(np.sum(pm))
        ov = int(np.sum((tm == 1) & (pm == 1)))

        teacher_ratios.append(n_tm / max(n, 1))
        physics_ratios.append(n_pm / max(n, 1))
        overlaps.append(ov)
        overlap_teacher.append(ov / max(n_tm, 1))
        overlap_physics.append(ov / max(n_pm, 1))
        ious.append(iou_binary(tm, pm))
        pos_ratios.append(int(np.sum(cl == 1)) / max(n, 1))
        neg_ratios.append(int(np.sum(cl == -1)) / max(n, 1))

    stats = {
        "total_files":         len(samples),
        "elements":            elems,
        "teacher_mark_ratio":  teacher_ratios,
        "physics_mark_ratio":  physics_ratios,
        "overlap_count":       overlaps,
        "overlap/teacher":     overlap_teacher,
        "overlap/physics":     overlap_physics,
        "IoU":                 ious,
        "correction_pos_ratio": pos_ratios,
        "correction_neg_ratio": neg_ratios,
    }

    print(f"  {'Metric':<30s} {'mean':>10s} {'min':>10s} {'max':>10s}")
    print(f"  {'-'*30} {'-'*10} {'-'*10} {'-'*10}")
    for key in [
        "elements", "teacher_mark_ratio", "physics_mark_ratio",
        "overlap_count", "overlap/teacher", "overlap/physics",
        "IoU", "correction_pos_ratio", "correction_neg_ratio",
    ]:
        arr = np.array(stats[key])
        print(f"  {key:<30s} {arr.mean():10.4f} {arr.min():10.4f} {arr.max():10.4f}")

    return stats


# ──────────────────────────────────────────────────────────────────────────────
# 3. Per-step trend analysis
# ──────────────────────────────────────────────────────────────────────────────


def check_per_step_trends(samples, issues):
    """Group by step, print trend table, flag anomalies."""
    print("\n" + "=" * 70)
    print("  [3/5]  Per-Step Trend Analysis")
    print("=" * 70)

    from collections import defaultdict
    step_data = defaultdict(lambda: {
        "elems": [], "tm_ratio": [], "pm_ratio": [],
        "overlap": [], "iou": [], "pos_ratio": [], "neg_ratio": [],
    })

    for s in samples:
        step = int(s["step"])
        n = int(s["elements"])
        tm = s["teacher_mark"]
        pm = s["physics_mark"]
        cl = s["correction_label"]
        n_tm = int(np.sum(tm))
        n_pm = int(np.sum(pm))
        ov = int(np.sum((tm == 1) & (pm == 1)))

        d = step_data[step]
        d["elems"].append(n)
        d["tm_ratio"].append(n_tm / max(n, 1))
        d["pm_ratio"].append(n_pm / max(n, 1))
        d["overlap"].append(ov)
        d["iou"].append(iou_binary(tm, pm))
        d["pos_ratio"].append(int(np.sum(cl == 1)) / max(n, 1))
        d["neg_ratio"].append(int(np.sum(cl == -1)) / max(n, 1))

    steps = sorted(step_data.keys())
    header = f"  {'step':>4s} {'n_files':>7s} {'mean_elem':>9s} {'tm_ratio':>9s} {'pm_ratio':>9s} {'IoU':>7s} {'pos%':>7s} {'neg%':>7s}"
    print(header)
    print(f"  {'-'*4} {'-'*7} {'-'*9} {'-'*9} {'-'*9} {'-'*7} {'-'*7} {'-'*7}")

    for st in steps:
        d = step_data[st]
        nf = len(d["elems"])
        me = np.mean(d["elems"])
        mt = np.mean(d["tm_ratio"])
        mp = np.mean(d["pm_ratio"])
        mi = np.mean(d["iou"])
        mpos = np.mean(d["pos_ratio"])
        mneg = np.mean(d["neg_ratio"])
        print(f"  {st:4d} {nf:7d} {me:9.1f} {mt:9.4f} {mp:9.4f} {mi:7.4f} {mpos:7.4f} {mneg:7.4f}")

    # Trend checks
    if len(steps) >= 3:
        # Check: correction signal should diminish at late steps (elements grow → marks shrink)
        late = steps[-1]
        early = steps[0]
        late_pos = np.mean(step_data[late]["pos_ratio"])
        early_pos = np.mean(step_data[early]["pos_ratio"])
        late_neg = np.mean(step_data[late]["neg_ratio"])
        early_neg = np.mean(step_data[early]["neg_ratio"])
        late_iou = np.mean(step_data[late]["iou"])
        early_iou = np.mean(step_data[early]["iou"])

        print(f"\n  Trend summary (step {early} → {late}):")
        print(f"    pos_ratio:  {early_pos:.4f} → {late_pos:.4f}")
        print(f"    neg_ratio:  {early_neg:.4f} → {late_neg:.4f}")
        print(f"    IoU:        {early_iou:.4f} → {late_iou:.4f}")

        # Warning: if late-step IoU is much lower than early, marks diverge
        if late_iou < early_iou * 0.5:
            issues.append(
                f"TREND: IoU drops significantly from step {early} ({early_iou:.3f}) "
                f"to step {late} ({late_iou:.3f}) — marks diverge at late steps."
            )

        # Warning: if correction signal vanishes (pos+neg → 0) at late steps
        late_corr = late_pos + late_neg
        early_corr = early_pos + early_neg
        if early_corr > 0.01 and late_corr < early_corr * 0.1:
            issues.append(
                f"TREND: Correction signal nearly vanishes at step {late} "
                f"({late_corr:.4f}) vs step {early} ({early_corr:.4f})."
            )

        # Warning: if IoU is very low overall (teacher and physics barely agree)
        all_ious = [iou for s in steps for iou in step_data[s]["iou"]]
        overall_iou = np.mean(all_ious)
        if overall_iou < 0.1:
            issues.append(
                f"LOW IoU: mean IoU across all steps is {overall_iou:.4f} — "
                f"teacher_mark and physics_mark barely overlap."
            )


# ──────────────────────────────────────────────────────────────────────────────
# 4. Spatial visualization
# ──────────────────────────────────────────────────────────────────────────────


def plot_spatial_samples(samples, output_dir: Path, n_plot: int = 5):
    """Randomly select n_plot files and save 4-panel spatial maps."""
    print("\n" + "=" * 70)
    print("  [4/5]  Spatial Visualization")
    print("=" * 70)

    output_dir.mkdir(parents=True, exist_ok=True)

    rng = np.random.RandomState(1234)
    indices = rng.choice(len(samples), size=min(n_plot, len(samples)), replace=False)
    selected = [samples[i] for i in sorted(indices)]

    for s in selected:
        name = s["_name"].replace(".npz", "")
        features = s["features"]
        x = features[:, 0]
        y = features[:, 1]

        physics_score = s["physics_score"]
        teacher_eta = s["teacher_eta"]
        teacher_mark = s["teacher_mark"]
        correction_label = s["correction_label"]

        fig, axes = plt.subplots(1, 4, figsize=(20, 5))
        fig.suptitle(
            f"{name}  (elem={int(s['elements'])}, k={float(s['k']):.2f}, "
            f"step={int(s['step'])})",
            fontsize=12,
        )

        # Panel 1: physics_score
        sc0 = axes[0].scatter(x, y, c=physics_score, s=8, cmap="viridis")
        axes[0].set_title("physics_score")
        axes[0].set_aspect("equal")
        plt.colorbar(sc0, ax=axes[0], shrink=0.8)

        # Panel 2: teacher_eta
        sc1 = axes[1].scatter(x, y, c=teacher_eta, s=8, cmap="magma")
        axes[1].set_title("teacher_eta")
        axes[1].set_aspect("equal")
        plt.colorbar(sc1, ax=axes[1], shrink=0.8)

        # Panel 3: teacher_mark
        sc2 = axes[2].scatter(x, y, c=teacher_mark, s=8, cmap="Reds", vmin=0, vmax=1)
        axes[2].set_title("teacher_mark")
        axes[2].set_aspect("equal")
        plt.colorbar(sc2, ax=axes[2], shrink=0.8)

        # Panel 4: correction_label
        sc3 = axes[3].scatter(
            x, y, c=correction_label, s=8, cmap="coolwarm", vmin=-1, vmax=1
        )
        axes[3].set_title("correction_label")
        axes[3].set_aspect("equal")
        plt.colorbar(sc3, ax=axes[3], shrink=0.8)

        for ax in axes:
            ax.set_xlabel("x")
            ax.set_ylabel("y")

        plt.tight_layout()
        out_path = output_dir / f"{name}.png"
        fig.savefig(out_path, dpi=150, bbox_inches="tight")
        plt.close(fig)
        print(f"  ✓ Saved: {out_path}")


# ──────────────────────────────────────────────────────────────────────────────
# 5. Verdict
# ──────────────────────────────────────────────────────────────────────────────


def print_verdict(issues):
    """Print final PASS / WARNING / FAIL."""
    print("\n" + "=" * 70)
    print("  [5/5]  Verdict")
    print("=" * 70)

    if not issues:
        print("\n  ✅  PASS — Data is consistent and ready for training.")
        return "PASS"

    # Classify issues
    errors = [i for i in issues if any(k in i.upper() for k in ["MISSING", "BAD ", "MISMATCH"])]
    warnings = [i for i in issues if i not in errors]

    if errors:
        print(f"\n  ❌  FAIL — {len(errors)} error(s) found:")
        for e in errors:
            print(f"    • {e}")
        if warnings:
            print(f"\n  ⚠️  Also {len(warnings)} warning(s):")
            for w in warnings:
                print(f"    • {w}")
        return "FAIL"

    print(f"\n  ⚠️  WARNING — {len(warnings)} suspicious pattern(s):")
    for w in warnings:
        print(f"    • {w}")
    return "WARNING"


# ──────────────────────────────────────────────────────────────────────────────
# Main
# ──────────────────────────────────────────────────────────────────────────────


def main():
    parser = argparse.ArgumentParser(
        description="Check step-wise correction dataset quality"
    )
    parser.add_argument(
        "--data-dir", type=str, required=True,
        help="Directory containing sampleXXXX_stepYYY.npz files",
    )
    parser.add_argument(
        "--output-dir", type=str, default="outlook/data_correction_check",
        help="Directory to save visualization plots",
    )
    args = parser.parse_args()

    data_dir = Path(args.data_dir)
    output_dir = Path(args.output_dir)

    print("=" * 70)
    print("  Correction Data Quality Check")
    print("=" * 70)
    print(f"  Data dir:    {data_dir}")
    print(f"  Output dir:  {output_dir}")

    # Load
    print("\n  Loading files...")
    samples = load_all_samples(data_dir)
    if not samples:
        print("  ✗ No sample files found!")
        sys.exit(1)
    print(f"  Loaded {len(samples)} files.")

    issues = []

    # 1. Field consistency
    check_fields(samples, issues)

    # 2. Global stats
    compute_global_stats(samples)

    # 3. Per-step trends
    check_per_step_trends(samples, issues)

    # 4. Visualization
    plot_spatial_samples(samples, output_dir)

    # 5. Verdict
    verdict = print_verdict(issues)

    print()
    return 0 if verdict == "PASS" else 1


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