#!/usr/bin/env python3
"""Offline benchmark of the Pac-Man problem in Python (NumPy and SciPy).

A port of the two best methods from the in-browser benchmark, the harmonic
basis and the graded-panel double layer, on the same problem, the same error
lattice, and the same timing protocol as the chapter driver.  Run

    python3 bench/python/bench_pacman.py

from anywhere; it writes bench-pacman-python.json into the current directory.
Load that file from the benchmark chapter's "load offline results" button.
Because it ran on the same machine as your browser, the timings share a clock
and can share an axis.

The point of the comparison is the language and its libraries, not the
author's fluency in each, so each port is written the way its own language
wants to be written.  For NumPy that means whole-array expressions, with two
details that matter more than they look:

  * the kernel matrix is assembled in place.  Written as the obvious sequence
    of broadcast expressions it allocates five n-by-n temporaries and spends
    more time on memory traffic than the dense solve itself spends on
    arithmetic; reusing two buffers cuts that stage by nearly half.  Writing
    it as a scalar double loop, which is the fastest shape in MATLAB and
    Julia, would be catastrophic here, and that asymmetry is the honest
    reason the ports do not look alike.

  * the least-squares fit goes through a QR driver.  numpy.linalg.lstsq is an
    SVD (LAPACK gelsd), while MATLAB's backslash and Julia's backslash are
    both pivoted QR, so using lstsq would have Python solving a harder
    problem than the others and calling it the same benchmark.
"""

import json
import platform
import sys
import time
from datetime import date

import numpy as np
from scipy.linalg import lstsq as qr_lstsq

KIND = "both"
Z0 = 1.6 + 1.3j


def exact(x, y, kind=KIND):
    """lp_lexact: the exact solution, singular and/or smooth part."""
    th = np.arctan2(y, x)
    th = np.where(th < -1e-13, th + 2 * np.pi, th)
    r = np.hypot(x, y)
    sing = r ** (2.0 / 3.0) * np.sin(2.0 * th / 3.0)
    smooth = 0.5 * np.log((x - Z0.real) ** 2 + (y - Z0.imag) ** 2)
    if kind == "sing":
        return sing
    if kind == "smooth":
        return smooth
    return sing + smooth


def lattice():
    """The driver's error lattice: spacing 1/8, margins from the boundary."""
    s = np.arange(-8, 9) / 8.0
    X, Y = np.meshgrid(s, s)
    R2 = X**2 + Y**2
    keep = (
        (R2 <= 0.7225 + 1e-9)
        & (R2 > 1e-12)
        & ~((X > 1e-12) & (Y < -1e-12))
        & ~((X > 1e-12) & (Y < 0.25 - 1e-9))
        & ~((Y < -1e-12) & (X > -0.25 + 1e-9))
    )
    return X[keep], Y[keep]


# ------------------------------------------- the method of particular sols --


def pac_basis(x, y, M, K):
    """Re z^m and Im z^m up to degree M, then K corner functions.

    Written into a preallocated Fortran-ordered array: LAPACK wants columns
    contiguous, and building a list of columns for column_stack would copy the
    whole matrix once more than it has to.
    """
    z = x + 1j * y
    A = np.empty((x.size, 2 * M + 1 + K), order="F")
    A[:, 0] = 1.0
    col = 1
    p = z.copy()
    for _ in range(M):
        A[:, col] = p.real
        A[:, col + 1] = p.imag
        col += 2
        p *= z
    th = np.arctan2(y, x)
    th = np.where(th < -1e-13, th + 2 * np.pi, th)
    r = np.abs(z)
    k, q = 0, 0
    while q < K:
        k += 1
        if k % 3 != 0:
            q += 1
            A[:, col] = r ** (2.0 * k / 3.0) * np.sin(2.0 * k * th / 3.0)
            col += 1
    return A


def pac_bpoints(ntot):
    ne = max(6, round(ntot / 6.7))
    na = max(12, ntot - 2 * ne)
    t = (np.arange(1, ne + 1) - 0.5) / ne
    c = t**2
    s = 1.5 * np.pi * (np.arange(1, na + 1) - 0.5) / na
    bx = np.concatenate([c, np.cos(s), np.zeros(ne)])
    by = np.concatenate([np.zeros(ne), np.sin(s), -c])
    return bx, by


def mps_solve(n, px, py):
    nb = 3 * n + 1
    bx, by = pac_bpoints(3 * nb)
    A = pac_basis(bx, by, n, n)
    # gelsy is pivoted QR: what MATLAB's and Julia's backslash do on a
    # rectangular system.  The default lstsq driver is an SVD and does more
    # work than either.
    c = qr_lstsq(A, exact(bx, by), lapack_driver="gelsy")[0]
    return pac_basis(px, py, n, n) @ c


# --------------------------------------------- boundary integral, panels ----


def graded_breaks(T, d, base):
    half = T / 2.0
    left = np.concatenate([[0.0], half * 2.0 ** np.arange(-d, 1)])
    bp0 = np.concatenate([left, T - left[-2::-1]])
    bp = [bp0[0]]
    for a, b in zip(bp0[:-1], bp0[1:]):
        ns = max(1, int(np.ceil((b - a) / base)))
        for q in range(1, ns + 1):
            bp.append(a + (b - a) * q / ns)
    return np.array(bp)


