"""New-source four-case repeat of 11901a7; frozen oracle and gates are unchanged."""

import argparse
import json
import math
from pathlib import Path

from assess import BASELINE_FACTOR, FLOORS
from oracle import LIMITS, TIMES, U, rectangles
from run_case import (INITIALIZATION, sample_steps, NUMERICAL_REVISION, SOURCE_SHA256,
                      BUILD_CAPTURE_SHA256, REFERENCE_HARNESS)

# Host-generated snapshots of a135252 at the frozen two local configurations.
# These describe expected inputs; actual sidecars are captured before cleanup.
BUILD_SOURCES = {
    128: "803177416b7599ab89cc5afd35cf4a28a673b2f4bb8e8a6a1abc0b6d77610873",
    256: "17ec5e60f87b840d3d13b8778bc8a3aecf485469c5a061af46c5cc19b73396cf",
}
CASES = ((100, 128, 1, "constant"), (100, 256, 1, "constant"),
         (100, 128, 2, "constant"), (100, 128, 2, "projected_linear"))
NUMERICS = dict(lattice_model="D2Q9", collision="SRT", store_precision="FP32",
                ddf_shifting=True, streaming="double_buffer", les_enabled=False,
                boundary="periodic", block_shape=[8, 8],
                sample_phase="postcollision G at root synchronization")


def finite(value):
    return isinstance(value, (int, float)) and not isinstance(value, bool) and math.isfinite(value)


def finite_tree(value):
    if isinstance(value, dict):
        return all(finite_tree(v) for v in value.values())
    if isinstance(value, list):
        return all(finite_tree(v) for v in value)
    return not isinstance(value, float) or math.isfinite(value)


def at_steps(rows, steps, n):
    return len(rows) == len(steps) and all(
        type(r.get("step")) is int and r["step"] == step and finite(r.get("tstar"))
        and abs(r["tstar"] - step * U / n) <= 3e-17
        for r, step in zip(rows, steps)
    )


def maximum(rows, metric):
    values = [r.get(metric) for r in rows]
    return max(values) if values and all(finite(v) and v >= 0 for v in values) else None


def valid_builds(record, n, depth):
    builds = record.get("actual_build_sidecars", [])
    if not isinstance(builds, list) or len(builds) != depth:
        return False
    for level, build in enumerate(builds):
        if not isinstance(build, dict):
            return False
        size = n * 2**level
        cfg = build.get("config", {})
        expected = dict(nx=10, ny=10, nz=1, lattice_model="D2Q9", collision="SRT",
                        store_precision="FP32", streaming="double_buffer", ddf_shifting=True,
                        les_enabled=False, velocity=0.0, viscosity=U*size/100, rho=1.0,
                        omega=1/(.5+3*U*size/100), threads_per_block=256)
        ptx = build.get("ptx_sha256", "")
        if (build.get("source_sha256") != BUILD_SOURCES[size]
                or build.get("arch") != "sm_70" or build.get("flags") != ["--use_fast_math"]
                or build.get("n_objects") != 0 or not isinstance(cfg, dict)
                or any(cfg.get(k) != value for k, value in expected.items())
                or not isinstance(ptx, str) or len(ptx) != 64
                or any(c not in "0123456789abcdef" for c in ptx)):
            return False
    return True


