import ctypes
import os
import platform

import numpy as np


class nfst_plan(ctypes.Structure):
    pass


# Flags (aligned with NFST.jl)
PRE_PHI_HUT = 1 << 0
FG_PSI = 1 << 1
PRE_LIN_PSI = 1 << 2
PRE_FG_PSI = 1 << 3
PRE_PSI = 1 << 4
PRE_FULL_PSI = 1 << 5
MALLOC_X = 1 << 6
MALLOC_F_HAT = 1 << 7
MALLOC_F = 1 << 8
FFT_OUT_OF_PLACE = 1 << 9
FFTW_INIT = 1 << 10
NFST_SORT_NODES = 1 << 11
NFST_OMP_BLOCKWISE_ADJOINT = 1 << 12

FFTW_MEASURE = 0
FFTW_DESTROY_INPUT = 1 << 0
FFTW_EXHAUSTIVE = 1 << 3
FFTW_PATIENT = 1 << 5
FFTW_ESTIMATE = 1 << 6


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, "libnfstjulia.dll")
    if platform.system() == "Darwin":
        return os.path.join(here, "libnfstjulia.dylib")

    return os.path.join(here, _glibc_subdir(), "libnfstjulia.so")


def _set_signatures(lib):
    lib.jnfst_alloc.restype = ctypes.POINTER(nfst_plan)
    lib.jnfst_finalize.argtypes = [ctypes.POINTER(nfst_plan)]

    lib.jnfst_init.argtypes = [
        ctypes.POINTER(nfst_plan),
        ctypes.c_int32,
        ctypes.POINTER(ctypes.c_int32),
        ctypes.c_int32,
        ctypes.POINTER(ctypes.c_int32),
        ctypes.c_int32,
        ctypes.c_uint32,
        ctypes.c_uint32,
    ]

    lib.jnfst_set_x.argtypes = [
        ctypes.POINTER(nfst_plan),
        np.ctypeslib.ndpointer(np.float64, flags="C"),
    ]
    lib.jnfst_set_x.restype = ctypes.POINTER(ctypes.c_double)

    lib.jnfst_set_f.argtypes = [
        ctypes.POINTER(nfst_plan),
        np.ctypeslib.ndpointer(np.float64, ndim=1, flags="C"),
    ]
    lib.jnfst_set_f.restype = ctypes.POINTER(ctypes.c_double)

    lib.jnfst_set_fhat.argtypes = [
        ctypes.POINTER(nfst_plan),
        np.ctypeslib.ndpointer(np.float64, ndim=1, flags="C"),
    ]
    lib.jnfst_set_fhat.restype = ctypes.POINTER(ctypes.c_double)

    lib.jnfst_trafo.argtypes = [ctypes.POINTER(nfst_plan)]
    lib.jnfst_trafo.restype = ctypes.POINTER(ctypes.c_double)

    lib.jnfst_adjoint.argtypes = [ctypes.POINTER(nfst_plan)]
    lib.jnfst_adjoint.restype = ctypes.POINTER(ctypes.c_double)

    lib.jnfst_trafo_direct.argtypes = [ctypes.POINTER(nfst_plan)]
    lib.jnfst_trafo_direct.restype = ctypes.POINTER(ctypes.c_double)

    lib.jnfst_adjoint_direct.argtypes = [ctypes.POINTER(nfst_plan)]
    lib.jnfst_adjoint_direct.restype = ctypes.POINTER(ctypes.c_double)

    lib.nfft_get_num_threads.restype = ctypes.c_int


