"""G0 fixed-density CPU noise measurement; no GPU or reference-program imports.

Run ``export`` once through vq to freeze direct-route SAD, mid-SCF, converged
and signed-perturbation densities. Run ``replay`` in separate processes at
OMP=1 and N, with BLAS threads pinned to one. ``summarize`` compares identical
input bytes, reports measured floors, and sets no scientific acceptance limit.
All fixtures and raw records belong in an explicitly selected external root.
"""
from __future__ import annotations

import argparse
import hashlib
import itertools
import json
import math
import os
from pathlib import Path
import time

import numpy as np

try:
    from examples.regression import gpu_profile as gp
except ModuleNotFoundError:
    # vq workspaces carry the two standalone harness files together.
    import importlib.util
    spec = importlib.util.spec_from_file_location(
        "gpu_profile", Path(__file__).with_name("gpu_profile.py"))
    gp = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(gp)

FIXTURE_SCHEMA = "vibeqc.gpu_g0.density_fixture/1"
REPLAY_SCHEMA = "vibeqc.gpu_g0.noise_replay/1"
DENSITIES = ("sad", "mid_scf", "converged", "perturb_plus", "perturb_minus")


def matrix_hash(value):
    a = np.asarray(value, dtype="<f8", order="C")
    return hashlib.sha256(str(a.shape).encode() + a.tobytes()).hexdigest()


def external_dir(value):
    root = Path(value).expanduser().resolve()
    # Also works in a standalone worker without the source checkout.
    for parent in Path(__file__).resolve().parents:
        if (parent / ".git").exists() and root.is_relative_to(parent):
            raise ValueError("raw evidence must be outside the checkout")
    root.mkdir(parents=True, exist_ok=True)
    return root


def write_json(path, data):
    path.write_text(json.dumps(data, indent=1, allow_nan=False) + "\n")


def system(case):
    import vibeqc as vq
    mol = vq.Molecule([vq.Atom(z, xyz) for z, xyz in case["atoms_bohr"]],
                      charge=case["charge"], multiplicity=case["multiplicity"])
    return mol, vq.BasisSet(mol, case["basis"])


def identity(prov):
    info = prov.get("build_info", {})
    if info.get("dirty") is not False or not info.get("sha_full"):
        raise ValueError("unknown or dirty source identity")
    if not prov.get("native_sha256"):
        raise ValueError("missing native identity")
    return info["sha_full"], prov["native_sha256"]


def require_awake(start):
    slept = gp.host_sleep_s(start, gp.sleep_clock_pair())
    if slept is None or slept > gp.HOST_SLEEP_TOL_S:
        raise ValueError("unavailable sleep clock or host suspension")
    return slept


def signed_perturbations(density, overlap, relative_scale=1e-7):
    """Small symmetric signed perturbation, preserving Tr(D S).

    No positive-semidefinite projection: replay must retain the density path.
    The generator and amplitude are fixture metadata, not admission limits.
    """
    rng = np.random.default_rng(20261028)
    a = rng.normal(size=density.shape)
    a = (a + a.T) / 2
    a -= np.eye(len(a)) * np.sum(a * overlap.T) / np.trace(overlap)
    a *= relative_scale * max(1., np.linalg.norm(density)) / np.linalg.norm(a)
    return density + a, density - a


