"""Small 2D radiation transport solvers used for the figures in the post on
neutrino transport approximations.

Everything is in slab geometry on the unit square with c = 1:

    M1          two-moment transport with the Levermore closure, explicit,
                second order, HLL fluxes
    FLD         flux-limited diffusion with the Levermore-Pomraning limiter,
                implicit in time
    sn_point    steady discrete-ordinates solution for a Gaussian point source
    mc_point    the same by Monte Carlo

    python transport2d.py       # runs the shadow test and prints the leak
"""
import numpy as np
import scipy.sparse as sp
import scipy.sparse.linalg as spla
from scipy.special import erf


def minmod(a, b):
    return np.where(a * b > 0, np.sign(a) * np.minimum(abs(a), abs(b)), 0.0)


def chi_levermore(f):
    """Eddington factor as a function of the flux factor f = |F| / (cE)."""
    return (3 + 4 * f**2) / (5 + 2 * np.sqrt(4 - 3 * f**2))


def lam_lp(R):
    """Levermore-Pomraning flux limiter, R = |grad E| / (kappa E)."""
    R = np.maximum(R, 1e-9)
    return (1 / np.tanh(R) - 1 / R) / R


class M1:
    """Two-moment radiation in 2D.  State U = (E, Fx, Fy) on an N x N grid.

    kappa_fn(X, Y) gives the absorption coefficient; inlet(y) gives
    (E, Fx, Fy) imposed in the ghost cells of the left boundary.  All other
    boundaries are vacuum.
    """

    def __init__(self, N, kappa_fn, inlet):
        self.N = N
        self.h = 1.0 / N
        x = (np.arange(N) + 0.5) * self.h
        self.X, self.Y = np.meshgrid(x, x, indexing="ij")
        self.kap = kappa_fn(self.X, self.Y)
        self.inlet = inlet
        self.U = np.zeros((3, N, N))

    def pressure(self, U):
        E, Fx, Fy = U
        Fm = np.sqrt(Fx**2 + Fy**2)
        f = np.clip(Fm / np.maximum(E, 1e-30), 0, 1)
        chi = chi_levermore(f)
        nx = np.where(Fm > 0, Fx / np.maximum(Fm, 1e-30), 0.0)
        ny = np.where(Fm > 0, Fy / np.maximum(Fm, 1e-30), 0.0)
        a, b = (1 - chi) / 2, (3 * chi - 1) / 2
        return E * (a + b * nx * nx), E * b * nx * ny, E * (a + b * ny * ny)

    def flux_x(self, U):
        Pxx, Pxy, _ = self.pressure(U)
        return np.array([U[1], Pxx, Pxy])

    def flux_y(self, U):
        _, Pxy, Pyy = self.pressure(U)
        return np.array([U[2], Pxy, Pyy])

    def padded(self, U):
        """Two ghost cells per side: vacuum everywhere except the inlet."""
        G = 2
        P = np.zeros((3, self.N + 2 * G, self.N + 2 * G))
        P[:, G:-G, G:-G] = U
        y = (np.arange(self.N) + 0.5) * self.h
        for k, v in enumerate(self.inlet(y)):
            P[k, :G, G:-G] = v
        return P

    def rhs(self, U):
        P = self.padded(U)
        h = self.h
        sx = minmod(P[:, 1:-1, :] - P[:, :-2, :], P[:, 2:, :] - P[:, 1:-1, :])[:, :, 1:-1]
        sy = minmod(P[:, :, 1:-1] - P[:, :, :-2], P[:, :, 2:] - P[:, :, 1:-1])[:, 1:-1, :]
        C = P[:, 1:-1, 1:-1]
        UL, UR = (C + 0.5 * sx)[:, :-1, :], (C - 0.5 * sx)[:, 1:, :]
        Fx = 0.5 * (self.flux_x(UL) + self.flux_x(UR)) - 0.5 * (UR - UL)   # HLL, wave speeds +-c
        UL, UR = (C + 0.5 * sy)[:, :, :-1], (C - 0.5 * sy)[:, :, 1:]
        Fy = 0.5 * (self.flux_y(UL) + self.flux_y(UR)) - 0.5 * (UR - UL)
        return (-(Fx[:, 1:, 1:-1] - Fx[:, :-1, 1:-1]) / h
                - (Fy[:, 1:-1, 1:] - Fy[:, 1:-1, :-1]) / h)

    def absorb(self, U, dt):
        return U / (1 + self.kap * dt)          # implicit absorption

    def run(self, T, cfl=0.4):
        dt, t = cfl * self.h, 0.0
        while t < T - 1e-12:
            dt = min(dt, T - t)
            U1 = self.absorb(self.U + dt * self.rhs(self.U), dt)
            U2 = self.absorb(U1 + dt * self.rhs(U1), dt)
            self.U = 0.5 * (self.U + U2)
            t += dt
        return self.U


