#!/usr/bin/env python3
"""Plot frozen FGMRES histories only after a completed full artifact audit.

  python plot_frozen_linear_rates.py linear_audit.json [followup_audit.json] \
      --output NEW_FIGURE_DIRECTORY

Inputs are full summarize_coupled_convergence.py audits of benchmark campaigns,
not the supplemental CSV-only comparison. Small summary/history CSVs are checked
against the audit hashes; large saved vectors and FEM actions are not reread.
Rendering must be scheduled outside timed numerical campaigns.
"""
from __future__ import annotations

import argparse
import csv
import hashlib
import io
import json
import math
from pathlib import Path
import sys


ROLES = ("baseline", "candidate")
VARIANT_LABELS = {
    "gravity-schur": "gravity Schur",
    "surface-border-schur": "surface/scalar Schur",
    "gravity-schur-surface-border-schur": "gravity + surface/scalar Schur",
    "stellar-ldu-gravity-schur": "stellar LDU + gravity Schur",
}


def require(condition, message):
    if not condition:
        raise ValueError(message)


def digest(raw):
    return hashlib.sha256(raw).hexdigest()


def number(value, name, positive=False):
    result = float(value)
    require(math.isfinite(result) and (result > 0 if positive else result >= 0),
            f"Invalid {name}: require finite {'positive' if positive else 'nonnegative'} value")
    return result


def near(left, right, name):
    require(math.isclose(left, right, rel_tol=3e-11, abs_tol=0.), f"Inconsistent {name}")


def flag(value, name):
    require(value in ("0", "1"), f"Invalid {name} flag")
    return value == "1"


def checked_csv(path, hashes, producer_hash, consumed):
    path = path.resolve()
    expected = hashes.get(str(path))
    require(isinstance(expected, str) and len(expected) == 64 and expected == producer_hash,
            f"Missing or inconsistent audit/producer input hash: {path}")
    raw = path.read_bytes()
    require(digest(raw) == expected, f"Input hash differs from the passed audit: {path}")
    consumed[str(path)] = expected
    return list(csv.DictReader(io.StringIO(raw.decode())))


def role_curve(row, samples):
    role, variant = row["role"], row["variant"]
    require(role in ROLES and bool(samples), "Missing linear history role")
    rhs = number(row["rhs_norm"], "RHS norm", positive=True)
    target = number(row["relative_tolerance"], "relative target", positive=True)
    absolute_target = number(row["absolute_target"], "absolute target", positive=True)
    near(absolute_target, target * rhs, "absolute/relative target")
    true_norm = number(row["true_residual_norm"], "true endpoint defect", positive=True)
    true_relative = number(row["true_relative"], "true relative endpoint defect", positive=True)
    near(true_relative, true_norm / rhs, "true relative endpoint defect")
    gate = flag(row["true_gate_pass"], "true gate")
    require(gate == (true_norm <= absolute_target), "True gate contradicts the independently checked endpoint")
    iterations, reported_iterations, maximum = (int(row[key]) for key in
        ("iterations", "reported_iterations", "maximum_iterations"))
    require(0 < iterations <= maximum and reported_iterations == iterations,
            "Reported iteration count differs from the completed preconditioner application count")
    x, values, final = [], [], []
    for index, sample in enumerate(samples):
        require(sample["role"] == role and sample["variant"] == variant, "Linear history role/variant mismatch")
        iteration = int(sample["iteration"])
        require(0 <= iteration <= iterations and (not x or iteration >= x[-1]),
                "Linear history has invalid or unordered iteration numbers")
        x.append(iteration)
        # Exact zero is a valid numerical outcome, but needs a different plotting
        # convention. Refuse it here rather than invent a small positive floor.
        values.append(number(sample["reported_norm"], "reported log residual", positive=True) / rhs)
        if flag(sample["final"], "final monitor"):
            final.append(index)
    require(set(x) == set(range(iterations + 1)), "Missing reported iteration samples")
    require(final == [len(samples) - 1] and x[-1] == reported_iterations,
            "Missing, repeated or misplaced final monitor sample")
    near(values[0], number(row["reported_initial_norm"], "reported initial norm", positive=True) / rhs,
         "initial monitor sample")
    near(values[0], 1., "cold-start relative norm")
    near(values[-1], number(row["reported_final_norm"], "reported final norm", positive=True) / rhs,
         "final monitor sample")
    # Repeated iteration numbers at restarts/final callbacks are meaningful.
    # Keep their order and choose the flagged endpoint, never the best residual.
    setup = math.fsum(number(row[key], key) for key in
                     ("construction_seconds", "normalization_binding_seconds", "krylov_setup_seconds"))
    solve, check = (number(row[key], key) for key in ("krylov_seconds", "true_check_seconds"))
    total = number(row["setup_solve_and_true_check_seconds"], "verified solve seconds", positive=True)
    near(number(row["setup_plus_solve_seconds"], "setup/solve seconds"), setup + solve, "setup/solve timing sum")
    near(total, setup + solve + check, "setup/solve/true-check timing sum")
    return {"role": role, "variant": variant, "iteration": x, "reported_relative": values,
            "final_reported_index": final[0], "completed_iterations": iterations,
            "maximum_iterations": maximum, "restart": int(row["restart"]),
            "rhs_norm": rhs, "relative_target": target, "absolute_target": absolute_target,
            "endpoint_true_relative": true_relative, "true_gate_pass": gate,
            "solver_recurrence_converged": flag(row["solver_converged"], "solver convergence"),
            "setup_seconds": setup, "solve_seconds": solve, "true_check_seconds": check,
            "setup_solve_true_check_seconds": total}


