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

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.
"""

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

import numpy as np

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):
    z = x + 1j * y
    cols = [np.ones_like(x)]
    p = z.copy()
    for _ in range(M):
        cols.append(p.real)
        cols.append(p.imag)
        p = 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
            cols.append(r ** (2.0 * k / 3.0) * np.sin(2.0 * k * th / 3.0))
    return np.column_stack(cols)


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)
    # QR-based least squares, the closest match to MATLAB's backslash
    c, *_ = np.linalg.lstsq(A, exact(bx, by), rcond=None)
    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
    DX = qx[:, None] - qx[None, :]
    DY = qy[:, None] - qy[None, :]
    R2 = DX**2 + DY**2
    np.fill_diagonal(R2, 1.0)
    K = (qnx[None, :] * DX + qny[None, :] * DY) / (2 * np.pi * R2)
    np.fill_diagonal(K, -qkap / (4 * np.pi))
    A = -0.5 * np.eye(n) + K * qw[None, :]
    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):
    E, T = [], []
    solve(ns[0], px[:1], py[:1])  # warm-up
    for n in ns:
        t0 = time.perf_counter()
        v = solve(n, px, py)
        el = time.perf_counter() - t0
        rep = 0
        while el < 0.005 and rep < 4:
            rep += 1
            t0 = time.perf_counter()
            v = solve(n, px, py)
            el = min(el, time.perf_counter() - t0)
        E.append(float(np.max(np.abs(v - ve))))
        T.append(el * 1000.0)
    return E, T


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")

    out = {
        "schema": "laplace-dirichlet-bench-1",
        "label": f"Python {sys.version_info.major}.{sys.version_info.minor} + NumPy {np.__version__}",
        "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()