import numpy as np
import tensorflow as tf

from ..._sparse_ops import (
    _bwd_segment_cmp,
    _csrmask,
    _csrmm,
    _csrsum,
    _gsddmm,
    _gspmm,
    _scatter_add,
    _segment_reduce,
)

from ...base import ALL, is_all
from ...heterograph_index import create_unitgraph_from_csr
from .tensor import asnumpy, context, copy_to, tensor, zerocopy_from_numpy

__all__ = [
    "gspmm",
    "gsddmm",
    "edge_softmax",
    "segment_reduce",
    "scatter_add",
    "csrmm",
    "csrsum",
    "csrmask",
]


def _scatter_nd(index, src, n_rows):
    assert index.shape == src.shape
    shp = index.shape
    ctx = context(src)
    ndim = index.ndim
    offsets = []
    stride = 1
    for i in reversed(range(1, ndim)):
        di = shp[i]
        offset_i = tf.range(di, dtype=index.dtype)
        offsets.append(
            tf.reshape(
                (stride * offset_i), (1,) * i + (di,) + (1,) * (ndim - 1 - i)
            )
        )
        stride *= di
    if ndim > 1:
        new_idx = index * stride + copy_to(sum(offsets), ctx)
    else:
        new_idx = index
    src = tf.reshape(src, (-1,))
    new_idx = tf.reshape(new_idx, (-1, 1))
    rst = tf.reshape(
        tf.scatter_nd(new_idx, src, (stride * n_rows,)), (n_rows, *shp[1:])
    )
    return rst


def _gather_nd(index, src):
    shp = index.shape
    ctx = context(src)
    ndim = index.ndim
    offsets = []
    stride = 1
    for i in reversed(range(1, ndim)):
        di = shp[i]
        offset_i = tf.range(di, dtype=index.dtype)
        offsets.append(
            tf.reshape(
                (stride * offset_i), (1,) * i + (di,) + (1,) * (ndim - 1 - i)
            )
        )
        stride *= di
    if ndim > 1:
        new_idx = index * stride + copy_to(sum(offsets), ctx)
    else:
        new_idx = index
    src = tf.reshape(src, (-1,))
    new_idx = tf.reshape(new_idx, (-1))
    rst = tf.reshape(tf.gather(src, new_idx), shp)
    return rst


def _reduce_grad(grad, shape):
    """Reduce gradient on the broadcast dimension
    If there is broadcast in forward pass, gradients need to be reduced on
    broadcast dimension. This function checks the input tensor shape and
    gradient shape and perform the reduction.
    Parameters
    ----------
    grad: Tensor
        Gradient tensor
    shape: tuple
        Shape of input tensor
    Returns
    -------
    Tensor
    """
    grad_shape = grad.shape[1:]
    in_shape = shape[1:]
    if in_shape == grad_shape:
        # no need to reduce
        return grad
    num_to_squeeze = len(grad_shape) - len(in_shape)
    # pad inshape
    in_shape = (1,) * num_to_squeeze + in_shape
    reduce_idx = np.asarray(
        np.nonzero(np.asarray(grad_shape) - np.asarray(in_shape))
    )
    reduce_idx += 1  # skip batch dim
    reduce_idx_tensor = tf.constant(
        tuple(reduce_idx.flatten().tolist()), dtype=tf.int32
    )
    grad = tf.reduce_sum(grad, axis=reduce_idx_tensor, keepdims=True)
    return tf.reshape(grad, shape)


def _need_reduce_last_dim(ufeat, efeat):
    """Indicates whether to reduce the last dimension on edges
    in the backward pass of spmm,
    if so, use dot instead of mul."""
    ushp = ufeat.shape
    eshp = efeat.shape
    return ushp[1:-1] == eshp[1:-1] and eshp[-1] == 1 and ushp[-1] > 1


def _muldiv(op, x):
    return 1.0 / x if op == "div" else x


def _addsub(op, x):
    return -x if op == "sub" else x