def load_cases(paths):
    cases, sources, consumed, seen = [], [], {}, set()
    for path in paths:
        path = Path(path).resolve()
        raw = path.read_bytes()
        audit = json.loads(raw)
        require(audit.get("complete") is True and audit.get("artifact_passed") is True,
                f"Require a completed, artifact-passed full linear audit: {path}")
        require(audit.get("auditor_sha256") and audit.get("benchmark_checker_sha256")
                and audit.get("consumed_sha256"), "CSV-only summaries are not a full linear artifact audit")
        sources.append({"path": str(path), "sha256": digest(raw), "auditor_sha256": audit["auditor_sha256"],
                        "benchmark_checker_sha256": audit["benchmark_checker_sha256"],
                        "scope": audit.get("scope"), "limitations": audit.get("limitations", [])})
        require(bool(audit["campaigns"]), "Empty audited campaign list")
        for campaign in audit["campaigns"]:
            require(campaign.get("complete") is True and campaign.get("artifact_passed") is True,
                    "Incomplete or failed linear campaign")
            require(bool(campaign["cases"]), "Empty linear campaign")
            for case in campaign["cases"]:
                name = case["id"]
                require(name not in seen, f"Duplicate case identifier across audits: {name}")
                seen.add(name)
                require(Path(name).name == name and name not in (".", ".."), "Invalid case identifier")
                require(case.get("complete") is True and case.get("artifact_passed") is True
                        and case.get("kind") == "benchmark", f"Missing complete linear benchmark: {name}")
                assessment, run = case["benchmark"], case["run"]
                require(assessment.get("passed") is True and assessment["comparison"].get("comparison_valid") is True,
                        "Frozen endpoint artifact audit did not pass")
                checks = {item["check"]: item["passed"] for item in assessment["checks"]}
                required_checks = {role + suffix for role in ROLES for suffix in
                                   ("_rhs_norm", "_defect_norm", "_relative_defect", "_honest_gate")}
                require(required_checks <= checks.keys() and all(value is True for value in checks.values()),
                        "Missing or failed saved-vector endpoint checks")
                require(run.get("finished_utc") and run.get("inputs_unchanged") is True
                        and run.get("exit_code") in (0, 3) and not run.get("timed_out"), "Unfinished linear run")
                folder = Path(campaign["path"]) / name
                tables = {stem: checked_csv(folder / (stem + ".csv"), audit["consumed_sha256"],
                          run["output_hashes"].get(stem + ".csv"), consumed)
                          for stem in ("preconditioner_summary", "preconditioner_history")}
                summary, history = tables["preconditioner_summary"], tables["preconditioner_history"]
                require(len(summary) == 2 and {row["role"] for row in summary} == set(ROLES), "Invalid paired summary")
                require({row["role"] for row in history} == set(ROLES), "Missing or unknown history role")
                curves = []
                for role in ROLES:
                    row = next(row for row in summary if row["role"] == role)
                    samples = [sample for sample in history if sample["role"] == role]
                    require(row == assessment["comparison"][role], "CSV summary differs from the audited endpoint record")
                    require(samples == case["linear_history"][role]["samples"], "CSV history differs from the audited history")
                    curves.append(role_curve(row, samples))
                baseline, candidate = curves
                require(baseline["variant"] == "independent", "Unexpected baseline preconditioner")
                for key in ("rhs_norm", "relative_target", "absolute_target", "restart", "maximum_iterations"):
                    require(baseline[key] == candidate[key], f"Unmatched baseline/candidate {key}")
                outcome = ("both_pass" if all(c["true_gate_pass"] for c in curves) else "baseline_only_pass"
                           if baseline["true_gate_pass"] else "candidate_only_pass" if candidate["true_gate_pass"] else "neither_passes")
                require(case["outcome"] == assessment["comparison"]["outcome"] == outcome, "Contradictory true-gate outcome")
                cases.append({"id": name, "campaign": campaign["path"], "audit_path": str(path),
                              "producer_commit": campaign["producer_commit"], "mesh_sha256": campaign["mesh_sha256"],
                              "metadata": case["metadata"], "source_input_hashes": run["input_state_hashes"],
                              "benchmark_audit_scope": case["benchmark_audit_scope"], "outcome": outcome, "curves": curves})
    require(bool(cases), "No complete linear cases")
    return cases, sources, consumed


