#!/usr/bin/env python3
"""Summarize completed density-order audits without opening campaign artifacts.

The input is audit_density_order.py JSON with artifact_passed=true. Tables keep
zero, unavailable and uncertain observations distinct. Matplotlib is imported
only when a figure is requested; --self-test performs tiny arithmetic tests.
"""
from __future__ import annotations

import argparse
from collections import defaultdict
import csv
import hashlib
import json
import math
import os
from pathlib import Path
import unittest


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


def finite(value, label, nonnegative=False):
    require(isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value),
            "Missing/nonfinite " + label)
    require(not nonnegative or value >= 0, "Negative " + label)
    return value


def ratio(numerator, denominator):
    return numerator / denominator if numerator is not None and denominator not in (None, 0) else None


def density_metrics(case, key):
    measure = case["measures"][key]
    fine = measure["summary"]["30"]
    error = case["unsliced_error_quadrature"][key]
    q30, q26 = error["30"], error["26"]
    for field in ("negative_mass", "negative_volume", "negative_mass_fraction", "negative_volume_fraction",
                  "negative_mass_uncertainty_scale", "resolved_negative_intervals", "uncertain_negative_intervals"):
        finite(fine[field], field, True)
    squared = q30["eos_squared_integral"]
    if q30["eos_defined_everywhere_sampled"]:
        finite(squared, "q30 EOS squared integral", True)
        finite(q30["volume"], "q30 volume", True)
        require(q30["volume"] > 0, "Zero q30 volume")
        l2, rms = math.sqrt(squared), math.sqrt(squared / q30["volume"])
    else:
        require(squared is None, "Undefined q30 EOS integral has a fabricated value")
        l2 = rms = None
    if fine["negative_mass"] == 0:
        classification = "zero_observed"
    elif fine["resolved_negative_intervals"]:
        classification = "resolved_witnesses_present"
    else:
        classification = "uncertain_witnesses_only"
    surface_h = measure["surface"].get("enthalpy", {})
    witness_h = measure.get("negative_interval_witness_enthalpy", {})
    return {
        "negative_mass_over_abs_mass": fine["negative_mass_fraction"],
        "negative_volume_fraction": fine["negative_volume_fraction"],
        "negative_mass": fine["negative_mass"], "mass_denominator": abs(finite(fine["mass"], "mass")),
        "minimum_density_sliced": finite(fine["minimum_density"], "minimum density"),
        "negative_classification": classification,
        "resolved_negative_intervals": fine["resolved_negative_intervals"],
        "uncertain_negative_intervals": fine["uncertain_negative_intervals"],
        "uncertain_negative_mass": fine["uncertain_negative_mass"],
        "negative_mass_uncertainty_scale": fine["negative_mass_uncertainty_scale"],
        "negative_mass_over_uncertainty_scale": ratio(fine["negative_mass"], fine["negative_mass_uncertainty_scale"]),
        "negative_mass_q30_minus_q22": measure["negative_mass_q30_minus_q22"],
        "negative_volume_q30_minus_q22": measure["negative_volume_q30_minus_q22"],
        "eos_l2_q30": l2, "eos_rms_q30": rms,
        "eos_squared_q26": q26["eos_squared_integral"], "eos_squared_q30": squared,
        "eos_squared_q30_minus_q26": error["eos_squared_q30_minus_q26"],
        "eos_squared_relative_q30_minus_q26": ratio(error["eos_squared_q30_minus_q26"], squared),
        "q30_negative_enthalpy_integration_nodes": q30["negative_enthalpy_samples"],
        "surface_minimum_enthalpy": surface_h.get("minimum_enthalpy"),
        "surface_negative_enthalpy_samples": surface_h.get("negative_enthalpy_samples"),
        "interval_witness_minimum_enthalpy": witness_h.get("minimum_enthalpy"),
        "interval_negative_enthalpy_witnesses": witness_h.get("negative_enthalpy_samples")}


