"""Helpers shared by the problems: the Green's function, boundary data, machine
information and saving results. No dependence on either backend."""

import json
import os
import platform
import socket
import subprocess

import numpy as np


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 impedance_data(xs, nrm, x0, k, lam):
    """du/dn + i k lambda u on the surface for u the field of a unit source at x0."""
    return green_dn(xs, nrm, x0, k) + 1j * k * lam * green(xs, x0, k)


def source_sum(x, x0s, ws, k, nrm=None, lam=0.0, chunk=64, xp=np):
    """Field sum_j ws[j] G(x, x0s[:, j]) at points x (3, n), or with nrm given its
    impedance data du/dn + i k lambda u. xp is numpy or cupy; the result is NumPy."""
    x, x0s, ws = xp.asarray(x), xp.asarray(x0s), xp.asarray(ws)
    nrm = None if nrm is None else xp.asarray(nrm)
    out = xp.zeros(x.shape[1], dtype=complex)
    for i in range(0, x0s.shape[1], chunk):
        d = x[:, None, :] - x0s[:, i:i + chunk, None]  # (3, chunk, n)
        r = xp.sqrt(xp.sum(d * d, axis=0))
        g = xp.exp(1j * k * r) / (4 * np.pi * r)
        if nrm is not None:
            g = g * ((1j * k - 1 / r) * xp.einsum("icn,in->cn", d, nrm) / r + 1j * k * lam)
        out += ws[i:i + chunk] @ g
    return out.get() if xp is not np else out


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)