def panel_shape(count):
    require(count > 0, "Need at least one panel")
    columns = min(3, count)
    return math.ceil(count / columns), columns


def timing_label(curve):
    status = "PASS" if curve["true_gate_pass"] else "FAIL"
    return (f"{curve['role'][0].upper()} {status}; true {curve['endpoint_true_relative']:.2e}; "
            f"{curve['setup_seconds']:.2f}+{curve['solve_seconds']:.2f}+{curve['true_check_seconds']:.2f}"
            f"={curve['setup_solve_true_check_seconds']:.2f} s")


def render(cases):
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib.lines import Line2D

    rows, columns = panel_shape(len(cases))
    height = 4.35 * rows + 1.2
    figure, axes = plt.subplots(rows, columns, squeeze=False, figsize=(max(10.5, 5.1 * columns), height))
    colors = ("#1769a8", "#c35b16")
    for axis, case in zip(axes.flat, cases):
        curves = case["curves"]
        for curve, color, style in zip(curves, colors, ("-", "--")):
            axis.plot(curve["iteration"], curve["reported_relative"], color=color, linestyle=style, linewidth=1.4)
            final = curve["final_reported_index"]
            axis.plot(curve["iteration"][final], curve["reported_relative"][final], marker="o", markersize=4,
                      markerfacecolor="white", markeredgecolor=color, linestyle="none", zorder=4)
            if curve["true_gate_pass"]:
                axis.scatter([curve["completed_iterations"]], [curve["endpoint_true_relative"]], marker="D", s=58,
                             facecolors="none", edgecolors=color, linewidths=1.5, zorder=5)
            else:
                axis.scatter([curve["completed_iterations"]], [curve["endpoint_true_relative"]], marker="x", s=58,
                             color=color, linewidths=1.7, zorder=5)
        target = curves[0]["relative_target"]
        axis.axhline(target, color="#555555", linestyle=":", linewidth=1.2)
        axis.text(.97, .96, f"target {target:g}", ha="right", va="top", fontsize=8, transform=axis.transAxes,
                  bbox={"facecolor": "white", "edgecolor": "none", "alpha": .8})
        variant = curves[1]["variant"]
        meta = case["metadata"]
        axis.set_title(f"n={float(meta['n']):g}, J={float(meta['J']):g} | {VARIANT_LABELS.get(variant, variant)}\n{case['id']}", fontsize=9)
        axis.set_yscale("log")
        axis.set_xlim(0, max(c["completed_iterations"] for c in curves) * 1.04)
        axis.set_ylabel("Reported relative residual /\nchecked endpoint true defect", fontsize=8)
        axis.set_xlabel("FGMRES iteration\n" + "\n".join(timing_label(c) for c in curves), fontsize=8)
        axis.tick_params(labelsize=8)
        axis.grid(True, which="major", alpha=.22)
    for axis in list(axes.flat)[len(cases):]:
        axis.set_visible(False)
    handles = [Line2D([], [], color=colors[0], label="Baseline (B): independent"),
               Line2D([], [], color=colors[1], linestyle="--", label="Candidate (C)"),
               Line2D([], [], color="#555555", linestyle=":", label="True-defect target"),
               Line2D([], [], color="#222222", marker="o", markerfacecolor="white", linestyle="none", label="Final reported norm"),
               Line2D([], [], color="#222222", marker="D", markerfacecolor="none", linestyle="none", label="Checked endpoint: PASS"),
               Line2D([], [], color="#222222", marker="x", linestyle="none", label="Checked endpoint: FAIL")]
    figure.suptitle("Frozen linear solves: reported FGMRES rates and checked endpoint defects", fontsize=14, y=1 - .10 / height)
    figure.legend(handles=handles, loc="upper center", bbox_to_anchor=(.5, 1 - .42 / height), ncol=3, fontsize=9, frameon=False)
    figure.text(.5, .12 / height, "Lines are reported recurrence norms / ||b||, not an independently checked true-residual history or a proven floor.\n"
                "Endpoints were checked from saved true-defect vectors. PASS/FAIL is linear accuracy, not nonlinear equilibrium.\n"
                "Times: setup + solve + true check = total seconds; no per-iteration time interpolation. Axes autoscale per panel.",
                ha="center", va="bottom", fontsize=8)
    figure.tight_layout(rect=(0, .72 / height, 1, 1 - 1. / height), h_pad=2.2, w_pad=1.5)
    return figure, matplotlib.__version__


