import json
import time
import numpy as np


def sigmoid(x):
    x = np.clip(x, -30.0, 30.0)
    return 1.0 / (1.0 + np.exp(-x))


def spectrum(T, r):
    j = np.arange(1, len(r) + 1)
    n = np.maximum(T - j + 1, 0)
    return float(np.sum(np.log1p(n * np.asarray(r) ** 2)))


def selection_objective(theta, rhat, t, lam):
    """Return C(h), h=0..H, using n_{t,j}=max(t-j+1,0)."""
    H = len(theta)
    j = np.arange(1, H + 1)
    n = np.maximum(t - j + 1, 0.0)
    price = lam * np.log1p(n * np.asarray(rhat) ** 2)
    tail = n * np.asarray(theta) ** 2
    # C(h)=sum_{j<=h} price_j + sum_{j>h} tail_j
    return np.r_[np.sum(tail), np.cumsum(price - tail) + np.sum(tail)]


def math_check(seed=7):
    rng = np.random.default_rng(seed)
    H, T = 18, 80
    theta = rng.normal(0, .12, H)
    rhat = np.abs(theta) + .2 / np.sqrt(np.arange(1, H + 1))
    lam, t = .35, 63
    C = selection_objective(theta, rhat, t, lam)
    j = np.arange(1, H + 1)
    n = np.maximum(t - j + 1, 0.)
    marginal = lam * np.log1p(n * rhat ** 2) - n * theta ** 2
    algebra_error = float(np.max(np.abs(np.diff(C) - marginal)))

    # Directly verify the paper's weighted tail scale against logistic KL.
    # For independent Rademacher inputs, compare source and truncated logits.
    B = .75
    th = rng.normal(0, .035, 7)
    th *= min(1., B / max(np.sum(np.abs(th)), 1e-12))
    hh, TT = 2, 120
    V = sum((TT-j+1) * th[j-1] ** 2 for j in range(hh + 1, len(th) + 1))
    kl_sum = 0.0
    reps = 4000
    for _ in range(reps):
        u = rng.choice([-1., 1.], size=len(th))
        a = float(np.dot(th, u))
        b = float(np.dot(th[:hh], u[:hh]))
        p, q = sigmoid(a), sigmoid(b)
        kl = p * np.log(p / q) + (1-p) * np.log((1-p) / (1-q))
        kl_sum += kl
    # Same expected KL is incurred on each affected round, hence V is its scale.
    empirical_ratio = (TT-hh) * kl_sum / reps / max(V, 1e-12)
    return {"objective_marginal_max_error": algebra_error,
            "tail_energy": float(V), "empirical_logistic_KL_over_tail_energy": float(empirical_ratio),
            "math_check_pass": bool(algebra_error < 1e-10 and .05 < empirical_ratio < .2)}


def make_data(T, H, seed):
    rng = np.random.default_rng(seed)
    u = rng.choice([-1., 1.], size=T + H + 1)
    theta = np.zeros((T, H))
    # Abruptly switch from short exponential memory to a longer polynomial-like tail.
    a = .40 * (.55 ** np.arange(H))
    b = .17 / (np.arange(1, H + 1) ** .72)
    theta[:T//2] = a
    theta[T//2:] = b
    y = np.zeros(T, dtype=np.int8)
    for t in range(T):
        x = u[H + t - np.arange(1, H + 1)]
        y[t] = rng.random() < sigmoid(np.dot(theta[t], x))
    return u, y, theta


def run_model(u, y, true_theta, mode, H, lam=.012, k=.08, lr=.035, warmup=100):
    T = len(y)
    est = np.zeros(H)
    count = np.zeros(H)
    active_h = H if mode == "fixed_long" else (4 if mode == "fixed_short" else 1)
    losses, hs, ops = [], [], []
    t0 = time.perf_counter()
    for t in range(T):
        x = u[H + t - np.arange(1, H + 1)]
        h_used = H if mode == "fixed_long" else (4 if mode == "fixed_short" else active_h)
        logit = float(np.dot(est[:h_used], x[:h_used]))
        p = float(sigmoid(logit))
        losses.append(-(y[t] * np.log(p + 1e-12) + (1-y[t]) * np.log(1-p + 1e-12)))
        err = p - y[t]
        # Full estimates are maintained as a transparent online pilot for unseen lags;
        # only the active prefix contributes to prediction and its measured inference ops.
        est -= lr * err * x
        count += 1
        rhat = np.abs(est) + k / np.sqrt(count)
        if mode == "adaptive" and t >= warmup and t % 5 == 0:
            C = selection_objective(est, rhat, t + 1, lam)
            best = np.min(C)
            candidates = np.flatnonzero(C <= best * 1.02 + 1e-9)
            active_h = int(candidates[0]) if len(candidates) else int(np.argmin(C))
            active_h = max(1, min(H, active_h))
        hs.append(h_used)
        ops.append(h_used)
    elapsed = time.perf_counter() - t0
    losses = np.asarray(losses)
    return {"cum_logloss": float(np.sum(losses)),
            "first_half_logloss": float(np.sum(losses[:T//2])),
            "second_half_logloss": float(np.sum(losses[T//2:])),
            "mean_active_h": float(np.mean(hs)),
            "second_half_active_h": float(np.mean(hs[T//2:])),
            "inference_ops_relative": float(np.sum(ops) / (T * H)),
            "wall_seconds": float(elapsed)}


def main():
    checks = math_check()
    T, H = 7000, 24
    u, y, truth = make_data(T, H, seed=19)
    results = {}
    for mode in ["fixed_long", "fixed_short", "adaptive"]:
        results[mode] = run_model(u, y, truth, mode, H)
    # Report the idea against the conventional fixed maximum context and a short fixed window.
    out = {"math": checks, "experiment": results,
           "settings": {"T": T, "H_max": H, "lambda": .012, "confidence_k": .08,
                        "seed": 19, "data": "short exponential -> long polynomial abrupt switch"}}
    with open("results.json", "w") as f:
        json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))


if __name__ == "__main__":
    main()