#!/usr/bin/env python3
"""Check saved native harmonic geometry/force algebra and campaign provenance.

No FEM residual or physical quadrature is recomputed. The force complement is
Euclidean in nodal coordinates, not the solver's dimensioned residual metric.
Artifact consistency, nonlinear convergence and physical accuracy are separate.
"""
import argparse
import csv
import hashlib
import json
import math
import os
from pathlib import Path
import tempfile

for _name in ("OPENBLAS_NUM_THREADS", "OMP_NUM_THREADS", "MKL_NUM_THREADS", "VECLIB_MAXIMUM_THREADS"):
    os.environ[_name] = "1"
import numpy as np

EPS = np.finfo(float).eps
RTOL = 5e-10
STAGES = ("initial", "final")
REQUIRED = {"metadata.txt", "layout.csv", "accepted_state.txt", "surface_basis.csv", "newton.csv"} | {
    f"{stage}_{suffix}.csv" for stage in STAGES
    for suffix in ("physical", "surface_coefficients", "surface_nodal_force", "surface_audit")
}


def digest(path):
    h = hashlib.sha256()
    with Path(path).open("rb") as stream:
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            h.update(chunk)
    return h.hexdigest()


def rows(path):
    with Path(path).open(newline="") as stream:
        reader = csv.DictReader(stream)
        if not reader.fieldnames or len(set(reader.fieldnames)) != len(reader.fieldnames):
            raise ValueError(f"Missing/duplicate CSV columns: {path}")
        data = list(reader)
    if any(None in row or any(v is None for v in row.values()) for row in data):
        raise ValueError(f"Malformed CSV row: {path}")
    return data


def finite(value):
    result = float(value)
    if not math.isfinite(result):
        raise ValueError(f"Nonfinite available algebraic value: {value}")
    return result


def integer(value):
    result = finite(value)
    if result != int(result):
        raise ValueError(f"Nonintegral index/count: {value}")
    return int(result)


def keyed(data, key):
    result = {key(row): row for row in data}
    if len(result) != len(data):
        raise ValueError("Duplicate observation key")
    return result


def contract(subscripts, *arrays):
    if any(not np.isfinite(a).all() for a in arrays):
        raise ValueError("Nonfinite contraction input")
    # Avoid the Accelerate matmul flag issue observed in the earlier auditors.
    result = np.einsum(subscripts, *arrays, optimize=False)
    if not np.isfinite(result).all():
        raise ValueError("Nonfinite contraction result")
    return result


def norm(a):
    return math.sqrt(float(contract("i,i->", np.ravel(a), np.ravel(a))))


class Audit:
    def __init__(self):
        self.checks = []

    def check(self, name, condition):
        self.checks.append({"check": name, "passed": bool(condition)})

    def close(self, name, actual, expected, scale=0.):
        a, b = np.asarray(actual, dtype=float), np.asarray(expected, dtype=float)
        valid = a.shape == b.shape and np.isfinite(a).all() and np.isfinite(b).all()
        self.check(name, valid and norm(a-b) <= RTOL * max(norm(a), norm(b)) + 256 * EPS * abs(scale))

    @property
    def passed(self):
        return all(c["passed"] for c in self.checks)


def modes(degree, maximum_order):
    if degree < 0 or maximum_order < -1 or maximum_order > degree:
        raise ValueError("Invalid harmonic degree/order")
    cutoff = degree if maximum_order == -1 else maximum_order
    return [(l, m) for l in range(degree+1) for m in range(-min(l, cutoff), min(l, cutoff)+1)]