def sidecar(cases, sources, consumed):
    return {"schema_version": 1, "source_audits": sources, "rechecked_csv_sha256": consumed,
            "plotter_sha256": digest(Path(__file__).read_bytes()), "cases": cases,
            "scope": ["Requires completed, artifact-passed full linear campaign audits; failed linear gates remain failed.",
                      "Summary/history CSV hashes are rechecked against both the full audit and producer record.",
                      "Large vectors, executable and mesh hashes, and FEM actions are not reevaluated during plotting.",
                      "Reported recurrence norms are not an independently checked true-residual history or a proven residual floor.",
                      "True endpoint norms/gates come from the saved-vector checks in the full audit; no intermediate true defects are inferred.",
                      "Repeated monitor iteration numbers, including a distinct final callback, are preserved in recorded order.",
                      "No time history is inferred. Times include inverse construction/binding/Krylov setup, solve and final true check; frozen-state preparation and other diagnostics are excluded.",
                      "Iteration counts are not equal work units, and linear success does not establish nonlinear equilibrium.",
                      "Log plots refuse exact zero residuals instead of silently replacing them with a positive floor."]}


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("audits", nargs="+", type=Path)
    parser.add_argument("--output", required=True, type=Path, help="New figure directory outside the source checkout")
    args = parser.parse_args()
    require(not args.output.exists(), "Refusing to overwrite an existing figure directory")
    require(not args.output.resolve().is_relative_to(Path(__file__).resolve().parents[1]),
            "Write figure artifacts outside the source checkout")
    cases, sources, consumed = load_cases(args.audits)
    report = sidecar(cases, sources, consumed)
    # Validate all inputs before starting Matplotlib or creating output, so an
    # incomplete ablation cannot quietly turn into a smaller, biased comparison.
    figure, version = render(cases)
    args.output.mkdir(parents=True)
    report.update(matplotlib_version=version, python_version=sys.version, figures=[])
    for extension in ("png", "pdf"):
        path = args.output / f"frozen_linear_rates.{extension}"
        figure.savefig(path, dpi=200)
        report["figures"].append({"path": str(path.resolve()), "sha256": digest(path.read_bytes())})
    from matplotlib import pyplot as plt
    plt.close(figure)
    path = args.output / "frozen_linear_rates.json"
    path.write_text(json.dumps(report, indent=2, allow_nan=False) + "\n")
    print(path)


if __name__ == "__main__":
    main()