def _expand(x, shape):
    return tf.broadcast_to(x, (x.shape[0], *shape))


def gspmm_real(gidx, op, reduce_op, X, Y):
    out, (argX, argY) = _gspmm(gidx, op, reduce_op, X, Y)

    def grad(dZ):
        dZ = tensor(dZ)
        if op != "copy_rhs":
            g_rev = gidx.reverse()
            if reduce_op == "sum":
                if op in ["mul", "div"]:
                    dX = _gspmm(g_rev, "mul", "sum", dZ, _muldiv(op, Y))[0]
                elif op in ["add", "sub"]:
                    dX = _gspmm(g_rev, "copy_lhs", "sum", dZ, Y)[0]
                elif op == "copy_lhs":
                    dX = _gspmm(g_rev, "copy_lhs", "sum", dZ, None)[0]
            else:
                if op in ["mul", "div"]:
                    dX = _scatter_nd(
                        argX,
                        _muldiv(op, _gather_nd(argY, _expand(Y, dZ.shape[1:])))
                        * dZ,
                        X.shape[0],
                    )
                elif op in ["add", "sub", "copy_lhs"]:
                    dX = _scatter_nd(argX, dZ, X.shape[0])
            dX = _reduce_grad(dX, X.shape)
        else:
            dX = tf.zeros_like(X)
        if op != "copy_lhs":
            if reduce_op == "sum":
                if op == "mul" and _need_reduce_last_dim(X, Y):
                    dY = _gsddmm(gidx, "dot", X, dZ)
                elif op in ["mul", "div"]:
                    dY = _gsddmm(gidx, "mul", X, dZ)
                    if op == "div":
                        dY = -dY / (Y**2)
                elif op in ["add", "sub", "copy_rhs"]:
                    dY = _gsddmm(gidx, "copy_rhs", X, _addsub(op, dZ))
            else:
                out_shp = (Y.shape[0],) + dZ.shape[1:]
                if op in ["mul", "div"]:
                    dY = _scatter_nd(
                        argY,
                        _gather_nd(argX, _expand(X, dZ.shape[1:])) * dZ,
                        Y.shape[0],
                    )
                    if op == "div":
                        dY = -dY / (Y**2)
                elif op in ["add", "sub", "copy_rhs"]:
                    dY = _scatter_nd(argY, _addsub(op, dZ), Y.shape[0])
            dY = _reduce_grad(dY, Y.shape)
        else:
            dY = tf.zeros_like(Y)
        return dX, dY

    return out, grad


def gspmm(gidx, op, reduce_op, X, Y):
    @tf.custom_gradient
    def _lambda(X, Y):
        return gspmm_real(gidx, op, reduce_op, X, Y)

    if X is None:
        X = tf.zeros(())
    if Y is None:
        Y = tf.zeros(())
    return _lambda(X, Y)


