"""Anderson's "FMM without multipoles": the outer- and inner-ring approximations in 2D.

Cleaned from the Scientific Computing 2 (TUM, 2020) worksheet 5 solution (fmm_anderson.ipynb, cell.py).
Note: ../../Fast_multipole_method.md   Source: Anderson (1992), SIAM J. Sci. Stat. Comput. 13, 923.

Idea: a potential that is harmonic outside (or inside) a circle of radius a is fixed by its values on the
circle. So store the potential at K = 2M+1 equally spaced points on a ring instead of series coefficients,
and rebuild it anywhere with a truncated Poisson integral (trapezoidal rule, exact for degree <= 2M):

  outer ring (target outside, |z| = r > a):
    phi(r,t) = kappa log r + 1/K sum_i [phi(a,s_i) - kappa log a] * P_M(a/r, t - s_i)
  inner ring (target inside, r < a):
    phi(r,t) =               1/K sum_i  phi(a,s_i)                * P_M(r/a, t - s_i)
  with  P_M(rho, u) = 1 + 2 sum_{n=1..M} rho^n cos(n u)   and kappa = total charge.

The only kernel the method ever calls is the potential itself (log r), which is why it carries over to
other kernels with little change.

Changes against the course code
  - The course used the closed-form Poisson kernel with weight h/(2 pi) = a/K, which is right only for a = 1.
    Here the weight is 1/K and the kernel is truncated at M terms, as in Anderson (1992), so M is the
    accuracy knob.
  - Charges q_i enter (the course summed unit charges); np.int / np.float removed.

Run:  python anderson_rings.py    (error of both approximations against the direct sum, for several M)
"""
import numpy as np


def direct(z_target, z_src, q):
    return np.log(np.abs(np.atleast_1d(z_target)[:, None] - z_src[None, :])) @ q


def poisson_truncated(rho, u, M):
    n = np.arange(1, M + 1)
    return 1 + 2 * (rho[..., None] ** n * np.cos(n * u[..., None])).sum(-1)


class Ring:
    def __init__(self, center: complex, a: float, M: int):
        self.c, self.a, self.M = center, a, M
        self.K = 2 * M + 1
        self.s = 2 * np.pi * np.arange(self.K) / self.K
        self.points = center + a * np.exp(1j * self.s)
        self.values = None        # potential at the ring points
        self.kappa = 0.0          # total charge inside (outer ring only)

    def outer(self, z):
        w = np.atleast_1d(z) - self.c
        r, t = np.abs(w), np.angle(w)
        f = self.values - self.kappa * np.log(self.a)
        P = poisson_truncated(self.a / r[:, None], t[:, None] - self.s[None, :], self.M)
        return self.kappa * np.log(r) + P @ f / self.K

    def inner(self, z):
        w = np.atleast_1d(z) - self.c
        r, t = np.abs(w), np.angle(w)
        P = poisson_truncated(r[:, None] / self.a, t[:, None] - self.s[None, :], self.M)
        return P @ self.values / self.K


def main():
    rng = np.random.default_rng(2)
    n = 50
    src = 0.5 + 0.5j + (rng.random(n) - 0.5) + 1j * (rng.random(n) - 0.5)   # unit box at (0.5, 0.5)
    q = np.ones(n)
    target = 5.0 + 0.5j + 0.3 * np.exp(0.7j)          # a point near the centre of a box at (5, 0.5)
    exact = direct(target, src, q)[0]
    print(f"direct potential {exact:.10f}")
    print(f"{'M':>3} {'outer ring':>11} {'inner ring':>11} {'outer->inner':>13}")
    for M in (2, 4, 6, 10, 14):
        src_ring = Ring(0.5 + 0.5j, a=1.0, M=M)        # ring radius = side length, as in the course
        src_ring.values = direct(src_ring.points, src, q)
        src_ring.kappa = q.sum()
        e_out = abs(src_ring.outer(target)[0] - exact)

        tgt_ring = Ring(5.0 + 0.5j, a=1.0, M=M)
        tgt_ring.values = direct(tgt_ring.points, src, q)          # exact values on the target ring
        e_in = abs(tgt_ring.inner(target)[0] - exact)
        tgt_ring.values = src_ring.outer(tgt_ring.points)          # the FMM translation: outer -> inner
        e_tr = abs(tgt_ring.inner(target)[0] - exact)
        print(f"{M:3d} {e_out:11.1e} {e_in:11.1e} {e_tr:13.1e}")


if __name__ == '__main__':
    main()