import ctypes

import numpy as np

def interp1d(src_start, src_step, src_val, dst_coords, stencil):
    idx_src   = ((dst_coords - src_start) / src_step - (stencil / 2 - 1)).astype(np.int)
    src_coord = np.linspace(src_start, src_start + src_step * (src_val.shape[0] - 1),
                            src_val.shape[0], dtype = src_val.dtype)

    fact = np.zeros((stencil, ) + idx_src.shape, dtype = src_val.dtype) + 1.0
    for i in range(stencil):
        for j in range(stencil):
            if i == j:
                continue

            fact[i] *= (dst_coords - src_coord[idx_src + j]) / (src_coord[idx_src + i] - src_coord[idx_src + j])

    ret = np.zeros_like(dst_coords, dtype = src_val.dtype)
    for i in range(stencil):
        ret += fact[i] * src_val[idx_src + i]

    return ret

def interp1d_irregular(src_coord, src_val, dst_coord, stencil):
    idx_src = np.empty(dst_coord.shape, dtype = np.int)
    for i, dc in enumerate(dst_coord):
        idx_src[i] = np.where(src_coord - dc < 0, src_coord, -np.inf).argmax() - stencil / 2 + 1

    fact = np.zeros((stencil, ) + idx_src.shape, dtype = src_val.dtype) + 1.0
    for i in range(stencil):
        for j in range(stencil):
            if i == j:
                continue

            fact[i] *= (dst_coord - src_coord[idx_src + j]) / (src_coord[idx_src + i] - src_coord[idx_src + j])

    ret = np.zeros_like(dst_coord, dtype = src_val.dtype)
    for i in range(stencil):
        ret += fact[i] * src_val[idx_src + i]

    return ret

def interp2d(src_start, src_step, src_val, dst_coords, stencil):
    ret = np.zeros_like(dst_coords[0], dtype = src_val.dtype)

    idx_src = []
    fact    = []
    for dim in range(len(src_start)):
        src_coord = np.linspace(src_start[dim], src_start[dim] + src_step[dim] * (src_val.shape[dim] - 1),
                                src_val.shape[dim], dtype = src_val.dtype)

        idx = ((dst_coords[dim] - src_start[dim]) / src_step[dim] - (stencil / 2 - 1)).astype(np.int)
        f   = np.zeros((stencil, ) + idx.shape, dtype = src_val.dtype) + 1.0


        for i in range(stencil):
            for j in range(stencil):
                if i == j:
                    continue

                f[i] *= (dst_coords[dim] - src_coord[idx + j]) / (src_coord[idx + i] - src_coord[idx + j])

        idx_src.append(idx)
        fact.append(f)

    for i in range(stencil):
        val = np.zeros_like(ret)
        for j in range(stencil):
            val += fact[1][j] * src_val[idx_src[0] + i, idx_src[1] + j]

        ret += fact[0][i] * val

    return ret

_doubleptr = ctypes.POINTER(ctypes.c_double)

class Interp2D_C:
    _lib    = None
    _interp_func = None

    _src_start_c  = None
    _src_step_c   = None
    _src_val_c    = None
    _src_stride_c = None
    _src_len_c    = None

    _x     = None
    _y     = None
    _ret   = None
    _x_c   = None
    _y_c   = None
    _ret_c = None

    def __init__(self, src_start, src_step, src_val, stencil):
        self._lib = ctypes.CDLL('lib_interp_c.so')


        func = getattr(self._lib, 'interp2d_%d' % stencil)
        func.argtypes = [_doubleptr, _doubleptr, _doubleptr,
            ctypes.POINTER(ctypes.c_int), ctypes.POINTER(ctypes.c_int),
            ctypes.c_int, _doubleptr, _doubleptr, _doubleptr]

        self._interp_func = func

        self._src_start_c  = (ctypes.c_double * 2)(src_start[0], src_start[1])
        self._src_step_c   = (ctypes.c_double * 2)(src_step[0],  src_step[1])
        self._src_val_c    = ctypes.cast(np.ctypeslib.as_ctypes(src_val), _doubleptr)
        self._src_stride_c = (ctypes.c_int * 2)(src_val.strides[0] // 8, src_val.strides[1] // 8)
        self._src_len_c    = (ctypes.c_int * 2)(src_val.shape[0], src_val.shape[1])

        self._setup_arrays((1,))

    def _setup_arrays(self, shape):
        self._x   = np.empty(shape)
        self._y   = np.empty(shape)
        self._ret = np.empty(shape)

        self._x_c   = ctypes.cast(np.ctypeslib.as_ctypes(self._x),   _doubleptr)
        self._y_c   = ctypes.cast(np.ctypeslib.as_ctypes(self._y),   _doubleptr)
        self._ret_c = ctypes.cast(np.ctypeslib.as_ctypes(self._ret), _doubleptr)

    def interp(self, x, y):
        if x.shape != self._x.shape:
            self._setup_arrays(x.shape)

        np.copyto(self._x,   x)
        np.copyto(self._y,   y)
        self._interp_func(self._src_start_c, self._src_step_c, self._src_val_c,
                          self._src_stride_c, self._src_len_c,
                          self._x.shape[0],
                          self._x_c, self._y_c, self._ret_c)
        return self._ret.copy()

    def __call__(self, x, y):
        return self.interp(x, y)