def surface_stage(folder, stage, basis, catalog, audit):
    label = stage + ":surface:"
    coefficients = rows(folder / f"{stage}_surface_coefficients.csv")
    forces = rows(folder / f"{stage}_surface_nodal_force.csv")
    summary = {k: finite(r["value"]) for k, r in keyed(
        rows(folder / f"{stage}_surface_audit.csv"), lambda r: r["metric"]).items()}
    n, k = basis.shape
    if len(coefficients) != k or len(forces) != n:
        raise ValueError("Surface coefficient/force counts disagree with basis")
    audit.check(label+"coefficient_indices", [integer(r["index"]) for r in coefficients] == list(range(k)))
    audit.check(label+"mode_catalog", [(integer(r["degree"]), integer(r["order"])) for r in coefficients] == catalog)
    audit.check(label+"node_indices", [integer(r["node"]) for r in forces] == list(range(n)))
    audit.check(label+"reported_dimensions", summary["parameter_count"] == k and summary["nodal_count"] == n)
    a = np.array([finite(r["radial_length_coefficient"]) for r in coefficients])
    radial, force, excluded = (np.array([finite(r[field]) for r in forces])
                               for field in ("radial_displacement", "force", "excluded_force"))
    audit.close(label+"radial_displacement", radial, contract("ij,j->i", basis, a), norm(basis)*norm(a))
    q, r = np.linalg.qr(basis, mode="reduced")
    if not np.isfinite(q).all() or not np.isfinite(r).all():
        raise ValueError("Nonfinite QR factors")
    relative_pivots = np.abs(np.diag(r)) / np.sqrt(contract("ij,ij->j", basis, basis))
    audit.check(label+"resolved_mode_rank", np.isfinite(relative_pivots).all() and np.all(relative_pivots > 1e-12))
    audit.close(label+"QR_reconstruction", contract("ij,jk->ik", q, r), basis)
    audit.close(label+"QR_orthogonality", contract("ij,ik->jk", q, q), np.eye(k))
    projection = contract("ij,j->i", q, contract("ij,i->j", q, force))
    expected_excluded = force-projection
    modal_force = contract("ij,i->j", basis, force)
    audit.close(label+"excluded_projection", excluded, expected_excluded, norm(force))
    audit.close(label+"excluded_annihilates_modes", contract("ij,i->j", basis, excluded), np.zeros(k),
                norm(basis)*norm(force))
    audit.close(label+"orthogonal_energy", norm(force)**2, norm(projection)**2 + norm(excluded)**2, norm(force)**2)
    audit.close(label+"virtual_work", float(contract("i,i->", a, modal_force)),
                float(contract("i,i->", radial, force)), norm(radial)*norm(force))
    for metric, actual in (("full_nodal_force_euclidean_norm", norm(force)),
                          ("excluded_nodal_force_euclidean_norm", norm(excluded)),
                          ("modal_force_euclidean_norm", norm(modal_force))):
        audit.close(label+metric, summary[metric], actual, norm(basis)*norm(force))
    audit.check(label+"producer_equivalence_errors_nonnegative", summary["volume_equivalence_error"] >= 0 and
                summary["force_pullback_equivalence_error"] >= 0)
    audit.check(label+"producer_pullback_bound", summary["force_pullback_equivalence_error"] <=
                1e-11 * (1 + norm(modal_force)))
    return a, {"summary": summary, "radial_displacement_euclidean_norm": norm(radial),
               "coefficient_euclidean_norm": norm(a), "minimum_relative_QR_pivot": float(min(relative_pivots)),
               "excluded_force_fraction": norm(excluded)/norm(force) if norm(force) > 0 else None,
               "basis_transpose_excluded_norm": norm(contract("ij,i->j", basis, excluded)),
               "volume_and_native_pullback_equivalence": "producer observations; underlying vectors are not saved"}


def physical_stage(folder, stage, metadata, audit):
    data = keyed(rows(folder / f"{stage}_physical.csv"), lambda r: (integer(r["quadrature"]), r["metric"]))
    audit.check(stage+":physical:quadratures", {q for q, _ in data} == {14, 18})
    result = {}
    for q in (14, 18):
        m = {name: float(row["value"]) for (order, name), row in data.items() if order == q}
        label = f"{stage}:physical:q{q}:"
        audit.check(label+"model", m["polytropic_index"] == finite(metadata["n"]) and
                    m["polytropic_constant"] == finite(metadata["K"]))
        for key in ("complete_volume_measurement", "pressure_from_enthalpy_defined_everywhere",
                    "pressure_from_density_defined_everywhere"):
            audit.check(label+key+":boolean", m[key] in (0., 1.))
        complete = m["complete_volume_measurement"] == 1
        if complete:
            audit.close(label+"mass_error", m["mass_absolute_error"], abs(m["mass"]-finite(metadata["M"])), abs(m["mass"]))
            audit.close(label+"J_error", m["angular_momentum_absolute_error"], abs(m["angular_momentum_z"]-finite(metadata["J"])), 1.)
        milestones = {}
        for prefix, pressure, valid in (("", "pressure_integral", "pressure_from_density_defined_everywhere"),
                                       ("enthalpy_", "enthalpy_pressure_integral", "pressure_from_enthalpy_defined_everywhere")):
            defined = complete and m[valid] == 1 and math.isfinite(m["binding_energy"]) and abs(m["binding_energy"]) > 0
            for force_prefix, binding in (("", "binding_energy"), ("force_", "force_binding_energy")):
                key = prefix+force_prefix+"virial_"
                signed, error = m[key+"signed"], m[key+"error"]
                if defined:
                    terms = (2*m["kinetic_energy"], m[binding], 3*m[pressure])
                    scale = math.fsum(abs(v) for v in terms)/abs(m["binding_energy"])
                    audit.close(label+key+"identity", signed, math.fsum(terms)/abs(m["binding_energy"]), scale)
                    audit.close(label+key+"absolute", error, abs(signed), scale)
                else:
                    audit.check(label+key+"undefined", math.isnan(signed) and math.isnan(error))
                milestones[key+"within_1e5"] = bool(error <= 1e-5) if defined and math.isfinite(error) else None
        audit.check(label+"finite_native_residual", math.isfinite(m["normalized_bordered_residual"]) and
                    m["normalized_bordered_residual"] >= 0)
        result[str(q)] = {"metrics": m, "virial_milestones": milestones,
                          "undefined_metrics": [name for name, value in m.items() if not math.isfinite(value)]}
    return result