def gsddmm_real(gidx, op, X, Y, lhs_target, rhs_target):
    out = _gsddmm(gidx, op, X, Y, lhs_target, rhs_target)

    def grad(dZ):
        if op != "copy_rhs":
            if lhs_target in ["u", "v"]:
                _gidx = gidx if lhs_target == "v" else gidx.reverse()
                if op in ["add", "sub", "copy_lhs"]:
                    dX = _gspmm(_gidx, "copy_rhs", "sum", None, dZ)[0]
                else:  # mul, div, dot
                    if rhs_target == lhs_target:
                        dX = _gspmm(_gidx, "copy_rhs", "sum", None, dZ)[
                            0
                        ] * _muldiv(op, Y)
                    elif rhs_target == "e":
                        dX = _gspmm(
                            _gidx, "copy_rhs", "sum", None, dZ * _muldiv(op, Y)
                        )[0]
                    else:  # rhs_target = !lhs_target
                        dX = _gspmm(_gidx, "mul", "sum", _muldiv(op, Y), dZ)[0]
            else:  # lhs_target == 'e'
                if op in ["add", "sub", "copy_lhs"]:
                    dX = dZ
                else:  # mul, div, dot
                    dX = _gsddmm(
                        gidx, "mul", dZ, _muldiv(op, Y), "e", rhs_target
                    )
            dX = _reduce_grad(dX, X.shape)
        else:
            dX = tf.zeros_like(X)
        if op != "copy_lhs":
            if rhs_target in ["u", "v"]:
                _gidx = gidx if rhs_target == "v" else gidx.reverse()
                if op in ["add", "sub", "copy_rhs"]:
                    dY = _gspmm(
                        _gidx, "copy_rhs", "sum", None, _addsub(op, dZ)
                    )[0]
                else:  # mul, div, dot
                    if lhs_target == rhs_target:
                        dY = _gspmm(_gidx, "copy_rhs", "sum", None, dZ)[0] * X
                    elif lhs_target == "e":
                        dY = _gspmm(_gidx, "copy_rhs", "sum", None, dZ * X)[0]
                    else:  # rhs_target = !lhs_target
                        dY = _gspmm(_gidx, "mul", "sum", X, dZ)[0]
                    if op == "div":
                        dY = -dY / (Y**2)
            else:
                if op in ["add", "sub", "copy_rhs"]:
                    dY = _addsub(op, dZ)
                else:  # mul, div, dot
                    dY = _gsddmm(gidx, "mul", dZ, X, "e", lhs_target)
                    if op == "div":
                        dY = -dY / (Y**2)
            dY = _reduce_grad(dY, Y.shape)
        else:
            dY = tf.zeros_like(Y)
        return dX, dY

    return out, grad


def gsddmm(gidx, op, X, Y, lhs_target="u", rhs_target="v"):
    @tf.custom_gradient
    def _lambda(X, Y):
        return gsddmm_real(gidx, op, X, Y, lhs_target, rhs_target)

    if X is None:
        X = tf.zeros(())
    if Y is None:
        Y = tf.zeros(())
    return _lambda(X, Y)


def edge_softmax_real(gidx, score, eids=ALL, norm_by="dst"):
    if not is_all(eids):
        gidx = gidx.edge_subgraph([eids], True).graph
    if norm_by == "src":
        gidx = gidx.reverse()
    score_max = _gspmm(gidx, "copy_rhs", "max", None, score)[0]
    score = tf.math.exp(_gsddmm(gidx, "sub", score, score_max, "e", "v"))
    score_sum = _gspmm(gidx, "copy_rhs", "sum", None, score)[0]
    out = _gsddmm(gidx, "div", score, score_sum, "e", "v")

    def edge_softmax_backward(grad_out):
        sds = out * grad_out
        accum = gspmm(gidx, "copy_rhs", "sum", None, sds)
        grad_score = sds - gsddmm(gidx, "mul", out, accum, "e", "v")
        return grad_score

    return out, edge_softmax_backward


def edge_softmax(gidx, logits, eids=ALL, norm_by="dst"):
    @tf.custom_gradient
    def _lambda(logits):
        return edge_softmax_real(gidx, logits, eids, norm_by)

    return _lambda(logits)


def segment_reduce_real(op, x, offsets):
    y, arg = _segment_reduce(op, x, offsets)

    def segment_reduce_backward(dy):
        m = x.shape[0]
        if op == "sum":
            offsets_np = asnumpy(offsets[1:])
            indices_np = np.zeros((m + 1,), dtype=offsets_np.dtype)
            np.add.at(indices_np, offsets_np, np.ones_like(offsets_np))
            indices_np = np.cumsum(indices_np, -1)[:-1]
            indices = zerocopy_from_numpy(indices_np)
            dx = tf.gather(dy, indices)
        else:
            dx = _bwd_segment_cmp(dy, arg, m)
        return dx

    return y, segment_reduce_backward


def segment_reduce(op, x, offsets):
    @tf.custom_gradient
    def _lambda(x):
        return segment_reduce_real(op, x, offsets)

    return _lambda(x)


