"""CPU-only 2D Taylor–Green reference and composite-leaf diagnostics."""

import numpy as np

CX = np.array([0, 1, -1, 0, 0, 1, -1, 1, -1])
CY = np.array([0, 0, 0, 1, -1, 1, -1, -1, 1])
W = np.array([4 / 9] + [1 / 9] * 4 + [1 / 36] * 4)
U = 0.03
HORIZON = 0.48  # t* = t U/L; exactly 16*N base steps
TIMES = [0.0, 0.12, 0.24, 0.36, 0.48]
# Frozen screening ceilings; baseline-relative qualification is in assess.py.
LIMITS = dict(
    velocity_l2=0.05, pressure_l2=0.15, energy_relative=0.05, mass_relative=1e-5
)


def exact(x, y, t, length, viscosity, amplitude=U):
    k = 2 * np.pi / length
    a = amplitude * np.exp(-2 * viscosity * k * k * t)
    # TensorMesh phase: negative cos(x)sin(y), positive sin(x)cos(y).
    ux = -a * np.cos(k * x) * np.sin(k * y)
    uy = a * np.sin(k * x) * np.cos(k * y)
    p = -0.25 * a * a * (np.cos(2 * k * x) + np.cos(2 * k * y))
    return ux, uy, p  # kinematic pressure, mean zero


def centers(shape, level=0):
    # Campaign deliberately uses level-wide union layouts, all origins zero.
    y, x = np.indices(shape, dtype=float)
    h = 2.0 ** (-level)
    return h * (x + 0.5), h * (y + 0.5)


def initial(shape, level, length, viscosity):
    x, y = centers(shape, level)
    ux, uy, p = exact(x, y, 0.0, length, viscosity)
    rho = 1 + 3 * p  # cs^2=1/3, rho0=1
    cu = CX[:, None, None] * ux + CY[:, None, None] * uy
    # Exact FP64 weights; public shifted initializer encodes before FP32 rounding.
    return (
        W[:, None, None]
        * rho
        * (1 + 3 * cu + 4.5 * cu * cu - 1.5 * (ux * ux + uy * uy))
    )


def rectangles(n, depth):
    if n < 32 or n % 32 or depth not in (2, 3):
        raise ValueError("AMR requires N >= 32 divisible by 32 and depth 2 or 3")
    # Broad off-axis rectangle cuts nonzero velocity and both gradients.
    out = [[(n // 8, n // 8, 7 * n // 8, 3 * n // 4)]]
    if depth == 3:
        # L1 index coordinates; global physical bounds (3N/8,N/4,3N/4,5N/8).
        out.append([(3 * n // 4, n // 2, 3 * n // 2, 5 * n // 4)])
    return out


def metrics(arrays, masks, t, length, viscosity, initial_mass=None):
    accum = []
    for level, (g, mask) in enumerate(zip(arrays, masks)):
        if g.shape != (9, *mask.shape):
            raise ValueError("population/mask shape mismatch")
        values = g[:, mask].astype(float)
        if not np.isfinite(values).all():
            raise FloatingPointError("nonfinite leaf populations")
        rho = values.sum(axis=0)
        if (rho <= 0).any():
            raise FloatingPointError("nonpositive leaf density")
        ux, uy = CX @ values / rho, CY @ values / rho
        x, y = centers(mask.shape, level)
        ex, ey, ep = exact(x[mask], y[mask], t, length, viscosity)
        area = np.full(rho.size, 0.25**level)
        accum.append((rho, ux, uy, (rho - 1) / 3, ex, ey, ep, area))
    rho, ux, uy, p, ex, ey, ep, area = [np.concatenate(v) for v in zip(*accum)]
    volume = area.sum()
    if not np.isclose(volume, length**2, rtol=0, atol=1e-10):
        raise ValueError("leaf partition does not cover domain once")
    mass = float(area @ rho)
    # Remove one global pressure gauge, never one gauge per level.
    p = p - area @ p / volume
    ep = ep - area @ ep / volume
    energy = float(area @ (ux * ux + uy * uy) / 2 / volume)
    exact_energy = U * U / 4 * np.exp(-4 * viscosity * (2 * np.pi / length) ** 2 * t)
    return dict(
        tstar=t * U / length,
        mass=mass,
        mass_relative=(
            0.0 if initial_mass is None else abs(mass - initial_mass) / initial_mass
        ),
        velocity_l2=float(
            np.sqrt(
                area @ ((ux - ex) ** 2 + (uy - ey) ** 2) / (area @ (ex * ex + ey * ey))
            )
        ),
        pressure_l2=float(np.sqrt(area @ ((p - ep) ** 2) / (area @ (ep * ep)))),
        energy=energy,
        exact_energy=exact_energy,
        energy_relative=abs(energy - exact_energy) / exact_energy,
        rho_min=float(rho.min()),
        rho_max=float(rho.max()),
        max_speed=float(np.max(np.hypot(ux, uy))),
        leaf_cells=int(len(rho)),
    )