from __future__ import absolute_import
import pycuda.gpuarray as gpuarray
import pycuda.elementwise as elementwise
import numpy as np
import warnings
from pycuda.driver import Stream


def _make_unary_array_func(name):
    def f(array, stream_or_out=None, **kwargs):

        if stream_or_out is not None:
            warnings.warn("please use 'out' or 'stream' keyword arguments", DeprecationWarning)
            if isinstance(stream_or_out, Stream):
                stream = stream_or_out
                out = None
            else:
                stream = None
                out = stream_or_out

        out, stream = None, None
        if 'out' in kwargs:
            out = kwargs['out']
        if 'stream' in kwargs:
            stream = kwargs['stream']

        if array.dtype == np.float32:
            func_name = name + "f"
        else:
            func_name = name

        if not array.flags.forc:
            raise RuntimeError("only contiguous arrays may "
                    "be used as arguments to this operation")

        if out is None:
            out = array._new_like_me()
        else:
            assert out.dtype == array.dtype
            assert out.strides == array.strides
            assert out.shape == array.shape

        func = elementwise.get_unary_func_kernel(func_name, array.dtype)
        func.prepared_async_call(array._grid, array._block, stream,
                array.gpudata, out.gpudata, array.mem_size)

        return out
    return f

fabs = _make_unary_array_func("fabs")
ceil = _make_unary_array_func("ceil")
floor = _make_unary_array_func("floor")
exp = _make_unary_array_func("exp")
log = _make_unary_array_func("log")
log10 = _make_unary_array_func("log10")
sqrt = _make_unary_array_func("sqrt")

sin = _make_unary_array_func("sin")
cos = _make_unary_array_func("cos")
tan = _make_unary_array_func("tan")
asin = _make_unary_array_func("asin")
acos = _make_unary_array_func("acos")
atan = _make_unary_array_func("atan")

sinh = _make_unary_array_func("sinh")
cosh = _make_unary_array_func("cosh")
tanh = _make_unary_array_func("tanh")

def fmod(arg, mod, stream=None):
    """Return the floating point remainder of the division `arg/mod`,
    for each element in `arg` and `mod`."""
    result = gpuarray.GPUArray(arg.shape, arg.dtype)

    if not arg.flags.forc or not mod.flags.forc:
        raise RuntimeError("only contiguous arrays may "
                "be used as arguments to this operation")

    func = elementwise.get_fmod_kernel()
    func.prepared_async_call(arg._grid, arg._block, stream,
            arg.gpudata, mod.gpudata, result.gpudata, arg.mem_size)

    return result

def frexp(arg, stream=None):
    """Return a tuple `(significands, exponents)` such that
    `arg == significand * 2**exponent`.
    """
    if not arg.flags.forc:
        raise RuntimeError("only contiguous arrays may "
                "be used as arguments to this operation")

    sig = gpuarray.GPUArray(arg.shape, arg.dtype)
    expt = gpuarray.GPUArray(arg.shape, arg.dtype)

    func = elementwise.get_frexp_kernel()
    func.prepared_async_call(arg._grid, arg._block, stream,
            arg.gpudata, sig.gpudata, expt.gpudata, arg.mem_size)

    return sig, expt

def ldexp(significand, exponent, stream=None):
    """Return a new array of floating point values composed from the
    entries of `significand` and `exponent`, paired together as
    `result = significand * 2**exponent`.
    """
    if not significand.flags.forc or not exponent.flags.forc:
        raise RuntimeError("only contiguous arrays may "
                "be used as arguments to this operation")

    result = gpuarray.GPUArray(significand.shape, significand.dtype)

    func = elementwise.get_ldexp_kernel()
    func.prepared_async_call(significand._grid, significand._block, stream,
            significand.gpudata, exponent.gpudata, result.gpudata,
            significand.mem_size)

    return result

def modf(arg, stream=None):
    """Return a tuple `(fracpart, intpart)` of arrays containing the
    integer and fractional parts of `arg`.
    """
    if not arg.flags.forc:
        raise RuntimeError("only contiguous arrays may "
                "be used as arguments to this operation")

    intpart = gpuarray.GPUArray(arg.shape, arg.dtype)
    fracpart = gpuarray.GPUArray(arg.shape, arg.dtype)

    func = elementwise.get_modf_kernel()
    func.prepared_async_call(arg._grid, arg._block, stream,
            arg.gpudata, intpart.gpudata, fracpart.gpudata,
            arg.mem_size)

    return fracpart, intpart