"""Short axis-shear orientation diagnostic; no physical acceptance gates."""

import json
import numpy as np

import shear_wave as sw
from shear_probes import sample, validate_row, TRACE_MAX_BYTES
from oracle import CX, CY, W, U, centers

ROOT_N = 128
MODE, E, PHI, ANGLE_DEG = (2, 0), np.array([0., 1.]), np.pi / 4, 0.
REFINEMENTS = [[(32, 32, 96, 96)]]
CASES = ((100, 128, 1, "constant"), (100, 256, 1, "constant"),
         (100, 128, 2, "constant"), (100, 128, 2, "projected_linear"))
STEPS = tuple(range(9))
INITIALIZATION_ID = "postg_first_gradient_v1"
INITIALIZATION = "physical FP64 axis-shear POST G=feq-(tau-1)A; first-gradient low-Mach CE"
GROUPS = dict(inner_left=(32, 64), inner_right=(95, 63),
              inner_bottom=(64, 32), inner_top=(63, 95),
              outer_left=(31, 64), outer_right=(96, 63),
              outer_bottom=(64, 31), outer_top=(63, 96))


def eps(n):
    return (4 * np.pi / n)**2


def exact(x, t, length, viscosity):
    """Transverse velocity and physical x derivative in the supplied length units."""
    k = 4 * np.pi / length
    amplitude = U * np.exp(-viscosity * k * k * t)
    return amplitude * np.sin(k * x + PHI), amplitude * k * np.cos(k * x + PHI)


def initial(shape, level, length, viscosity):
    x, _ = centers(shape, level)
    uy, gradient = exact(x, 0, length, viscosity)
    tau = .5 + 3 * viscosity * 2**level
    a = (3 * W * CX * CY)[:, None, None] * gradient * 2.**-level
    return sw.equilibrium(np.ones_like(uy), np.zeros_like(uy), uy) - (tau - 1) * a


def metrics(arrays, masks, t, length, viscosity, initial_mass=None):
    rho, ux, uy, x, _, w, _ = sw._leaves(arrays, masks)
    if not np.isclose(w.sum(), length**2, rtol=0, atol=1e-10):
        raise ValueError("leaf partition does not cover domain once")
    ey, _ = exact(x, t, length, viscosity)
    p = (rho - 1) / 3
    p -= w @ p / w.sum()  # one domain gauge, never the probe mean
    mass = float(w @ rho)
    return dict(tstar=t * U / length, mass=mass,
                mass_relative=0. if initial_mass is None else abs(mass-initial_mass)/initial_mass,
                velocity_l2=float(np.sqrt(w @ (ux**2 + (uy-ey)**2) / (w @ ey**2))),
                spurious_pressure_rms=float(np.sqrt(w @ p**2 / w.sum())),
                rho_min=float(rho.min()), rho_max=float(rho.max()),
                max_speed=float(np.hypot(ux, uy).max()), leaf_cells=len(rho))


def case_steps(n):
    return [r * (n // ROOT_N) for r in STEPS]


def select(masks, n=ROOT_N):
    """Exact native cells filling eight unit-area groups; no interpolated samples."""
    if (n, len(masks)) not in ((128, 1), (256, 1), (128, 2)):
        raise ValueError("unsupported square probe layout")
    probes = []
    for name, (x, y) in GROUPS.items():
        level = int(len(masks) == 2 and name.startswith("inner"))
        factor = (n // ROOT_N) * 2**level
        for dy in range(factor):
            for dx in range(factor):
                i, j = factor*x + dx, factor*y + dy
                if not masks[level][j, i]:
                    raise ValueError("square probe is not ACTIVE")
                probes.append(dict(group=name, level=level, i=i, j=j,
                                   center=[(i+.5)/factor, (j+.5)/factor],
                                   width=1/factor, area=1/factor**2, parity=[i % 2, j % 2]))
    if len({(p["level"], p["i"], p["j"]) for p in probes}) != len(probes):
        raise ValueError("duplicate square probes")
    return probes


def new_trace(probes, case_viscosity, case):
    return dict(schema=1, kind="axis-square orientation 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/100, case_viscosity=case_viscosity,
                mode=list(MODE), phase=PHI, amplitude=U, fields=["rho", "ux", "uy"],
                probes=probes, steps=[], values=[])


def group_moments(trace):
    """Signed unit-area means; normalize by each group's mean actual-center derivative.

    Momentum is rho*u. Raw density is ungauged. Per-cell parity stays in the trace.
    U128 samples one center, whereas U256/refined-inner use four different centers.
    """
    out = []
    for step, row in zip(trace["steps"], trace["values"]):
        values, groups = np.asarray(row).reshape(-1, 3), {}
        for name in GROUPS:
            indices = [i for i, p in enumerate(trace["probes"]) if p["group"] == name]
            probes = [trace["probes"][i] for i in indices]
            weights = np.array([p["area"] for p in probes])
            if weights.sum() != 1:
                raise ValueError("probe group does not have unit area")
            rho, ux, uy = values[indices].T
            ey, gradient = exact(np.array([p["center"][0] for p in probes]), step,
                                 ROOT_N, trace["viscosity_root"])
            drho, djx, djy = (float(weights @ v) for v in (rho-1, rho*ux, rho*uy-ey))
            derivative = float(weights @ gradient)
            groups[name] = dict(drho=drho, djx=djx, djy=djy, mean_duy_dx=derivative,
                                djx_over_duy_dx=djx/derivative, djy_over_duy_dx=djy/derivative)
        out.append(dict(root_step=step, groups=groups))
    return out


def validate(trace):
    if trace["steps"] != list(STEPS) or len(trace["values"]) != len(STEPS):
        raise ValueError("square trace schedule mismatch")
    _, n, depth, scheme = trace["case"]
    if (100, n, depth, scheme) not in CASES or len(trace["probes"]) != (
            20 if depth == 2 else 8*(n//ROOT_N)**2):
        raise ValueError("square trace case/count mismatch")
    for row in trace["values"]:
        validate_row(row, len(trace["probes"]))
    trace["group_moments"] = group_moments(trace)
    size = 0
    while True:
        final = len(json.dumps(dict(trace, bytes=size), allow_nan=False).encode())
        if final == size:
            break
        size = final
    if size > TRACE_MAX_BYTES:
        raise ValueError("square trace exceeds byte cap")
    return size