def build_rows(audits):
    """Deduplicate overlapping audited case records, never average repeat runs."""
    records = {}
    for audit_path, audit in audits:
        require(audit.get("artifact_passed") is True, "Refusing an incomplete/failed audit: " + audit_path)
        for campaign, run in audit["campaigns"].items():
            for name, case in run["cases"].items():
                require(case.get("artifact_passed") is True, "Case artifact audit failed: " + name)
                key = campaign, name
                payload = run["producer_commit"], run["mesh_sha256"], case
                if key in records:
                    require(records[key][1] == payload, "Overlapping audit records disagree: " + str(key))
                else:
                    records[key] = audit_path, payload
    frozen, coupled = [], []
    for (campaign, name), (audit_path, (commit, mesh, case)) in sorted(records.items()):
        common = {"audit_json": audit_path, "campaign": campaign, "case": name, "producer_commit": commit,
                  "mesh_sha256": mesh, "n": case["n"], "J": case["J"], "K": case["K"],
                  "source_density_degree": case["density_degree"], "source_state_dofs": case["state_size"]}
        if case["kind"] == "frozen":
            for degree in (2, 3, 4):
                key = f"Q{degree}/eos_l2_q26"
                row = {**common, "density_degree": degree, "control": "eos_l2_q26",
                       "density_dofs": case["projection"][f"Q{degree}"]["active_density_dofs"],
                       **density_metrics(case, key)}
                frozen.append(row)
        else:
            require(case["kind"] == "coupled", "Unknown case kind")
            newton = case["newton"]
            require(isinstance(newton["converged"], bool), "Invalid convergence observation")
            initial = finite(newton["initial_residual"], "initial residual", True)
            final = finite(newton["final_residual"], "final residual", True)
            reduction = ratio(final, initial)
            row = {**common, "density_degree": case["density_degree"], "control": "stored_density",
                   "density_dofs": case["block_sizes"]["value/density"], **density_metrics(case, f"Q{case['density_degree']}/stored_density"),
                   "converged": newton["converged"], "failure": newton["failure"],
                   "initial_residual": initial, "final_residual": final, "final_over_initial_residual": reduction,
                   "residual_reduction_percent": 100 * (1 - reduction) if reduction is not None else None,
                   "nonlinear_threshold": newton["threshold"], "accepted_steps": newton["accepted_steps"],
                   "iteration_records": len(newton["iterations"]), "krylov_iterations": newton["linear_iterations"],
                   "newton_iteration_seconds": finite(newton["iteration_seconds"], "Newton time", True),
                   "linear_solve_seconds": math.fsum(finite(r["linear_seconds"], "linear time", True) for r in newton["iterations"]),
                   "observer_seconds": case["observer_timings"]["observer_total"],
                   "elapsed_seconds_including_observers": case["run"]["seconds"],
                   "utc_minus_monotonic_seconds": case["run"]["utc_minus_monotonic_seconds"]}
            physical18 = case["physical"]["18"]["metrics"]
            # A bordered solve can meet its reported tolerance while the
            # unbordered equations still carry the auxiliary border action.
            for key in ("normalized_unbordered_residual", "central_border", "normalized_border_action"):
                row["final_" + key] = finite(physical18[key], "final " + key, nonnegative=key != "central_border")
            for q in (14, 18):
                physical = case["physical"][str(q)]["metrics"]
                for source, label in (("enthalpy_virial_error", "h_energy_virial"),
                                      ("enthalpy_force_virial_error", "h_force_virial"),
                                      ("virial_error", "density_energy_virial"),
                                      ("force_virial_error", "density_force_virial")):
                    value = physical[source]
                    if value is not None:
                        finite(value, source, True)
                    row[f"{label}_q{q}"] = value
                    row[f"{label}_q{q}_within_1e5"] = value <= 1e-5 if value is not None else None
            coupled.append(row)
    # Repeated same-model/same-degree runs require explicit selection upstream;
    # their averaging would conceal convergence outcomes and timing variability.
    seen = set()
    for row in coupled:
        key = row["case"], row["n"], row["J"], row["mesh_sha256"], row["density_degree"]
        require(key not in seen, "Duplicate coupled model/degree; select one audited run: " + str(key))
        seen.add(key)
    return frozen, coupled