def inspect_case(folder, record, metadata, audit):
    if metadata["surface_prescription"] != "harmonics":
        raise ValueError("The native harmonic auditor requires harmonic cases")
    catalog = modes(integer(metadata["surface_degree"]), integer(metadata["surface_maximum_order"]))
    with (folder / "surface_basis.csv").open(newline="") as stream:
        basis = np.array([[finite(value) for value in row] for row in csv.reader(stream)])
    if basis.ndim != 2 or basis.shape[1] != len(catalog) or basis.shape[0] < basis.shape[1]:
        raise ValueError("Invalid tall harmonic basis dimensions")
    audit.close("basis:monopole_normalization", basis[:, 0], np.full(basis.shape[0], 1/math.sqrt(4*math.pi)))
    surface, coefficients = {}, {}
    for stage in STAGES:
        coefficients[stage], surface[stage] = surface_stage(folder, stage, basis, catalog, audit)
    layout = keyed(rows(folder / "layout.csv"), lambda r: (r["kind"], r["block"]))
    totals = {}
    for kind in ("value", "residual"):
        blocks = sorted((r for (k, _), r in layout.items() if k == kind), key=lambda r: integer(r["offset"]))
        offset = 0
        for row in blocks:
            audit.check("layout:"+kind+":"+row["block"], integer(row["offset"]) == offset and integer(row["size"]) > 0 and finite(row["scale"]) > 0)
            offset += integer(row["size"])
        totals[kind] = offset
    audit.check("layout:square", totals["value"] == totals["residual"] > 0)
    shape = layout["value", "surface_deformation"]
    audit.check("layout:shape_dimensions", integer(shape["size"]) == len(catalog) and
                integer(layout["residual", "surface_shape_balance"]["size"]) == len(catalog))
    with (folder / "accepted_state.txt").open() as stream:
        size = integer(stream.readline())
        state = np.loadtxt(stream, ndmin=1)
    if state.ndim != 1 or len(state) != size or size != totals["value"] or not np.isfinite(state).all():
        raise ValueError("Invalid accepted physical state")
    start = integer(shape["offset"])
    audit.close("final:coefficients_match_accepted_state", coefficients["final"], state[start:start+len(catalog)])
    physics = {stage: physical_stage(folder, stage, metadata, audit) for stage in STAGES}
    history = rows(folder / "newton.csv")
    audit.check("newton:ordered_iterations", [integer(row["iteration"]) for row in history] == list(range(len(history))))
    accepted = 0
    before = physics["initial"]["18"]["metrics"]["normalized_bordered_residual"]
    for row in history:
        label = "newton:"+row["iteration"]+":"
        flag, residual, alpha = integer(row["accepted"]), finite(row["F"]), finite(row["alpha"])
        audit.check(label+"accepted_boolean", flag in (0, 1))
        audit.check(label+"nonnegative_norm", residual >= 0)
        for name in ("seconds", "linear_seconds", "prepare_seconds", "refresh_seconds"):
            audit.check(label+name, finite(row[name]) >= 0)
        if flag == 1:
            accepted += 1
            audit.check(label+"accepted_alpha", 0 < alpha <= 1)
            audit.check(label+"recorded_true_linear_gate", 0 <= finite(row["linear_true"]) <= finite(metadata["linear_tol"])*(1+RTOL))
            audit.check(label+"residual_decrease", residual <= before*(1+RTOL))
        else:
            audit.close(label+"rejected_state_preserved", residual, before)
        before = residual
    audit.check("newton:accepted_count", integer(metadata["accepted_steps"]) == accepted)
    converged = integer(metadata["converged"])
    audit.check("newton:converged_boolean", converged in (0, 1))
    audit.check("newton:exit_matches_outcome", record["exit_code"] == (0 if converged else 3))
    for q in ("14", "18"):
        audit.close("newton:final_residual:q"+q, physics["final"][q]["metrics"]["normalized_bordered_residual"], before)
    sensitivities = {stage: {name: physics[stage]["18"]["metrics"][name]-value
                            for name, value in physics[stage]["14"]["metrics"].items()
                            if math.isfinite(value) and math.isfinite(physics[stage]["18"]["metrics"].get(name, math.nan))}
                     for stage in STAGES}
    return {"metadata": metadata, "modes": catalog, "basis_shape": list(basis.shape), "surface": surface,
            "physical": physics, "q18_minus_q14": sensitivities, "newton": history,
            "accepted_steps": accepted, "converged": bool(converged), "failure": metadata.get("failure"),
            "initial_residual": physics["initial"]["18"]["metrics"]["normalized_bordered_residual"],
            "final_residual": before, "summed_iteration_seconds": math.fsum(finite(r["seconds"]) for r in history),
            "basis_values_independently_reconstructed": False,
            "basis_limitation": "Reference node directions are not saved; mode labels and monopole are checked, not all Y_lm values."}


