import json, time
import numpy as np


def make_dictionary(m, n, d, seed=0):
    rng = np.random.default_rng(seed)
    rows = np.empty(d*n, dtype=np.int64)
    cols = np.repeat(np.arange(n), d)
    weights = np.empty(d*n, dtype=np.float64)
    for j in range(n):
        # Sampling without replacement gives exactly d active coordinates per column.
        rr = rng.choice(m, size=d, replace=False)
        sl = slice(j*d, (j+1)*d)
        rows[sl] = rr
        weights[sl] = rng.normal(size=d)
        weights[sl] /= np.linalg.norm(weights[sl])
    W = np.zeros((m, n), dtype=np.float64)
    W[rows, cols] = weights
    return W, rows, cols, weights


def edge_correlations(R, rows, cols, weights):
    # R is [batch,m]. This is the direct gather/multiply/reduce realization.
    C = np.zeros((R.shape[0], int(cols.max()) + 1), dtype=R.dtype)
    for e in range(len(rows)):
        C[:, cols[e]] += R[:, rows[e]] * weights[e]
    return C


def dense_omp(H, W, k, edge=None, lam=1e-9):
    B, m = H.shape
    n = W.shape[1]
    R = H.copy()
    X = np.zeros((B, n))
    supports = [[] for _ in range(B)]
    for _ in range(k):
        C = W.T @ R.T
        C = C.T
        for b in range(B):
            if supports[b]: C[b, supports[b]] = 0.0
            j = int(np.argmax(np.abs(C[b])))
            supports[b].append(j)
            S = supports[b]
            WS = W[:, S]
            # Ridge only protects the intentionally tiny toy cases from singularity.
            G = WS.T @ WS + lam*np.eye(len(S))
            coef = np.linalg.solve(G, WS.T @ H[b])
            X[b, S] = coef
            R[b] = H[b] - WS @ coef
    return X, R, supports


def edge_omp(H, W, rows, cols, weights, k, lam=1e-9):
    B = H.shape[0]
    n = W.shape[1]
    R = H.copy()
    X = np.zeros((B, n))
    supports = [[] for _ in range(B)]
    for _ in range(k):
        C = edge_correlations(R, rows, cols, weights)
        for b in range(B):
            if supports[b]: C[b, supports[b]] = 0.0
            j = int(np.argmax(np.abs(C[b])))
            supports[b].append(j)
            S = supports[b]
            WS = W[:, S]
            coef = np.linalg.solve(WS.T @ WS + lam*np.eye(len(S)), WS.T @ H[b])
            X[b, S] = coef
            R[b] = H[b] - WS @ coef
    return X, R, supports


def timed(fn, repeats=3):
    vals = []
    for _ in range(repeats):
        t = time.perf_counter(); fn(); vals.append(time.perf_counter()-t)
    return float(np.median(vals))


def main():
    rng = np.random.default_rng(123)
    # Core algebra check, including batched residuals.
    W, rows, cols, weights = make_dictionary(48, 96, 5, 4)
    R = rng.normal(size=(7, 48))
    cd = R @ W
    ce = edge_correlations(R, rows, cols, weights)
    corr_err = float(np.max(np.abs(cd-ce)))

    # Same observations and exact same refit rule for both implementations.
    m, n, d, B, k = 96, 384, 6, 24, 6
    W, rows, cols, weights = make_dictionary(m, n, d, 9)
    true_x = np.zeros((B, n))
    for b in range(B):
        true_x[b, rng.choice(n, k, replace=False)] = rng.normal(size=k)
    H = true_x @ W.T + 0.01*rng.normal(size=(B,m))
    xd, rd, sd = dense_omp(H, W, k)
    xe, re, se = edge_omp(H, W, rows, cols, weights, k)
    dense_mse = float(np.mean(rd**2)); edge_mse = float(np.mean(re**2))
    support_equal = float(np.mean([set(a)==set(b) for a,b in zip(sd,se)]))
    xdiff = float(np.max(np.abs(xd-xe)))

    # Correlation-only scaling/cost comparison. Dense BLAS is the control;
    # edge implementation is intentionally transparent Python gather-reduce.
    scaling = []
    for nn in (128, 256, 512, 1024):
        mm, dd, bb = 128, 8, 32
        WW, rr, cc, ww = make_dictionary(mm, nn, dd, nn+1)
        RR = rng.normal(size=(bb, mm))
        td = timed(lambda: RR @ WW)
        te = timed(lambda: edge_correlations(RR, rr, cc, ww))
        scaling.append({"n":nn, "dense_sec":td, "edge_sec":te,
                        "edge_over_dense":te/td,
                        "dense_flops_proxy":mm*nn*bb,
                        "edge_flops_proxy":dd*nn*bb})
    out = {"max_correlation_abs_error":corr_err,
           "omp_max_coefficient_abs_diff":xdiff,
           "support_set_agreement":support_equal,
           "dense_reconstruction_mse":dense_mse,
           "edge_reconstruction_mse":edge_mse,
           "theoretical_correlation_work_ratio_d_over_m":d/m,
           "scaling":scaling}
    print(json.dumps(out, indent=2))

if __name__ == '__main__': main()