def comparisons(frozen, coupled):
    result = []
    for kind, table in (("frozen", frozen), ("coupled", coupled)):
        groups = defaultdict(list)
        for row in table:
            key = (row["campaign"], row["case"]) if kind == "frozen" else (row["case"], row["n"], row["J"], row["mesh_sha256"])
            groups[key].append(row)
        for group in groups.values():
            baseline = next((r for r in group if r["density_degree"] == 2), None)
            if baseline is None:
                continue
            keys = ["negative_mass_over_abs_mass", "negative_volume_fraction", "eos_l2_q30"]
            if kind == "coupled":
                keys += ["newton_iteration_seconds", "krylov_iterations", "final_residual", "h_energy_virial_q18", "h_force_virial_q18"]
            for row in group:
                if row["density_degree"] != 2:
                    result.append({"kind": kind, "case": row["case"], "n": row["n"], "J": row["J"],
                                   "baseline_campaign": baseline["campaign"], "campaign": row["campaign"],
                                   "density_degree": row["density_degree"],
                                   **{key + "_over_Q2": ratio(row[key], baseline[key]) for key in keys}})
    return result


def write_csv(path, table, fields=None):
    fields = fields or list(dict.fromkeys(key for row in table for key in row))
    with path.open("w", newline="") as stream:
        if fields:
            writer = csv.DictWriter(stream, fieldnames=fields, extrasaction="ignore")
            writer.writeheader()
            writer.writerows(table)


