"""TopEdu project 10: The race between solvers

The test problem: -laplacian(u) = f on the unit square, u = 0 on the edge, with f chosen so the exact answer is
u = g(x) g(y), g(s) = s (1 - s) e^s. Knowing the answer lets you measure the real error of every method.
(The lab uses u = sin(pi x) sin(pi y). Try it here, in problem(): CG then finishes in one iteration. Why?)

This file races three ways of solving the same grid equations (plain second-order finite differences):
  Jacobi:               each point becomes the average the equation wants. Simple, and very slow on big grids.
  Conjugate gradients:  the standard method for this kind of system; its iterations grow with the grid.
  Multigrid (V-cycle):  fixes smooth errors on coarser grids; its iterations should stay flat.

Stage 1: run it. It checks that the error falls like h^2 (rate near 2). The CHECK line should say PASS.
Stage 2: time everything as the grid grows, fit the power laws, and find your own machine's break-even point.
The online lab uses higher-order SBP-SAT methods; sbp-sat-suite has a mentor-run multigrid and GPU version.

Run:  python 10_solver_race.py  [--show]
Needs: numpy and matplotlib
"""
import os
import sys
import time

import numpy as np
import matplotlib

if "--show" not in sys.argv:
    matplotlib.use("Agg")
import matplotlib.pyplot as plt

OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "output")
os.makedirs(OUT, exist_ok=True)


def problem(n):
    """Grid spacing h, the forcing f and the exact answer u on an n by n grid."""
    x = np.linspace(0, 1, n)
    X, Y = np.meshgrid(x, x, indexing="ij")
    g = lambda s: s * (1 - s) * np.exp(s)
    minus_g2 = lambda s: s * (s + 3) * np.exp(s)          # -g''(s)
    return x[1] - x[0], minus_g2(X) * g(Y) + g(X) * minus_g2(Y), g(X) * g(Y)


def apply(u, h):
    """-laplacian(u) at the interior points (the edge stays 0)."""
    out = np.zeros_like(u)
    out[1:-1, 1:-1] = (4 * u[1:-1, 1:-1] - u[2:, 1:-1] - u[:-2, 1:-1] - u[1:-1, 2:] - u[1:-1, :-2]) / h ** 2
    return out


def jacobi(f, h, tolerance=1e-8, limit=200_000):
    u = np.zeros_like(f)
    stop = tolerance * np.linalg.norm(f[1:-1, 1:-1])
    for iteration in range(1, limit + 1):
        u[1:-1, 1:-1] = (u[2:, 1:-1] + u[:-2, 1:-1] + u[1:-1, 2:] + u[1:-1, :-2] + h * h * f[1:-1, 1:-1]) / 4
        if iteration % 50 == 0 and np.linalg.norm((f - apply(u, h))[1:-1, 1:-1]) < stop:
            return u, iteration
    return u, limit


def conjugate_gradients(f, h, tolerance=1e-8, precondition=None):
    """CG, or multigrid-preconditioned CG when precondition is given (one V-cycle per iteration)."""
    u = np.zeros_like(f)
    r = f.copy(); r[0, :] = r[-1, :] = r[:, 0] = r[:, -1] = 0
    z = precondition(r, h) if precondition else r
    p = z.copy(); rz = np.vdot(r, z)
    stop = tolerance * np.linalg.norm(r)
    for iteration in range(1, 100_000):
        Ap = apply(p, h)
        alpha = rz / np.vdot(p, Ap)
        u += alpha * p; r -= alpha * Ap
        if np.linalg.norm(r) < stop:
            return u, iteration
        z = precondition(r, h) if precondition else r
        rz, rz_old = np.vdot(r, z), rz
        p = z + (rz / rz_old) * p
    return u, iteration