def scatter_add_real(x, idx, m):
    y = _scatter_add(x, idx, m)

    def scatter_add_backward(dy):
        return tf.gather(dy, idx)

    return y, scatter_add_backward


def scatter_add(x, idx, m):
    @tf.custom_gradient
    def _lambda(x):
        return scatter_add_real(x, idx, m)

    return _lambda(x)


def csrmm_real(gidxA, A_weights, gidxB, B_weights, num_vtypes):
    gidxC, C_weights = _csrmm(gidxA, A_weights, gidxB, B_weights, num_vtypes)
    nrows, ncols, C_indptr, C_indices, C_eids = gidxC.adjacency_matrix_tensors(
        0, False, "csr"
    )

    def grad(dnrows, dncols, dC_indptr, dC_indices, dC_eids, dC_weights):
        # Only the last argument is meaningful.
        dgidxA, dA_weights = _csrmm(
            gidxC,
            dC_weights,
            gidxB.reverse(),
            B_weights,
            gidxA.number_of_ntypes(),
        )
        dgidxB, dB_weights = _csrmm(
            gidxA.reverse(),
            A_weights,
            gidxC,
            dC_weights,
            gidxB.number_of_ntypes(),
        )
        dA_weights = _csrmask(dgidxA, dA_weights, gidxA)
        dB_weights = _csrmask(dgidxB, dB_weights, gidxB)
        return dA_weights, dB_weights

    return (
        tf.constant(nrows),
        tf.constant(ncols),
        C_indptr,
        C_indices,
        C_eids,
        C_weights,
    ), grad


def csrmm(gidxA, A_weights, gidxB, B_weights, num_vtypes):
    @tf.custom_gradient
    def _lambda(A_weights, B_weights):
        return csrmm_real(gidxA, A_weights, gidxB, B_weights, num_vtypes)

    nrows, ncols, C_indptr, C_indices, C_eids, C_weights = _lambda(
        A_weights, B_weights
    )
    gidxC = create_unitgraph_from_csr(
        num_vtypes,
        nrows.numpy(),
        ncols.numpy(),
        C_indptr,
        C_indices,
        C_eids,
        ["coo", "csr", "csc"],
    )
    return gidxC, C_weights


def csrsum_real(gidxs, weights):
    gidxC, C_weights = _csrsum(gidxs, weights)
    nrows, ncols, C_indptr, C_indices, C_eids = gidxC.adjacency_matrix_tensors(
        0, False, "csr"
    )

    def grad(dnrows, dncols, dC_indptr, dC_indices, dC_eids, dC_weights):
        # Only the last argument is meaningful.
        return tuple(_csrmask(gidxC, dC_weights, gidx) for gidx in gidxs)

    return (
        tf.constant(nrows),
        tf.constant(ncols),
        C_indptr,
        C_indices,
        C_eids,
        C_weights,
    ), grad


def csrsum(gidxs, weights):
    @tf.custom_gradient
    def _lambda(*weights):
        return csrsum_real(gidxs, weights)

    nrows, ncols, C_indptr, C_indices, C_eids, C_weights = _lambda(*weights)
    num_vtypes = gidxs[0].number_of_ntypes()
    gidxC = create_unitgraph_from_csr(
        num_vtypes,
        nrows.numpy(),
        ncols.numpy(),
        C_indptr,
        C_indices,
        C_eids,
        ["coo", "csr", "csc"],
    )
    return gidxC, C_weights


def csrmask_real(gidxA, A_weights, gidxB):
    B_weights = _csrmask(gidxA, A_weights, gidxB)

    def grad(dB_weights):
        return _csrmask(gidxB, dB_weights, gidxA)

    return B_weights, grad


def csrmask(gidxA, A_weights, gidxB):
    @tf.custom_gradient
    def _lambda(A_weights):
        return csrmask_real(gidxA, A_weights, gidxB)

    return _lambda(A_weights)