def audit_campaign(campaign):
    campaign = campaign.resolve()
    audit, results = Audit(), {}
    provenance = json.loads((campaign / "provenance.json").read_text())
    cases = provenance["cases"]
    audit.check("campaign:manifest", cases == json.loads((campaign / "cases.json").read_text()) and
                len(cases) > 0 and len({c["id"] for c in cases}) == len(cases))
    audit.check("campaign:binary", digest(campaign / "plateau_diagnostic") == provenance["binary_sha256"])
    audit.check("campaign:mesh", digest(provenance["mesh"]) == provenance["mesh_sha256"])
    for case in cases:
        name, arguments = case["id"], case["arguments"]
        local = Audit()
        try:
            if not name or any(c not in "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789_-" for c in name):
                raise ValueError("Invalid case directory name")
            folder = campaign / name
            record_path = campaign / (name+"_run.json")
            record = json.loads(record_path.read_text())
            local.check("provenance:completed", bool(record.get("finished_utc")) and not record.get("timed_out", False) and record.get("exit_code") in (0, 3))
            local.check("provenance:command", record["id"] == name and record["command"] ==
                        [str(campaign / "plateau_diagnostic"), *arguments, "--mesh", provenance["mesh"], "--output", str(folder)])
            hashes = record["output_hashes"]
            local.check("provenance:required_hashes", REQUIRED <= set(hashes))
            observed = {str(p.relative_to(folder)) for p in folder.rglob("*") if p.is_file() and p.suffix in (".txt", ".csv", ".json")}
            local.check("provenance:complete_output_manifest", set(hashes) == observed)
            for relative, expected in hashes.items():
                path = folder / relative
                if not path.resolve().is_relative_to(folder.resolve()):
                    raise ValueError("Output hash path escapes case directory")
                local.check("provenance:output:"+relative, digest(path) == expected)
            for path, expected in record.get("input_state_hashes", {}).items():
                local.check("provenance:input:"+path, digest(path) == expected)
            entries = [line.split("=", 1) for line in (folder / "metadata.txt").read_text().splitlines() if "=" in line]
            metadata = dict(entries)
            if len(metadata) != len(entries):
                raise ValueError("Duplicate metadata key")
            for option, key in (("--index", "n"), ("--angular-momentum", "J"), ("--surface-degree", "surface_degree"),
                                ("--surface-maximum-order", "surface_maximum_order"), ("--max-newton", "max_newton"),
                                ("--linear-tol", "linear_tol"), ("--max-linear", "max_linear"), ("--restart", "restart")):
                if option in arguments:
                    local.check("requested:"+key, finite(metadata[key]) == finite(arguments[arguments.index(option)+1]))
            for option, key in (("--surface-prescription", "surface_prescription"), ("--strategy", "strategy"),
                                ("--preconditioner-variant", "preconditioner_variant")):
                if option in arguments:
                    local.check("requested:"+key, metadata[key] == arguments[arguments.index(option)+1])
            if "--no-central" in arguments:
                local.check("requested:no_central", metadata["central_phase"] == "0")
            result = inspect_case(folder, record, metadata, local)
            result["run_sha256"] = digest(record_path)
            result["timing"] = {key: record.get(key) for key in ("seconds", "utc_elapsed_seconds", "utc_minus_monotonic_seconds", "child_cpu_seconds")}
        except (OSError, ValueError, KeyError, TypeError, IndexError, OverflowError, np.linalg.LinAlgError) as error:
            local.check("readable_complete_artifacts:"+str(error), False)
            result = {"error": str(error)}
        result.update(artifact_passed=local.passed, checks=local.checks)
        results[name] = result
        audit.check("case:"+name, local.passed)
    return {"schema_version": 1, "artifact_passed": audit.passed, "campaign": str(campaign),
            "source_commit": provenance["commit"], "provenance_sha256": digest(campaign / "provenance.json"),
            "auditor_sha256": digest(__file__), "tolerances": {"relative": RTOL, "arithmetic_allowance_eps": 256},
            "scope": "Saved algebra, hashes and reported scalar identities; no fresh FEM, quadrature or model validity certificate.",
            "checks": audit.checks, "cases": results}


