"""Minimal ctypes interface to the fmm3dbie shared library.

Only the routines needed for the interior Helmholtz impedance problem are
wrapped. All Fortran arguments are passed by reference; integers are
integer*8, and arrays are Fortran (column-major) ordered.
"""

import contextlib
import ctypes
import os
import sys

# OpenBLAS threads compete with the library's OpenMP threads
os.environ.setdefault("OPENBLAS_NUM_THREADS", "1")
# write Fortran console output immediately so it can be redirected to a log
os.environ.setdefault("GFORTRAN_UNBUFFERED_PRECONNECTED", "y")

import numpy as np

LIBPATH = os.environ.get(
    "FMM3DBIE_LIB", os.path.expanduser("~/src/fmm3dbie/lib/libfmm3dbie.so")
)
_lib = ctypes.CDLL(LIBPATH)

_i8 = ctypes.c_int64
_f8 = ctypes.c_double


@contextlib.contextmanager
def fortran_output_to(path):
    """Redirect the library's console output (Fortran unit 6, i.e. fd 1) to a file."""
    sys.stdout.flush()
    saved = os.dup(1)
    with open(path, "a") as f:
        os.dup2(f.fileno(), 1)
        try:
            yield
        finally:
            _libc.fflush(None)
            os.dup2(saved, 1)
            os.close(saved)


_libc = ctypes.CDLL(None)


def _ref_i8(v):
    return ctypes.byref(_i8(int(v)))


def _ref_f8(v):
    return ctypes.byref(_f8(float(v)))


def _ptr(a):
    return a.ctypes.data_as(ctypes.c_void_p)


def _farr(a, dtype):
    return np.asfortranarray(a, dtype=dtype)


class Surface:
    """Discretized surface in the fmm3dbie format (all patches share one order and type)."""

    def __init__(self, norder, iptype, srcvals):
        self.norder = norder
        self.iptype = iptype
        self.srcvals = _farr(srcvals, np.float64)  # (12, npts)
        self.npts = self.srcvals.shape[1]
        self.npols = npols_for(norder, iptype)
        self.npatches = self.npts // self.npols
        self.norders = np.full(self.npatches, norder, dtype=np.int64)
        self.iptypes = np.full(self.npatches, iptype, dtype=np.int64)
        self.ixyzs = (np.arange(self.npatches + 1, dtype=np.int64) * self.npols + 1)
        self.srccoefs = surf_vals_to_coefs(self, self.srcvals[0:9, :])
        self.wts = get_qwts(self)

    def args(self):
        return (_ref_i8(self.npatches), _ptr(self.norders), _ptr(self.ixyzs),
                _ptr(self.iptypes), _ref_i8(self.npts), _ptr(self.srccoefs),
                _ptr(self.srcvals))


def npols_for(norder, iptype):
    if iptype == 1:
        return (norder + 1) * (norder + 2) // 2
    return (norder + 1) ** 2


def get_disc_nodes(norder, iptype):
    npols = npols_for(norder, iptype)
    uvs = np.zeros((2, npols), order="F")
    _lib.get_disc_nodes_(_ref_i8(norder), _ref_i8(npols), _ref_i8(iptype), _ptr(uvs))
    return uvs


def surf_vals_to_coefs(surf, vals):
    vals = _farr(vals, np.float64)
    nd = vals.shape[0]
    coefs = np.zeros_like(vals, order="F")
    _lib.surf_vals_to_coefs_(_ref_i8(nd), _ref_i8(surf.npatches), _ptr(surf.norders),
                             _ptr(surf.ixyzs), _ptr(surf.iptypes), _ref_i8(surf.npts),
                             _ptr(vals), _ptr(coefs))
    return coefs


def get_qwts(surf):
    wts = np.zeros(surf.npts)
    _lib.get_qwts_(_ref_i8(surf.npatches), _ptr(surf.norders), _ptr(surf.ixyzs),
                   _ptr(surf.iptypes), _ref_i8(surf.npts), _ptr(surf.srcvals), _ptr(wts))
    return wts