class NFSTplan:
    def __init__(self, N, M, n=None, m=8, f1=None, f2=None):
        N = np.asarray(N, dtype=np.int32)
        if N.ndim != 1:
            raise ValueError("N has to be one-dimensional")
        if np.any(N <= 0) or np.any(N % 2 != 0):
            raise ValueError("Every entry of N has to be an even, positive integer")
        if M <= 0:
            raise ValueError("M has to be a positive integer")

        self.N = N
        self.M = int(M)
        self.D = len(N)
        self.Ns = int(np.prod(N - 1))

        if n is None:
            n = (2 ** (np.ceil(np.log(N) / np.log(2)) + 1)).astype(np.int32)
        else:
            n = np.asarray(n, dtype=np.int32)

        if np.any(n <= N) or np.any(n % 2 != 0):
            raise ValueError("Every entry of n has to be even and larger than N")

        self.n = n
        self.m = int(m)

        if f1 is None:
            if self.D > 1:
                self.f1 = np.uint32(
                    PRE_PHI_HUT
                    | PRE_PSI
                    | MALLOC_X
                    | MALLOC_F_HAT
                    | MALLOC_F
                    | FFTW_INIT
                    | FFT_OUT_OF_PLACE
                    | NFST_SORT_NODES
                    | NFST_OMP_BLOCKWISE_ADJOINT
                )
            else:
                self.f1 = np.uint32(
                    PRE_PHI_HUT
                    | PRE_PSI
                    | MALLOC_X
                    | MALLOC_F_HAT
                    | MALLOC_F
                    | FFTW_INIT
                    | FFT_OUT_OF_PLACE
                )
        else:
            self.f1 = np.uint32(f1)

        self.f2 = np.uint32(FFTW_ESTIMATE | FFTW_DESTROY_INPUT) if f2 is None else np.uint32(f2)

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

        self._plan = self.lib.jnfst_alloc()
        self.lib.jnfst_init(
            self._plan,
            ctypes.c_int32(self.D),
            self.N.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
            ctypes.c_int32(self.M),
            self.n.ctypes.data_as(ctypes.POINTER(ctypes.c_int32)),
            ctypes.c_int32(self.m),
            ctypes.c_uint32(int(self.f1)),
            ctypes.c_uint32(int(self.f2)),
        )

        self._x = None
        self._f = None
        self._fhat = None
        self.finalized = False

    def __del__(self):
        try:
            self.finalize()
        except Exception:
            pass

    @property
    def num_threads(self):
        return int(self.lib.nfft_get_num_threads())

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

    @x.setter
    def x(self, value):
        arr = np.asarray(value, dtype=np.float64, order="C")
        if self.D == 1:
            if arr.ndim != 1 or arr.shape[0] != self.M:
                raise ValueError("x has to be a Float64 vector of length M")
            ptr = self.lib.jnfst_set_x(self._plan, arr)
            self._x = np.ctypeslib.as_array(ptr, shape=(self.M,))
        else:
            if arr.ndim != 2 or arr.shape != (self.M, self.D):
                raise ValueError("x has to be a Float64 matrix of size MxD")
            ptr = self.lib.jnfst_set_x(self._plan, arr)
            self._x = np.ctypeslib.as_array(ptr, shape=(self.M * self.D,)).reshape((self.M, self.D))

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

    @f.setter
    def f(self, value):
        arr = np.asarray(value, dtype=np.float64, order="C")
        if arr.ndim != 1 or arr.shape[0] != self.M:
            raise ValueError("f has to be a Float64 vector of size M")
        ptr = self.lib.jnfst_set_f(self._plan, arr)
        self._f = np.ctypeslib.as_array(ptr, shape=(self.M,))

    @property
    def fhat(self):
        return self._fhat

    @fhat.setter
    def fhat(self, value):
        arr = np.asarray(value, dtype=np.float64, order="C")
        if arr.ndim != 1 or arr.shape[0] != self.Ns:
            raise ValueError("fhat has to be a Float64 vector of size prod(N-1)")
        ptr = self.lib.jnfst_set_fhat(self._plan, arr)
        self._fhat = np.ctypeslib.as_array(ptr, shape=(self.Ns,))

    def trafo(self):
        ptr = self.lib.jnfst_trafo(self._plan)
        self._f = np.ctypeslib.as_array(ptr, shape=(self.M,))
        return self._f

    def adjoint(self):
        ptr = self.lib.jnfst_adjoint(self._plan)
        self._fhat = np.ctypeslib.as_array(ptr, shape=(self.Ns,))
        return self._fhat

    def trafo_direct(self):
        ptr = self.lib.jnfst_trafo_direct(self._plan)
        self._f = np.ctypeslib.as_array(ptr, shape=(self.M,))
        return self._f

    def adjoint_direct(self):
        ptr = self.lib.jnfst_adjoint_direct(self._plan)
        self._fhat = np.ctypeslib.as_array(ptr, shape=(self.Ns,))
        return self._fhat

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