"""Test-only straight-face centered witness; private CUDA snapshot, frozen source intact."""
from contextlib import contextmanager
import hashlib
import json
from pathlib import Path
import numpy as np

import shear_axis_square as axis
from oracle import CX, CY, W

DIAGNOSTIC = "centered_face_square_n128"
VARIANT = "centered_face_v1"
BUILD_SHA256 = "81e880017495190209bd867d4b523a2c243be496ed4681892e4cd789b7ba1f21"
CASES = ((128, 1, "current"), (256, 1, "current"),
         (128, 2, "current"), (128, 2, VARIANT))
# Local arithmetic-derived diagnostic deadbands, NOT accumulated trajectory bounds.
POP_ENVELOPE, OP_ATOL, OP_RTOL = 1/32, 6e-8, 2e-7
MARGIN = 2 * np.array([9, 6, 6]) * (OP_ATOL + OP_RTOL * POP_ENVELOPE)


def routing(hierarchy):
    """Square-only mask: full donor quartets, one missing normal axis, legal tangent."""
    from CelerisLab.lbm.blocks import BlockLayout2D
    from CelerisLab.lbm.refinement_tables import projected_explosion_tables
    if hierarchy.depth != 2 or hierarchy.kinds[0].shape != (128, 128):
        raise ValueError("centered witness requires the fixed two-level square")
    pair = hierarchy.pairs[0]
    covered = np.zeros((128, 128), bool); covered[32:96, 32:96] = True
    if pair.fine_origin != (0., 0.) or not np.array_equal(pair.covered, covered):
        raise ValueError("wrong centered-witness geometry")
    layout = BlockLayout2D.build(128, 128, 8, 8, periodic=(True, True))
    neighbors, offsets = projected_explosion_tables(pair, layout, hierarchy.kinds[0])
    donors = layout.owner.ravel()[pair.ghost_coarse]
    off = neighbors.reshape(-1, 2, 2) == donors[:, None, None]
    disabled = off.all(axis=2)
    y, x = np.divmod(pair.ghost_coarse, 128)
    vertical = np.isin(x, [30, 31, 96, 97]) & (y >= 36) & (y < 92)
    horizontal = np.isin(y, [30, 31, 96, 97]) & (x >= 36) & (x < 92)
    mask = ((vertical & disabled[:, 0] & ~off[:, 1].any(axis=1)) |
            (horizontal & disabled[:, 1] & ~off[:, 0].any(axis=1))).astype(np.uint8)
    for donor in np.unique(pair.ghost_coarse):
        rows = pair.ghost_coarse == donor
        if rows.sum() != 4 or len(set(mask[rows])) != 1 or sorted(map(tuple, offsets[rows])) != [(-.25,-.25), (-.25,.25), (.25,-.25), (.25,.25)]:
            raise ValueError("routing must preserve complete +/- quarter donor quartets")
    return mask


def metadata(mask):
    return dict(variant=VARIANT, mask_sha256=hashlib.sha256(mask.tobytes()).hexdigest(),
                ghost_rows=len(mask), selected_rows=int(mask.sum()),
                rule="diagonal projected correction x2; audited straight faces only")


def patch_source(source, mask):
    """Change only the projected function in the private generated build directory."""
    anchor = 'extern "C" __global__ void BlockExplodeProjected('
    position = "        to[index_f(t % TOTAL_CELLS, q)] ="
    before, kernel = source.split(anchor)
    body, after = kernel.split("// Targets are receiver-local SOURCE slots", 1)
    if body.count(position) != 1:
        raise ValueError("unexpected projected kernel source")
    flag = "__device__ __constant__ unsigned char centered_face_mask[] = {" + ",".join(map(str, mask)) + "};\n"
    edit = "        if (centered_face_mask[row] && ex != 0.f && ey != 0.f) { gx *= 2.f; gy *= 2.f; }\n"
    return before + flag + anchor + body.replace(position, edit + position) + "// Targets are receiver-local SOURCE slots" + after


def source_digest(build):
    digest = hashlib.sha256()
    for path in sorted(Path(build).rglob("*")):
        if path.is_file() and path.name != ".celeris-build":
            digest.update(str(path.relative_to(build)).encode()); digest.update(path.read_bytes())
    return digest.hexdigest()


