import json
import numpy as np
from invariant_reconstruction import pressure, admissible, limited_state, density_theta_bound, GAMMA


def state(rho, vx, vy, p):
    return np.array([rho, rho*vx, rho*vy,
                     p/(GAMMA-1.0) + 0.5*rho*(vx*vx+vy*vy)])


def pressure_root(w0, wr, pf):
    # The pressure-floor condition after multiplying by rho is quadratic in theta.
    d = wr - w0
    def f(t):
        w = w0 + t*d
        return w[3]*w[0] - .5*np.dot(w[1:3], w[1:3]) - pf/(GAMMA-1.0)*w[0]
    y0, y1, y2 = f(0.), f(1.), f(2.)
    a = .5*(y2 - 2*y1 + y0)
    b = y1 - y0 - a
    if abs(a) < 1e-13:
        roots = [-y0/b] if abs(b) > 1e-13 else []
    else:
        roots = np.roots([a, b, y0])
    good = [float(np.real(x)) for x in roots
            if abs(np.imag(x)) < 1e-8 and np.real(x) >= 0]
    return min(good) if good else np.inf


def main():
    rng = np.random.default_rng(1178)
    rf, pf = 1e-6, 1e-6
    w0 = state(1.0, .7, -.2, 1.0)
    out = {"predictions": {}, "comparison": {}}

    # Prediction 1: if pressure is safely high, the exact affine density bound is used.
    density_rows = []
    for rr in np.geomspace(1e-9, 0.9*rf, 8):
        wr = state(rr, 0., 0., 1.0)
        _, got = limited_state(w0, wr, rf, pf)
        expected = density_theta_bound(w0, wr, rf)
        density_rows.append([float(rr), float(got), float(expected), abs(got-expected)])
    out["predictions"]["density_affine_bound"] = {
        "prediction": "safe high-pressure endpoints satisfy theta=(rho0-rhofloor)/(rho0-rhoraw)",
        "max_abs_error": float(max(r[-1] for r in density_rows)), "rows": density_rows}

    # Prediction 2: pressure positivity is enforced at the first analytic quadratic root.
    pressure_rows = []
    for mach_momentum in [2., 4., 8., 16., 32.]:
        # Negative thermodynamic pressure makes the raw endpoint inadmissible.
        wr = state(1.0, mach_momentum, 0., -0.1)
        _, got = limited_state(w0, wr, rf, pf)
        expected = min(1.0, pressure_root(w0, wr, pf))
        pressure_rows.append([mach_momentum, float(got), float(expected), abs(got-expected),
                              float(pressure(wr))])
    out["predictions"]["pressure_root"] = {
        "prediction": "accepted theta equals the first pressure-floor crossing on a bad segment",
        "max_abs_error": float(max(r[3] for r in pressure_rows)), "rows": pressure_rows}

    # Prediction 3: scaling a fixed bad ray by Lambda gives theta proportional to 1/Lambda.
    direction = state(0.2, 18., 0., -0.1) - w0
    scales = np.array([1., 2., 4., 8., 16., 32.])
    theta = np.array([limited_state(w0, w0 + lam*direction, rf, pf)[1]
                      for lam in scales])
    scaled = scales * theta
    out["predictions"]["inverse_strength_scaling"] = {
        "prediction": "for Lambda beyond the first crossing, Lambda*theta is constant",
        "scales": scales.tolist(), "theta": theta.tolist(),
        "lambda_theta": scaled.tolist(),
        "tail_ratio": float(max(scaled[-3:])/min(scaled[-3:]))}

    # Secondary fixed-seed comparison against first-order (cell-average) reconstruction.
    n = 3000
    cells, raws = [], []
    for _ in range(n):
        c = state(rng.uniform(.5, 1.5), rng.uniform(-1, 1), rng.uniform(-1, 1),
                  rng.uniform(.3, 2.))
        raws.append(c + rng.normal(0, 1.8, 4)); cells.append(c)
    raw_invalid = sum(not admissible(r, rf, pf) for r in raws)
    safe, ts = [], []
    for c, r in zip(cells, raws):
        s, t = limited_state(c, r, rf, pf); safe.append(s); ts.append(t)
    out["comparison"] = {
        "samples": n,
        "raw_invalid_rate": raw_invalid/n,
        "invariant_invalid_rate": sum(not admissible(s,rf,pf) for s in safe)/n,
        "baseline_first_order_l2_to_raw": float(np.mean([np.linalg.norm(c-r) for c,r in zip(cells,raws)])),
        "idea_l2_to_raw": float(np.mean([np.linalg.norm(s-r) for s,r in zip(safe,raws)])),
        "idea_l2_to_cell": float(np.mean([np.linalg.norm(s-c) for s,c in zip(safe,cells)])),
        "activation_rate": float(np.mean(np.array(ts) < .999999))}
    print(json.dumps(out, indent=2))
    with open("results.json", "w") as f: json.dump(out, f, indent=2)

if __name__ == "__main__":
    main()