def panel_nodes(bp, gx, gw):
    a = bp[:-1]
    b = bp[1:]
    t = ((a + b) / 2)[:, None] + ((b - a) / 2)[:, None] * gx[None, :]
    w = ((b - a) / 2)[:, None] * gw[None, :]
    return t.ravel(), w.ravel()


def pac_panels(d):
    gx, gw = np.polynomial.legendre.leggauss(12)
    be = graded_breaks(1.0, d, 0.35)
    ba = graded_breaks(1.5 * np.pi, d, 0.35)
    t1, w1 = panel_nodes(be, gx, gw)
    t2, w2 = panel_nodes(ba, gx, gw)
    t3, w3 = panel_nodes(be, gx, gw)
    z1 = np.zeros_like(t1)
    z3 = np.zeros_like(t3)
    qx = np.concatenate([t1, np.cos(t2), z3])
    qy = np.concatenate([z1, np.sin(t2), t3 - 1.0])
    qnx = np.concatenate([z1, np.cos(t2), np.ones_like(t3)])
    qny = np.concatenate([-np.ones_like(t1), np.sin(t2), z3])
    qw = np.concatenate([w1, w2, w3])
    qkap = np.concatenate([z1, np.ones_like(t2), z3])
    return qx, qy, qnx, qny, qw, qkap


def bie_solve(d, px, py):
    qx, qy, qnx, qny, qw, qkap = pac_panels(d)
    n = qx.size

    # A is built in two buffers and finished in place: every operation below
    # writes into memory that already exists.  The obvious version, one
    # broadcast expression per line, needs five n-by-n arrays and is about
    # 1.6 times slower for that reason alone.
    A = qx[:, None] - qx[None, :]              # dx
    B = qy[:, None] - qy[None, :]              # dy
    R2 = A * A
    R2 += B * B
    np.fill_diagonal(R2, 1.0)                  # keep 0/0 off the diagonal
    A *= qnx[None, :]
    B *= qny[None, :]
    A += B
    A /= R2
    A *= qw[None, :] / (2 * np.pi)
    # the diagonal is the smooth limit of the kernel, plus the jump relation
    diag = np.arange(n)
    A[diag, diag] = -0.5 - qkap * qw / (4 * np.pi)

    sigma = np.linalg.solve(A, exact(qx, qy))

    EX = px[:, None] - qx[None, :]
    EY = py[:, None] - qy[None, :]
    KE = (qnx[None, :] * EX + qny[None, :] * EY) / (2 * np.pi * (EX**2 + EY**2))
    return (KE * qw[None, :]) @ sigma


# ------------------------------------------------------------- the sweep ----


def sweep(solve, ns, px, py, ve):
    """Warm up, then time each resolution the way the chapter driver does.

    Keep the fastest of several runs. Three runs at least whenever a solve
    is cheap enough to allow it, many more when it is very cheap, and one
    when it is expensive: the budget below caps the wall clock per point.
    The minimum over runs is what makes these numbers comparable at all,
    here and in the offline scripts, which apply the same rule. A single
    measurement of a fifty-millisecond solve reports whatever else the
    machine happened to be doing at the time.
    """
    E, T = [], []
    solve(ns[0], px[:1], py[:1])  # warm-up
    for n in ns:
        el, tot, rep = float("inf"), 0.0, 0
        while rep == 0 or (rep < 50 and tot < 0.6 and (rep < 3 or tot < 0.1)):
            t0 = time.perf_counter()
            v = solve(n, px, py)
            dt = time.perf_counter() - t0
            el = min(el, dt)
            tot += dt
            rep += 1
        E.append(float(np.max(np.abs(v - ve))))
        T.append(el * 1000.0)
    return E, T


def threads():
    """How many threads the BLAS behind NumPy will use, if we can find out."""
    try:
        from threadpoolctl import threadpool_info

        n = max((p["num_threads"] for p in threadpool_info()), default=None)
        if n:
            return n
    except ImportError:
        pass
    return None


def main():
    px, py = lattice()
    ve = exact(px, py)
    methods = []
    for name, slot, solve, ns in [
        ("harmonic basis, least squares", "mps", mps_solve, [4, 8, 16, 32, 48]),
        ("double layer, graded panels", "bie", bie_solve, [4, 8, 12, 16, 20, 24]),
    ]:
        E, T = sweep(solve, ns, px, py, ve)
        methods.append({"name": name, "slot": slot, "ns": ns, "E": E, "T": T})
        print(name)
        for n, e, t in zip(ns, E, T):
            print(f"  n={n:3d}  err={e:.3e}  {t:8.2f} ms")

    nt = threads()
    label = f"Python {sys.version_info.major}.{sys.version_info.minor} + NumPy {np.__version__}"
    if nt:
        label += f", {nt} threads"
    out = {
        "schema": "laplace-dirichlet-bench-1",
        "label": label,
        "kind": KIND,
        "date": date.today().isoformat(),
        "machine": platform.machine(),
        "methods": methods,
    }
    with open("bench-pacman-python.json", "w") as f:
        json.dump(out, f)
    print("wrote bench-pacman-python.json")


if __name__ == "__main__":
    main()