import ctypes
import os
import platform

import numpy as np


class fastsum_plan(ctypes.Structure):
    pass


def _glibc_subdir() -> str:
    if os.name == "nt" or platform.system() == "Darwin":
        return ""

    libc = ctypes.CDLL(None)
    libc.gnu_get_libc_version.restype = ctypes.c_char_p
    version = libc.gnu_get_libc_version().decode("ascii")

    parts = version.split(".")
    major = int(parts[0])
    minor = int(parts[1]) if len(parts) > 1 else 0

    if (major, minor) < (2, 36):
        return "glibc2.22"
    return "glibc2.40"


def _lib_path() -> str:
    here = os.path.dirname(__file__)
    if os.name == "nt":
        return os.path.join(here, "libfastsumjulia.dll")
    if platform.system() == "Darwin":
        return os.path.join(here, "libfastsumjulia.dylib")
    return os.path.join(here, _glibc_subdir(), "libfastsumjulia.so")


def _set_signatures(lib):
    lib.jfastsum_alloc.restype = ctypes.POINTER(fastsum_plan)

    lib.jfastsum_init.argtypes = [
        ctypes.POINTER(fastsum_plan),
        ctypes.c_int,
        ctypes.c_char_p,
        np.ctypeslib.ndpointer(np.float64, flags="C"),
        ctypes.c_uint32,
        ctypes.c_int,
        ctypes.c_int,
        ctypes.c_double,
        ctypes.c_float,
        ctypes.c_int,
        ctypes.c_int,
        ctypes.c_int,
        ctypes.c_int,
        ctypes.c_int,
        ctypes.c_int,
    ]

    lib.jfastsum_finalize.argtypes = [ctypes.POINTER(fastsum_plan)]

    lib.jfastsum_set_x.argtypes = [
        ctypes.POINTER(fastsum_plan),
        np.ctypeslib.ndpointer(np.float64, flags="F"),
    ]
    lib.jfastsum_set_x.restype = ctypes.POINTER(ctypes.c_double)

    lib.jfastsum_set_y.argtypes = [
        ctypes.POINTER(fastsum_plan),
        np.ctypeslib.ndpointer(np.float64, flags="F"),
    ]
    lib.jfastsum_set_y.restype = ctypes.POINTER(ctypes.c_double)

    lib.jfastsum_set_alpha.argtypes = [
        ctypes.POINTER(fastsum_plan),
        np.ctypeslib.ndpointer(np.complex128, flags="F"),
    ]
    lib.jfastsum_set_alpha.restype = ctypes.POINTER(ctypes.c_double)

    lib.jfastsum_trafo.argtypes = [ctypes.POINTER(fastsum_plan)]
    lib.jfastsum_trafo.restype = ctypes.POINTER(ctypes.c_double)

    lib.jfastsum_exact.argtypes = [ctypes.POINTER(fastsum_plan)]
    lib.jfastsum_exact.restype = ctypes.POINTER(ctypes.c_double)


class FASTSUM:
    def __init__(
        self,
        d,
        N,
        M,
        kernel,
        c,
        n=256,
        p=8,
        eps_I=None,
        eps_B=None,
        nn=None,
        m=None,
    ):
        self.d = int(d)
        self.N = int(N)
        self.M = int(M)
        self.n = int(n)
        self.p = int(p)
        self.kernel_name = str(kernel)

        self.c = (
            np.array([c], dtype=np.float64)
            if isinstance(c, (int, float))
            else np.asarray(c, dtype=np.float64)
        )

        self.eps_I = float(self.p / self.n) if eps_I is None else float(eps_I)
        self.eps_B = float(max(1.0 / 16.0, self.p / self.n)) if eps_B is None else float(eps_B)
        self.nn_x = int(2 * self.n) if nn is None else int(nn)
        self.nn_y = self.nn_x
        self.m_x = int(self.p) if m is None else int(m)
        self.m_y = self.m_x
        self.flags = 0

        self.lib_path = _lib_path()
        self.lib = ctypes.CDLL(self.lib_path)
        _set_signatures(self.lib)

        self._plan = self.lib.jfastsum_alloc()

        code = self.lib.jfastsum_init(
            self._plan,
            ctypes.c_int(self.d),
            ctypes.create_string_buffer(self.kernel_name.encode("utf-8")),
            np.ascontiguousarray(self.c, dtype=np.float64),
            ctypes.c_uint32(self.flags),
            ctypes.c_int(self.n),
            ctypes.c_int(self.p),
            ctypes.c_double(self.eps_I),
            ctypes.c_float(self.eps_B),
            ctypes.c_int(self.N),
            ctypes.c_int(self.M),
            ctypes.c_int(self.nn_x),
            ctypes.c_int(self.nn_y),
            ctypes.c_int(self.m_x),
            ctypes.c_int(self.m_y),
        )
        if code == 1:
            raise RuntimeError("Unknown kernel")

        self._x = None
        self._y = None
        self._alpha = None
        self._f = None
        self.finalized = False

    @property
    def x(self):
        return self._x

    @x.setter
    def x(self, value):
        arr = np.asarray(value, dtype=np.float64)
        if arr.ndim != 2 or arr.shape != (self.N, self.d):
            raise ValueError(f"x must have shape ({self.N}, {self.d})")
        arrf = np.asfortranarray(arr)
        ptr = self.lib.jfastsum_set_x(self._plan, arrf)
        self._x = np.ctypeslib.as_array(ptr, shape=(self.N * self.d,)).reshape((self.N, self.d))

    @property
    def y(self):
        return self._y

    @y.setter
    def y(self, value):
        arr = np.asarray(value, dtype=np.float64)
        if arr.ndim != 2 or arr.shape != (self.M, self.d):
            raise ValueError(f"y must have shape ({self.M}, {self.d})")
        arrf = np.asfortranarray(arr)
        ptr = self.lib.jfastsum_set_y(self._plan, arrf)
        self._y = np.ctypeslib.as_array(ptr, shape=(self.M * self.d,)).reshape((self.M, self.d))

    @property
    def alpha(self):
        return self._alpha

    @alpha.setter
    def alpha(self, value):
        arr = np.asarray(value, dtype=np.complex128)
        if arr.ndim != 1 or arr.shape[0] != self.N:
            raise ValueError(f"alpha must have shape ({self.N},)")
        arrf = np.asfortranarray(arr)
        ptr = self.lib.jfastsum_set_alpha(self._plan, arrf)
        self._alpha = np.ctypeslib.as_array(ptr, shape=(self.N * 2,)).view(np.complex128)

    @property
    def f(self):
        return self._f

    def trafo(self):
        ptr = self.lib.jfastsum_trafo(self._plan)
        self._f = np.ctypeslib.as_array(ptr, shape=(self.M * 2,)).view(np.complex128)
        return self._f

    def trafo_exact(self):
        ptr = self.lib.jfastsum_exact(self._plan)
        self._f = np.ctypeslib.as_array(ptr, shape=(self.M * 2,)).view(np.complex128)
        return self._f

    def finalize(self):
        if not self.finalized:
            self.lib.jfastsum_finalize(self._plan)
            self.finalized = True