class FLD:
    """Flux-limited diffusion in 2D, implicit in time with the limiter lagged.

    The inlet is a Dirichlet value of E on the left boundary; every other
    boundary has E = 0.
    """

    def __init__(self, N, kappa_fn, inlet):
        self.N = N
        self.h = 1.0 / N
        x = (np.arange(N) + 0.5) * self.h
        self.X, self.Y = np.meshgrid(x, x, indexing="ij")
        self.kap = kappa_fn(self.X, self.Y)
        self.Ein = inlet(x)[0]
        self.E = np.zeros((N, N))

    def diffusion(self, E):
        h = self.h
        Ep = np.pad(E, 1, mode="edge")
        Ep[0, 1:-1] = np.where(self.Ein > 0, self.Ein, Ep[1, 1:-1])
        gx = (Ep[2:, 1:-1] - Ep[:-2, 1:-1]) / (2 * h)
        gy = (Ep[1:-1, 2:] - Ep[1:-1, :-2]) / (2 * h)
        R = np.sqrt(gx**2 + gy**2) / (self.kap * np.maximum(E, 1e-12))
        return lam_lp(R) / self.kap

    def step(self, dt):
        N, h = self.N, self.h
        D = self.diffusion(self.E)
        Dp = np.pad(D, 1, mode="edge")
        hm = lambda a, b: 2 * a * b / (a + b + 1e-30)          # harmonic mean on faces
        Dw, De = hm(Dp[:-2, 1:-1], D), hm(Dp[2:, 1:-1], D)
        Ds, Dn = hm(Dp[1:-1, :-2], D), hm(Dp[1:-1, 2:], D)
        idx = np.arange(N * N).reshape(N, N)
        c = dt / h**2
        diag = 1 + self.kap * dt + c * (Dw + De + Ds + Dn)
        rows, cols, vals = [idx.ravel()], [idx.ravel()], [diag.ravel()]
        rhs = self.E.copy()
        for shift, Dface, m in [(-N, Dw, np.s_[1:, :]), (N, De, np.s_[:-1, :]),
                                (-1, Ds, np.s_[:, 1:]), (1, Dn, np.s_[:, :-1])]:
            mask = np.zeros((N, N), bool)
            mask[m] = True
            rows.append(idx[mask]); cols.append(idx[mask] + shift); vals.append(-c * Dface[mask])
        rhs[0, :] += c * Dw[0, :] * self.Ein                     # Dirichlet inlet
        A = sp.csc_matrix((np.concatenate(vals), (np.concatenate(rows), np.concatenate(cols))),
                          shape=(N * N, N * N))
        self.E = spla.spsolve(A, rhs.ravel()).reshape(N, N)

    def run(self, T, dt):
        t = 0.0
        while t < T - 1e-12:
            self.step(min(dt, T - t))
            t += dt
        return self.E


def sn_point(X, Y, ndir, sig=0.03):
    """Steady energy density from an isotropic Gaussian source in vacuum,
    using ndir discrete directions.  The line integral along each ray is
    done exactly, so the only error is the angular discretisation."""
    E = np.zeros_like(X)
    for k in range(ndir):
        th = 2 * np.pi * (k + 0.5) / ndir
        a = X * np.cos(th) + Y * np.sin(th)
        b = -X * np.sin(th) + Y * np.cos(th)
        E += np.exp(-b**2 / (2 * sig**2)) * sig * np.sqrt(np.pi / 2) * (1 + erf(a / (sig * np.sqrt(2))))
    return E / ndir


def mc_point(N, L, npk, sig=0.03, seed=1):
    """The same by Monte Carlo: npk packets, isotropic, path length tallied
    in each cell of an N x N grid on [-L, L]^2."""
    rng = np.random.default_rng(seed)
    E = np.zeros((N, N))
    ds = 0.5 * (2 * L / N)
    x0 = rng.normal(0, sig, (npk, 2))
    th = rng.uniform(0, 2 * np.pi, npk)
    n = np.stack([np.cos(th), np.sin(th)], 1)
    for j in range(int(2 * np.sqrt(2) * L / ds) + 2):
        p = x0 + (j + 0.5) * ds * n
        i = ((p + L) / (2 * L) * N).astype(int)
        ok = (i[:, 0] >= 0) & (i[:, 0] < N) & (i[:, 1] >= 0) & (i[:, 1] < N)
        np.add.at(E, (i[ok, 0], i[ok, 1]), ds)
    return E * (2 * np.pi * sig**2 / npk) * (N / (2 * L)) ** 2


if __name__ == "__main__":
    cx, cy, rad, halfw = 0.3, 0.5, 0.09, 0.2
    kap = lambda X, Y: np.where((X - cx) ** 2 + (Y - cy) ** 2 < rad**2, 2000.0, 0.3)
    inlet = lambda y: (np.where(abs(y - 0.5) < halfw, 1.0, 0.0),
                       np.where(abs(y - 0.5) < halfw, 1.0, 0.0), np.zeros_like(y))
    N = 160
    m1 = M1(N, kap, inlet)
    E_m1 = m1.run(2.5)[0]
    E_fld = FLD(N, kap, inlet).run(6.0, 0.02)
    X, Y = m1.X, m1.Y
    shadow = (X > 0.6) & (X < 0.9) & (abs(Y - cy) < 0.6 * rad)
    lit = np.exp(-0.3 * X)[shadow].mean()
    print(f"energy density in the shadow, relative to the unobstructed beam:")
    print(f"  flux-limited diffusion: {E_fld[shadow].mean() / lit:.3f}")
    print(f"  M1:                     {E_m1[shadow].mean() / lit:.3f}")