"""CPU-only oblique shear-wave reference and composite-leaf diagnostics.

Convention: integer mode MODE=(2,1) on the L x L periodic root domain, so the
wavevector k = (2*pi/L)*(2,1) points 26.565 deg from +x (mode (1,2) would be
63.435 deg). Velocity is polarized along E = (-1,2)/sqrt(5), perpendicular to k.
The frozen Taylor-Green oracle is reused for constants and geometry only.
"""

import math

import numpy as np
from oracle import CX, CY, W, U, TIMES, centers

MODE = (2, 1)
PHI = 0.0
E = np.array([-MODE[1], MODE[0]], dtype=float) / math.hypot(*MODE)
ANGLE_DEG = math.degrees(math.atan2(MODE[1], MODE[0]))
RE = 100  # Re = U L / nu, the same nu as Taylor-Green Re100 at equal N
FIT_FROM_TSTAR = 0.12  # late-time samples only; reduces startup sensitivity
INITIALIZATION = (
    "physical FP64 equilibrium G at actual cell centers; rho=1 (exact p const)"
)
# Per-row screens. phase_drift_abs, residual_l2 and the run-level rate screen
# use eps(n) and are preregistered diagnostic choices, not error bounds.
LIMITS = dict(velocity_l2=0.05, spurious_pressure_scaled=0.02, mass_relative=1e-5)
# The comparative pressure budget max(2*control, 0.02) is never below the 0.02
# absolute ceiling, so it is always redundant; raw control ratios are kept.
FLOORS = dict(rate_error=0.002, residual_l2=0.002, spurious_pressure_scaled=0.02)
BASELINE_FACTOR = 2.0


def wavevector(length):
    return 2 * np.pi / length * np.array(MODE, dtype=float)


def decay_rate(length, viscosity):
    """Exact amplitude decay rate per root step: nu |k|^2."""
    k = wavevector(length)
    return viscosity * float(k @ k)


def eps(n):
    """(|k| h_root)^2 with h_root = 1: dimensionless, coefficient one."""
    k = wavevector(n)
    return float(k @ k)


def exact(x, y, t, length, viscosity, amplitude=U):
    k = wavevector(length)
    a = amplitude * np.exp(-decay_rate(length, viscosity) * t)
    s = np.sin(k[0] * x + k[1] * y + PHI)
    p = np.zeros_like(s)
    return a * E[0] * s, a * E[1] * s, p  # kinematic pressure, constant


def equilibrium(rho, ux, uy):
    # Same second-order D2Q9 equilibrium as oracle.initial.
    cu = CX[:, None, None] * ux + CY[:, None, None] * uy
    return (
        W[:, None, None]
        * rho
        * (1 + 3 * cu + 4.5 * cu * cu - 1.5 * (ux * ux + uy * uy))
    )


def initial(shape, level, length, viscosity):
    x, y = centers(shape, level)
    ux, uy, p = exact(x, y, 0.0, length, viscosity)
    return equilibrium(1 + 3 * p, ux, uy)


def _leaves(arrays, masks):
    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")
        x, y = centers(mask.shape, level)
        area = np.full(rho.size, 0.25**level)
        lev = np.full(rho.size, level)
        accum.append((rho, CX @ values / rho, CY @ values / rho, x[mask], y[mask], area, lev))
    return [np.concatenate(v) for v in zip(*accum)]


def metrics(arrays, masks, t, length, viscosity, initial_mass=None):
    rho, ux, uy, x, y, w, lev = _leaves(arrays, masks)
    volume = w.sum()
    if not np.isclose(volume, length**2, rtol=0, atol=1e-10):
        raise ValueError("leaf partition does not cover domain once")
    k = wavevector(length)
    theta = k[0] * x + k[1] * y + PHI
    s, c = np.sin(theta), np.cos(theta)
    v = ux * E[0] + uy * E[1]
    # Weighted least squares on the actual composite quadrature (no
    # continuous-orthogonality assumption).
    gram = np.array([[w @ (s * s), w @ (s * c)], [w @ (s * c), w @ (c * c)]])
    a_hat, b_hat = np.linalg.solve(gram, [w @ (v * s), w @ (v * c)])
    ex, ey, _ = exact(x, y, t, length, viscosity)
    exact_norm = w @ (ex * ex + ey * ey)
    fit = a_hat * s + b_hat * c
    longitudinal = (ux * k[0] + uy * k[1]) / math.sqrt(k @ k)
    amplitude = U * math.exp(-decay_rate(length, viscosity) * t)
    scale = amplitude * amplitude / 4  # continuum reference scale, not the TG norm
    modal = math.hypot(a_hat, b_hat)
    p = (rho - 1) / 3
    p = p - w @ p / volume  # one global gauge
    mass = float(w @ rho)
    return dict(
        tstar=t * U / length,
        mass=mass,
        mass_relative=(
            0.0 if initial_mass is None else abs(mass - initial_mass) / initial_mass
        ),
        a_hat=float(a_hat),
        b_hat=float(b_hat),
        amplitude_exact=amplitude,
        amplitude_ratio=float(a_hat / amplitude),
        modal_amplitude=modal,
        amplitude_error=modal / amplitude - 1,
        phase_drift=float(math.atan2(b_hat, a_hat)),
        phase_drift_abs=abs(float(math.atan2(b_hat, a_hat))),
        velocity_l2=float(
            np.sqrt(w @ ((ux - ex) ** 2 + (uy - ey) ** 2) / exact_norm)
        ),
        residual_l2=float(
            np.sqrt(w @ ((ux - fit * E[0]) ** 2 + (uy - fit * E[1]) ** 2) / exact_norm)
        ),
        longitudinal_l2=float(np.sqrt(w @ (longitudinal**2) / exact_norm)),
        spurious_pressure_rms=float(np.sqrt(w @ (p * p) / volume)),
        spurious_pressure_scaled=float(np.sqrt(w @ (p * p) / volume) / scale),
        spurious_pressure_max_scaled=float(np.max(np.abs(p)) / scale),
        spurious_pressure_scaled_by_level=[
            float(np.sqrt(w[lev == i] @ (p[lev == i] ** 2) / w[lev == i].sum()) / scale)
            for i in range(len(masks))
        ],
        rho_min=float(rho.min()),
        rho_max=float(rho.max()),
        max_speed=float(np.max(np.hypot(ux, uy))),
        leaf_cells=int(len(rho)),
    )


