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

SEED = 1303

def block_matrix(n_left, h, delta=1.0):
    """Uniformly sampled discretization of 1/(s+r), one dyadic r block."""
    s = np.arange(1, n_left + 1, dtype=float)[:, None] + delta - 1.0
    r = h + np.arange(max(1, int(h)), dtype=float) + delta
    return 1.0 / (s + r[None, :])


def svd_observations():
    rows = []
    for h in [4, 8, 16, 32, 64]:
        H = block_matrix(4 * h, h)
        sv = np.linalg.svd(H / np.linalg.norm(H, 2), compute_uv=False)
        # The paper predicts s_(N+1) <= 4^-N. This normalized check is
        # conservative: normalization removes irrelevant scale factors.
        bounds = [{"N": n, "observed": float(sv[n]), "predicted_bound": 4.0 ** (-n),
                   "satisfies": bool(sv[n] <= 4.0 ** (-n) * (1 + 1e-10))}
                  for n in range(min(5, len(sv)))]
        rows.append({"h": h, "singular_values": [float(x) for x in sv[:6]],
                     "geometric_bound_check": bounds,
                     "ratios": [float(sv[i+1] / sv[i]) for i in range(min(4, len(sv)-1))]})
    return rows


def rank_scaling():
    rows = []
    for h in [8, 16, 32, 64]:
        sv = np.linalg.svd(block_matrix(4*h, h), compute_uv=False)
        sv /= sv[0]
        for eps in [1e-2, 1e-3, 1e-4]:
            observed = int(np.sum(sv > eps))
            predicted = int(math.ceil(math.log(1 / eps, 4)))
            rows.append({"h": h, "epsilon": eps, "observed_rank": observed,
                         "predicted_rank_bound": predicted,
                         "bound_holds": observed <= predicted})
    return rows


def application_scaling():
    """Compare cached SVD application with dense multiplication over sizes."""
    rng = np.random.default_rng(SEED)
    rows = []
    for h in [64, 128, 256, 512]:
        nl, nr, d, rank = 4*h, h, 16, 4
        H = block_matrix(nl, h)
        U, sv, VT = np.linalg.svd(H, full_matrices=False)
        A, B = U[:, :rank] * sv[:rank][None, :], VT[:rank, :].T
        V = rng.standard_normal((nr, d))
        dense = H @ V
        approx = A @ (B.T @ V)
        err = np.linalg.norm(dense - approx) / np.linalg.norm(dense)
        # Small fixed repetitions; warmup avoids first-call allocation bias.
        H @ V; A @ (B.T @ V)
        reps = 8 if h <= 256 else 4
        t0 = time.perf_counter()
        for _ in range(reps): H @ V
        td = (time.perf_counter() - t0) / reps
        t0 = time.perf_counter()
        for _ in range(reps): A @ (B.T @ V)
        tf = (time.perf_counter() - t0) / reps
        dense_work = nl * nr * d
        fact_work = (nl + nr) * rank * d
        rows.append({"h": h, "n_left": nl, "n_right": nr, "rank": rank,
                     "relative_output_error": float(err), "dense_ms": td*1000,
                     "factorized_ms": tf*1000, "measured_speedup": td/max(tf,1e-12),
                     "ideal_work_ratio": dense_work / fact_work})
    return rows


def main():
    out = {"seed": SEED, "singular_decay": svd_observations(),
           "rank_scaling": rank_scaling(), "application_scaling": application_scaling()}
    Path("results.json").write_text(json.dumps(out, indent=2))
    print(json.dumps(out, indent=2))

if __name__ == "__main__": main()