def plot_figure(frozen, coupled, output):
    os.environ.setdefault("MPLCONFIGDIR", str(output / "matplotlib_cache"))
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt
    from matplotlib.lines import Line2D

    plt.rcParams.update({"font.size": 9, "axes.spines.top": False, "axes.spines.right": False,
                         "savefig.dpi": 190, "pdf.fonttype": 42})
    row_count = bool(frozen) + 2 * bool(coupled)
    require(row_count > 0, "No audited rows to plot")
    fig, axes = plt.subplots(row_count, 3, figsize=(13.8, 3.9 * row_count + .7), squeeze=False)

    def groups(table, frozen_case):
        result = defaultdict(list)
        for row in table:
            key = (row["campaign"], row["case"]) if frozen_case else (row["case"], row["n"], row["J"], row["mesh_sha256"])
            result[key].append(row)
        return [sorted(group, key=lambda r: r["density_degree"]) for group in result.values()]

    def panel(ax, grouped, field, title, ylabel, negative=False, log=True, second=None, legend=False):
        annotations, positive = [], False
        for index, group in enumerate(grouped):
            color = plt.get_cmap("tab10")(index % 10)
            label = group[0]["case"]
            x = [r["density_degree"] for r in group]
            series = [(field, "-")] + ([(second, "--")] if second else [])
            for metric, style in series:
                values = []
                for row in group:
                    value = row[metric]
                    if value is None or (log and value == 0):
                        annotations.append(f"{label} Q{row['density_degree']}: {'undefined' if value is None else 'observed zero'}"
                                           + (" (force)" if metric == second else ""))
                        values.append(math.nan)
                        continue
                    finite(value, "plotted " + metric, log)
                    positive = True
                    values.append(value)
                    not_converged = row.get("converged") is False
                    uncertain = negative and row["negative_classification"] == "uncertain_witnesses_only"
                    ax.plot(row["density_degree"], value, marker="x" if not_converged else "o", markersize=5,
                            color=color, markerfacecolor="none" if uncertain else color, linestyle="none")
                ax.plot(x, values, style, color=color, linewidth=1.15, label=label if metric == field else None)
        if log and positive:
            ax.set_yscale("log")
        if not positive:
            ax.text(.5, .5, "No positive finite observations" if log else "No finite observations",
                    ha="center", va="center", transform=ax.transAxes)
        ax.set_xticks([2, 3, 4], ["Q2", "Q3", "Q4"])
        ax.set_xlim(1.85, 4.15)
        ax.set(title=title, ylabel=ylabel)
        ax.grid(alpha=.2, which="both")
        if annotations:
            ax.text(0, -.19, "\n".join(annotations), fontsize=6.5, transform=ax.transAxes, va="top")
        if legend:
            ax.legend(fontsize=7, loc="best")
        if second:
            handles = [Line2D([], [], color=".3", linestyle="-", label="Energy virial"),
                       Line2D([], [], color=".3", linestyle="--", label="Force virial")]
            ax.legend(handles=handles, fontsize=7)
            ax.axhline(1e-5, color=".4", linestyle=":", linewidth=.9)

    row = 0
    if frozen:
        grouped = groups(frozen, True)
        panel(axes[row, 0], grouped, "negative_mass_over_abs_mass", "Frozen L2 projection: negative mass", r"$M_- / |M|$", negative=True, legend=True)
        panel(axes[row, 1], grouped, "negative_volume_fraction", "Frozen L2 projection: negative volume", r"$V_- / V$", negative=True)
        panel(axes[row, 2], grouped, "eos_l2_q30", "Frozen L2 projection: EOS error", r"$\|\rho-\rho_{\rm EOS}(h)\|_{L^2}$ (q30)")
        row += 1
    if coupled:
        grouped = groups(coupled, False)
        panel(axes[row, 0], grouped, "negative_mass_over_abs_mass", "Coupled endpoint: negative mass", r"$M_- / |M|$", negative=True, legend=True)
        panel(axes[row, 1], grouped, "eos_l2_q30", "Coupled endpoint: EOS error", r"$\|\rho-\rho_{\rm EOS}(h)\|_{L^2}$ (q30)")
        panel(axes[row, 2], grouped, "h_energy_virial_q18", "Coupled endpoint: P(h) virials", "Absolute virial residual / |W| (q18)", second="h_force_virial_q18")
        row += 1
        panel(axes[row, 0], grouped, "final_over_initial_residual", "Reported nonlinear residual reduction", r"$\|F_{\rm final}\| / \|F_{\rm initial}\|$")
        panel(axes[row, 1], grouped, "newton_iteration_seconds", "Newton iteration cost", "Seconds (excludes observer and setup)", log=False)
        panel(axes[row, 2], grouped, "krylov_iterations", "Total reported Krylov work", "Iterations", log=False)
    fig.suptitle("Density approximation order — audited observations", fontsize=14, y=.985)
    fig.text(.025, .018, "Open markers: uncertain-only negative witnesses; ×: nonlinear solve did not converge. Zero/undefined values are listed below panels, not placed on a log floor.\n"
             "Negative regions use root-sliced quadrature; EOS norms use unsliced q30. P(h) virials do not certify density-EOS validity. See CSV/JSON for all outcomes and timing scope.", fontsize=8)
    fig.subplots_adjust(left=.075, right=.985, top=.92 if row_count == 1 else .95, bottom=.15 if row_count == 1 else .085,
                        wspace=.35, hspace=.72)
    fig.savefig(output / "density_order_comparison.png")
    fig.savefig(output / "density_order_comparison.pdf")
    plt.close(fig)


