"""Benchmarks for 'Matrix Multiplication, From Definition to Cache Lines'.

Protocol: the two slow loops run once; everything else is best-of-5.
Single-threaded BLAS except the explicit multithread runs at the end.
The timed functions are EXACTLY the code shown in the essay.
"""
import os
os.environ.setdefault('OPENBLAS_NUM_THREADS', '1')
os.environ.setdefault('OMP_NUM_THREADS', '1')

import json
import subprocess
import sys
import time

import numpy as np

R = {}


def best_of(fn, repeat=5):
    best = float('inf')
    for _ in range(repeat):
        t0 = time.perf_counter()
        fn()
        best = min(best, time.perf_counter() - t0)
    return best


# ── the essay's naive matmul, verbatim ──────────────────────────────────
def matmul_naive(A, B, n):
    C = [[0.0] * n for _ in range(n)]
    for i in range(n):
        for j in range(n):
            acc = 0.0
            for p in range(n):
                acc += A[i][p] * B[p][j]
            C[i][j] = acc
    return C


# ── same loops, NumPy arrays with scalar indexing ───────────────────────
def matmul_numpy_scalar(A, B, n):
    C = np.zeros((n, n))
    for i in range(n):
        for j in range(n):
            acc = 0.0
            for p in range(n):
                acc += A[i, p] * B[p, j]
            C[i, j] = acc
    return C


rng = np.random.default_rng(0)
n = 512
A = rng.standard_normal((n, n))
B = rng.standard_normal((n, n))

t0 = time.perf_counter()
C_list = matmul_naive(A.tolist(), B.tolist(), n)
R['naive_python_512'] = time.perf_counter() - t0

t0 = time.perf_counter()
C_scalar = matmul_numpy_scalar(A, B, n)
R['numpy_scalar_512'] = time.perf_counter() - t0

# ── A @ B across sizes, single thread ───────────────────────────────────
R['blas_single'] = {}
for m in [256, 512, 1024, 2048, 4096]:
    X = rng.standard_normal((m, m))
    Y = rng.standard_normal((m, m))
    X @ Y  # warm up
    t = best_of(lambda: X @ Y)
    R['blas_single'][m] = {'sec': t, 'gflops': 2 * m**3 / t / 1e9}

# ── row-major vs column-major traversal (cache-line demo) ───────────────
m = 8192
M = rng.standard_normal((m, m))


def sum_by_rows():
    s = 0.0
    for i in range(m):
        s += M[i, :].sum()
    return s


def sum_by_cols():
    s = 0.0
    for j in range(m):
        s += M[:, j].sum()
    return s


sum_by_rows(); sum_by_cols()  # warm up
R['traverse_rows_8192'] = best_of(sum_by_rows)
R['traverse_cols_8192'] = best_of(sum_by_cols)

# ── multithreaded A @ B at 4096: OpenBLAS default vs 4 threads ──────────
MT_CODE = (
    "import time, numpy as np\n"
    "rng = np.random.default_rng(0)\n"
    "X = rng.standard_normal((4096, 4096)); Y = rng.standard_normal((4096, 4096))\n"
    "X @ Y\n"
    "best = float('inf')\n"
    "for _ in range(5):\n"
    "    t0 = time.perf_counter(); X @ Y; best = min(best, time.perf_counter() - t0)\n"
    "print(best)\n"
)
R['blas_multi_4096'] = {}
for label, extra in [('default', {}), ('threads4', {'OPENBLAS_NUM_THREADS': '4'})]:
    env = {k: v for k, v in os.environ.items()
           if k not in ('OPENBLAS_NUM_THREADS', 'OMP_NUM_THREADS')}
    env.update(extra)
    out = subprocess.run([sys.executable, '-c', MT_CODE], env=env,
                         capture_output=True, text=True, check=True)
    t = float(out.stdout.strip())
    R['blas_multi_4096'][label] = {'sec': t, 'gflops': 2 * 4096**3 / t / 1e9}

# ── correctness cross-check ─────────────────────────────────────────────
ref = A @ B
R['naive_max_abs_err'] = float(np.max(np.abs(np.array(C_list) - ref)))
R['numpy_scalar_max_abs_err'] = float(np.max(np.abs(C_scalar - ref)))

with open(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'matmul_bench2.json'), 'w') as f:
    json.dump(R, f, indent=1)
print(json.dumps(R, indent=1))