def fit_rate(rows, length, viscosity):
    """Free-intercept least squares of ln R, R = hypot(a_hat, b_hat), late samples.

    The intercept absorbs only a constant amplitude prefactor once a single
    exponential dominates; per-interval rates and the fit residual are kept so
    a lingering transient stays visible. Rates are per root step of the record
    (uniform 2N uses its own nu and doubled step counts).
    """
    late = [r for r in rows if r["tstar"] >= FIT_FROM_TSTAR - 1e-12]
    amp = np.array([math.hypot(r["a_hat"], r["b_hat"]) for r in late])
    if len(late) < 3 or not np.isfinite(amp).all():
        raise ValueError("need >= 3 late samples with finite R")
    exact_rate = decay_rate(length, viscosity)
    to_tstar = length / U  # root steps per unit t*
    if not (amp > 0).all():
        # Finite modal extinction: a physical failure, not malformed data. No
        # clipping and no invented rate; fit quantities are unavailable.
        return dict(rate_fit_valid=False, rate=None, exact_rate=exact_rate,
                    rate_tstar=None, exact_rate_tstar=exact_rate * to_tstar,
                    rate_error=None, interval_rate_errors=None, log_fit_residual_max=None)
    t = np.array([r["step"] for r in late], dtype=float)
    (slope, intercept) = np.polyfit(t, np.log(amp), 1)
    residual = np.log(amp) - (slope * t + intercept)
    return dict(
        rate_fit_valid=True, rate=float(-slope), exact_rate=exact_rate,
        rate_tstar=float(-slope * to_tstar),
        exact_rate_tstar=exact_rate * to_tstar,  # = 5 (2 pi)^2 / Re
        rate_error=abs(-slope / exact_rate - 1),
        interval_rate_errors=[
            float(-np.log(amp[i + 1] / amp[i]) / (t[i + 1] - t[i]) / exact_rate - 1)
            for i in range(len(t) - 1)
        ],
        log_fit_residual_max=float(np.max(np.abs(residual))),
    )


NONNEGATIVE = ("tstar", "mass", "mass_relative", "amplitude_exact", "modal_amplitude",
               "phase_drift_abs", "velocity_l2", "residual_l2", "longitudinal_l2",
               "spurious_pressure_rms", "spurious_pressure_scaled",
               "spurious_pressure_max_scaled", "rho_min", "rho_max", "max_speed")
SIGNED = ("a_hat", "b_hat", "amplitude_ratio", "amplitude_error", "phase_drift")


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


def validate_rows(record):
    """Exact schedule, finite numbers (incl. lists) and nonnegative error fields.

    Signed quantities (amplitudes, amplitude/phase errors) keep their sign.
    Raises ValueError so callers can report an invalid run before any screen.
    """
    n, rows = record["n"], record["rows"]
    steps = [round(tstar * n / U) for tstar in TIMES]
    if not isinstance(rows, list) or len(rows) != len(steps):
        raise ValueError("rows do not match the scheduled samples")
    for row, step in zip(rows, steps):
        if type(row.get("step")) is not int or row["step"] != step:
            raise ValueError(f"unscheduled or duplicate step {row.get('step')!r}")
        if not _number(row.get("tstar")) or abs(row["tstar"] - step * U / n) > 3e-17:
            raise ValueError(f"wrong sample time at step {step}")
        if not all(_number(row.get(key)) for key in (*NONNEGATIVE, *SIGNED)):
            raise ValueError(f"missing or nonfinite diagnostic at step {step}")
        if any(row[key] < 0 for key in NONNEGATIVE) or row["rho_min"] <= 0:
            raise ValueError(f"negative error quantity or density at step {step}")
        levels = row.get("spurious_pressure_scaled_by_level")
        if (not isinstance(levels, list) or len(levels) != record["depth"]
                or not all(_number(v) and v >= 0 for v in levels)):
            raise ValueError(f"invalid per-level pressure at step {step}")


