"""The room problem.

A point source at x0 inside a rounded cube (the superellipsoid
|x|^p + |y|^p + |z|^p = a^p) with impedance walls. With time dependence
exp(-i omega t) the pressure amplitude u satisfies

    (Delta + k^2) u = -delta(x - x0)   in the room,
    du/dn + i k lambda u = 0            on the walls,

with n the outward normal and lambda = -beta, where beta > 0 is the wall
admittance. We write u = G + v with G the free-space field of the source, and
find v from the right-preconditioned combined field integral equation of
fmm3dbie (helm_rpcomb_imp_solver, interior mode).

This module has no dependence on either backend.
"""

import time

import numpy as np

from .common import green

NAME = "room"
IFINOUT = 0  # interior problem

# Named cases. Both backends discretize the surface identically, so the same
# case gives the same unknowns on both.
CASES = {
    "small": dict(k=10.0, beta=0.2, m=4, norder=6),     # 3.2 wavelengths, 4,704 points
    "medium": dict(k=10.0, beta=0.2, m=6, norder=8),    # 3.2 wavelengths, 17,496 points
    "large": dict(k=40.0, beta=0.05, m=10, norder=8),   # 12.7 wavelengths, 48,600 points
}
A, P = 1.0, 8                          # room half-width and superellipsoid exponent
SOURCE = np.array([-0.35, 0.2, 0.1])   # the point source
EXTERIOR_SOURCE = np.array([1.7, 0.6, -0.5])  # for the self-test (outside the room)
ALPHA = 1.0
NGRID = 101                   # default grid size of the field slices


def add_arguments(ap):
    ap.add_argument("--k", type=float, help="override the case's wavenumber")
    ap.add_argument("--beta", type=float, help="override the case's wall admittance")
    ap.add_argument("--m", type=int, help="override the patches per face edge")
    ap.add_argument("--norder", type=int, help="override the order per patch")


def make_case(name, args):
    case = dict(CASES[name])
    for key in ["k", "beta", "m", "norder"]:
        if getattr(args, key, None) is not None:
            case[key] = getattr(args, key)
    case["lam"] = -case["beta"]
    return case


def results_name(case_name, backend):
    return f"{case_name}-{backend}"


def describe(case):
    return (f"k = {case['k']}, beta = {case['beta']}, "
            f"room width {2 * A * case['k'] / (2 * np.pi):.2f} wavelengths")


def gpu_surface(case, args):
    from fmm3dbie_gpu import surface
    return surface.rounded_cube(a=A, p=P, m=case["m"], norder=case["norder"],
                                exact=not args.interp_geometry)


def fmm3dbie_surface(case):
    import rounded_cube as rc
    return rc.rounded_cube(a=A, p=P, m=case["m"], norder=case["norder"])


def source(case):
    return SOURCE


def selftest_sources(case):
    """A source outside the room, (positions (3, n), strengths (n,)); its field is the
    exact solution inside."""
    return EXTERIOR_SOURCE[:, None], np.ones(1)


def selftest_points(case, n=50):
    return np.random.default_rng(0).uniform(-0.8 * A, 0.8 * A, size=(3, n))


def inside(x, a=A, p=P, margin=0.0):
    return np.sum(np.abs(x / a) ** p, axis=0) < (1 - margin) ** p


def slice_points(ngrid):
    """Points on three planes through the source; returns grid, {name: (pts, mask, shape)}."""
    g = np.linspace(-A, A, ngrid)
    out = {}
    for name, (i, j, fixed) in {"xy": (0, 1, 2), "xz": (0, 2, 1), "yz": (1, 2, 0)}.items():
        Ag, Bg = np.meshgrid(g, g, indexing="xy")
        pts = np.zeros((3, Ag.size))
        pts[i], pts[j], pts[fixed] = Ag.ravel(), Bg.ravel(), SOURCE[fixed]
        out[name] = (pts, inside(pts, margin=1e-3), Ag.shape)
    return g, out


def energy_balance(k, lam, v_at_source, wts, u_wall):
    """Power radiated by the source, k/(4 pi) + Im v(x0), and power absorbed by the
    walls, -k Re(lambda) int |u|^2 dS. They must agree."""
    p_src = k / (4 * np.pi) + v_at_source.imag
    p_wall = -k * np.real(lam) * np.sum(wts * np.abs(u_wall) ** 2)
    return p_src, p_wall


def report(be, state, case, args, summary, arrays):
    """Energy balance and field slices of the room solution."""
    k, lam = case["k"], case["lam"]
    xs, wts = np.asarray(be.surf.srcvals)[0:3], np.asarray(be.surf.wts)
    v0 = complex(be.eval_points(state, SOURCE[:, None])[0])
    u_wall = green(xs, SOURCE, k) + be.eval_wall(state)
    p_src, p_wall = energy_balance(k, lam, v0, wts, u_wall)
    mismatch = abs(p_src - p_wall) / p_src
    print(f"  energy balance: source {p_src:.8f}, walls {p_wall:.8f}, relative mismatch {mismatch:.2e}")
    summary.update(p_src=p_src, p_wall=p_wall, energy_mismatch=mismatch)
    arrays.update(x0=SOURCE, k=k, lam=lam, a=A, p=P, xs=xs, wts=wts, u_wall=u_wall)
    if args.ngrid > 0:
        g, sl = slice_points(args.ngrid)
        t0 = time.time()
        for name, (pts, mask, shape) in sl.items():
            u = np.full(mask.size, np.nan + 1j * np.nan)
            u[mask] = green(pts[:, mask], SOURCE, k) + be.eval_points(state, pts[:, mask])
            arrays[f"u_{name}"] = u.reshape(shape)
        be.sync()
        summary["slices_s"] = time.time() - t0
        arrays["grid"] = g
        print(f"  field on three {args.ngrid} x {args.ngrid} slices: {summary['slices_s']:.2f} s")