import json, time
import numpy as np

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


def softmax(a):
    a = a - np.max(a, axis=-1, keepdims=True)
    z = np.exp(a)
    return z / z.sum(axis=-1, keepdims=True)


def make_frozen(d, dh, m, seed):
    r = np.random.default_rng(seed)
    # Variance-normalized frozen Gaussian maps.
    Q = r.normal(0, 1 / np.sqrt(dh), size=(dh, dh))
    K = r.normal(0, 1 / np.sqrt(dh), size=(dh, dh))
    V = r.normal(0, 1 / np.sqrt(dh), size=(m, dh))
    A = (Q.T @ K)[: d + 1] / np.sqrt(dh)
    return Q, K, V, A


def targets(Xs, sigma):
    # b_i(x), modulo the query-only term -||x||^2/(2 sigma^2).
    return np.column_stack((Xs / sigma**2, -np.sum(Xs * Xs, axis=1) / (2 * sigma**2)))


def construct(Xs, ys, dh, sigma, ridge=0.0, seed=0):
    d = Xs.shape[1]
    m = ys.shape[1]
    _, _, V, A = make_frozen(d, dh, m, seed)
    C = targets(Xs, sigma)
    B = np.vstack((A, V))
    rhs = np.hstack((C, ys)).T  # (d+1+m, N)
    if ridge == 0:
        P = np.linalg.lstsq(B, rhs, rcond=None)[0]
    else:
        P = np.linalg.solve(B.T @ B + ridge * np.eye(dh), B.T @ rhs)
    return P.T, A, V, C


def predict(Xq, P, A, V, d):
    E = np.column_stack((Xq, np.ones(len(Xq)), np.zeros((len(Xq), A.shape[1] - d - 1))))
    logits = E[:, : d + 1] @ (A @ P.T)
    W = softmax(logits)
    vals = (V @ P.T).T
    return W @ vals, W, logits


def nw(Xq, Xs, ys, sigma):
    b = -((Xq[:, None, :] - Xs[None, :, :]) ** 2).sum(axis=2) / (2 * sigma**2)
    return softmax(b) @ ys, softmax(b)


def knn(Xq, Xs, ys, k=3):
    dist = ((Xq[:, None, :] - Xs[None, :, :]) ** 2).sum(axis=2)
    ix = np.argpartition(dist, k - 1, axis=1)[:, :k]
    return ys[ix].mean(axis=1)


def holder(x):
    return (np.sin(2.3 * x[:, :1]) + 0.6 * np.cos(3.1 * x[:, 1:2]) +
            0.25 * np.sin(5 * (x[:, :1] + x[:, 1:2])))


def math_verification():
    d, m, N, sigma = 2, 1, 24, 0.35
    Xs = rng.uniform(-1, 1, (N, d)); ys = holder(Xs)
    rows = []
    # Prediction 1: generic full row rank appears at dh >= d+1+m = 4.
    for dh in [3, 4, 5, 8, 16]:
        P, A, V, C = construct(Xs, ys, dh, sigma, 0, 11)
        B = np.vstack((A, V)); rhs = np.vstack((C.T, ys.T))
        residual = np.linalg.norm((B @ P.T - rhs).T) / np.sqrt(N)
        rank = int(np.linalg.matrix_rank(B, tol=1e-10))
        rows.append({'dh': dh, 'joint_rank': rank, 'joint_rows': d+1+m,
                     'rms_joint_residual': float(residual),
                     'mean_prompt_norm': float(np.linalg.norm(P, axis=1).mean())})
    # Prediction 2: ridge regularization trades residual for prompt norm monotonically.
    ridge_rows = []
    for lam in [0, 1e-5, 1e-3, 1e-1, 1.0]:
        P, A, V, C = construct(Xs, ys, 16, sigma, lam, 11)
        B = np.vstack((A, V)); rhs = np.vstack((C.T, ys.T))
        ridge_rows.append({'lambda': lam,
            'rms_residual': float(np.linalg.norm((B @ P.T-rhs).T)/np.sqrt(N)),
            'mean_prompt_norm': float(np.linalg.norm(P,axis=1).mean())})
    # Prediction 3: target coefficient and prompt norm scale as sigma^-2, while
    # normalized kernel agreement remains exact in the full-rank regime.
    sigma_rows = []
    Xq = rng.uniform(-1, 1, (100, d))
    for s in [0.25, 0.35, 0.5, 0.7]:
        P, A, V, C = construct(Xs, ys, 16, s, 0, 11)
        pred, W, logits = predict(Xq, P, A, V, d)
        ref, Wref = nw(Xq, Xs, ys, s)
        sigma_rows.append({'sigma': s, 'mean_prompt_norm': float(np.linalg.norm(P,axis=1).mean()),
                           'weight_max_abs_error': float(np.max(np.abs(W-Wref))),
                           'prediction_mse_to_nw': float(np.mean((pred-ref)**2))})
    return {'rank_prediction': rows, 'ridge_prediction': ridge_rows,
            'bandwidth_prediction': sigma_rows}


def mini_experiment():
    d, m, dh, sigma = 2, 1, 16, 0.38
    r = np.random.default_rng(77)
    Xtest = r.uniform(-1, 1, (1000, d)); ytest = holder(Xtest)
    result = []
    for N in [32, 64, 128, 256]:
        Xs = r.uniform(-1, 1, (N, d)); ys = holder(Xs)
        t0 = time.perf_counter()
        P, A, V, C = construct(Xs, ys, dh, sigma, 0, 19)
        adaptation_ms = (time.perf_counter()-t0)*1000
        pred, W, logits = predict(Xtest, P, A, V, d)
        ref, Wref = nw(Xtest, Xs, ys, sigma)
        result.append({'N': N, 'analytic_mse': float(np.mean((pred[:,0]-ytest[:,0])**2)),
          'nw_mse': float(np.mean((ref[:,0]-ytest[:,0])**2)),
          'knn3_mse': float(np.mean((knn(Xtest,Xs,ys)[:,0]-ytest[:,0])**2)),
          'attention_weight_error': float(np.max(np.abs(W-Wref))),
          'logit_rms_residual_mod_common': float(np.sqrt(np.mean((logits - (Xtest @ Xs.T / sigma**2 - (Xs**2).sum(axis=1)[None,:] / (2*sigma**2)))**2))),
          'prompt_memory_floats': int(P.size), 'adaptation_ms': adaptation_ms})
    return result


def main():
    out = {'seed': SEED, 'math_verification': math_verification(),
           'mini_experiment': mini_experiment()}
    with open('results.json', 'w') as f: json.dump(out, f, indent=2)
    print(json.dumps(out, indent=2))

if __name__ == '__main__':
    main()