class Tests(unittest.TestCase):
    def test_zero_and_undefined_are_not_floored(self):
        fine = {k: 0. for k in ("negative_mass", "negative_volume", "negative_mass_fraction", "negative_volume_fraction",
                "negative_mass_uncertainty_scale", "resolved_negative_intervals", "uncertain_negative_intervals", "uncertain_negative_mass")}
        fine.update(mass=1., minimum_density=0.)
        measure = {"summary": {"30": fine}, "surface": {}, "negative_mass_q30_minus_q22": 0., "negative_volume_q30_minus_q22": 0.}
        q = {"eos_squared_integral": 0., "eos_defined_everywhere_sampled": True, "volume": 1., "negative_enthalpy_samples": 0}
        error = {"30": q, "26": q.copy(), "eos_squared_q30_minus_q26": 0.}
        case = {"measures": {"Q2/stored_density": measure}, "unsliced_error_quadrature": {"Q2/stored_density": error}}
        row = density_metrics(case, "Q2/stored_density")
        self.assertEqual(row["negative_classification"], "zero_observed")
        self.assertEqual(row["eos_l2_q30"], 0.)
        self.assertIsNone(row["negative_mass_over_uncertainty_scale"])
        q.update(eos_squared_integral=None, eos_defined_everywhere_sampled=False, negative_enthalpy_samples=1)
        self.assertIsNone(density_metrics(case, "Q2/stored_density")["eos_l2_q30"])
        fine.update(negative_mass=1e-20, negative_mass_fraction=1e-20, uncertain_negative_intervals=1)
        self.assertEqual(density_metrics(case, "Q2/stored_density")["negative_classification"], "uncertain_witnesses_only")

    def test_failed_audit_rejected(self):
        with self.assertRaises(ValueError):
            build_rows([("bad.json", {"artifact_passed": False})])

    def test_zero_baseline_ratio_undefined(self):
        self.assertIsNone(ratio(1., 0.))
        self.assertEqual(ratio(0., 1.), 0.)
        self.assertIsNone(ratio(None, 1.))

    def test_audited_coupled_outcome_keeps_iteration_cost(self):
        fine = {k: 0. for k in ("negative_mass", "negative_volume", "negative_mass_fraction", "negative_volume_fraction",
                "negative_mass_uncertainty_scale", "resolved_negative_intervals", "uncertain_negative_intervals", "uncertain_negative_mass")}
        fine.update(mass=1., minimum_density=0.)
        measure = {"summary": {"30": fine}, "surface": {}, "negative_mass_q30_minus_q22": 0., "negative_volume_q30_minus_q22": 0.}
        q = {"eos_squared_integral": 4e-8, "eos_defined_everywhere_sampled": True, "volume": 1., "negative_enthalpy_samples": 0}
        metrics = {"enthalpy_virial_error": 2e-6, "enthalpy_force_virial_error": 3e-5,
                   "normalized_unbordered_residual": .04, "central_border": -.125, "normalized_border_action": .03,
                   "virial_error": None, "force_virial_error": None}
        case = {"artifact_passed": True, "kind": "coupled", "n": 1.5, "J": 0., "K": 1., "density_degree": 2,
                "state_size": 50, "block_sizes": {"value/density": 27},
                "measures": {"Q2/stored_density": measure}, "unsliced_error_quadrature": {"Q2/stored_density": {
                    "30": q, "26": q, "eos_squared_q30_minus_q26": 0.}},
                "newton": {"converged": False, "failure": "iteration limit", "initial_residual": .1,
                    "final_residual": .01, "threshold": 1e-8, "accepted_steps": 1, "linear_iterations": 7,
                    "iteration_seconds": 10., "iterations": [{"linear_seconds": 6.}]},
                "observer_timings": {"observer_total": 200.}, "run": {"seconds": 215., "utc_minus_monotonic_seconds": .01},
                "physical": {"14": {"metrics": metrics}, "18": {"metrics": metrics}}}
        audit = {"artifact_passed": True, "campaigns": {"/recorded/campaign": {
            "producer_commit": "test", "mesh_sha256": "test", "cases": {"sample": case}}}}
        frozen, coupled = build_rows([("test.json", audit)])
        self.assertEqual(frozen, [])
        row = coupled[0]
        self.assertEqual(row["newton_iteration_seconds"], 10.)
        self.assertEqual(row["observer_seconds"], 200.)
        self.assertAlmostEqual(row["residual_reduction_percent"], 90.)
        self.assertFalse(row["converged"])
        self.assertTrue(row["h_energy_virial_q18_within_1e5"])
        self.assertFalse(row["h_force_virial_q18_within_1e5"])
        self.assertIsNone(row["density_energy_virial_q18_within_1e5"])
        self.assertEqual(row["final_normalized_unbordered_residual"], .04)
        self.assertEqual(row["final_central_border"], -.125)
        self.assertEqual(row["final_normalized_border_action"], .03)


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("audits", nargs="*", type=Path)
    parser.add_argument("--output", type=Path, help="Fresh output directory")
    parser.add_argument("--tables-only", action="store_true")
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()
    if args.self_test:
        test = unittest.TextTestRunner(verbosity=2).run(unittest.defaultTestLoader.loadTestsFromTestCase(Tests))
        if not test.wasSuccessful():
            raise SystemExit(1)
        if not args.audits:
            return
    if not args.audits or args.output is None:
        parser.error("Supply completed audit JSON and --output")
    inputs = [(str(path.resolve()), json.loads(path.read_text())) for path in args.audits]
    frozen, coupled = build_rows(inputs)
    require(frozen or coupled, "No cases in input audits")
    relative = comparisons(frozen, coupled)
    args.output.mkdir(parents=True, exist_ok=False)
    write_csv(args.output / "frozen_projection.csv", frozen, ["case", "n", "J", "density_degree", "density_dofs",
        "negative_mass_over_abs_mass", "negative_volume_fraction", "negative_classification", "negative_mass_over_uncertainty_scale",
        "eos_l2_q30", "eos_rms_q30", "eos_squared_relative_q30_minus_q26", "negative_mass_q30_minus_q22",
        "surface_minimum_enthalpy", "surface_negative_enthalpy_samples"])
    write_csv(args.output / "coupled_outcomes.csv", coupled, ["case", "n", "J", "density_degree", "source_state_dofs", "density_dofs",
        "converged", "failure", "accepted_steps", "krylov_iterations", "newton_iteration_seconds", "linear_solve_seconds",
        "initial_residual", "final_residual", "final_over_initial_residual", "h_energy_virial_q18", "h_force_virial_q18",
        "final_normalized_unbordered_residual", "final_central_border", "final_normalized_border_action",
        "h_energy_virial_q14", "h_force_virial_q14", "negative_mass_over_abs_mass", "negative_volume_fraction",
        "negative_classification", "eos_l2_q30", "eos_rms_q30"])
    write_csv(args.output / "relative_to_Q2.csv", relative)
    if not args.tables_only:
        plot_figure(frozen, coupled, args.output)
    digest = lambda path: hashlib.sha256(path.read_bytes()).hexdigest()
    summary = {"schema": "density_order_summary_v1", "generator_sha256": digest(Path(__file__)),
               "input_audit_sha256": {str(path.resolve()): digest(path) for path in args.audits},
               "scope": "Presentation of audited observations only; no campaign files reopened and no independent FEM work.",
               "notes": ["Frozen controls are fixed-q26 L2 projections; EOS error norms use separate unsliced q30.",
                         "Negative-mass denominator is the absolute measured density mass, not a substituted exact target mass.",
                         "Uncertainty scales are reconstruction indicators, not rigorous bounds; zero observed negativity is not a global positivity certificate.",
                         "Newton seconds sum reported nonlinear-iteration times, excluding observer and setup costs; total elapsed time is kept separately.",
                         "Final unbordered residual, signed central border and normalized border action come from q18 physical diagnostics; bordered convergence does not imply that these vanish.",
                         "Ratios to Q2 compare recorded endpoints and do not imply equal convergence or error tolerances were attained.",
                         "Undefined values are null/blank; zero is kept exactly. No logarithmic plot floor is inserted.",
                         "P(h) virial milestones do not establish pointwise density validity; density-based virials are retained separately."],
               "frozen_projection": frozen, "coupled_outcomes": coupled, "relative_to_Q2": relative,
               "artifacts_sha256": {path.name: digest(path) for path in args.output.iterdir() if path.is_file()}}
    (args.output / "summary.json").write_text(json.dumps(summary, indent=2, allow_nan=False) + "\n")
    print("Wrote", args.output)


if __name__ == "__main__":
    main()