import json, math
from pathlib import Path
import numpy as np

# Ordered Diffusion Message Passing MVP.
# The toy grid makes the local Gaussian calculation transparent; the forecasting
# task uses the same row-normalized edge operator.
SEED = 7
rng = np.random.default_rng(SEED)


def kernel_on_offsets(offsets, eps, beta):
    """One row of K/P at x_i=0, with s(x)=x and f evaluated on offsets."""
    logits = -offsets**2 / (4.0 * eps) + beta * offsets
    logits -= logits.max()  # numerically stable softmax
    p = np.exp(logits)
    p /= p.sum()
    return p


def core_sweeps():
    # Fine grid and broad enough window make truncation error negligible.
    dx = 0.005
    offsets = np.arange(-6.0, 6.0 + dx / 2, dx)
    drift_rows = []
    for eps in [0.015, 0.03, 0.06, 0.12]:
        for beta in [0.0, 0.25, 0.5, 1.0]:
            p = kernel_on_offsets(offsets, eps, beta)
            mean = float(p @ offsets)
            second = float(p @ (offsets**2))
            drift_rows.append({
                "eps": eps, "beta": beta,
                "observed_mean": mean,
                "predicted_mean": 2 * beta * eps,
                "relative_drift_error": abs(mean - 2*beta*eps) / max(abs(2*beta*eps), 1e-12),
                "observed_generator_x": mean / eps,
                "predicted_generator_x": 2 * beta,
                "observed_generator_x2": second / eps,
                "predicted_generator_x2": 2 + 4 * beta * beta * eps,
            })

    # Prediction 1: drift is zero at beta=0 and linear in beta.
    linear = [r for r in drift_rows if r["eps"] == 0.06]
    b = np.array([r["beta"] for r in linear])
    m = np.array([r["observed_mean"] for r in linear])
    slope = float(np.polyfit(b, m, 1)[0])
    # Prediction 2: at fixed beta, displacement scales linearly in epsilon.
    beta_rows = [r for r in drift_rows if r["beta"] == 0.5]
    e = np.array([r["eps"] for r in beta_rows])
    mm = np.array([r["observed_mean"] for r in beta_rows])
    eps_slope = float(np.polyfit(e, mm, 1)[0])
    # Prediction 3: generator of x^2 tends to 2 as eps -> 0 (at beta fixed).
    small = [r for r in drift_rows if r["beta"] == 0.5 and r["eps"] <= 0.06]
    gen_errors = [abs(r["observed_generator_x2"] - r["predicted_generator_x2"]) for r in small]
    # Exact row-stochastic/nonexpansive check on a finite random row and repeated rows.
    p = kernel_on_offsets(offsets, 0.06, 1.0)
    v = rng.normal(size=len(p))
    stability = {
        "row_sum": float(p.sum()),
        "max_entry": float(p.max()),
        "inf_norm_input": float(np.max(np.abs(v))),
        "inf_norm_one_step": float(abs(p @ v)),
        "inf_norm_bound_holds": bool(abs(p @ v) <= np.max(np.abs(v)) + 1e-12),
    }
    summary = {
        "drift_beta_slope_observed": slope,
        "drift_beta_slope_predicted": 2 * 0.06,
        "drift_eps_slope_observed_beta_0.5": eps_slope,
        "drift_eps_slope_predicted_beta_0.5": 1.0,
        "x2_generator_max_abs_error_small_eps": max(gen_errors),
        "x2_generator_prediction_tolerance": 0.02,
        "rows": drift_rows,
        "stability": stability,
    }
    return summary


def directed_forecast():
    """Advect a smooth signal on a randomly sampled 1-D point cloud.

    h(x)=sin(x), target is h(x+shift). A positive ordered tilt shifts the
    local average forward, while symmetric diffusion only blurs h.
    """
    n = 700
    x = np.sort(rng.uniform(-5.0, 5.0, size=n))
    h = np.sin(x)
    target = np.sin(x + 0.22)
    k = 35
    # kNN in 1-D, same neighbors for both methods.
    dist = np.abs(x[:, None] - x[None, :])
    nn = np.argpartition(dist, kth=k, axis=1)[:, :k]
    eps = float(np.median(np.sort(dist, axis=1)[:, k])**2 / 4.0)
    eps = max(eps, 1e-5)
    delta = x[nn] - x[:, None]
    base = -delta**2 / (4 * eps)

    def aggregate(beta):
        logits = base + beta * delta
        logits -= logits.max(axis=1, keepdims=True)
        p = np.exp(logits)
        p /= p.sum(axis=1, keepdims=True)
        return np.sum(p * h[nn], axis=1), p

    pred0, p0 = aggregate(0.0)
    # Oracle scalar ordering s(x)=x is the simplest learned-ordering target.
    # Sweep beta and report the best value without changing neighbors or width.
    candidates = np.linspace(0, 1.8, 19)
    mses = []
    entropies = []
    for beta in candidates:
        pred, pp = aggregate(float(beta))
        mses.append(float(np.mean((pred - target)**2)))
        entropies.append(float(np.mean(-np.sum(pp * np.log(pp + 1e-12), axis=1))))
    ib = int(np.argmin(mses))
    best_beta = float(candidates[ib])
    pred, pb = aggregate(best_beta)
    return {
        "n": n, "k": k, "epsilon": eps, "shift": 0.22,
        "symmetric_mse": float(np.mean((pred0-target)**2)),
        "ordered_best_mse": float(mses[ib]),
        "ordered_best_beta": best_beta,
        "ordered_improvement_fraction": float(1 - mses[ib]/max(mses[0], 1e-12)),
        "symmetric_entropy": float(np.mean(-np.sum(p0*np.log(p0+1e-12), axis=1))),
        "ordered_entropy": float(entropies[ib]),
        "all_beta_mse": {str(float(b)): float(m) for b,m in zip(candidates, mses)},
    }


def main():
    out = {"seed": SEED, "core_verification": core_sweeps(), "forecast": directed_forecast()}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__":
    main()