library(microbenchmark)

dyn.load("wip.so")


## rowSum* .call --------------------------------------------------------------
rowSums.c <- function(M) {
    stopifnot(
        is.matrix(M),
        is.numeric(M)
    )
    if (!is.double(M)) {
        M <- matrix(as.double(M), nrow = nrow(M))
    }

    .Call('R_rowSums', PACKAGE = 'wip', M)
}
colSums.c <- function(M) {
    stopifnot(
        is.matrix(M),
        is.numeric(M)
    )
    if (!is.double(M)) {
        M <- matrix(as.double(M), nrow = nrow(M))
    }

    .Call('R_colSums', PACKAGE = 'wip', M)
}
rowSquareSums.c <- function(M) {
    stopifnot(
        is.matrix(M),
        is.numeric(M)
    )
    if (!is.double(M)) {
        M <- matrix(as.double(M), nrow = nrow(M))
    }

    .Call('R_rowSquareSums', PACKAGE = 'wip', M)
}
rowSumsSymVec.c <- function(vecA, nrow, diag = 0.0) {
    stopifnot(
        is.vector(vecA),
        is.numeric(vecA),
        is.numeric(diag),
        nrow * (nrow - 1) == length(vecA) * 2
    )
    if (!is.double(vecA)) {
        vecA <- as.double(vecA)
    }
    .Call('R_rowSumsSymVec', PACKAGE = 'wip',
          vecA, as.integer(nrow), as.double(diag))
}
rowSweep.c <- function(A, v, op = '-') {
    stopifnot(
        is.matrix(A),
        is.numeric(v)
    )
    if (!is.double(A)) {
        A <- matrix(as.double(A), nrow = nrow(A))
    }
    if (!is.vector(v) || !is.double(v)) {
        v <- as.double(v)
    }
    stopifnot(
        nrow(A) == length(v),
        op %in% c('+', '-', '*', '/')
    )

    .Call('R_rowSweep', PACKAGE = 'wip', A, v, op)
}

## row*, col* tests ------------------------------------------------------------
n <- 3000
M <- matrix(runif(n * 12), n, 12)
stopifnot(
    all.equal(rowSums(M^2), rowSums.c(M^2)),
    all.equal(colSums(M), colSums.c(M))
)
microbenchmark(
    rowSums     = rowSums(M^2),
    rowSums.c   = rowSums.c(M^2),
    rowSqSums.c = rowSquareSums.c(M)
)
microbenchmark(
    colSums   = colSums(M),
    colSums.c = colSums.c(M)
)

sum = rowSums(M)
stopifnot(all.equal(
    sweep(M, 1, sum, FUN = `/`),
    rowSweep.c(M, sum, '/') # Col-Normalize)
), all.equal(
    sweep(M, 1, sum, FUN = `/`),
    M / sum
))
microbenchmark(
    sweep = sweep(M, 1, sum, FUN = `/`),
    M / sum,
    rowSweep.c = rowSweep.c(M, sum, '/') # Col-Normalize)
)

# Create symmetric matrix with constant diagonal entries.
nrow <- 200
diag <- 1.0
Sym <- tcrossprod(runif(nrow))
diag(Sym) <- diag
# Get vectorized lower triangular part of `Sym` matrix.
SymVec <- Sym[lower.tri(Sym)]
stopifnot(all.equal(
    rowSums(Sym),
    rowSumsSymVec.c(SymVec, nrow, diag)
))
microbenchmark(
    rowSums = rowSums(Sym),
    rowSums.c = rowSums.c(Sym),
    rowSumsSymVec.c = rowSumsSymVec.c(SymVec, nrow, diag)
)

## Matrix-Matrix opperation .call ---------------------------------------------
transpose.c <- function(A) {
    stopifnot(
        is.matrix(A), is.numeric(A)
    )
    if (!is.double(A)) {
        A <- matrix(as.double(A), nrow(A), ncol(A))
    }

    .Call('R_transpose', PACKAGE = 'wip', A)
}

matrixprod.c <- function(A, B) {
    stopifnot(
        is.matrix(A), is.numeric(A),
        is.matrix(B), is.numeric(B),
        ncol(A) == nrow(B)
    )
    if (!is.double(A)) {
        A <- matrix(as.double(A), nrow = nrow(A))
    }
    if (!is.double(B)) {
        B <- matrix(as.double(B), nrow = nrow(B))
    }

    .Call('R_matrixprod', PACKAGE = 'wip', A, B)
}
crossprod.c <- function(A, B) {
    stopifnot(
        is.matrix(A), is.numeric(A),
        is.matrix(B), is.numeric(B),
        nrow(A) == nrow(B)
    )
    if (!is.double(A)) {
        A <- matrix(as.double(A), nrow = nrow(A))
    }
    if (!is.double(B)) {
        B <- matrix(as.double(B), nrow = nrow(B))
    }

    .Call('R_crossprod', PACKAGE = 'wip', A, B)
}
skewSymRank2k.c <- function(A, B, alpha = 1, beta = 0) {
    stopifnot(
        is.matrix(A), is.numeric(A),
        is.matrix(B), is.numeric(B),
        all(dim(A) == dim(B)),
        is.numeric(alpha), length(alpha) == 1L,
        is.numeric(beta), length(beta) == 1L
    )
    if (!is.double(A)) {
        A <- matrix(as.double(A), nrow = nrow(A))
    }
    if (!is.double(B)) {
        B <- matrix(as.double(B), nrow = nrow(B))
    }

    .Call('R_skewSymRank2k', PACKAGE = 'wip', A, B,
          as.double(alpha), as.double(beta))
}

## Matrix-Matrix opperation tests ---------------------------------------------
n <- 200
k <- 100
m <- 300