def v_cycle(r, h, sweeps=2):
    """One multigrid V-cycle for -laplacian(e) = r, on a grid of 2^k + 1 points. Symmetric, so it can precondition CG:
    weighted-Jacobi smoothing, restrict the leftover to a grid half as fine, solve there, bring the fix back."""
    n = r.shape[0]
    e = np.zeros_like(r)
    if n <= 3:
        e[1, 1] = r[1, 1] * h * h / 4
        return e
    for _ in range(sweeps):
        e[1:-1, 1:-1] += 0.8 * (h * h / 4) * (r - apply(e, h))[1:-1, 1:-1]
    left = r - apply(e, h)
    coarse = np.zeros(((n + 1) // 2, (n + 1) // 2))          # full-weighting restriction
    c = left
    coarse[1:-1, 1:-1] = (4 * c[2:-2:2, 2:-2:2] + 2 * (c[1:-3:2, 2:-2:2] + c[3:-1:2, 2:-2:2] + c[2:-2:2, 1:-3:2] + c[2:-2:2, 3:-1:2])
                          + c[1:-3:2, 1:-3:2] + c[3:-1:2, 1:-3:2] + c[1:-3:2, 3:-1:2] + c[3:-1:2, 3:-1:2]) / 16
    fix = v_cycle(coarse, 2 * h, sweeps)
    fine = np.zeros_like(r)                                   # bilinear interpolation back up
    fine[::2, ::2] = fix
    fine[1::2, ::2] = (fix[:-1, :] + fix[1:, :]) / 2
    fine[:, 1::2] = (fine[:, :-1:2] + fine[:, 2::2]) / 2
    e += fine
    for _ in range(sweeps):
        e[1:-1, 1:-1] += 0.8 * (h * h / 4) * (r - apply(e, h))[1:-1, 1:-1]
    return e


def race(n, methods):
    h, f, exact = problem(n)
    results = {}
    for name, run in methods.items():
        start = time.perf_counter()
        u, iterations = run(f, h)
        results[name] = (iterations, time.perf_counter() - start, np.abs(u - exact).max())
    return results


def main():
    methods = {
        "CG": conjugate_gradients,
        "multigrid CG": lambda f, h: conjugate_gradients(f, h, precondition=v_cycle),
    }
    sizes = [17, 33, 65, 129, 257]
    table = {}
    print("points   method          iterations   seconds   max error")
    for n in sizes:
        table[n] = race(n, methods if n > 65 else {**methods, "Jacobi": jacobi})
        for name, (iterations, seconds, error) in table[n].items():
            print(f"{n:6d}   {name:14s} {iterations:10d} {seconds:9.3f}   {error:.2e}")

    errors = [table[n]["CG"][2] for n in sizes]
    hs = [1 / (n - 1) for n in sizes]
    rate = np.polyfit(np.log(hs[-3:]), np.log(errors[-3:]), 1)[0]
    print(f"measured convergence rate: {rate:.2f} (second-order differences should give 2)")
    print(f"CHECK convergence rate near 2 -> {'PASS' if abs(rate - 2) < 0.1 else 'FAIL'} (limit 2 +/- 0.1)")

    fig, (left, right) = plt.subplots(1, 2, figsize=(9, 3.6))
    for name in methods:
        left.loglog(sizes, [table[n][name][0] for n in sizes], "o-", label=name)
        right.loglog(sizes, [table[n][name][1] for n in sizes], "o-", label=name)
    small = [n for n in sizes if "Jacobi" in table[n]]
    left.loglog(small, [table[n]["Jacobi"][0] for n in small], "o-", label="Jacobi")
    right.loglog(small, [table[n]["Jacobi"][1] for n in small], "o-", label="Jacobi")
    left.set_xlabel("points per side"); left.set_ylabel("iterations"); left.legend()
    right.set_xlabel("points per side"); right.set_ylabel("seconds (this machine)")
    fig.tight_layout(); fig.savefig(os.path.join(OUT, "10_race.png"), dpi=150)
    print(f"plots saved in {OUT}")

    # ---- YOUR TURN ----
    # 1. Predict first: if CG's iterations grow like n^p, what is p? Fit it from the table (polyfit on the logs).
    #    Do the same for time. Why does time grow faster than iterations?
    # 2. Add 513 points. From what size does multigrid win on time here? Compiled code (sbp-sat-suite) shifts that
    #    point, because each method's cost per iteration changes. Your machine's break-even point is your own result.
    # 3. In the online lab, run orders 2, 4 and 6 at 17, 33, 65 and 129 points and measure each order's rate.
    #    Order 4 measures about 3 in our tests. Why might the boundary be the reason?
    # 4. With a mentor: race CPU against GPU in sbp-sat-suite. Report where your numbers differ from its README.
    if "--show" in sys.argv:
        plt.show()


if __name__ == "__main__":
    main()