### =========================================================================
### 'Arith' operations on SparseArray objects
### -------------------------------------------------------------------------
###
### 'Arith' operations: "+", "-", "*", "/", "^", "%%", "%/%"
###
### See '?S4groupGeneric' for more information.
###
### We also implement unary "+" and "-" for SparseArray objects.
###


ARITH_INPUT_TYPES <- c("integer", "double", "complex")

check_Arith_input_type <- function(type, what)
{
    if (!(type %in% ARITH_INPUT_TYPES))
        stop(wmsg("arithmetic operation not supported ",
                  "on ", what, " of type() \"", type , "\""))
}

check_vector_operand_length <- function(y_len, x_dim, recycle.along,
                                        what_x, side=c("right", "left"))
{
    stopifnot(isSingleInteger(y_len),
              is.integer(x_dim),
              isSingleInteger(recycle.along),
              isSingleString(what_x))
    side <- match.arg(side)

    dim1 <- x_dim[[recycle.along]]
    if (y_len == dim1 || y_len == 1L)
        return()
    operand <- paste(side, "vector")
    if (recycle.along == 1L) {
        which_dim <- "first (a.k.a. innermost) dimension"
    } else if (recycle.along == length(x_dim)) {
        which_dim <- "last (a.k.a. outermost) dimension"
    } else {
        th <- switch(as.character(recycle.along %% 10L),
                     `2`="nd", `3`="rd", "th")
        which_dim <- paste0(recycle.along, th, " dimension")
    }
    what <- paste(which_dim, " of ", what_x)
    if (y_len > dim1)
        stop(wmsg(operand, " is longer than ", what))
    if (y_len == 0L)
        stop(wmsg(operand, " length cannot be 0 unless ", what, " is 0"))
    if (dim1 %% y_len != 0L)
        warning(wmsg(what, " is not a multiple of ", operand, " length"))
}

op_is_commutative <- function(op)
    (op %in% c("+", "*", "==", "!=", "&", "|"))

error_on_left_sparsity_not_preserved <- function(op, when)
{
    flipped_op <- flip_Compare_op(op)
    show_flipped_op <- flipped_op != op || op_is_commutative(op)
    if (show_flipped_op) {
        msg <- c("'x ", op, " y' and 'y ", flipped_op, " x': operations")
    } else {
        msg <- c("x ", op, " y: operation")
    }
    stop(wmsg(msg, " not supported on SparseArray object x ",
              "when ", when, " (result wouldn't be sparse)"))
}

get_Arith_output_type <- function(op, x_type, y_type)
{
    if (op %in% c("/", "^") && x_type == "integer" && y_type == "integer")
        return("double")
    output_type <- type(c(vector(x_type), vector(y_type)))
    if (output_type == "complex") {
        ## temporary
        stop(wmsg("'x ", op, " y' is not supported yet when 'x' or 'y' ",
                  "is an SVT_SparseArray or NaArray object ",
                  "of type() \"complex\""))
    }
    if (op %in% c("%%", "%/%") && output_type == "complex")
        stop(wmsg("unimplemented complex operation"))
    output_type
}


### - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
### Unary "+" and "-"
###

.unary_plus_SparseArray <- function(x)
{
    check_Arith_input_type(type(x), "SparseArray object")
    x  # no-op
}

.unary_minus_SparseArray <- function(x)
{
    check_Arith_input_type(type(x), "SparseArray object")
    if (is(x, "COO_SparseArray")) {
        ans <- BiocGenerics:::replaceSlots(x, nzdata=-x@nzdata, check=FALSE)
    } else if (is(x, "SVT_SparseArray")) {
        check_svt_version(x)
        new_SVT <- SparseArray.Call("C_unary_minus_SVT", x@dim, x@type, x@SVT)
        ans <- BiocGenerics:::replaceSlots(x, SVT=new_SVT, check=FALSE)
    } else {
        stop(wmsg("unary \"-\" is not supported on ", class(x), " objects"))
    }
    ans
}

setMethod("+", c("SparseArray", "missing"),
    function(e1, e2) .unary_plus_SparseArray(e1)
)

setMethod("-", c("SparseArray", "missing"),
    function(e1, e2) .unary_minus_SparseArray(e1)
)


### - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - - -
### 'Arith' group
###

.check_right_vector_for_Arith_SVT1_v2 <- function(y, x_dim, op,
                                                  recycle.along=1L)
{
    check_vector_operand_length(length(y), x_dim, recycle.along,
                                "SparseArray object")
    if (anyNA(y))
        error_on_left_sparsity_not_preserved(op,
                 "y contains NA or NaN values")
    if (op == "*" && any(is.infinite(y)))
        error_on_left_sparsity_not_preserved(op,
                 "y contains infinite values")
    if (op == "^" && any(y <= 0))
        error_on_left_sparsity_not_preserved(op,
                 "y contains non-positive values")
    if (op != "*" && any(y == 0))
        error_on_left_sparsity_not_preserved(op,
                 "y contains zeros")
}

