library(tidyverse)
library(Matrix)

set.seed(1234)

m = 10
n = 8
d = 5

f = \(x,y) { sqrt(sum((x - y)^2)) }
g = \(x,y) { f(x, y) <= 3 }

# ----- Test outer1, outer1_triplet, and outer1_matvec -----
x = rnorm(n*d)
X = matrix(x, n, d)

# Expected result
E = matrix(0, n, n)
for (i in 1:n) {
	for (j in 1:n) {
		E[i,j] = f(X[i,], X[j,])
	}
}
a = rnorm(n)

expect_error( outer1(x, f) )

A = outer1(X, f)
expect_equal(dim(A), c(n,n))
expect_equal(A, E)

sp = outer1_triplet(X, f, g = \(x,y) { TRUE })
B = sparseMatrix(sp$i, sp$j, x = sp$x, dims = c(sp$m, sp$n), symmetric = T)
expect_true(all(sp$i >= 1))
expect_true(all(sp$i <= m))
expect_true(all(sp$j >= 1))
expect_true(all(sp$j <= n))
expect_equal(as.matrix(B), E)

sp = outer1_triplet(X, f, g)
B = sparseMatrix(sp$i, sp$j, x = sp$x, dims = c(sp$m, sp$n), symmetric = T)
idx0 = which(E > 3, arr.ind = T)
idx1 = which(E <= 3, arr.ind = T)
expect_equal(B[idx1], E[idx1])
expect_true(all(B[idx0] == 0))

expect_true(all.equal(outer1_matvec(X, f, a), E %*% a,
	check.attributes = FALSE, check.class = FALSE))

# ----- Test outer2, outer2_triplet, and outer2_matvec -----
x = rnorm(m*d)
y = rnorm(n*d)
X = matrix(x, m, d)
Y = matrix(y, n, d)

# Expected result
E = matrix(0, m, n)
for (i in 1:m) {
	for (j in 1:n) {
		E[i,j] = f(X[i,], Y[j,])
	}
}
a = rnorm(n)

expect_error( outer2(x, y, f) )
expect_error( outer2(X, y, f) )
expect_error( outer2(x, Y, f) )

A = outer2(X, Y, f)
expect_equal(dim(A), c(m,n))
expect_equal(A, E)

sp = outer2_triplet(X, Y, f, g = \(x,y) { TRUE })
B = sparseMatrix(sp$i, sp$j, x = sp$x, dims = c(sp$m, sp$n), symmetric = F)
expect_true(all(sp$i >= 1))
expect_true(all(sp$i <= m))
expect_true(all(sp$j >= 1))
expect_true(all(sp$j <= n))
expect_equal(as.matrix(B), E)

sp = outer2_triplet(X, Y, f, g)
B = sparseMatrix(sp$i, sp$j, x = sp$x, dims = c(sp$m, sp$n), symmetric = F)
idx0 = which(E > 3, arr.ind = T)
idx1 = which(E <= 3, arr.ind = T)
expect_equal(B[idx1], E[idx1])
expect_true(all(B[idx0] == 0))

expect_true(all.equal(outer2_matvec(X, Y, f, a), E %*% a,
	check.attributes = FALSE, check.class = FALSE))
