import json, time
import numpy as np
from scipy.linalg import null_space
from scipy.spatial.distance import cdist

SEED = 2128

def phs(X, Y=None, k=2):
    if Y is None: Y = X
    r = cdist(X, Y)
    if k % 2: return r ** k
    return r ** k * np.log(np.maximum(r, 1e-15))

def basis(X):
    return np.column_stack([np.ones(len(X)), X[:, 0], X[:, 1]])

def head(X, w, beta, Xq, k=2, chunk=512):
    out = basis(Xq) @ beta
    for a in range(0, len(Xq), chunk):
        out[a:a+chunk] += phs(Xq[a:a+chunk], X, k) @ w
    return out

def vecchia_precision(X, m=8, ell=None):
    n = len(X)
    D = cdist(X, X); D[np.diag_indices(n)] = np.inf
    if ell is None: ell = 1.5 * np.median(np.min(D, axis=1))
    C = np.exp(-cdist(X, X) / max(ell, 1e-8))
    T = np.zeros((n, n))
    order = np.argsort(X[:, 0] + .173 * X[:, 1])
    for pos, i in enumerate(order):
        if pos == 0:
            T[i, i] = 1.0
            continue
        earlier = order[:pos]
        nn = earlier[np.argsort(np.linalg.norm(X[earlier] - X[i], axis=1))[:m]]
        S = C[np.ix_(nn, nn)] + 1e-8*np.eye(len(nn))
        c = C[i, nn]
        a = np.linalg.solve(S, c)
        var = max(C[i, i] - a @ c, 1e-8)
        T[i, i] = 1.0 / np.sqrt(var)
        T[i, nn] = -a / np.sqrt(var)
    return T

def projected_pcg(M, Q, rhs, pre=None, tol=1e-8, maxit=300):
    A = lambda z: Q.T @ (M @ (Q @ z))
    z = np.zeros_like(rhs)
    r = rhs - A(z)
    r0 = np.linalg.norm(r)
    if pre is None: solve = lambda x: x
    else:
        # P_vecchia is the approximate inverse/preconditioning operator; apply it directly
        G = Q.T @ pre @ Q
        G = (G + G.T)/2 + 1e-8*np.eye(G.shape[0])
        solve = lambda x: G @ x
    q = solve(r); p = q.copy(); rz = r @ q; it = 0
    for it in range(1, maxit+1):
        Ap = A(p); den = p @ Ap
        if abs(den) < 1e-14: break
        alpha = rz / den; z += alpha*p; r -= alpha*Ap
        if np.linalg.norm(r) <= tol * max(r0, 1e-15): break
        q = solve(r); rz_new = r @ q
        p = q + (rz_new/rz)*p; rz = rz_new
    return z, it, np.linalg.norm(r)/max(r0,1e-15)

def main():
    rng = np.random.default_rng(SEED)
    n, nq = 180, 300
    X = rng.uniform(-1,1,(n,2)); Xq = rng.uniform(-1,1,(nq,2))
    y = np.sin(3*X[:,0]) * np.cos(2*X[:,1]) + .2*X[:,0]**2
    B = basis(X); Q = null_space(B.T)
    M = phs(X, k=2)
    # Exact augmented solution is the reference interpolation.
    aug = np.block([[M, B], [B.T, np.zeros((3,3))]])
    sol = np.linalg.solve(aug, np.r_[y, np.zeros(3)])
    w0, beta0 = sol[:n], sol[n:]
    pred = head(X, w0, beta0, Xq)
    ref_err = np.max(np.abs(head(X,w0,beta0,X)-y))
    constraint = np.linalg.norm(B.T @ w0)

    # Prediction 1: geometric scaling M(rho) projected scales as rho^2 for k=2.
    scales = [.5, 1., 2., 4.]
    scale_ratios = []
    for rho in scales:
        Mr = phs(rho*X, k=2)
        scale_ratios.append(np.linalg.norm(Q.T@Mr@Q) / np.linalg.norm(Q.T@M@Q))
    predicted = [rho*rho for rho in scales]

    # Prediction 2: projected polynomial terms are annihilated exactly.
    D2 = cdist(X,X)**2
    proj_poly = np.linalg.norm(Q.T @ D2 @ Q) / max(np.linalg.norm(Q.T@M@Q),1e-15)
    # Random coefficients expose constraint preservation.
    wr = Q @ rng.normal(size=n-3)
    constraint_random = np.linalg.norm(B.T @ wr)

    # Prediction 3: Vecchia quality improves as neighbor set grows; report PCG residual/iters.
    rhs = Q.T @ y
    rows = []
    z, it, rr = projected_pcg(M, Q, rhs, pre=None)
    rows.append({'m':'none', 'iterations':it, 'relative_residual':float(rr),
                 'constraint':float(np.linalg.norm(B.T@(Q@z)))})
    for m in [2,4,8,16,32]:
        T = vecchia_precision(X, m=m)
        P = T @ T.T
        z, it, rr = projected_pcg(M,Q,rhs,pre=P)
        rows.append({'m':m, 'iterations':it, 'relative_residual':float(rr),
                     'constraint':float(np.linalg.norm(B.T@(Q@z)))})

    # Dense versus chunked head evaluation (same result, lower peak interaction memory).
    t0=time.perf_counter(); dense = basis(Xq)@beta0 + phs(Xq,X)@w0; td=time.perf_counter()-t0
    t0=time.perf_counter(); chunked=head(X,w0,beta0,Xq,chunk=32); tc=time.perf_counter()-t0
    result = {'seed':SEED, 'n':n, 'queries':nq,
      'math_checks': {
        'scale_rho':scales, 'predicted_rho2':predicted, 'observed_projected_norm_ratios':scale_ratios,
        'scale_max_abs_error':float(max(abs(a-b) for a,b in zip(scale_ratios,predicted))),
        'projected_squared_distance_ratio':float(proj_poly),
        'predicted_projected_polynomial_ratio':0.0,
        'random_nullspace_constraint_norm':float(constraint_random),
        'exact_interpolation_max_error':float(ref_err), 'exact_solution_constraint_norm':float(constraint)},
      'pcg_vecchia':rows,
      'evaluation': {'dense_seconds':td, 'chunked_seconds':tc,
                     'chunked_vs_dense_max_error':float(np.max(np.abs(chunked-dense))),
                     'reference_query_rmse':float(np.sqrt(np.mean((pred-(np.sin(3*Xq[:,0])*np.cos(2*Xq[:,1])+.2*Xq[:,0]**2))**2)))}}
    print(json.dumps(result, indent=2))

if __name__ == '__main__': main()