### Supports: "*", "/", "^", "%%", "%/%"
### Returns an SVT_SparseArray object.
.Arith_SVT1_v2 <- function(op, x, y, recycle.along=1L)
{
    stopifnot(isSingleString(op),
              is(x, "SVT_SparseArray"),
              isSingleInteger(recycle.along))
    check_svt_version(x)

    ## Check types.
    check_Arith_input_type(type(x), "SparseArray object")
    if (!is.atomic(y))
        stop(wmsg("arithmetic operations between SparseArray objects ",
                  "and non-atomic vectors are not supported"))
    if (!(type(y) %in% ARITH_INPUT_TYPES))
        stop(wmsg("arithmetic operations between SparseArray objects ",
                  "and ", class(y), " vectors are not supported"))

    ## Check 'op'.
    if (!(op %in% c("*", "/", "^", "%%", "%/%")))
        stop(wmsg("\"", op, "\" is not supported between a SparseArray ",
                  "object and a ", class(y), " vector (result wouldn't ",
                  "be sparse in general)"))

    ## Check 'y'.
    .check_right_vector_for_Arith_SVT1_v2(y, dim(x), op,
                                          recycle.along=recycle.along)

    ## Compute 'ans_type'.
    if (type(x) == "double" && type(y) == "integer" || op %in% c("/", "^"))
        type(y) <- "double"  # cheap
    ans_type <- get_Arith_output_type(op, type(x), type(y))

    new_SVT <- SparseArray.Call("C_Arith_SVT1_v2",
                                x@dim, x@type, x@SVT, FALSE,
                                y, recycle.along, op, ans_type)
    BiocGenerics:::replaceSlots(x, type=ans_type, SVT=new_SVT, check=FALSE)
}

setMethod("Arith", c("SVT_SparseArray", "vector"),
    function(e1, e2) .Arith_SVT1_v2(.Generic, e1, e2)
)

setMethod("Arith", c("vector", "SVT_SparseArray"),
    function(e1, e2) {
        if (.Generic != "*")
            stop(wmsg("\"", .Generic, "\" is not supported between ",
                      "a ", class(e1), " vector on the left and an ",
                      "SVT_SparseArray object on the right (result ",
                      "wouldn't be sparse in general)"))
        .Arith_SVT1_v2(.Generic, e2, e1)
    }
)

### Supports: "+", "-", "*"
### Returns an SVT_SparseArray object.
.Arith_SVT1_SVT2 <- function(op, x, y)
{
    stopifnot(isSingleString(op),
              is(x, "SVT_SparseArray"),
              is(y, "SVT_SparseArray"))
    check_svt_version(x)
    check_svt_version(y)

    ## Check types.
    check_Arith_input_type(type(x), "SparseArray object")
    check_Arith_input_type(type(y), "SparseArray object")

    ## Check 'op'.
    if (!(op %in% c("+", "-", "*")))
        stop(wmsg("\"", op, "\" is not supported between SparseArray ",
                  "objects (result wouldn't be sparse in general)"))

    ## Check array conformability.
    x_dim <- dim(x)
    y_dim <- dim(y)
    if (!identical(x_dim, y_dim))
        stop(wmsg("non-conformable arrays"))

    ## Compute 'ans_dimnames'.
    ans_dimnames <- S4Arrays:::get_first_non_NULL_dimnames(list(x, y))

    ## Compute 'ans_type'.
    ans_type <- get_Arith_output_type(op, type(x), type(y))

    ans_SVT <- SparseArray.Call("C_Arith_SVT1_SVT2",
                                x_dim, x@type, x@SVT, FALSE,
                                y_dim, y@type, y@SVT, FALSE, op, ans_type)
    new_SVT_SparseArray(x_dim, ans_dimnames, ans_type, ans_SVT, check=FALSE)
}

setMethod("Arith", c("SVT_SparseArray", "SVT_SparseArray"),
    function(e1, e2) .Arith_SVT1_SVT2(.Generic, e1, e2)
)

### We could either:
###  (a) Turn 'e2' into an SVT_SparseArray object and return an
###      SVT_SparseArray object. This is the dgCMatrix approach.
###  (b) Turn 'e1' into an ordinary array and return an ordinary array.
### We choose to do (a). Note that there's no best choice in general as it
### really depends on the 'Arith' operation (i.e. "+", "-", or "*") and
### whether 'e2' has a lot of zeros or not. If it has no or very little
### zeros, then (a) will tend to be less memory efficient than (b) when
### doing "+" or "-". When doing "*", (a) should be always more memory
### efficient than (b), no matter what.
### The cautious user would typically make that choice upfront anyway, by
### coercing one or the other object before calling the 'Arith' op on them.
setMethod("Arith", c("SVT_SparseArray", "array"),
    function(e1, e2) .Arith_SVT1_SVT2(.Generic, e1, as(e2, "SVT_SparseArray"))
)

setMethod("Arith", c("array", "SVT_SparseArray"),
    function(e1, e2) .Arith_SVT1_SVT2(.Generic, as(e1, "SVT_SparseArray"), e2)
)