def helm_rpcomb_imp_solver(surf, zk, alpha, zlams, rhs, ifinout, eps=1e-6,
                           eps_gmres=1e-6, numit=400):
    """Solve n.grad(u) + i k lambda u = rhs with u = S_k[s] + i alpha D_k S_{i|k|}[s].

    ifinout = 0 for the interior problem, 1 for the exterior problem.
    Returns (sigma, siksigma, niter, errs, rres).
    """
    zpars = np.array([zk, alpha], dtype=np.complex128)
    zlams = _farr(zlams, np.complex128)
    rhs = _farr(rhs, np.complex128)
    sigma = np.zeros(surf.npts, dtype=np.complex128)
    siksigma = np.zeros(surf.npts, dtype=np.complex128)
    errs = np.zeros(numit + 1)
    niter = _i8(0)
    rres = _f8(0.0)
    _lib.helm_rpcomb_imp_solver_(*surf.args(), _ref_f8(eps), _ptr(zpars), _ptr(zlams),
                                 _ref_i8(numit), _ref_i8(ifinout), _ptr(rhs),
                                 _ref_f8(eps_gmres), ctypes.byref(niter), _ptr(errs),
                                 ctypes.byref(rres), _ptr(sigma), _ptr(siksigma))
    return sigma, siksigma, niter.value, errs[: niter.value], rres.value


def helm_rpcomb_eval(surf, zk, alpha, sigma, siksigma, targs, eps=1e-6):
    """Evaluate u = S_k[s] + i alpha D_k S_{i|k|}[s] at off-surface targets (3, ntarg)."""
    targs = _farr(targs, np.float64)
    ntarg = targs.shape[1]
    zpars = np.array([zk, alpha], dtype=np.complex128)
    ipatch_id = np.full(ntarg, -1, dtype=np.int64)
    uvs_targ = np.zeros((2, ntarg), order="F")
    pot = np.zeros(ntarg, dtype=np.complex128)
    _lib.helm_rpcomb_eval_(*surf.args(), _ref_i8(3), _ref_i8(ntarg), _ptr(targs),
                           _ptr(ipatch_id), _ptr(uvs_targ), _ref_f8(eps), _ptr(zpars),
                           _ptr(_farr(sigma, np.complex128)),
                           _ptr(_farr(siksigma, np.complex128)), _ptr(pot))
    return pot


def helm_rpcomb_eval_boundary(surf, zk, alpha, sigma, siksigma, eps=1e-6, ifinout=0):
    """Boundary limit of u = S_k[s] + i alpha D_k S_{i|k|}[s] at the nodes, from the
    inside (ifinout = 0) or the outside (ifinout = 1).

    The on-surface evaluation gives the principal value of D_k; with outward
    normals the interior limit adds -1/2 times the double-layer density and the
    exterior limit +1/2 times it.
    """
    uvs = get_disc_nodes(surf.norder, surf.iptype)
    ipatch_id = np.repeat(np.arange(1, surf.npatches + 1), surf.npols).astype(np.int64)
    uvs_targ = np.asfortranarray(np.tile(uvs, (1, surf.npatches)))
    zpars = np.array([zk, alpha], dtype=np.complex128)
    sigma = _farr(sigma, np.complex128)
    siksigma = _farr(siksigma, np.complex128)
    pot = np.zeros(surf.npts, dtype=np.complex128)
    _lib.helm_rpcomb_eval_(*surf.args(), _ref_i8(12), _ref_i8(surf.npts), _ptr(surf.srcvals),
                           _ptr(ipatch_id), _ptr(uvs_targ), _ref_f8(eps), _ptr(zpars),
                           _ptr(sigma), _ptr(siksigma), _ptr(pot))
    return pot + (-1) ** (ifinout + 1) * 0.5j * alpha * siksigma