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

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


def mechanism_checks():
    rows = []
    # Prediction 1: nested cumulative sums reproduce products exactly in log space.
    for K in (4, 8, 32, 128):
        for scale in (0.05, 0.2, 0.8):
            alpha = rng.normal(0, scale, K)
            later_sum = np.cumsum(alpha[::-1])[::-1] - alpha
            direct_log_product = np.array([np.sum(alpha[j + 1:]) for j in range(K)])
            err = float(np.max(np.abs(later_sum - direct_log_product)))
            rows.append({"check": "product_identity", "K": K, "scale": scale, "max_abs_log_error": err})

    # Prediction 2: shell work is O(K), pairwise work is O(K^2), with fitted exponents.
    Ks = np.array([8, 16, 32, 64, 128, 256, 512], dtype=float)
    shell_work = Ks
    pair_work = Ks ** 2
    shell_exp = float(np.polyfit(np.log(Ks), np.log(shell_work), 1)[0])
    pair_exp = float(np.polyfit(np.log(Ks), np.log(pair_work), 1)[0])
    rows.append({"check": "complexity_exponent", "shell_exponent": shell_exp, "pairwise_exponent": pair_exp})

    # Prediction 3: bounded shell coefficients have no K-dependent amplification,
    # while an unnormalized dense sum amplifies a constant signal linearly in K.
    amp_rows = []
    for K in (8, 16, 32, 64, 128, 256):
        beta = np.ones(K) / K
        shell_amp = float(np.sum(np.cumsum(beta[::-1])[::-1]) / K)
        dense_amp = float(np.sum(np.ones((K, K)) / K))
        amp_rows.append({"K": K, "shell_normalized_amplification": shell_amp, "dense_amplification": dense_amp})
    return {"rows": rows, "complexity": {"shell_exponent": shell_exp, "pairwise_exponent": pair_exp}, "amplification": amp_rows}


def make_data(n=1600, T=24, K=32, regime="ballistic"):
    # Predict a vector of future values from a short observed context.
    # Ballistic: almost deterministic shifted latent wave; mixing: independent future noise.
    x = rng.normal(size=(n, T + K + 1)).astype(np.float32)
    if regime == "ballistic":
        latent = rng.normal(size=n).astype(np.float32)
        for t in range(T + K + 1):
            x[:, t] = latent + 0.03 * rng.normal(size=n)
        # small trend makes horizons distinguishable but memory remains strong
        y = x[:, T + 1:T + K + 1] + 0.01 * np.arange(1, K + 1)[None, :]
    else:
        y = rng.normal(size=(n, K)).astype(np.float32)
    context = x[:, :T]
    return context, y.astype(np.float32)


def fit_models(context, y, K, epochs=80):
    # Numpy least-squares implementation isolates decoder structure and makes
    # the O(K) vs O(K^2) operations explicit and reproducible.
    n, T = context.shape
    split = int(.75 * n)
    Xtr, Xva, Ytr, Yva = context[:split], context[split:], y[:split], y[split:]
    htr = np.c_[Xtr.mean(1), Xtr[:, -1], np.arange(1, K + 1)[None, :].repeat(len(Xtr), 0) / K]
    hva = np.c_[Xva.mean(1), Xva[:, -1], np.arange(1, K + 1)[None, :].repeat(len(Xva), 0) / K]
    beta = np.linspace(1.0, 0.1, K, dtype=np.float32)
    B = np.cumsum(beta[::-1])[::-1] / K
    # Shell features: shared state plus current horizon and later cumulative coefficient.
    shell_tr = np.stack([np.c_[htr[:, :2], np.full((len(Xtr), 1), B[j]), htr[:, 2]] for j in range(K)], 1)
    shell_va = np.stack([np.c_[hva[:, :2], np.full((len(Xva), 1), B[j]), hva[:, 2]] for j in range(K)], 1)
    # Pairwise baseline: each horizon can use all horizon coefficients (K features).
    pair_tr = np.concatenate([htr[:, :2], np.tile(beta, (len(Xtr), 1))], 1)
    pair_va = np.concatenate([hva[:, :2], np.tile(beta, (len(Xva), 1))], 1)
    def ridge(A, target, lam=1e-3):
        return np.linalg.solve(A.T @ A + lam * np.eye(A.shape[1]), A.T @ target)
    t0 = time.perf_counter()
    Ws = np.stack([ridge(np.c_[shell_tr[:, j, :], np.ones(len(Xtr))], Ytr[:, j]) for j in range(K)])
    shell_time = time.perf_counter() - t0
    t0 = time.perf_counter()
    Wp = ridge(np.c_[pair_tr, np.ones(len(Xtr))], Ytr)
    pair_time = time.perf_counter() - t0
    ps = np.stack([np.c_[shell_va[:, j, :], np.ones(len(Xva))] @ Ws[j] for j in range(K)], 1)
    pp = np.c_[pair_va, np.ones(len(Xva))] @ Wp
    return {"shell_mse": float(np.mean((ps-Yva)**2)), "pairwise_mse": float(np.mean((pp-Yva)**2)), "shell_seconds": shell_time, "pairwise_seconds": pair_time, "shell_parameters": int(Ws.size), "pairwise_parameters": int(Wp.size)}


def main():
    checks = mechanism_checks()
    experiments = {}
    for regime in ("ballistic", "mixing"):
        context, y = make_data(regime=regime)
        experiments[regime] = fit_models(context, y, 32)
    result = {"seed": SEED, "mechanism_checks": checks, "experiments": experiments}
    Path("results.json").write_text(json.dumps(result, indent=2))
    print(json.dumps(result, indent=2))

if __name__ == "__main__":
    main()