A <- matrix(runif(n * k), n, k)
B <- matrix(runif(k * m), k, m)
stopifnot(
    all.equal(t(A), transpose.c(A))
)
microbenchmark(
    t(A),
    transpose.c(A)
)

stopifnot(
    all.equal(A %*% B, matrixprod.c(A, B))
)
microbenchmark(
    "%*%" = A %*% B,
    matrixprod.c = matrixprod.c(A, B)
)

A <- matrix(runif(k * n), k, n)
B <- matrix(runif(k * m), k, m)
stopifnot(
    all.equal(crossprod(A, B), crossprod.c(A, B))
)
microbenchmark(
    crossprod = crossprod(A, B),
    crossprod.c = crossprod.c(A, B)
)

n <- 12
k <- 11
A <- matrix(runif(n * k), n, k)
B <- matrix(runif(n * k), n, k)
stopifnot(all.equal(
    A %*% t(B) - B %*% t(A), skewSymRank2k.c(A, B)
))
microbenchmark(
    A %*% t(B) - B %*% t(A),
    skewSymRank2k.c(A, B)
)

## Orthogonal projection onto null space .Call --------------------------------
nullProj.c <- function(B) {
    stopifnot(
        is.matrix(B), is.numeric(B)
    )
    if (!is.double(B)) {
        B <- matrix(as.double(B), nrow = nrow(B))
    }

    .Call('R_nullProj', PACKAGE = 'wip', B)
}
## Orthogonal projection onto null space tests --------------------------------
p <- 12
q <- 10
V <- qr.Q(qr(matrix(rnorm(p * q, 0, 1), p, q)))

# Projection matrix onto `span(V)`
Q <- diag(1, p) - tcrossprod(V, V)
stopifnot(
    all.equal(Q, nullProj.c(V))
)
microbenchmark(
    nullProj = diag(1, p) - tcrossprod(V, V),
    nullProj.c = nullProj.c(V)
)


# ## WIP for gradient. ----------------------------------------------------------

gradient.c <- function(X, X_diff, Y, V, h) {
    stopifnot(
        is.matrix(X), is.double(X),
        is.matrix(X_diff), is.double(X_diff),
            ncol(X_diff) == ncol(X), nrow(X_diff) == nrow(X) * (nrow(X) - 1) / 2,
        is.vector(Y) || (is.matrix(Y) && pmin(dim(Y)) == 1L), is.double(Y),
            length(Y) == nrow(X),
        is.matrix(V), is.double(V),
            nrow(V) == ncol(X),
        is.vector(h), is.numeric(h), length(h) == 1
    )

    .Call('R_gradient', PACKAGE = 'wip',
          X, X_diff, as.double(Y), V, as.double(h));
}

elem.pairs <- function(elements) {
    # Number of elements to match.
    n <- length(elements)
    # Create all combinations.
    pairs <- rbind(rep(elements, each=n), rep(elements, n))
    # Select unique combinations without self interaction.
    return(pairs[, pairs[1, ] < pairs[2, ]])
}
grad <- function(X, Y, V, h, persistent = TRUE) {
    n <- nrow(X)
    p <- ncol(X)

    if (!persistent) {
        pair.index <- elem.pairs(seq(n))
        i <- pair.index[, 1] # `i` indices of `(i, j)` pairs
        j <- pair.index[, 2] # `j` indices of `(i, j)` pairs
        lower <- ((i - 1) * n) + j
        upper <- ((j - 1) * n) + i
        X_diff <- X[i, , drop = F] - X[j, , drop = F]
    }

    # Projection matrix onto `span(V)`
    Q <- diag(1, p) - tcrossprod(V, V)
    # Vectorized distance matrix `D`.
    vecD <- rowSums((X_diff %*% Q)^2)


    # Weight matrix `W` (dnorm ... gaussean density function)
    W <- matrix(1, n, n) # `exp(0) == 1`
    W[lower] <- exp((-0.5 / h) * vecD^2) # Set lower tri. part
    W[upper] <- t.default(W)[upper] # Mirror lower tri. to upper
    W <- sweep(W, 2, colSums(W), FUN = `/`) # Col-Normalize

    # Weighted `Y` momentums
    y1 <- Y %*% W # Result is 1D -> transposition irrelevant
    y2 <- Y^2 %*% W

    # Per example loss `L(V, X_i)`
    L <- y2 - y1^2

    # Vectorized Weights with forced symmetry
    vecS <- (L[i] - (Y[j] - y1[i])^2) * W[lower]
    vecS <- vecS + ((L[j] - (Y[i] - y1[j])^2) * W[upper])
    # Compute scaling of `X` row differences
    vecS <- vecS * vecD

    G <- crossprod(X_diff, X_diff * vecS) %*% V
    G <- (-2 / (n * h^2)) * G
    return(G)
}
rStiefl <- function(p, q) {
    return(qr.Q(qr(matrix(rnorm(p * q, 0, 1), p, q))))
}

n <- 200
p <- 12
q <- 10

X <- matrix(runif(n * p), n, p)
Y <- runif(n)
V <- rStiefl(p, q)
h <- 0.1

pair.index <- elem.pairs(seq(n))
i <- pair.index[1, ] # `i` indices of `(i, j)` pairs
j <- pair.index[2, ] # `j` indices of `(i, j)` pairs
lower <- ((i - 1) * n) + j
upper <- ((j - 1) * n) + i
X_diff <- X[i, , drop = F] - X[j, , drop = F]

stopifnot(all.equal(
    grad(X, Y, V, h),
    gradient.c(X, X_diff, Y, V, h)
))
microbenchmark(
    grad = grad(X, Y, V, h),
    gradient.c = gradient.c(X, X_diff, Y, V, h)
)