def summarize(record):
    validate_rows(record)
    rows = record["rows"]
    n = record["n"]
    out = fit_rate(rows, n, record["viscosity"])
    out.update(
        eps=eps(n),
        complete=record["completed"] is True,
        residual_l2=max(r["residual_l2"] for r in rows),
        phase_drift_abs=max(r["phase_drift_abs"] for r in rows),
        spurious_pressure_scaled=max(r["spurious_pressure_scaled"] for r in rows),
        spurious_pressure_rms=max(r["spurious_pressure_rms"] for r in rows),
        amplitude_error_abs=max(abs(r["amplitude_error"]) for r in rows),
        velocity_l2=max(r["velocity_l2"] for r in rows),
        mass_relative=max(r["mass_relative"] for r in rows),
    )
    if out["rate_fit_valid"] and not all(
            _number(v) for v in (out["rate"], out["rate_error"], *out["interval_rate_errors"],
                                 out["log_fit_residual_max"])):
        raise ValueError("nonfinite rate fit")
    out["own_screens_pass"] = out["complete"] and out["rate_fit_valid"] and all(
        out[key] <= out["eps"] for key in ("rate_error", "phase_drift_abs", "residual_l2")
    ) and all(out[key] <= limit for key, limit in LIMITS.items())
    return out


def _usable(record):
    try:
        out = summarize(record)
    except (KeyError, TypeError, ValueError):
        return None
    return out if out["complete"] else None


def assess(base_control, fine_control, refined):
    """Uniform root-spacing and finest-spacing controls versus refined runs.

    INVALID_CONTROL: missing, incomplete, invalid-row or mixed-source controls.
    INVALID_RUN: a refined record has invalid rows (schedule, NaN, sign).
    CONTROL_SCREEN_FAIL: complete valid controls that fail their own screens;
    refined own-screen results are kept, but no controlled comparison passes.
    PASS/FAIL: every refined run's controlled_pass (own screens and budgets).
    """
    if base_control is None or fine_control is None or not refined:
        return dict(status="INVALID_CONTROL", reason="missing control or refined run")
    n = base_control["n"]
    if (base_control["depth"], fine_control["depth"], fine_control["n"]) != (1, 1, 2 * n):
        raise ValueError("controls must be uniform N and uniform 2N")
    if any(r["depth"] != 2 or r["n"] != n for r in refined):
        raise ValueError("refined runs must be depth 2 at root N")
    if len({r.get("explosion") for r in refined}) != len(refined):
        raise ValueError("duplicate refined explosion scheme")
    sources = {repr(r.get("source")) for r in [base_control, fine_control, *refined]}
    base, fine = _usable(base_control), _usable(fine_control)
    if base is None or fine is None or len(sources) != 1:
        return dict(status="INVALID_CONTROL",
                    reason="incomplete/invalid control or mixed source")
    controls_pass = base["own_screens_pass"] and fine["own_screens_pass"]
    report = dict(base_control=base, fine_control=fine, refined={},
                  controls_pass=controls_pass)
    statuses = []
    for r in refined:
        try:
            s = summarize(r)
        except (KeyError, TypeError, ValueError) as error:
            report["refined"][r["explosion"]] = dict(valid=False, reason=str(error))
            statuses.append("INVALID_RUN")
            continue
        s["valid"] = True
        # rate_error is None when a modal fit is unavailable; never a pass.
        s["budget"] = {
            key: (None if base[key] is None else max(BASELINE_FACTOR * base[key], floor))
            for key, floor in FLOORS.items()
        }
        s["budget_pass"] = all(
            s[key] is not None and s["budget"][key] is not None and s[key] <= s["budget"][key]
            for key in FLOORS)
        s["control_ratios"] = {
            name: {key: (s[key] / c[key] if s[key] is not None and c[key] else None)
                   for key in (*FLOORS, "spurious_pressure_rms")}
            for name, c in (("root", base), ("finest", fine))
        }
        s["controlled_pass"] = controls_pass and s["own_screens_pass"] and s["budget_pass"]
        report["refined"][r["explosion"]] = s
        statuses.append("PASS" if s["controlled_pass"] else "FAIL")
    if "INVALID_RUN" in statuses:
        report["status"] = "INVALID_RUN"
    elif not controls_pass:
        report["status"] = "CONTROL_SCREEN_FAIL"
    else:
        report["status"] = "PASS" if all(v == "PASS" for v in statuses) else "FAIL"
    return report