#!/usr/bin/env julia
# Offline benchmark of the Pac-Man problem in Julia.
#
# 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
#
#     julia bench/julia/bench_pacman.jl
#
# from anywhere; it writes bench-pacman-julia.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.  Uses only the standard library.

using LinearAlgebra
using Dates
using Printf

const KIND = "both"
const Z0 = 1.6 + 1.3im

function exact(x, y)
    th = atan(y, x)
    th < -1e-13 && (th += 2pi)
    r = hypot(x, y)
    sing = r^(2 / 3) * sin(2 * th / 3)
    smooth = 0.5 * log((x - real(Z0))^2 + (y - imag(Z0))^2)
    KIND == "sing" && return sing
    KIND == "smooth" && return smooth
    return sing + smooth
end
exact(x::AbstractVector, y::AbstractVector) = [exact(x[k], y[k]) for k in eachindex(x)]

"""The driver's error lattice: spacing 1/8, margins from the boundary."""
function lattice()
    px = Float64[]
    py = Float64[]
    for x in -1.0:0.125:1.0, y in -1.0:0.125:1.0
        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)
        keep || continue
        push!(px, x)
        push!(py, y)
    end
    return px, py
end

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

function pac_basis(x, y, M, K)
    np = length(x)
    A = zeros(np, 2M + 1 + K)
    z = complex.(x, y)
    col = 1
    p = ones(ComplexF64, np)
    for mdeg in 0:M
        A[:, col] = real.(p)
        col += 1
        if mdeg > 0
            A[:, col] = imag.(p)
            col += 1
        end
        p .*= z
    end
    th = [t < -1e-13 ? t + 2pi : t for t in atan.(y, x)]
    r = abs.(z)
    k, q = 0, 0
    while q < K
        k += 1
        k % 3 == 0 && continue
        q += 1
        A[:, col] = r .^ (2k / 3) .* sin.(2k .* th ./ 3)
        col += 1
    end
    return A
end

function pac_bpoints(ntot)
    ne = max(6, round(Int, ntot / 6.7))
    na = max(12, ntot - 2ne)
    t = ((1:ne) .- 0.5) ./ ne
    c = t .^ 2
    s = 1.5pi .* ((1:na) .- 0.5) ./ na
    bx = vcat(c, cos.(s), zeros(ne))
    by = vcat(zeros(ne), sin.(s), -c)
    return bx, by
end

function mps_solve(n, px, py)
    nb = 3n + 1
    bx, by = pac_bpoints(3nb)
    A = pac_basis(bx, by, n, n)
    c = A \ exact(bx, by)          # pivoted QR least squares
    return pac_basis(px, py, n, n) * c
end

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

"""Gauss-Legendre nodes and weights on [-1, 1], by Newton on P_p."""
function gauss(p)
    x = [cos(pi * (k - 0.25) / (p + 0.5)) for k in 1:p]
    dp = ones(p)
    for _ in 1:50
        p0 = ones(p)
        p1 = copy(x)
        for k in 2:p
            p2 = ((2k - 1) .* x .* p1 .- (k - 1) .* p0) ./ k
            p0 = p1
            p1 = p2
        end
        dp = p .* (x .* p1 .- p0) ./ (x .^ 2 .- 1)
        dx = p1 ./ dp
        x .-= dx
        maximum(abs, dx) < 1e-15 && break
    end
    w = 2 ./ ((1 .- x .^ 2) .* dp .^ 2)
    return x, w
end

function graded_breaks(T, d, base)
    half = T / 2
    left = vcat(0.0, half .* 2.0 .^ (-d:0))
    bp0 = vcat(left, T .- reverse(left[1:end-1]))
    bp = [bp0[1]]
    for k in 2:length(bp0)
        a, b = bp0[k-1], bp0[k]
        ns = max(1, ceil(Int, (b - a) / base))
        for q in 1:ns
            push!(bp, a + (b - a) * q / ns)
        end
    end
    return bp