def export_fixture(case, out_dir):
    import vibeqc as vq
    out = external_dir(out_dir)
    if case["route"] != "direct":
        raise ValueError("freeze one common direct-route density set per method/system")
    before = gp.runtime_provenance()
    identity(before)
    clocks = gp.sleep_clock_pair()
    vq.set_num_threads(1)
    mol, basis = system(case)
    runner, opts = gp.make_options(case, component_timing=True, stability_check=False)
    full = runner(mol, basis, opts)
    trace = list(full.scf_trace)
    if not full.converged or not trace or trace[-1].grad_norm > gp.RECIPE["conv_tol_grad"]:
        raise ValueError("direct fixture SCF did not meet the pinned residual target")
    mid_iter = max(1, int(full.n_iter) // 2)
    opts.max_iter = mid_iter
    mid = runner(mol, basis, opts)
    overlap = vq.compute_overlap(basis)
    normalize = vq._vibeqc_core._normalize_guess_density
    density = np.asarray(normalize(full.density, overlap, mol.n_electrons())).real.copy()
    plus, minus = signed_perturbations(density, overlap)
    arrays = dict(zip(DENSITIES, (vq.sad_density(mol, basis), mid.density,
                                 density, plus, minus)))
    preparation = {}
    for name, value in arrays.items():
        prepared = np.asarray(normalize(value, overlap, mol.n_electrons())).real.copy()
        preparation[name] = float(np.max(np.abs(prepared - value)))
        arrays[name] = prepared
    arrays["overlap"] = overlap
    for value in arrays.values():
        if not np.isfinite(value).all():
            raise ValueError("non-finite fixture")
    after = gp.runtime_provenance()
    if identity(before) != identity(after):
        raise ValueError("runtime changed during fixture export")
    meta = {"schema": FIXTURE_SCHEMA, "case": case,
            "basis_sha256": gp.basis_sha256(basis),
            "harness_sha256": gp._sha256_file(__file__),
            "matrix_sha256": {k: matrix_hash(v) for k, v in arrays.items()},
            "matrix_layout": {k: {"shape": list(np.asarray(v).shape),
                                   "strides_bytes": list(np.asarray(v).strides),
                                   "dtype": str(np.asarray(v).dtype)}
                              for k, v in arrays.items()},
            "converged_energy_eh": float(full.energy),
            "converged_iterations": int(full.n_iter), "mid_scf_iteration_cap": mid_iter,
            "mid_scf_iterations": int(mid.n_iter),
            "converged_grad_norm": float(trace[-1].grad_norm),
            "perturbation_relative_frobenius": 1e-7,
            "density_convention": "restricted total D = 2 C_occ C_occ^T",
            "density_preparation": "native restart normalization before freezing",
            "density_preparation_max_abs": preparation,
            "electron_counts": {k: float(np.sum(v * overlap.T)) for k, v in arrays.items()
                                if k != "overlap"},
            "host_sleep_s": require_awake(clocks),
            "provenance_before": before, "provenance": after}
    np.savez(out / "densities.npz", **arrays)
    meta["archive_sha256"] = gp._sha256_file(str(out / "densities.npz"))
    write_json(out / "fixture.json", meta)
    return meta


def load_fixture(path):
    path = Path(path).resolve()
    meta = json.loads(path.read_text())
    if meta["schema"] != FIXTURE_SCHEMA:
        raise ValueError("unknown fixture schema")
    archive = path.with_name("densities.npz")
    if gp._sha256_file(str(archive)) != meta["archive_sha256"]:
        raise ValueError("fixture archive changed")
    with np.load(archive, allow_pickle=False) as data:
        arrays = {key: np.array(data[key], copy=True) for key in data.files}
    if set(arrays) != set(DENSITIES) | {"overlap"}:
        raise ValueError("incomplete fixture density set")
    for key, value in arrays.items():
        if (matrix_hash(value) != meta["matrix_sha256"][key]
                or value.shape != arrays["overlap"].shape
                or value.ndim != 2 or value.shape[0] != value.shape[1]
                or not np.isfinite(value).all()):
            raise ValueError("invalid fixture matrix")
    return meta, arrays


def xc_evaluator(case, mol, basis, overlap):
    """Use the production unpolarised XC path without J/K subtraction.

    Prefer the direct restricted evaluator, which accepts signed perturbations.
    Older runtimes use NOITER with H=J=K=0 and explicitly refuse inputs that
    the SCF restart interface rejects or changes.
    """
    import vibeqc as vq
    core = vq._vibeqc_core
    _, opts = gp.make_options(case, component_timing=False, stability_check=False)
    opts.max_iter = 0
    grid = vq.build_grid(mol, opts.grid)
    metadata = {"points_sha256": matrix_hash(grid.points),
                "weights_sha256": matrix_hash(grid.weights),
                "n_points": len(grid.weights), "options": gp.serialize(opts.grid)}
    if hasattr(core, "evaluate_rks_xc_potential"):
        functional = vq.Functional(opts.functional, 1)

        def evaluate(density):
            value, energy, workers = core.evaluate_rks_xc_potential(
                functional, basis, grid, density)
            return np.array(value), float(energy), int(workers)

        return evaluate, dict(metadata, evaluator="restricted_fixed_density")

    zero = np.zeros(overlap.shape)

    class ZeroJK(core.JKBuilder):
        def __init__(self):
            super().__init__()

        def build_J(self, density):
            return zero

        def build_K(self, density):
            return zero

    jk = ZeroJK()

    def evaluate(density):
        result = core.run_rks_scf_with_jk(
            basis, mol.n_electrons(), overlap, zero, 0., jk, grid, opts, density)
        if result.n_iter != 0 or not np.array_equal(result.density, density):
            raise ValueError("NOITER did not preserve the replay density")
        return np.array(result.fock), float(result.e_xc), int(result.xc_batch_workers_used)

    return evaluate, dict(metadata, evaluator="restricted_noiter_fallback")


def replay(fixture_path, route, threads, repeats, out_dir):
    import vibeqc as vq
    core = vq._vibeqc_core
    if threads < 1 or repeats < 2:
        raise ValueError("positive threads and at least two repeats required")
    for name in ("OPENBLAS_NUM_THREADS", "MKL_NUM_THREADS", "VECLIB_MAXIMUM_THREADS",
                 "BLIS_NUM_THREADS"):
        if os.environ.get(name) != "1":
            raise ValueError(f"pin {name}=1 before starting the process")
    out = external_dir(out_dir)
    meta, arrays = load_fixture(fixture_path)
    before = gp.runtime_provenance()
    if identity(before) != identity(meta["provenance"]):
        raise ValueError("fixture and replay runtime identities differ")
    vq.set_num_threads(threads)
    case = dict(meta["case"], route=route, request=gp.route_request(route, meta["case"]["method"]))
    mol, basis = system(case)
    if gp.basis_sha256(basis) != meta["basis_sha256"]:
        raise ValueError("resolved basis changed")
    _, opts = gp.make_options(case, component_timing=False, stability_check=False)
    resolved = gp.resolve_settings(case, mol, basis, opts)
    grids = {}
    aux = None
    if route in ("direct", "auto"):
        jk = (core.make_four_index_jk_builder(basis)
              if resolved["scf_mode_resolved"] == "CONVENTIONAL"
              else core.make_direct_jk_builder(basis, opts.schwarz_threshold, False))
    elif route == "df":
        aux = vq.BasisSet(mol, case["request"]["aux_basis"])
        jk = core.make_df_jk_builder(basis, aux)
    elif route == "rijcosx":
        if resolved.get("cosx_variant") != "FITTED":
            raise ValueError("the exposed COSX builder requires the FITTED variant")
        aux = vq.BasisSet(mol, case["request"]["aux_basis"])
        stage = vq.cosx_grid_stages_for_level(resolved["cosx_grid_level"])[-1]
        grid = vq.build_grid(mol, stage)
        jk = core.make_cosx_jk_builder(basis, aux, grid)
        grids["cosx"] = {"points_sha256": matrix_hash(grid.points),
                         "weights_sha256": matrix_hash(grid.weights),
                         "n_points": len(grid.weights), "stage": "final",
                         "one_center_correction": "excluded; iterated K only"}
    else:
        raise ValueError("replay route must be direct, df, rijcosx or auto")
    evaluate_xc = None
    if case["method"] != "RHF":
        evaluate_xc, grids["xc"] = xc_evaluator(case, mol, basis, arrays["overlap"])
    ax = gp.EXCHANGE_FRACTION[case["method"]]
    values, observations = {}, []
    for repeat in range(repeats):
        for label in DENSITIES:
            density = arrays[label]
            operations = ["J"] + (["K"] if ax else []) + (["Vxc"] if evaluate_xc else [])
            for op in operations:
                jk.reset_state()
                clock = gp.sleep_clock_pair()
                t0 = time.perf_counter()
                key = f"{label}__{op}__rep{repeat}"
                if op == "Vxc":
                    try:
                        value, energy, workers = evaluate_xc(density)
                    except ValueError as exc:
                        if (not label.startswith("perturb_") or str(exc) !=
                                "initial guess: density must be positive semidefinite"):
                            raise
                        observations.append({"key": key, "density": label, "operator": op,
                                             "repeat": repeat, "status": "REFUSED",
                                             "reason": "SCF restart rejects indefinite density",
                                             "host_sleep_s": require_awake(clock)})
                        continue
                else:
                    value = np.array(jk.build_J(density) if op == "J" else jk.build_K(density))
                    factor = 0.5 if op == "J" else -0.25 * ax
                    energy = factor * float(np.sum(density * value.T))
                    workers = None  # Not all builders instrument their actual team.
                elapsed = time.perf_counter() - t0
                slept = require_awake(clock)
                if not np.isfinite(value).all() or not math.isfinite(energy):
                    raise ValueError("non-finite replay observable")
                values[key] = value
                observations.append({"key": key, "density": label, "operator": op,
                                     "repeat": repeat, "energy_eh": energy,
                                     "wall_s": elapsed, "host_sleep_s": slept,
                                     "xc_workers_observed": workers,
                                     "matrix_sha256": matrix_hash(value)})
    after = gp.runtime_provenance()
    if identity(before) != identity(after):
        raise ValueError("runtime changed during replay")
    stem = f"{route}__t{threads}"
    np.savez(out / f"{stem}.npz", **values)
    record = {"schema": REPLAY_SCHEMA, "fixture": meta, "route": route,
              "threads": threads, "threads_observed": vq.get_num_threads(),
              "repeats": repeats, "grids": grids, "resolved": resolved,
              "screening": {"incremental": False, "schwarz_threshold": opts.schwarz_threshold},
              "observations": observations, "provenance_before": before, "provenance": after,
              "archive_sha256": gp._sha256_file(str(out / f"{stem}.npz")),
              "harness_sha256": gp._sha256_file(__file__)}
    write_json(out / f"{stem}.json", record)
    from vibeqc.output import OutputChannel, write
    with OutputChannel.to_file(out / f"{stem}.out"):
        write(f"G0 fixed-density replay: {case['system']} {case['method']} {route}\n")
        write(f"Threads {threads}; repeats {repeats}; observables {len(observations)}\n")
    return record


def error_metrics(candidate, reference, density, factor):
    """FP64 matrix errors; zero reference norm gives None unless both zero."""
    delta = candidate - reference
    norm = float(np.linalg.norm(reference))
    diff = float(np.linalg.norm(delta))
    return {"max_abs_eh": float(np.max(np.abs(delta))),
            "frobenius_relative": diff / norm if norm else (0. if diff == 0 else None),
            "symmetry_defect_eh": float(np.max(np.abs(candidate - candidate.T))),
            "contraction_error_eh": abs(factor * float(np.sum(density * delta.T)))}


def summarize(raw_dir, threads, expected_source_sha, expected_native_sha256):
    """All-pairs repeat and cross-thread floors, separated by frozen input."""
    if len(threads) != 2 or 1 not in threads or min(threads) < 1 or len(set(threads)) != 2:
        raise ValueError("request exactly OMP=1 and one N > 1")
    grouped = {}
    for path in Path(raw_dir).rglob("*.json"):
        record = json.loads(path.read_text())
        if not isinstance(record, dict) or record.get("schema") != REPLAY_SCHEMA:
            continue
        for prov in (record["provenance_before"], record["provenance"],
                     record["fixture"]["provenance_before"], record["fixture"]["provenance"]):
            if identity(prov) != (expected_source_sha, expected_native_sha256):
                raise ValueError("unexpected noise runtime identity")
        operators = ["J"]
        if gp.EXCHANGE_FRACTION[record["fixture"]["case"]["method"]]:
            operators.append("K")
        if record["fixture"]["case"]["method"] != "RHF":
            operators.append("Vxc")
        expected = {(d, op, rep) for d in DENSITIES for op in operators
                    for rep in range(record["repeats"])}
        observed = [(o["density"], o["operator"], o["repeat"])
                    for o in record["observations"]]
        if len(observed) != len(set(observed)) or set(observed) != expected:
            raise ValueError("missing or duplicate replay observable")
        if any(o.get("host_sleep_s") is None or o["host_sleep_s"] > gp.HOST_SLEEP_TOL_S
               or (o.get("status") != "REFUSED" and not math.isfinite(o["energy_eh"]))
               for o in record["observations"]):
            raise ValueError("invalid replay clock or energy")
        archive = path.with_suffix(".npz")
        if gp._sha256_file(str(archive)) != record["archive_sha256"]:
            raise ValueError("replay archive changed")
        key = (record["fixture"]["archive_sha256"], record["route"])
        grouped.setdefault(key, []).append((record, archive))
    rows = []
    for (fixture_hash, route), entries in sorted(grouped.items()):
        first = entries[0][0]
        case = first["fixture"]["case"]
        matrices, observations, refusals = {}, {}, {}
        configs = set()
        for record, archive in entries:
            t = record["threads"]
            if t in configs or record["threads_observed"] != t:
                raise ValueError("duplicate replay or wrong thread count")
            configs.add(t)
            if (record["grids"] != first["grids"] or record["screening"] != first["screening"]
                    or record["resolved"] != first["resolved"]
                    or record["harness_sha256"] != first["harness_sha256"]
                    or any(record["provenance"].get(k) != first["provenance"].get(k)
                           for k in ("cpu_model", "machine", "platform", "native_build_info",
                                     "blas_info", "numpy"))):
                raise ValueError("replay settings differ across threads")
            with np.load(archive, allow_pickle=False) as data:
                for obs in record["observations"]:
                    k = (obs["density"], obs["operator"])
                    matrices.setdefault(k, [])
                    observations.setdefault(k, [])
                    if obs.get("status") == "REFUSED":
                        refusals.setdefault(k, []).append(
                            {"threads": t, "repeat": obs["repeat"], "reason": obs["reason"]})
                        continue
                    value = np.array(data[obs["key"]])
                    if not np.isfinite(value).all() or matrix_hash(value) != obs["matrix_sha256"]:
                        raise ValueError("replay matrix changed")
                    k = (obs["density"], obs["operator"])
                    matrices.setdefault(k, []).append((t, value))
                    observations.setdefault(k, []).append(obs)
        # Density arrays are deliberately external: archive path is a raw-only field.
        matches = [p for p in Path(raw_dir).rglob("densities.npz")
                   if gp._sha256_file(str(p)) == fixture_hash]
        if len(matches) != 1:
            raise ValueError("exactly one matching frozen density archive required")
        with np.load(matches[0], allow_pickle=False) as data:
            densities = {k: np.array(data[k]) for k in DENSITIES}
        for (label, op), values in sorted(matrices.items()):
            obs = observations[label, op]
            factor = 0.5 if op == "J" else (-0.25 * gp.EXCHANGE_FRACTION[case["method"]]
                                            if op == "K" else 1.)
            categories = {"repeat": [], "omp_1_vs_n": []}
            for i, j in itertools.combinations(range(len(values)), 2):
                ta, a = values[i]
                tb, b = values[j]
                if ta == 1 and tb != 1:
                    a, b = b, a  # OMP=1 is the cross-thread reference.
                kind = "repeat" if ta == tb else "omp_1_vs_n"
                m = error_metrics(a, b, densities[label], factor)
                m["symmetry_defect_eh"] = max(
                    m["symmetry_defect_eh"], float(np.max(np.abs(b - b.T))))
                m["energy_error_eh"] = abs(obs[i]["energy_eh"] - obs[j]["energy_eh"])
                categories[kind].append(m)
            floors = {}
            for kind, errors in categories.items():
                floors[kind] = {k: (None if any(e[k] is None for e in errors)
                                    else max(e[k] for e in errors))
                                for k in errors[0]} if errors else None
            counts = {t: sum(ta == t for ta, _ in values) for t in configs}
            complete = (configs == set(threads) and all(n >= 3 for n in counts.values())
                        and not refusals.get((label, op)))
            rows.append({"system": case["system"], "method": case["method"],
                         "route": route, "density": label, "operator": op,
                         "fixture_sha256": fixture_hash,
                         "density_sha256": first["fixture"]["matrix_sha256"][label],
                         "basis_sha256": first["fixture"]["basis_sha256"],
                         "aux_basis_sha256": first["resolved"].get("aux_basis_sha256"),
                         "cpu_model": first["provenance"].get("cpu_model"),
                         "compiler": (first["provenance"].get("native_build_info") or {}).get("compiler_version"),
                         "blas": gp.redact((first["provenance"].get("blas_info") or {}).get("libraries")),
                         "grids": first["grids"], "screening": first["screening"],
                         "repeats_by_threads": counts, "noise_floors": floors,
                         "refusals": refusals.get((label, op), []),
                         "status": "MEASURED" if complete else "INCOMPLETE"})
    if not rows:
        raise ValueError("no noise records found")
    result = {"schema": "vibeqc.gpu_g0.noise_summary/1", "rows": rows,
              "source_sha": expected_source_sha, "native_sha256": expected_native_sha256,
              "scope": "real restricted fixed-density CPU arithmetic; no scientific tolerance assigned",
              "vxc_contraction": "Tr(D delta Vxc), a diagnostic; energy_error uses nonlinear E_xc",
              "cosx_scope": "FITTED final-grid iterated K; one-center correction excluded",
              "zero_denominator": "relative error is zero if both matrices vanish, otherwise null"}
    gp.assert_public_safe(result)
    return result


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    sub = ap.add_subparsers(dest="cmd", required=True)
    ex = sub.add_parser("export")
    ex.add_argument("case_json"); ex.add_argument("--out-dir", required=True)
    rp = sub.add_parser("replay")
    rp.add_argument("fixture_json"); rp.add_argument("--out-dir", required=True)
    rp.add_argument("--route", choices=["direct", "df", "rijcosx", "auto"], required=True)
    rp.add_argument("--threads", type=int, required=True)
    rp.add_argument("--repeats", type=int, default=3)
    sm = sub.add_parser("summarize")
    sm.add_argument("raw_dir"); sm.add_argument("--json-out", required=True)
    sm.add_argument("--threads", default="1,2")
    sm.add_argument("--expected-source-sha", required=True)
    sm.add_argument("--expected-native-sha256", required=True)
    args = ap.parse_args()
    if args.cmd == "export":
        export_fixture(json.loads(Path(args.case_json).read_text()), args.out_dir)
    elif args.cmd == "replay":
        replay(args.fixture_json, args.route, args.threads, args.repeats, args.out_dir)
    else:
        result = summarize(args.raw_dir, [int(t) for t in args.threads.split(",")],
                           args.expected_source_sha, args.expected_native_sha256)
        write_json(Path(args.json_out), result)


if __name__ == "__main__":
    main()