#######################################################
# Copyright (c) 2015, ArrayFire
# All rights reserved.
#
# This file is distributed under 3-clause BSD license.
# The complete license agreement can be obtained at:
# http://arrayfire.com/licenses/BSD-3-Clause
########################################################

import inspect
from .library import *
from .util import *
from .broadcast import *
from .base import *
from .index import *

def create_array(buf, numdims, idims, dtype):
    out_arr = ct.c_longlong(0)
    ct.c_dims = dim4(idims[0], idims[1], idims[2], idims[3])
    safe_call(clib.af_create_array(ct.pointer(out_arr), ct.c_longlong(buf),\
                                   numdims, ct.pointer(ct.c_dims), dtype))
    return out_arr

def constant_array(val, d0, d1=None, d2=None, d3=None, dtype=f32):

    if not isinstance(dtype, ct.c_int):
        if isinstance(dtype, int):
            dtype = ct.c_int(dtype)
        else:
            raise TypeError("Invalid dtype")

    out = ct.c_longlong(0)
    dims = dim4(d0, d1, d2, d3)

    if isinstance(val, complex):
        c_real = ct.c_double(val.real)
        c_imag = ct.c_double(val.imag)

        if (dtype != c32 and dtype != c64):
            dtype = c32

        safe_call(clib.af_constant_complex(ct.pointer(out), c_real, c_imag,\
                                           4, ct.pointer(dims), dtype))
    elif dtype == s64:
        c_val = ct.c_longlong(val.real)
        safe_call(clib.af_constant_long(ct.pointer(out), c_val, 4, ct.pointer(dims)))
    elif dtype == u64:
        c_val = ct.c_ulonglong(val.real)
        safe_call(clib.af_constant_ulong(ct.pointer(out), c_val, 4, ct.pointer(dims)))
    else:
        c_val = ct.c_double(val)
        safe_call(clib.af_constant(ct.pointer(out), c_val, 4, ct.pointer(dims), dtype))

    return out


def binary_func(lhs, rhs, c_func):
    out = array()
    other = rhs

    if (is_number(rhs)):
        ldims = dim4_tuple(lhs.dims())
        lty = lhs.type()
        other = array()
        other.arr = constant_array(rhs, ldims[0], ldims[1], ldims[2], ldims[3], lty)
    elif not isinstance(rhs, array):
        raise TypeError("Invalid parameter to binary function")

    safe_call(c_func(ct.pointer(out.arr), lhs.arr, other.arr, bcast.get()))

    return out

def binary_funcr(lhs, rhs, c_func):
    out = array()
    other = lhs

    if (is_number(lhs)):
        rdims = dim4_tuple(rhs.dims())
        rty = rhs.type()
        other = array()
        other.arr = constant_array(lhs, rdims[0], rdims[1], rdims[2], rdims[3], rty)
    elif not isinstance(lhs, array):
        raise TypeError("Invalid parameter to binary function")

    c_func(ct.pointer(out.arr), other.arr, rhs.arr, bcast.get())

    return out

def transpose(a, conj=False):
    out = array()
    safe_call(clib.af_transpose(ct.pointer(out.arr), a.arr, conj))
    return out

def transpose_inplace(a, conj=False):
    safe_call(clib.af_transpose_inplace(a.arr, conj))

def ctype_to_lists(ctype_arr, dim, shape, offset=0):
    if (dim == 0):
        return list(ctype_arr[offset : offset + shape[0]])
    else:
        dim_len = shape[dim]
        res = [[]] * dim_len
        for n in range(dim_len):
            res[n] = ctype_to_lists(ctype_arr, dim - 1, shape, offset)
            offset += shape[0]
        return res

def get_info(dims, buf_len):
    elements = 1
    numdims = len(dims)
    idims = [1]*4

    for i in range(numdims):
        elements *= dims[i]
        idims[i] = dims[i]

    if (elements == 0):
        if (buf_len != 0):
            idims = [buf_len, 1, 1, 1]
            numdims = 1
        else:
            raise RuntimeError("Invalid size")

    return numdims, idims