end

function panel_nodes(bp, gx, gw)
    t = Float64[]
    w = Float64[]
    for k in 1:length(bp)-1
        a, b = bp[k], bp[k+1]
        append!(t, (a + b) / 2 .+ (b - a) / 2 .* gx)
        append!(w, (b - a) / 2 .* gw)
    end
    return t, w
end

function pac_panels(d)
    gx, gw = gauss(12)
    be = graded_breaks(1.0, d, 0.35)
    ba = graded_breaks(1.5pi, d, 0.35)
    t1, w1 = panel_nodes(be, gx, gw)
    t2, w2 = panel_nodes(ba, gx, gw)
    t3, w3 = panel_nodes(be, gx, gw)
    n1, n3 = length(t1), length(t3)
    qx = vcat(t1, cos.(t2), zeros(n3))
    qy = vcat(zeros(n1), sin.(t2), t3 .- 1)
    qnx = vcat(zeros(n1), cos.(t2), ones(n3))
    qny = vcat(-ones(n1), sin.(t2), zeros(n3))
    qw = vcat(w1, w2, w3)
    qkap = vcat(zeros(n1), ones(length(t2)), zeros(n3))
    return qx, qy, qnx, qny, qw, qkap
end

function bie_solve(d, px, py)
    qx, qy, qnx, qny, qw, qkap = pac_panels(d)
    n = length(qx)
    A = zeros(n, n)
    for j in 1:n, i in 1:n
        if i == j
            A[i, j] = -0.5 - qkap[j] / (4pi) * qw[j]
        else
            dx = qx[i] - qx[j]
            dy = qy[i] - qy[j]
            A[i, j] = (qnx[j] * dx + qny[j] * dy) / (2pi * (dx^2 + dy^2)) * qw[j]
        end
    end
    sigma = A \ exact(qx, qy)
    v = zeros(length(px))
    for k in eachindex(px)
        s = 0.0
        for j in 1:n
            dx = px[k] - qx[j]
            dy = py[k] - qy[j]
            s += (qnx[j] * dx + qny[j] * dy) / (2pi * (dx^2 + dy^2)) * qw[j] * sigma[j]
        end
        v[k] = s
    end
    return v
end

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

function sweep(solve, ns, px, py, ve)
    E = Float64[]
    T = Float64[]
    solve(ns[1], px[1:1], py[1:1])   # warm-up, which in Julia is compilation
    for n in ns
        el = @elapsed v = solve(n, px, py)
        # a second timed run always: stray compilation and GC pauses
        # otherwise land in single measurements
        el = min(el, @elapsed v = solve(n, px, py))
        rep = 0
        while el < 0.005 && rep < 4
            rep += 1
            el = min(el, @elapsed v = solve(n, px, py))
        end
        push!(E, maximum(abs.(v .- ve)))
        push!(T, el * 1000)
    end
    return E, T
end

function json_num(v)
    isfinite(v) || return "null"
    return string(v)
end
json_arr(v) = "[" * join(json_num.(float.(v)), ",") * "]"

function main()
    px, py = lattice()
    ve = exact(px, py)
    entries = String[]
    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)
        println(name)
        for k in eachindex(ns)
            @printf("  n=%3d  err=%.3e  %8.2f ms\n", ns[k], E[k], T[k])
        end
        push!(entries,
            """{"name":"$(name)","slot":"$(slot)","ns":$(json_arr(ns)),"E":$(json_arr(E)),"T":$(json_arr(T))}""")
    end
    label = "Julia $(VERSION.major).$(VERSION.minor)"
    open("bench-pacman-julia.json", "w") do f
        write(f,
            """{"schema":"laplace-dirichlet-bench-1","label":"$(label)","kind":"$(KIND)","date":"$(Dates.today())","methods":[$(join(entries, ","))]}""")
    end
    println("wrote bench-pacman-julia.json")
end

main()