import json
from pathlib import Path
import numpy as np

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


def nearest(x, h):
    # Round-to-nearest, ties are irrelevant for the chosen probes.
    return np.floor(x / h + 0.5) * h


def full_state(deltas, h):
    z = 0.0
    qs = []
    for delta in deltas:
        q = nearest(z + delta, h) - z
        z += q
        qs.append(q)
    return z, np.asarray(qs)


def feedback(deltas, h):
    z, c = 0.0, 0.0
    qs, carries, sat = [], [], 0
    for delta in deltas:
        u = delta + c
        q = nearest(u, h)
        c = u - q
        z += q
        qs.append(q)
        carries.append(c)
        sat += int(abs(u) >= 127.5 * h)  # diagnostic only; no clipping in this toy
    return z, np.asarray(qs), np.asarray(carries), sat


def stochastic_state(deltas, h, local_rng):
    z = 0.0
    qs = []
    for delta in deltas:
        x = (z + delta) / h
        lo = np.floor(x)
        p = x - lo
        q = (lo + (local_rng.random() < p)) * h
        z = q
        qs.append(q)
    return z, np.asarray(qs)


def identity_check():
    h = 0.7
    deltas = rng.normal(0, 0.31, 97)
    z, qs, carries, _ = feedback(deltas, h)
    identity_err = abs(z - (deltas.sum() - carries[-1]))
    return {"identity_abs_error": float(identity_err),
            "max_abs_carry_over_h": float(np.max(abs(carries)) / h),
            "predicted_carry_bound": 0.5}


def depth_sweep():
    h = 1.0
    delta = 0.49 * h
    rows = []
    for d in [1, 2, 4, 8, 16, 32, 64, 128, 256, 512]:
        deltas = np.full(d, delta)
        target = deltas.sum()
        zb, _ = full_state(deltas, h)
        zi, _, c, _ = feedback(deltas, h)
        rows.append({"depth": d,
                     "baseline_abs_error": abs(zb-target),
                     "idea_abs_error": abs(zi-target),
                     "idea_max_carry_over_h": float(np.max(abs(c))/h),
                     "predicted_idea_bound": 0.5})
    # Fit baseline error slope in units h/depth after the transient.
    ds = np.array([r["depth"] for r in rows], float)
    eb = np.array([r["baseline_abs_error"] for r in rows])
    ei = np.array([r["idea_abs_error"] for r in rows])
    slope = np.polyfit(ds[2:], eb[2:], 1)[0]
    return rows, {"baseline_error_slope_per_depth": float(slope),
                  "predicted_baseline_slope": 0.49,
                  "idea_max_error": float(ei.max()),
                  "predicted_idea_error_bound": 0.5}


def scale_sweep():
    rows = []
    d = 257
    for h in [0.125, 0.25, 0.5, 1.0, 2.0, 4.0]:
        # Relative increment is fixed, so normalized carry should be invariant.
        deltas = np.full(d, 0.49*h)
        target = deltas.sum()
        zi, _, c, _ = feedback(deltas, h)
        rows.append({"h": h, "final_abs_error": abs(zi-target),
                     "max_abs_carry": float(np.max(abs(c))),
                     "max_abs_carry_over_h": float(np.max(abs(c))/h),
                     "predicted_max_abs_carry": h/2})
    return rows


def random_increment_comparison():
    # Mimics a residual stream with varying proposals, without conflating the
    # conservation claim with a trained-network effect.
    h = 0.5
    out = []
    for d in [8, 32, 128, 512]:
        errors_b, errors_i, errors_s = [], [], []
        for trial in range(200):
            deltas = rng.uniform(-0.49*h, 0.49*h, d)
            target = deltas.sum()
            zb, _ = full_state(deltas, h)
            zi, _, _, _ = feedback(deltas, h)
            zs, _ = stochastic_state(deltas, h, np.random.default_rng(SEED + trial + d))
            errors_b.append(abs(zb-target)); errors_i.append(abs(zi-target)); errors_s.append(abs(zs-target))
        out.append({"depth": d, "baseline_mean_abs_error": float(np.mean(errors_b)),
                    "idea_mean_abs_error": float(np.mean(errors_i)),
                    "stochastic_mean_abs_error": float(np.mean(errors_s))})
    return out


def main():
    ident = identity_check()
    depth, depth_summary = depth_sweep()
    scales = scale_sweep()
    random_cmp = random_increment_comparison()

    # Quantitative mechanism checks (not merely qualitative sanity checks).
    assert ident["identity_abs_error"] < 1e-12
    assert ident["max_abs_carry_over_h"] <= 0.5 + 1e-12
    assert abs(depth_summary["baseline_error_slope_per_depth"] - 0.49) < 1e-10
    assert depth_summary["idea_max_error"] <= 0.5 + 1e-10
    for row in scales:
        assert abs(row["max_abs_carry_over_h"] - 0.5) < 1e-10
        assert abs(row["max_abs_carry"] - row["predicted_max_abs_carry"]) < 1e-10
    assert random_cmp[-1]["idea_mean_abs_error"] < random_cmp[-1]["baseline_mean_abs_error"]

    result = {"seed": SEED, "identity": ident, "depth_sweep": depth,
              "depth_summary": depth_summary, "scale_sweep": scales,
              "random_comparison": random_cmp,
              "assertions": "passed"}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()