"""The showcase problem, shared by both backends.

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 json
import os
import platform
import socket
import subprocess

import numpy as np

# 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


def green(x, x0, k):
    d = x - x0[:, None]
    r = np.linalg.norm(d, axis=0)
    return np.exp(1j * k * r) / (4 * np.pi * r)


def green_dn(x, nrm, x0, k):
    d = x - x0[:, None]
    r = np.linalg.norm(d, axis=0)
    g = np.exp(1j * k * r) / (4 * np.pi * r)
    return g * (1j * k - 1 / r) * np.sum(d * nrm, axis=0) / r


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


def room_rhs(xs, nrm, k, lam):
    """Impedance data for v = u - G: -(dG/dn + i k lambda G)."""
    return -(green_dn(xs, nrm, SOURCE, k) + 1j * k * lam * green(xs, SOURCE, k))


def selftest_rhs(xs, nrm, k, lam):
    """Impedance data of the field of a source outside the room."""
    return green_dn(xs, nrm, EXTERIOR_SOURCE, k) + 1j * k * lam * green(xs, EXTERIOR_SOURCE, k)


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


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 machine_info():
    info = dict(host=socket.gethostname(), platform=platform.platform(), cpu_count=os.cpu_count())
    try:
        with open("/proc/cpuinfo") as f:
            for line in f:
                if line.startswith("model name"):
                    info["cpu"] = line.split(":", 1)[1].strip()
                    break
    except OSError:
        pass
    try:
        gpu = subprocess.run(["nvidia-smi", "--query-gpu=name,memory.total", "--format=csv,noheader"],
                             capture_output=True, text=True, timeout=10)
        if gpu.returncode == 0:
            info["gpu"] = gpu.stdout.strip().splitlines()[0]
    except (OSError, subprocess.SubprocessError):
        pass
    return info


def save(outdir, summary, **arrays):
    os.makedirs(outdir, exist_ok=True)
    np.savez(os.path.join(outdir, "field.npz"), **arrays)
    with open(os.path.join(outdir, "summary.json"), "w") as f:
        json.dump(summary, f, indent=2)