"""The cymbal problem: a reciprocity solve for a listener near a rigid cymbal.

A point source at the listener's position L radiates in free space; the only
obstacle is a stationary, sound-hard cymbal. With time dependence
exp(-i omega t) the pressure amplitude u satisfies

    (Delta + k^2) u = -delta(x - L)   outside the cymbal,
    du/dn = 0                          on the cymbal,

and the Sommerfeld radiation condition. By reciprocity, u(x) is also the
pressure at the listener due to a unit point source at x, so one solve gives
the listener's response to a source anywhere; we evaluate it on a 3D grid in a
bounding box, and on three planes through the box. The box only selects where
the field is evaluated: the integral representation satisfies the radiation
condition, so the problem needs no truncation.

The cymbal is a thin closed body (fmm3dbie_gpu/cymbal.py): fmm3dbie's integral
equations need a closed surface. We write u = G + v, with G the free-space
field of the source, and find v from fmm3dbie's right-preconditioned combined
field equation for the impedance problem in exterior mode with lambda = 0
(helm_rpcomb_imp_solver, ifinout = 1), the same solver as the room problem.
Lengths are in meters and the speed of sound is 343 m/s.

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

import time

import numpy as np

from fmm3dbie_gpu.cymbal import CymbalShape, cymbal, cymbal_srcvals

from .common import green, impedance_data

NAME = "cymbal"
IFINOUT = 1  # exterior problem
ALPHA = 1.0
C_SOUND = 343.0

# h is the target patch size in meters; the mesh is also refined toward the rim
# and across the bell, independently of the frequency.
CASES = {
    "small": dict(freq=1000.0, h=0.06, norder=6),
    "medium": dict(freq=3000.0, h=0.04, norder=8),
    "large": dict(freq=8000.0, h=0.03, norder=8),
}
RADIUS = 0.2                      # a 40 cm cymbal
T_RIM = 2 * RADIUS / 40           # rim thickness diameter/40; twice that at the center
LISTENER = np.array([0.25, 0.0, 0.45])       # above the cymbal and off its axis
BOX = np.array([[-0.5, 0.5], [-0.5, 0.5], [-0.35, 0.65]])  # evaluation region (1 m cube)
RECIPROCITY_POINT = np.array([-0.15, 0.05, -0.12])        # below the cymbal
SLICE_BELOW = -0.1                # height of the horizontal slice, below the cymbal
NGRID = 201                       # default grid size of the slices, 5 mm spacing


def add_arguments(ap):
    ap.add_argument("--freq", type=float, help="override the case's frequency (Hz)")
    ap.add_argument("--h", type=float, help="override the case's target patch size (m)")
    ap.add_argument("--norder", type=int, help="override the order per patch")
    ap.add_argument("--t-rim", type=float, help=f"rim thickness in m (default {T_RIM})")
    ap.add_argument("--listener", type=lambda s: [float(v) for v in s.split(",")],
                    help="listener (= source) position x,y,z in m")
    ap.add_argument("--reciprocity", action="store_true",
                    help="also solve with the source at a point below the cymbal and check "
                         "that the field there due to the listener equals the field at the "
                         "listener due to it (a second solve)")
    ap.add_argument("--n3d", type=int, default=-1,
                    help="grid points per axis of the 3D grid in the box (default: a quarter "
                         "wavelength spacing, 41 to 121 points; 0 to skip)")


def make_case(name, args):
    case = dict(CASES[name])
    for key in ["freq", "h", "norder", "t_rim"]:
        if getattr(args, key, None) is not None:
            case[key] = getattr(args, key)
    case.setdefault("t_rim", T_RIM)
    case["listener"] = list(args.listener) if getattr(args, "listener", None) else list(LISTENER)
    case["k"] = 2 * np.pi * case["freq"] / C_SOUND
    case["lam"] = 0.0
    return case


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


def shape(case):
    return CymbalShape(radius=RADIUS, t_rim=case["t_rim"])


def describe(case):
    wl = C_SOUND / case["freq"]
    return (f"f = {case['freq']:g} Hz, k = {case['k']:.2f}/m, cymbal diameter "
            f"{2 * RADIUS / wl:.2f} wavelengths, rim thickness {1000 * case['t_rim']:.1f} mm "
            f"(diameter/{2 * RADIUS / case['t_rim']:.0f})")


def gpu_surface(case, args):
    return cymbal(shape(case), case["h"], case["norder"])


def fmm3dbie_surface(case):
    import fmm3dbie_ctypes as fb
    srcvals, _ = cymbal_srcvals(shape(case), case["h"], case["norder"])
    return fb.Surface(case["norder"], 11, srcvals)


def source(case):
    return np.array(case["listener"])


def selftest_sources(case, r_max=0.12, spacing=0.003):
    """Point sources inside the cymbal, (positions (3, n), strengths (n,)); their field
    is the exact solution outside. A single source would be at most T/2 = 1 cm from
    the surface, and its data would vary on that scale, which the main problem's
    data does not. Instead we spread sources 3 mm apart over the midsurface out to
    r = 12 cm, where they are at least 7 mm from the surface, with strengths that
    taper smoothly to zero at r_max. The data is then smooth on the scale of the
    patches, apart from a ripple of relative size about exp(-2 pi 7 / 3) ~ 4e-7."""
    sh = shape(case)
    g = np.arange(-r_max, r_max + spacing / 2, spacing)
    X, Y = np.meshgrid(g, g, indexing="ij")
    rho = (X * X + Y * Y).ravel()
    m = rho < r_max ** 2
    x0s = np.stack([X.ravel()[m], Y.ravel()[m], sh.F(rho[m])])
    ws = (1 - rho[m] / r_max ** 2) ** 3 * spacing ** 2 / (np.pi * r_max ** 2 / 4)  # sum ~ 1
    return x0s, ws


def selftest_points(case, n=50):
    """Points around the cymbal, at least 5 mm from its surface."""
    sh = shape(case)
    pts = np.random.default_rng(0).uniform([-0.3, -0.3, -0.15], [0.3, 0.3, 0.2], size=(4 * n, 3)).T
    return pts[:, ~sh.inside(pts, margin=0.005)][:, :n]


def far_field_power(k, x0, xs, nrm, wts, sigma, phi1, alpha=ALPHA, xp=np):
    """Power radiated to infinity, k int |u_inf|^2 dOmega, for u = G(., x0) + S_k sigma
    + i alpha D_k phi1, with u ~ u_inf exp(ikr)/r. The far field pattern is computed
    with the smooth rule at the surface nodes, and the sphere integral with Gauss-
    Legendre in cos(theta) and the trapezoid rule in phi, exact for the band limit.
    xp is numpy or cupy."""
    rmax = max(np.linalg.norm(xs, axis=0).max(), np.linalg.norm(x0))
    L = int(np.ceil(k * rmax)) + 20
    ct, wct = np.polynomial.legendre.leggauss(L + 1)
    nphi = 2 * L + 2
    ph = 2 * np.pi * np.arange(nphi) / nphi
    st = np.sqrt(1 - ct ** 2)
    dirs = np.stack([np.outer(st, np.cos(ph)), np.outer(st, np.sin(ph)),
                     np.outer(ct, np.ones(nphi))]).reshape(3, -1)
    wdir = np.outer(wct, np.full(nphi, 2 * np.pi / nphi)).ravel()
    xs, nrm, dirs = xp.asarray(xs), xp.asarray(nrm), xp.asarray(dirs)
    qs, qd = xp.asarray(wts * sigma), xp.asarray(1j * alpha * wts * phi1)
    uinf = xp.exp(-1j * k * (dirs.T @ xp.asarray(x0)))
    step = 64 if xp is np else 1024
    for i in range(0, dirs.shape[1], step):
        d = dirs[:, i:i + step]
        e = xp.exp(-1j * k * (d.T @ xs))
        uinf[i:i + step] += e @ qs + ((-1j * k * (d.T @ nrm)) * e) @ qd
    uinf = uinf.get() if xp is not np else uinf
    uinf /= 4 * np.pi
    return k * np.sum(wdir * np.abs(uinf) ** 2)


def box_grid(n):
    axes = [np.linspace(lo, hi, n) for lo, hi in BOX]
    X, Y, Z = np.meshgrid(*axes, indexing="ij")
    return axes, np.stack([X.ravel(), Y.ravel(), Z.ravel()])


def slice_points(ngrid, x0):
    """Three planes through the box: xz through the listener, yz through the cymbal's
    axis, and xy below the cymbal. Returns {name: (pts, (a-axis, b-axis), fixed)}."""
    out = {}
    for name, (i, j, fixed, val) in {"xz": (0, 2, 1, x0[1]), "yz": (1, 2, 0, 0.0),
                                     "xy": (0, 1, 2, SLICE_BELOW)}.items():
        ga, gb = np.linspace(*BOX[i], ngrid), np.linspace(*BOX[j], ngrid)
        Ag, Bg = np.meshgrid(ga, gb, indexing="xy")
        pts = np.zeros((3, Ag.size))
        pts[i], pts[j], pts[fixed] = Ag.ravel(), Bg.ravel(), val
        out[name] = (pts, (ga, gb), (fixed, val))
    return out


def eval_total(be, state, sh, x0, k, pts, spacing):
    """u = G + v at the points, NaN inside the cymbal and at the source."""
    mask = ~sh.inside(pts, margin=1e-3) & (np.linalg.norm(pts - x0[:, None], axis=0) > 1e-3 * spacing)
    u = np.full(mask.size, np.nan + 1j * np.nan)
    u[mask] = green(pts[:, mask], x0, k) + be.eval_points(state, pts[:, mask])
    return u


def report(be, state, case, args, summary, arrays):
    k, x0, sh = case["k"], source(case), shape(case)
    srcvals, wts = np.asarray(be.surf.srcvals), np.asarray(be.surf.wts)
    xs, nrm = srcvals[0:3], srcvals[9:12]
    v0 = complex(be.eval_points(state, x0[:, None])[0])
    u_wall = green(xs, x0, k) + be.eval_wall(state)

    t0 = time.time()
    sigma, phi1 = (np.asarray(be.to_numpy(q)) for q in state)
    p_src = k / (4 * np.pi) + v0.imag
    p_far = far_field_power(k, x0, xs, nrm, wts, sigma, phi1, xp=be.xp)
    mismatch = abs(p_src - p_far) / p_src
    print(f"  energy balance: source {p_src:.8f}, far field {p_far:.8f}, relative mismatch "
          f"{mismatch:.2e} ({time.time() - t0:.1f} s)")
    summary.update(p_src=p_src, p_far=p_far, energy_mismatch=mismatch, v_listener=[v0.real, v0.imag])

    if args.reciprocity:
        xb = RECIPROCITY_POINT
        (state_b, niter_b, _), t = be.timed_solve(-impedance_data(xs, nrm, xb, k, 0.0))
        v_lb = complex(be.eval_points(state, xb[:, None])[0])
        v_bl = complex(be.eval_points(state_b, x0[:, None])[0])
        rdiff = abs(v_lb - v_bl) / abs(v_lb)
        print(f"  reciprocity: scattered field at {xb} due to the listener {v_lb:.6e}, at the "
              f"listener due to {xb} {v_bl:.6e}, relative difference {rdiff:.2e} "
              f"({niter_b} iterations, {t:.2f} s)")
        summary.update(reciprocity_diff=rdiff, reciprocity_solve_s=t)

    arrays.update(x0=x0, k=k, freq=case["freq"], xs=xs, wts=wts, u_wall=u_wall, box=BOX,
                  outline=sh.cross_section(), radius=RADIUS)
    wl = 2 * np.pi / k
    if args.ngrid > 0:
        t0 = time.time()
        for name, (pts, (ga, gb), (fixed, val)) in slice_points(args.ngrid, x0).items():
            u = eval_total(be, state, sh, x0, k, pts, (ga[1] - ga[0]))
            arrays[f"u_{name}"] = u.reshape(gb.size, ga.size)
            arrays[f"axes_{name}"] = np.stack([ga, gb])
            arrays[f"plane_{name}"] = np.array([fixed, val])
        be.sync()
        summary["slices_s"] = time.time() - t0
        print(f"  field on three {args.ngrid} x {args.ngrid} slices: {summary['slices_s']:.2f} s")
    n3d = args.n3d
    if n3d < 0:
        n3d = int(np.clip(np.ceil((BOX[0, 1] - BOX[0, 0]) / (wl / 4)) + 1, 41, 121))
    if n3d > 0:
        t0 = time.time()
        axes, pts = box_grid(n3d)
        u = eval_total(be, state, sh, x0, k, pts, axes[0][1] - axes[0][0])
        arrays["u_grid"] = u.reshape(n3d, n3d, n3d)
        arrays["grid_axes"] = np.stack(axes)
        be.sync()
        summary.update(grid_s=time.time() - t0, n3d=n3d)
        print(f"  field on a {n3d}^3 grid in the box: {summary['grid_s']:.2f} s")