@contextmanager
def private_build(hierarchy, report, variant):
    """No kernel upload/launch monkey-patch; same production function and arguments."""
    if variant == "current":
        yield
        return
    from CelerisLab.cuda import compiler_v2 as compiler
    mask = routing(hierarchy)
    report["reconstruction"] = metadata(mask)
    original, calls = compiler.generate_config, []
    def generate(cfg, *args, **kwargs):
        build = Path(original(cfg, *args, **kwargs))
        calls.append(cfg.viscosity)
        if len(calls) == 2:
            if calls != [.0384, .0768]:
                raise ValueError("unexpected candidate level configuration")
            path = build / "blocks.cu"
            path.write_text(patch_source(path.read_text(), mask))
            sidecar = build / ".celeris-build"
            data = json.loads(sidecar.read_text())
            data.update(source_sha256=source_digest(build), experimental_reconstruction=metadata(mask))
            sidecar.write_text(json.dumps(data, sort_keys=True))
        elif len(calls) > 2:
            raise ValueError("candidate is one adjacent pair only")
        return str(build)
    compiler.generate_config = generate
    try:
        yield
        if calls != [.0384, .0768]:
            raise ValueError("candidate did not build its exact two levels")
    finally:
        compiler.generate_config = original


def observe(arrays):
    """Bound observed synchronized ACTIVE values only, not intermediate GPU states."""
    maximum = max(float(np.nanmax(abs(a - W[:, None, None]))) for a in arrays)
    if not np.isfinite(maximum) or maximum > POP_ENVELOPE:
        raise ValueError("observed shifted-population envelope exceeded")
    return maximum


