"""Signed interface-probe traces for the nested/twin shear presets (mechanism diagnostic).

No pass gates. Native-validated 24-target layout through box 1, in root-128
physical coordinates, crossing L0->L1 and (nested only) L1->L2 away from corners:
  horizontal line y = 56 across the vertical faces x = 16 and x = 22;
  vertical line x = 38 across the horizontal faces y = 16 and y = 22;
  normal coordinates NORMAL_COORDS on both (nested: 3 L0 + 5 L1 + 4 L2 each).
Each target maps to the nearest ACTIVE leaf center (tie rule: distance, then
finer level, then smaller j, then smaller i). Grids differ between cases
(U128, U512, nested, twin), so actual centers, levels and displacements are
stored and analytic fields are evaluated at the actual centers.
Samples are POST-collision G at root synchronization only (root-step units);
no substep content is observed, and grid-scale alternation can alias.
"""

import math

import numpy as np

import shear_wave as sw
from oracle import CX, CY, U

ROOT_N = 128
NORMAL_COORDS = (12.5, 14.5, 15.5, 16.25, 17.25, 19.25, 20.75, 21.75,
                 22.125, 23.125, 24.125, 26.125)
TRANSECTS = {"horizontal_y56": dict(axis=0, other=56.0),   # x varies, y fixed
             "vertical_x38": dict(axis=1, other=38.0)}     # y varies, x fixed
STEPS = tuple(range(0, 129)) + tuple(range(1920, 2049))  # root steps, both inclusive
TRACE_MAX_BYTES = 600_000


def points():
    out = []
    for name, t in TRANSECTS.items():
        for along in NORMAL_COORDS:
            xy = (along, t["other"]) if t["axis"] == 0 else (t["other"], along)
            out.append((name, along, xy))
    return out


def case_steps(n):
    """Case lattice steps at the fixed root-physical sample times (U512: 4 per root step)."""
    return [s * (n // ROOT_N) for s in STEPS]


def select(masks, n=ROOT_N):
    """Nearest ACTIVE leaf center for each target (deterministic tie rule).

    Targets and stored centers are in root-physical units (root cell = 1); a case
    with root N = n has cell size h = ROOT_N / (n * 2**level) there (U512: 0.25).
    """
    if n % ROOT_N:
        raise ValueError("case root N must be a multiple of 128")
    probes = []
    for name, along, (x, y) in points():
        best = None
        for level, mask in enumerate(masks):
            h = ROOT_N / (n * 2 ** level)
            i0, j0 = int(math.floor(x / h)), int(math.floor(y / h))
            for j in range(j0 - 2, j0 + 3):
                for i in range(i0 - 2, i0 + 3):
                    if 0 <= j < mask.shape[0] and 0 <= i < mask.shape[1] and mask[j, i]:
                        cx, cy = h * (i + 0.5), h * (j + 0.5)
                        key = (math.hypot(cx - x, cy - y), -level, j, i)
                        if best is None or key < best[0]:
                            best = (key, level, i, j, cx, cy)
        if best is None:
            raise ValueError(f"no ACTIVE leaf near target {(x, y)}")
        _, level, i, j, cx, cy = best
        probes.append(dict(transect=name, normal=along, point=[x, y], level=level, i=i, j=j,
                           center=[cx, cy], displacement=[cx - x, cy - y]))
    if len({(p["level"], p["i"], p["j"]) for p in probes}) != len(probes):
        raise ValueError("probe targets map to duplicate cells")
    return probes


def sample(arrays, probes):
    """Flat [rho, ux, uy] per probe from FP64 physical populations."""
    row = []
    for p in probes:
        f = np.asarray(arrays[p["level"]][:, p["j"], p["i"]], dtype=float)
        rho = float(f.sum())
        row += [rho, float(CX @ f) / rho, float(CY @ f) / rho]
    return row


def new_trace(probes, case_viscosity, case):
    """Times/coordinates are root-physical (N=128); nu_root = U*128/Re for all cases."""
    return dict(schema=1, kind="mechanism diagnostic, no pass gates", case=case,
                time_unit="root step (N=128)", sample_phase="postcollision G at root synchronization",
                root_n=ROOT_N, viscosity_root=U * ROOT_N / sw.RE, case_viscosity=case_viscosity,
                mode=list(sw.MODE), phase=sw.PHI,
                amplitude=U, fields=["rho", "ux", "uy"], probes=probes, steps=[], values=[])


def residuals(trace):
    """Signed (drho, dux, duy) per step and probe, analytic field at each actual center."""
    out = []
    for step, row in zip(trace["steps"], trace["values"]):
        res = []
        for n, p in enumerate(trace["probes"]):
            x, y = p["center"]
            ux, uy, _ = sw.exact(x, y, step, trace["root_n"], trace["viscosity_root"])
            rho, vx, vy = row[3 * n: 3 * n + 3]
            res.append((rho - 1.0, vx - float(ux), vy - float(uy)))
        out.append(res)
    return out


def validate_row(row, probe_count):
    """Reject an invalid optional sample before it can enter a scalar receipt."""
    if len(row) != 3 * probe_count or not all(math.isfinite(v) for v in row):
        raise ValueError("malformed or nonfinite trace row")


def validate(trace):
    """Exact schedule, finite values, fixed probe count and byte cap."""
    import json
    if trace.get("steps") != list(STEPS) or len(trace["values"]) != len(STEPS):
        raise ValueError("trace steps differ from the fixed schedule")
    if len(trace["probes"]) != len(NORMAL_COORDS) * len(TRANSECTS):
        raise ValueError("unexpected probe count")
    for row in trace["values"]:
        validate_row(row, len(trace["probes"]))
    # Include the final byte-count field itself, including its number of digits.
    size = 0
    while True:
        final_size = len(json.dumps(dict(trace, bytes=size), allow_nan=False).encode())
        if final_size == size:
            break
        size = final_size
    if size > TRACE_MAX_BYTES:
        raise ValueError(f"trace exceeds {TRACE_MAX_BYTES} bytes ({size})")
    return size