def self_test():
    """Tiny real-mode fixture; deliberate corruptions must fail independently."""
    constant, linear = 1/math.sqrt(4*math.pi), math.sqrt(3/(4*math.pi))
    points = np.array([[1, 0, 0], [-1, 0, 0], [0, 1, 0], [0, -1, 0], [0, 0, 1], [0, 0, -1]])
    basis = np.column_stack((np.full(6, constant), -linear*points[:, 1], linear*points[:, 2], -linear*points[:, 0]))
    coefficients = np.array([.2, .01, .02, .03])
    force = np.array([.3, -.4, .2, .7, -.3, .9])
    modal = contract("ij,i->j", basis, force)
    excluded = force-contract("ij,j->i", basis, modal/contract("ij,ij->j", basis, basis))
    radial = contract("ij,j->i", basis, coefficients)
    outcomes = []
    for corruption in (None, "catalog", "radial", "excluded", "norm", "nonfinite", "coefficient_count"):
        with tempfile.TemporaryDirectory() as directory:
            folder = Path(directory)
            mode_rows = [[i, l, m, coefficients[i]] for i, (l, m) in enumerate(modes(1, -1))]
            force_rows = [[i, radial[i], force[i], excluded[i]] for i in range(6)]
            summary = {"parameter_count": 4, "nodal_count": 6, "full_nodal_force_euclidean_norm": norm(force),
                       "excluded_nodal_force_euclidean_norm": norm(excluded), "modal_force_euclidean_norm": norm(modal),
                       "volume_equivalence_error": 0., "force_pullback_equivalence_error": 0.}
            if corruption == "catalog": mode_rows[1][2] = 0
            if corruption == "radial": force_rows[0][1] += .01
            if corruption == "excluded": force_rows[1][3] += .01
            if corruption == "norm": summary["modal_force_euclidean_norm"] += .01
            if corruption == "nonfinite": force_rows[2][2] = math.nan
            if corruption == "coefficient_count": mode_rows.pop()
            for name, header, data in (("coefficients", ("index", "degree", "order", "radial_length_coefficient"), mode_rows),
                                       ("nodal_force", ("node", "radial_displacement", "force", "excluded_force"), force_rows),
                                       ("audit", ("metric", "value"), summary.items())):
                with (folder / ("initial_surface_"+name+".csv")).open("w", newline="") as stream:
                    writer = csv.writer(stream); writer.writerow(header); writer.writerows(data)
            check = Audit()
            try:
                surface_stage(folder, "initial", basis, modes(1, -1), check)
                passed = check.passed
            except (ValueError, KeyError):
                passed = False
            outcomes.append({"fixture": corruption or "valid", "passed": passed == (corruption is None)})
    return {"self_test_passed": all(r["passed"] for r in outcomes), "fixtures": outcomes}


def safe(value):
    if isinstance(value, float) and not math.isfinite(value): return None
    if isinstance(value, dict): return {k: safe(v) for k, v in value.items()}
    if isinstance(value, (list, tuple)): return [safe(v) for v in value]
    return value


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("campaign", type=Path, nargs="?")
    parser.add_argument("--output", type=Path)
    parser.add_argument("--self-test", action="store_true")
    args = parser.parse_args()
    if args.self_test:
        result = self_test()
        print(json.dumps(result, indent=2))
        return 0 if result["self_test_passed"] else 1
    if args.campaign is None or args.output is None:
        parser.error("A campaign and --output are required unless --self-test is selected")
    try:
        result = audit_campaign(args.campaign)
    except (OSError, ValueError, KeyError, TypeError) as error:
        result = {"artifact_passed": False, "error": str(error), "auditor_sha256": digest(__file__)}
    args.output.write_text(json.dumps(safe(result), indent=2, allow_nan=False)+"\n")
    print("PASS" if result["artifact_passed"] else "FAIL", args.output)
    return 0 if result["artifact_passed"] else 1


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