def compare(records, expected_source):
    """Four fresh receipts; diagnostic screens only, never a physical qualification."""
    import run_case as rc
    from assess_projected_pairs import BUILD_SOURCES, NUMERICS, valid_builds, finite_tree
    from CelerisLab.lbm.hierarchy import build_hierarchy
    from CelerisLab.lbm.refinement import ACTIVE
    if not finite_tree(records) or not expected_source.get("revision"):
        raise ValueError("nonfinite receipt or missing exact source identity")
    keys = [(r['n'], r['depth'], r.get('test_reconstruction')) for r in records]
    if keys != list(CASES):
        raise ValueError("expected ordered fresh U128/U256/current/candidate cases")
    for r in records:
        for key in ('source', 'runtime_identity', 'runtime_source', 'cuda_visible_devices'):
            if not r.get(key) or r[key] != records[0][key]:
                raise ValueError(f"mixed or missing {key}")
        if (r['source'] != expected_source or r.get('numerical_revision') != rc.NUMERICAL_REVISION
                or r.get('numerics') != NUMERICS or r.get('re_L') != 100 or r.get('U') != .03
                or r.get('polarization') != [0., 1.] or r.get('initialization_id') != axis.INITIALIZATION_ID
                or r['source'].get('dirty') is not False or r['source'].get('source_tree_sha256') != rc.SOURCE_SHA256
                or r.get('diagnostic') != DIAGNOSTIC or not r.get('completed')
                or r.get('qualification_eligible') is not False or r.get('initialization') != axis.INITIALIZATION
                or r.get('mode') != [2, 0] or r.get('phase') != axis.PHI
                or r.get('explosion') != ('constant' if r['depth'] == 1 else 'projected_linear')
                or not 0 <= r.get('observed_shifted_max', np.inf) <= POP_ENVELOPE):
            raise ValueError("wrong candidate receipt identity or observed envelope")
        h = build_hierarchy(rc.config(r['n'], 100), [] if r['depth'] == 1 else axis.REFINEMENTS)
        if r.get('refinements') != ([] if r['depth'] == 1 else [[list(axis.REFINEMENTS[0][0])]]):
            raise ValueError("wrong geometry")
        if (r.get('viscosity') != .03*r['n']/100
                or r.get('tau_by_level') != [.5+3*c.viscosity for c in h.configs]
                or r.get('workload') != rc.workload([k == ACTIVE for k in h.kinds], 8*(r['n']//128))):
            raise ValueError("wrong physical scales or workload")
        builds = r.get('actual_build_sidecars', [])
        canonical = []
        for level, build in enumerate(builds):
            modified = r['test_reconstruction'] == VARIANT and level == 1
            expected = BUILD_SHA256 if modified else BUILD_SOURCES[r['n']*2**level]
            if (build.get('source_sha256') != expected
                    or build.get('experimental_reconstruction') != (metadata(routing(h)) if modified else None)):
                raise ValueError("wrong actual build or reconstruction identity")
            # Reuse the frozen build-metadata checker after checking the private-source hash above.
            canonical.append(dict(build, source_sha256=BUILD_SOURCES[r['n']*2**level]))
        if not valid_builds(dict(actual_build_sidecars=canonical), r['n'], r['depth']):
            raise ValueError("invalid architecture, flags, config or PTX digest")
        if r.get('reconstruction') != (metadata(routing(h)) if r['test_reconstruction'] == VARIANT else dict(variant='current')):
            raise ValueError("wrong receipt mask identity")
        trace = r['probe_trace']
        declared = axis.new_trace(trace['probes'], r['viscosity'], [DIAGNOSTIC, r['n'], r['depth'], r['explosion']])
        if any(trace.get(k) != v for k,v in declared.items() if k not in ('steps','values')):
            raise ValueError("wrong trace phase, clock or physical identity")
        if trace['probes'] != axis.select([k == ACTIVE for k in h.kinds], r['n']):
            raise ValueError("wrong actual probe cells")
        if trace.get("bytes") != axis.validate(trace):
            raise ValueError("trace serialized byte count mismatch")
        if np.min(np.asarray(trace["values"])[:,0::3]) <= 0:
            raise ValueError("nonpositive sampled density")
        if [row['step'] for row in r['rows']] != [0, 8*(r['n']//128)] or not all(rc.finite_row(v) for v in r['rows']):
            raise ValueError("invalid scalar rows")
        if min(row['rho_min'] for row in r['rows']) <= 0:
            raise ValueError("nonpositive density")
    current, candidate = records[2:]
    if current['probe_trace']['values'][0] != candidate['probe_trace']['values'][0]:
        raise ValueError("candidate/current initial samples differ; cancellation invalid")
    vectors = []
    for r in records:
        groups = axis.group_moments(r['probe_trace'])
        vectors.append(np.array([[[row['groups'][name][key] for key in ('drho','djx','djy')]
                                  for name in axis.GROUPS] for row in groups]))
    # Inner quadrature matches U256; outer quadrature matches U128.
    uniform = vectors[1].copy(); uniform[:, 4:] = vectors[0][:, 4:]
    residual = [v-v[0]-uniform+uniform[0] for v in vectors[2:]]
    old, new = residual
    checks = dict(bottom_step1_clear=abs(old[1,2,1])-abs(new[1,2,1]) > MARGIN[1],
                  top_step1_nonworse=abs(new[1,3,1]) <= abs(old[1,3,1])+MARGIN[1])
    for face, indices in [('bottom',[2,6]), ('top',[3,7])]:
        rms = lambda a: float(np.sqrt(np.mean(a[8,indices,1]**2)))
        checks[face+'_step8_nonworse'] = rms(new) <= rms(old)+MARGIN[1]
    checks['vertical_null'] = bool(np.all(abs(new[:,[0,1,4,5]]-old[:,[0,1,4,5]]) <= MARGIN))
    acoustic = lambda a: np.sqrt(a[...,0]**2/3 + a[...,1]**2 + a[...,2]**2)
    acoustic_margin = float(np.sqrt(MARGIN[0]**2/3 + MARGIN[1]**2 + MARGIN[2]**2))
    checks['acoustic_nonworse'] = bool(np.all(acoustic(new) <= acoustic(old)+acoustic_margin))
    raw_pressure = lambda r: float(np.max(abs(np.asarray(r['probe_trace']['values'])[:,0::3]-1))/3)
    checks['sampled_pressure_nonworse'] = raw_pressure(candidate) <= raw_pressure(current)+MARGIN[0]/3
    checks['step8_pressure_nonworse'] = candidate['rows'][-1]['spurious_pressure_rms'] <= current['rows'][-1]['spurious_pressure_rms']+MARGIN[0]/3
    return dict(qualified=False, advance_to_review=all(checks.values()), checks=checks,
                moment_deadband=MARGIN.tolist(), deadband_scope="local arithmetic-derived diagnostic; no accumulated error guarantee",
                groups=list(axis.GROUPS), components=['drho','djx','djy'],
                current_signed_residual=old.tolist(), candidate_signed_residual=new.tolist(),
                raw_current=vectors[2].tolist(), raw_candidate=vectors[3].tolist(),
                group_moments=[axis.group_moments(r['probe_trace']) for r in records],
                significant_sign_flips=((old*new < 0) & (abs(old)>MARGIN) & (abs(new)>MARGIN)).tolist(),
                top_step1_below_deadband=bool(abs(old[1,3,1]) <= MARGIN[1]),
                note="Absolute reduction with sign retained; top below deadband may be inconclusive.")


def main():
    import argparse
    p = argparse.ArgumentParser(description=__doc__)
    p.add_argument("--repo", required=True, type=Path)
    p.add_argument("--expected-revision", required=True)
    p.add_argument("receipts", nargs=4, type=Path, help="U128 U256 current candidate, in order")
    p.add_argument("--output", required=True, type=Path)
    args = p.parse_args()
    import run_case as rc
    source = rc.verify_source(args.repo.resolve(), args.expected_revision, "shear_centered", probes=True)
    result = compare([json.loads(path.read_text()) for path in args.receipts], source)
    with args.output.open("x") as stream:
        json.dump(result, stream, indent=2, allow_nan=False)
    return 0 if result["advance_to_review"] else 1


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