def assess(records, expected_revision=None, current=False):
    """current=True assesses current-checkout receipts (run_case --current-checkout) with
    the same gates; it never accepts them as the frozen a135252 qualification, and the
    frozen mode never accepts current-checkout receipts."""
    indexed, duplicates = {}, []
    for record in records:
        key = tuple(record.get(k) for k in ("re_L", "n", "depth", "explosion"))
        if key in indexed:
            duplicates.append(str(key))
        else:
            indexed[key] = record
    missing = [str(k) for k in CASES if k not in indexed]
    checks = dict(complete_matrix=not missing, no_extra_cases=set(indexed) <= set(CASES),
                  no_duplicate_cases=not duplicates, finite_receipts=finite_tree(records))
    sources = [r.get("source", {}) for r in records]
    checks["same_source_and_harness"] = bool(sources) and len({
        json.dumps(s, sort_keys=True) for s in sources}) == 1
    if current:
        checks["clean_current_checkout_source"] = bool(sources) and expected_revision is not None and all(
            s.get("dirty") is False and s.get("source_mode") == "current_checkout"
            and bool(s.get("source_tree_sha256")) and bool(s.get("harness_sha256"))
            and s.get("build_capture_sha256") == BUILD_CAPTURE_SHA256
            and s.get("revision") == expected_revision for s in sources)
    else:
        checks["clean_a135252_numerical_source"] = bool(sources) and all(
            s.get("dirty") is False and s.get("source_tree_sha256") == SOURCE_SHA256
            and "source_mode" not in s
            and bool(s.get("revision")) and bool(s.get("harness_sha256"))
            and s.get("build_capture_sha256") == BUILD_CAPTURE_SHA256
            and (expected_revision is None or s.get("revision") == expected_revision) for s in sources)
    runtimes = [r.get("runtime_identity") for r in records]
    checks["same_recorded_runtime"] = bool(runtimes) and all(
        isinstance(v, dict) and all(v.get(k) for k in
            ("device", "nvcc_path", "nvcc_version", "python", "numpy", "pycuda"))
        and isinstance(v["device"], dict) and v["device"].get("capability") == [7, 0]
        and all(v["device"].get(k) for k in ("name", "pci_bus_id", "driver", "runtime"))
        for v in runtimes
    ) and len({json.dumps(v, sort_keys=True) for v in runtimes}) == 1
    devices = [r.get("cuda_visible_devices") for r in records]
    checks["same_explicit_device"] = bool(devices) and all(
        isinstance(d, str) and d and "," not in d and d != "-1" for d in devices
    ) and len(set(devices)) == 1
    enabled = [r.get("early_pressure", {}).get("enabled") for r in records]
    checks["same_early_sampling"] = bool(enabled) and all(type(v) is bool for v in enabled) and len(set(enabled)) == 1
    shared_valid = all(checks.values())
    complete, case_valid, early = {}, {}, []
    for key in CASES:
        if key not in indexed:
            continue
        r = indexed[key]
        _, n, depth, scheme = key
        label = f"{n}/{depth}/{scheme}"
        rows = r.get("rows", [])
        main_steps, early_steps = sample_steps(n, True)
        complete[key] = r.get("completed") is True and at_steps(rows, main_steps, n)
        checks[label + " complete_times"] = complete[key]
        checks[label + " scales"] = (
            r.get("U") == U and finite(r.get("viscosity"))
            and abs(r["viscosity"] - U * n / 100) < 1e-15
            and r.get("tau_by_level") == [0.5 + 3 * U * n / 100 * 2**l for l in range(depth)]
            and r.get("mach") == U * math.sqrt(3)
        )
        checks[label + " actual_builds"] = valid_builds(r, n, depth)
        if current:
            checks[label + " reference"] = (r.get("numerical_revision") is None
                                             and r.get("source_mode") == "current_checkout"
                                             and r.get("reference_harness") == REFERENCE_HARNESS)
        else:
            checks[label + " reference"] = (r.get("numerical_revision") == NUMERICAL_REVISION
                                             and "source_mode" not in r
                                             and r.get("reference_harness") == REFERENCE_HARNESS)
        checks[label + " numerics"] = r.get("numerics") == NUMERICS and r.get("initialization") == INITIALIZATION
        runtime = r.get("runtime_source", {})
        package = Path(runtime.get("checkout", "")) / "src/CelerisLab"
        checks[label + " runtime_source"] = bool(runtime.get("checkout")) and all(
            runtime.get(k) == str(package / tail) for k, tail in
            (("module", "__init__.py"), ("compiler", "cuda/compiler_v2.py"), ("kernel_root", "lbm/kernels")))
        expected_patch = [] if depth == 1 else rectangles(n, depth)
        checks[label + " patch"] = json.dumps(r.get("refinements")) == json.dumps(expected_patch)
        cells = [n*n] if depth == 1 else [8704, 30720]
        hops = [16*n * 2**l for l in range(depth)]
        work = r.get("workload", {})
        checks[label + " workload"] = (
            work.get("active_cells_by_level") == cells and work.get("level_hops") == hops
            and work.get("root_intervals") == 16*n
            and work.get("active_collision_updates_by_level") == [c*h for c, h in zip(cells, hops)]
            and work.get("active_collision_updates") == sum(c*h for c, h in zip(cells, hops))
        )
        checks[label + " absolute"] = bool(rows) and all(
            finite(row.get(metric)) and 0 <= row[metric] <= limit
            for row in rows for metric, limit in LIMITS.items())
        trace = r.get("early_pressure", {})
        trace_rows = trace.get("rows", [])
        peak = maximum(trace_rows, "pressure_l2")
        trace_complete = at_steps(trace_rows, early_steps, n) and peak is not None
        checks[label + " early_trace"] = (
            trace_complete and trace.get("sampled_max_pressure_l2") == peak
            and len(rows) >= 2 and all(trace_rows[i]["pressure_l2"] == rows[j].get("pressure_l2")
                                      for i, j in ((0, 0), (-1, 1)))
            if trace.get("enabled") is True else
            trace.get("enabled") is False and trace_rows == [] and trace.get("sampled_max_pressure_l2") is None
        )
        early.append(dict(n=n, depth=depth, explosion=scheme, enabled=trace.get("enabled"),
                          samples=len(trace_rows), complete=trace_complete,
                          sampled_max_pressure_l2=peak))
        case_valid[key] = all(value for name, value in checks.items() if name.startswith(label + " "))

    def worst(key, metric):
        return maximum(indexed[key]["rows"][1:], metric) if complete.get(key) else None

    comparisons, scheme_results = [], {}
    for scheme in ("constant", "projected_linear"):
        key = (100, 128, 2, scheme)
        for metric, floor in FLOORS.items():
            amr, coarse, fine = (worst(k, metric) for k in (key, CASES[0], CASES[1]))
            budget = None if coarse is None else max(BASELINE_FACTOR * coarse, floor)
            budget = budget if finite(budget) else None
            ratio = None if amr is None or fine is None else amr / max(fine, 1e-30)
            passed = amr is not None and budget is not None and amr <= budget
            checks[f"baseline {scheme}/{metric}"] = passed
            comparisons.append(dict(explosion=scheme, metric=metric, amr=amr,
                                    uniform128=coarse, baseline_budget=budget, passed=passed,
                                    uniform256=fine, amr_to_uniform256=ratio if finite(ratio) else None))
        baseline_passed = all(checks[f"baseline {scheme}/{metric}"] for metric in FLOORS)
        valid = shared_valid and all(case_valid.get(k, False) for k in (CASES[0], CASES[1], key))
        scheme_results[scheme] = dict(passed=valid and baseline_passed,
                                     valid_case_and_controls=valid, baseline_passed=baseline_passed)
    return dict(passed=all(checks.values()), missing=missing, duplicates=duplicates, checks=checks,
                comparisons=comparisons, scheme_results=scheme_results, early_pressure=early,
                source_mode="current_checkout" if current else "frozen",
                numerical_revision=None if current else NUMERICAL_REVISION,
                source_tree_sha256=(sources[0].get("source_tree_sha256") if current and sources
                                    else None if current else SOURCE_SHA256),
                reference_harness=REFERENCE_HARNESS, expected_revision=expected_revision,
                floors=FLOORS, baseline_factor=BASELINE_FACTOR,
                main_statistic="maximum of four post-initial samples at t*=.12,.24,.36,.48",
                early_statistic="sampled pressure maximum only, t*=k*.001875, k=0..64; not an all-step or continuous maximum",
                finest_control="uniform256 and AMR/ uniform256 ratios are descriptive, not acceptance gates",
                scope="2D periodic Re100 smooth-decay diagnostic; no new algorithm, convergence-order or general AMR qualification")


if __name__ == "__main__":
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("files", type=Path, nargs="+")
    p.add_argument("--output", type=Path, required=True)
    p.add_argument("--expected-revision", required=True)
    p.add_argument("--current-checkout", action="store_true",
                   help="assess current-checkout receipts; never the frozen qualification")
    args = p.parse_args()
    if args.output.exists():
        p.error("Output already exists; choose a new path to retain prior evidence")
    result = assess([json.loads(path.read_text()) for path in args.files], args.expected_revision,
                    current=args.current_checkout)
    args.output.write_text(json.dumps(result, indent=2, allow_nan=False) + "\n")
    print(json.dumps(result, indent=2, allow_nan=False))
    raise SystemExit(0 if result["passed"] else 1)