class array(base_array):

    def __init__(self, src=None, dims=(0,), type_char=None):

        super(array, self).__init__()

        buf=None
        buf_len=0
        _type_char='f'
        dtype = f32

        if src is not None:

            if (isinstance(src, array)):
                safe_call(clib.af_retain_array(ct.pointer(self.arr), src.arr))
                return

            host = __import__("array")

            if isinstance(src, host.array):
                buf,buf_len = src.buffer_info()
                _type_char = src.typecode
                numdims, idims = get_info(dims, buf_len)
            elif isinstance(src, list):
                tmp = host.array('f', src)
                buf,buf_len = tmp.buffer_info()
                _type_char = tmp.typecode
                numdims, idims = get_info(dims, buf_len)
            elif isinstance(src, int) or isinstance(src, ct.c_ulonglong):
                buf = src
                numdims, idims = get_info(dims, buf_len)

                elements = 1
                for dim in idims:
                    elements *= dim

                if (elements == 0):
                    raise RuntimeError("Expected dims when src is data pointer")

                if (type_char is None):
                    raise TypeError("Expected type_char when src is data pointer")

                _type_char = type_char

            else:
                raise TypeError("src is an object of unsupported class")

            if (type_char is not None and
                type_char != _type_char):
                raise TypeError("Can not create array of requested type from input data type")

            self.arr = create_array(buf, numdims, idims, to_dtype[_type_char])

    def copy(self):
        out = array()
        safe_call(clib.af_retain_array(ct.pointer(out.arr), self.arr))
        return out

    def __del__(self):
        if (self.arr.value != 0):
            clib.af_release_array(self.arr)

    def elements(self):
        num = ct.c_ulonglong(0)
        safe_call(clib.af_get_elements(ct.pointer(num), self.arr))
        return num.value

    def type(self):
        dty = ct.c_int(f32.value)
        safe_call(clib.af_get_type(ct.pointer(dty), self.arr))
        return dty.value

    def dims(self):
        d0 = ct.c_longlong(0)
        d1 = ct.c_longlong(0)
        d2 = ct.c_longlong(0)
        d3 = ct.c_longlong(0)
        safe_call(clib.af_get_dims(ct.pointer(d0), ct.pointer(d1),\
                                   ct.pointer(d2), ct.pointer(d3), self.arr))
        dims = (d0.value,d1.value,d2.value,d3.value)
        return dims[:self.numdims()]

    def numdims(self):
        nd = ct.c_uint(0)
        safe_call(clib.af_get_numdims(ct.pointer(nd), self.arr))
        return nd.value

    def is_empty(self):
        res = ct.c_bool(False)
        safe_call(clib.af_is_empty(ct.pointer(res), self.arr))
        return res.value

    def is_scalar(self):
        res = ct.c_bool(False)
        safe_call(clib.af_is_scalar(ct.pointer(res), self.arr))
        return res.value

    def is_row(self):
        res = ct.c_bool(False)
        safe_call(clib.af_is_row(ct.pointer(res), self.arr))
        return res.value

    def is_column(self):
        res = ct.c_bool(False)
        safe_call(clib.af_is_column(ct.pointer(res), self.arr))
        return res.value

    def is_vector(self):
        res = ct.c_bool(False)
        safe_call(clib.af_is_vector(ct.pointer(res), self.arr))
        return res.value

    def is_complex(self):
        res = ct.c_bool(False)
        safe_call(clib.af_is_complex(ct.pointer(res), self.arr))
        return res.value

    def is_real(self):
        res = ct.c_bool(False)
        safe_call(clib.af_is_real(ct.pointer(res), self.arr))
        return res.value

    def is_double(self):
        res = ct.c_bool(False)
        safe_call(clib.af_is_double(ct.pointer(res), self.arr))
        return res.value

    def is_single(self):
        res = ct.c_bool(False)
        safe_call(clib.af_is_single(ct.pointer(res), self.arr))
        return res.value

    def is_real_floating(self):
        res = ct.c_bool(False)
        safe_call(clib.af_is_realfloating(ct.pointer(res), self.arr))
        return res.value

    def is_floating(self):
        res = ct.c_bool(False)
        safe_call(clib.af_is_floating(ct.pointer(res), self.arr))
        return res.value

    def is_integer(self):
        res = ct.c_bool(False)
        safe_call(clib.af_is_integer(ct.pointer(res), self.arr))
        return res.value

    def is_bool(self):
        res = ct.c_bool(False)
        safe_call(clib.af_is_bool(ct.pointer(res), self.arr))
        return res.value

    def __add__(self, other):
        return binary_func(self, other, clib.af_add)

    def __iadd__(self, other):
        self = binary_func(self, other, clib.af_add)
        return self

    def __radd__(self, other):
        return binary_funcr(other, self, clib.af_add)

    def __sub__(self, other):
        return binary_func(self, other, clib.af_sub)

    def __isub__(self, other):
        self = binary_func(self, other, clib.af_sub)
        return self

    def __rsub__(self, other):
        return binary_funcr(other, self, clib.af_sub)

    def __mul__(self, other):
        return binary_func(self, other, clib.af_mul)

    def __imul__(self, other):
        self = binary_func(self, other, clib.af_mul)
        return self

    def __rmul__(self, other):
        return binary_funcr(other, self, clib.af_mul)

    # Necessary for python3
    def __truediv__(self, other):
        return binary_func(self, other, clib.af_div)

    def __itruediv__(self, other):
        self =  binary_func(self, other, clib.af_div)
        return self

    def __rtruediv__(self, other):
        return binary_funcr(other, self, clib.af_div)

    # Necessary for python2
    def __div__(self, other):
        return binary_func(self, other, clib.af_div)

    def __idiv__(self, other):
        self =  binary_func(self, other, clib.af_div)
        return self

    def __rdiv__(self, other):
        return binary_funcr(other, self, clib.af_div)

    def __mod__(self, other):
        return binary_func(self, other, clib.af_mod)

    def __imod__(self, other):
        self =  binary_func(self, other, clib.af_mod)
        return self

    def __rmod__(self, other):
        return binary_funcr(other, self, clib.af_mod)

    def __pow__(self, other):
        return binary_func(self, other, clib.af_pow)

    def __ipow__(self, other):
        self =  binary_func(self, other, clib.af_pow)
        return self

    def __rpow__(self, other):
        return binary_funcr(other, self, clib.af_pow)

    def __lt__(self, other):
        return binary_func(self, other, clib.af_lt)

    def __gt__(self, other):
        return binary_func(self, other, clib.af_gt)

    def __le__(self, other):
        return binary_func(self, other, clib.af_le)

    def __ge__(self, other):
        return binary_func(self, other, clib.af_ge)

    def __eq__(self, other):
        return binary_func(self, other, clib.af_eq)

    def __ne__(self, other):
        return binary_func(self, other, clib.af_neq)

    def __and__(self, other):
        return binary_func(self, other, clib.af_bitand)

    def __iand__(self, other):
        self = binary_func(self, other, clib.af_bitand)
        return self

    def __or__(self, other):
        return binary_func(self, other, clib.af_bitor)

    def __ior__(self, other):
        self = binary_func(self, other, clib.af_bitor)
        return self

    def __xor__(self, other):
        return binary_func(self, other, clib.af_bitxor)

    def __ixor__(self, other):
        self = binary_func(self, other, clib.af_bitxor)
        return self

    def __lshift__(self, other):
        return binary_func(self, other, clib.af_bitshiftl)

    def __ilshift__(self, other):
        self = binary_func(self, other, clib.af_bitshiftl)
        return self

    def __rshift__(self, other):
        return binary_func(self, other, clib.af_bitshiftr)

    def __irshift__(self, other):
        self = binary_func(self, other, clib.af_bitshiftr)
        return self

    def __neg__(self):
        return 0 - self

    def __pos__(self):
        return self

    def __invert__(self):
        return self == 0

    def __nonzero__(self):
        return self != 0

    # TODO:
    # def __abs__(self):
    #     return self

    def __getitem__(self, key):
        try:
            out = array()
            n_dims = self.numdims()
            inds = get_indices(key, n_dims)

            safe_call(clib.af_index_gen(ct.pointer(out.arr),\
                                        self.arr, ct.c_longlong(n_dims), ct.pointer(inds)))
            return out
        except RuntimeError as e:
            raise IndexError(str(e))


    def __setitem__(self, key, val):
        try:
            n_dims = self.numdims()

            if (is_number(val)):
                tdims = get_assign_dims(key, self.dims())
                other_arr = constant_array(val, tdims[0], tdims[1], tdims[2], tdims[3])
            else:
                other_arr = val.arr

            out_arr = ct.c_longlong(0)
            inds  = get_indices(key, n_dims)

            safe_call(clib.af_assign_gen(ct.pointer(out_arr),\
                                         self.arr, ct.c_longlong(n_dims), ct.pointer(inds),\
                                         other_arr))
            safe_call(clib.af_release_array(self.arr))
            self.arr = out_arr

        except RuntimeError as e:
            raise IndexError(str(e))

    def to_ctype(self, row_major=False, return_shape=False):

        if (self.arr.value == 0):
            raise RuntimeError("Can not call to_ctype on empty array")

        tmp = transpose(self) if row_major else self
        ctype_type = to_c_type[self.type()] * self.elements()
        res = ctype_type()
        safe_call(clib.af_get_data_ptr(ct.pointer(res), self.arr))
        if (return_shape):
            return res, self.dims()
        else:
            return res

    def to_array(self, row_major=False, return_shape=False):

        if (self.arr.value == 0):
            raise RuntimeError("Can not call to_array on empty array")

        res = self.to_ctype(row_major, return_shape)

        host = __import__("array")
        h_type = to_typecode[self.type()]

        if (return_shape):
            return host.array(h_type, res[0]), res[1]
        else:
            return host.array(h_type, res)

    def to_list(self, row_major=False):
        ct_array, shape = self.to_ctype(row_major, True)
        return ctype_to_lists(ct_array, len(shape) - 1, shape)

def display(a):
    expr = inspect.stack()[1][-2]
    if (expr is not None):
        print('%s' % expr[0].split('display(')[1][:-2])
    safe_call(clib.af_print_array(a.arr))