"""Toy verification of Rankine--Hugoniot front tokens on Burgers' equation.

Run with the prescribed Python interpreter. Outputs a concise JSON report and
an optional CSV-like result file in this directory.
"""
import json, math
from pathlib import Path
import numpy as np

SEED = 7
rng = np.random.default_rng(SEED)


def flux(u):
    return 0.5 * u * u


def rh_speed(ul, ur):
    den = ul - ur
    return (flux(ul) - flux(ur)) / den if abs(den) > 1e-12 else ul


def godunov_flux(ul, ur):
    # Exact scalar Burgers Godunov flux for the Riemann problem.
    if ul <= ur:  # rarefaction
        if ul >= 0: return flux(ul)
        if ur <= 0: return flux(ur)
        return 0.0
    # shock
    s = rh_speed(ul, ur)
    return flux(ul) if s >= 0 else flux(ur)


def godunov_step(u, dx, dt):
    # Fixed exterior states equal to the nearest interior value (shock stays away).
    ext = np.concatenate(([u[0]], u, [u[-1]]))
    f = np.array([godunov_flux(ext[i], ext[i+1]) for i in range(len(ext)-1)])
    return u - dt / dx * (f[1:] - f[:-1])


def exact_shock(x, xfront, ul=1.0, ur=0.0):
    return np.where(x < xfront, ul, ur)


def detect_front(u, x, threshold=0.35):
    jumps = np.abs(np.diff(u))
    ids = np.flatnonzero(jumps >= threshold)
    if len(ids) == 0:
        return []
    # For this toy, adjacent flagged edges are one cluster; code also handles many.
    clusters = []
    start = prev = int(ids[0])
    for q in ids[1:]:
        q = int(q)
        if q > prev + 1:
            clusters.append((start, prev)); start = q
        prev = q
    clusters.append((start, prev))
    out = []
    for a, b in clusters:
        # traces are pooled from the cells immediately outside the flagged cluster
        li = max(0, a)
        ri = min(len(u)-1, b+1)
        out.append({"x": float(0.5*(x[a]+x[b+1])),
                    "ul": float(u[li]), "ur": float(u[ri])})
    return out


def token_rollout(u0, x, dt, nsteps):
    """Propagate detected fronts analytically and render the piecewise field."""
    tokens = detect_front(u0, x)
    # In the intended module this is where attention/residual prediction enters;
    # for a clean shock the exact traces are already the sufficient local state.
    xs = []
    fields = []
    u = u0.copy()
    for _ in range(nsteps + 1):
        if tokens:
            t = tokens[0]
            fields.append(exact_shock(x, t["x"], t["ul"], t["ur"]))
            xs.append(t["x"])
        else:
            fields.append(u.copy()); xs.append(float("nan"))
        if tokens:
            for t in tokens:
                t["x"] += dt * rh_speed(t["ul"], t["ur"])
    return np.asarray(fields), np.asarray(xs)


def crossing_position(u, x, level=0.5):
    # Linear interpolation of the numerical shock crossing.
    k = np.flatnonzero((u[:-1] >= level) & (u[1:] < level))
    if len(k) == 0: return float("nan")
    i = int(k[0]); den = u[i] - u[i+1]
    return float(x[i] + (level-u[i]) * (x[i+1]-x[i]) / den) if den else float(x[i])


def main():
    # Core mathematical check over random traces and x-independent flux.
    pairs = rng.uniform(-1.0, 1.0, (100, 2))
    pairs = pairs[np.abs(pairs[:,0]-pairs[:,1]) > .1]
    speed_err = []
    for ul, ur in pairs:
        s = rh_speed(ul, ur)
        speed_err.append(abs((flux(ul)-flux(ur))/(ul-ur)-s))
    math_check = {"max_rh_identity_error": float(max(speed_err)),
                  "mean_rh_identity_error": float(np.mean(speed_err)),
                  "identity_pass": bool(max(speed_err) < 1e-12)}

    # Clean shock benchmark at three resolutions, same physical horizon.
    records = []
    for N in (64, 128, 256):
        dx = 1.0/N; x = (np.arange(N)+.5)*dx
        x0 = .25; ul, ur = 1., 0.; s = rh_speed(ul, ur)
        u0 = exact_shock(x, x0, ul, ur)
        # Conservative baseline has a standard CFL timestep.
        dt = .35*dx/max(abs(ul), abs(ur)); steps = int(.20/dt)
        dt = .20/steps
        ub = u0.copy(); base_pos=[]
        for _ in range(steps):
            ub = godunov_step(ub, dx, dt)
            base_pos.append(crossing_position(ub, x))
        tf, token_pos = token_rollout(u0, x, dt, steps)
        times = np.arange(steps+1)*dt
        true_pos = x0 + s*times
        # compare only positions available at the grid output times
        bpos = np.asarray(base_pos)
        bpos_err = float(np.nanmean(np.abs(bpos-true_pos[1:])))
        tpos_err = float(np.mean(np.abs(token_pos-true_pos)))
        bfield = float(np.mean(np.abs(ub-exact_shock(x,true_pos[-1],ul,ur))))
        tfield = float(np.mean(np.abs(tf[-1]-exact_shock(x,true_pos[-1],ul,ur))))
        records.append({"N":N,"dt":dt,"steps":steps,
                        "baseline_front_L1":bpos_err,"idea_front_L1":tpos_err,
                        "baseline_field_L1":bfield,"idea_field_L1":tfield})

    # One-step propagation error vs dt: discretization of the initial location is
    # held fixed, so the analytic kinematic update should have no time integration error.
    N=128; dx=1/N; x=(np.arange(N)+.5)*dx; x0=.25; u0=exact_shock(x,x0)
    tok=detect_front(u0,x)[0]; xdet=tok['x']; s=rh_speed(tok['ul'],tok['ur'])
    scales=[]
    for dt in (0.002, 0.004, 0.008):
        pred=xdet+dt*s; truth=x0+dt*s
        scales.append({"dt":dt,"abs_position_error":abs(pred-truth)})
    report={"seed":SEED,"math_check":math_check,"records":records,
            "one_step_scaling":scales,
            "note":"Front-token field uses exact local traces on a clean Riemann shock; no learned residual is needed in this deliberately minimal case."}
    Path("results.json").write_text(json.dumps(report, indent=2))
    print(json.dumps(report, indent=2))

if __name__ == '__main__': main()