library(tensorPredictors)

plot.tsir.vs.gmlm <- function(fit.gmlm, fit.tsir) {
    par.reset <- par(mfrow = c(2, 2))

    betas.gmlm <- fit.gmlm$betas
    Omegas <- fit.gmlm$Omegas

    betas.tsir <- fit.tsir$betas

    b1.gmlm <- (length(betas.gmlm[[1]]) / norm(betas.gmlm[[1]])) * betas.gmlm[[1]]
    b1.tsir <- (length(betas.tsir[[1]]) / norm(betas.tsir[[1]])) * betas.tsir[[1]]
    if (sum(abs(b1.gmlm + b1.tsir)) < sum(abs(b1.gmlm - b1.tsir))) {
        b1.tsir <- -b1.tsir
    }
    time <- seq(0, 1, len = 256)
    plot(range(time), range(b1.gmlm, b1.tsir), type = "n",
        main = "Time", xlab = "Time [s]", ylab = expression(beta[1])
    )
    lines(time, b1.gmlm, col = "blue")
    lines(time, b1.tsir, col = "red")

    b2.gmlm <- (length(betas.gmlm[[2]]) / norm(betas.gmlm[[2]])) * betas.gmlm[[2]]
    b2.tsir <- (length(betas.tsir[[2]]) / norm(betas.tsir[[2]])) * betas.tsir[[2]]
    if (sum(abs(b2.gmlm + b2.tsir)) < sum(abs(b2.gmlm - b2.tsir))) {
        b2.tsir <- -b2.tsir
    }
    plot(c(1, length(b2.gmlm)), range(b2.gmlm, b2.tsir), type = "n",
        main = "Sensors", xlab = "Sensor Index", ylab = expression(beta[2])
    )
    lines(b2.gmlm, lty = 3, col = "blue")
    lines(b2.tsir, lty = 3, col = "red")
    points(b2.gmlm, pch = 16, col = "blue")
    points(b2.tsir, pch = 16, col = "red")


    gammas.gmlm <- Map(solve, Omegas, betas.gmlm)
    gammas.tsir <- fit.tsir$Gammas

    g1.gmlm <- (length(gammas.gmlm[[1]]) / norm(gammas.gmlm[[1]])) * gammas.gmlm[[1]]
    g1.tsir <- (length(gammas.tsir[[1]]) / norm(gammas.tsir[[1]])) * gammas.tsir[[1]]
    if (sum(abs(g1.gmlm + g1.tsir)) < sum(abs(g1.gmlm - g1.tsir))) {
        g1.tsir <- -g1.tsir
    }
    time <- seq(0, 1, len = 256)
    plot(range(time), range(g1.gmlm, g1.tsir), type = "n",
        main = "Time", xlab = "Time [s]", ylab = expression(Gamma[1])
    )
    lines(time, g1.gmlm, col = "blue")
    lines(time, g1.tsir, col = "red")

    g2.gmlm <- (length(gammas.gmlm[[2]]) / norm(gammas.gmlm[[2]])) * gammas.gmlm[[2]]
    g2.tsir <- (length(gammas.tsir[[2]]) / norm(gammas.tsir[[2]])) * gammas.tsir[[2]]
    if (sum(abs(g2.gmlm + g2.tsir)) < sum(abs(g2.gmlm - g2.tsir))) {
        g2.tsir <- -g2.tsir
    }
    plot(c(1, length(g2.gmlm)), range(g2.gmlm, g2.tsir), type = "n",
        main = "Sensors", xlab = "Sensor Index", ylab = expression(Gamma[2])
    )
    lines(g2.gmlm, lty = 3, col = "blue")
    lines(g2.tsir, lty = 3, col = "red")
    points(g2.gmlm, pch = 16, col = "blue")
    points(g2.tsir, pch = 16, col = "red")

    par(par.reset)
}



# Load 2D (S1 stimulus only) EEG dataset of all subjects
c(X, y) %<-% readRDS("eeg_data_2d.rds")

# fit a tensor normal model to the full data set
fit.gmlm <- gmlm_tensor_normal(X, y, sample.axis = 3L,
    logger = function(iter, betas, Omegas, resid, loss) {
        cat(sprintf("GMLM - iter: %d - loss: %.3f\n", iter, loss))
    }
)

fit.tsir <- TSIR(X, y, sample.axis = 3L,
    logger = function(iter, Gammas, Sigmas, loss) {
        cat(sprintf("TSIR - iter: %d - loss: %.3f\n", iter, loss))
    }
)

plot.tsir.vs.gmlm(fit.gmlm, fit.tsir)

fit.tsir.reg <- TSIR(X, y, sample.axis = 3L, reg.threshold = 25,
    logger = function(iter, Gammas, Sigmas, loss) {
        cat(sprintf("TSIR - iter: %d - loss: %.3f\n", iter, loss))
    }
)

# Save estimates to csv for nicer plots
fit.gmlm$Gammas <- with(fit.gmlm, Map(solve, Omegas, betas))
standardize <- function(x, y) {
    x <- (length(x) / sqrt(sum(x^2))) * x
    if (!missing(y)) {
        y <- (length(y) / sqrt(sum(y^2))) * y
        if (sum((x + y)^2) < sum(x - y)^2) {
            x <- -x
        }
    }
    x
}
write.table(data.frame(
    time = seq(0, 1, len = length(fit.gmlm$betas[[1]])),
    gmlm.beta      = standardize(fit.gmlm$betas[[1]]),
    gmlm.Gamma     = standardize(fit.gmlm$Gammas[[1]]),
    tsir.beta      = standardize(fit.tsir$betas[[1]], fit.gmlm$betas[[1]]),
    tsir.Gamma     = standardize(fit.tsir$Gammas[[1]], fit.gmlm$Gammas[[1]]),
    tsir.reg.beta  = standardize(fit.tsir.reg$betas[[1]], fit.gmlm$betas[[1]]),
    tsir.reg.Gamma = standardize(fit.tsir.reg$Gammas[[1]], fit.gmlm$Gammas[[1]]),
    check.names = FALSE
), file = "tsir-vs-gmlm-1.csv", row.names = FALSE, quote = FALSE)
write.table(data.frame(
    sensor = seq_along(fit.gmlm$betas[[2]]),
    gmlm.beta      = standardize(fit.gmlm$betas[[2]]),
    gmlm.Gamma     = standardize(fit.gmlm$Gammas[[2]]),
    tsir.beta      = standardize(fit.tsir$betas[[2]], fit.gmlm$betas[[2]]),
    tsir.Gamma     = standardize(fit.tsir$Gammas[[2]], fit.gmlm$Gammas[[2]]),
    tsir.reg.beta  = standardize(fit.tsir.reg$betas[[2]], fit.gmlm$betas[[2]]),
    tsir.reg.Gamma = standardize(fit.tsir.reg$Gammas[[2]], fit.gmlm$Gammas[[2]]),
    check.names = FALSE
), file = "tsir-vs-gmlm-2.csv", row.names = FALSE, quote = FALSE)




fit.gmlm <- gmlm_tensor_normal(X, y, sample.axis = 3L, reg.threshold = 25,
    logger = function(iter, betas, Omegas, resid, loss) {
        cat(sprintf("GMLM - iter: %d - loss: %.3f\n", iter, loss))
    }
)

fit.tsir <- TSIR(X, y, sample.axis = 3L, reg.threshold = 25,
    logger = function(iter, Gammas, Sigmas, loss) {
        cat(sprintf("TSIR - iter: %d - loss: %.3f\n", iter, loss))
    }
)

plot.tsir.vs.gmlm(fit.gmlm, fit.tsir)


fit.gmlm <- gmlm_tensor_normal(X, y, sample.axis = 3L, reg.threshold = Inf,
    logger = function(iter, betas, Omegas, resid, loss) {
        cat(sprintf("GMLM - iter: %d - loss: %.3f\n", iter, loss))
    }
)

fit.tsir <- TSIR(X, y, sample.axis = 3L, reg.threshold = 25,
    logger = function(iter, Gammas, Sigmas, loss) {
        cat(sprintf("TSIR - iter: %d - loss: %.3f\n", iter, loss))
    }
)

plot.tsir.vs.gmlm(fit.gmlm, fit.tsir)






# Load 3D EEG dataset of all subjects
c(X, y) %<-% readRDS("eeg_data_3d.rds")

# fit a tensor normal model to the full data set
fit.gmlm <- gmlm_tensor_normal(X, y, sample.axis = 4L, reg.threshold = Inf,
    logger = function(iter, betas, Omegas, resid, loss) {
        cat(sprintf("GMLM - iter: %d - loss: %.3f\n", iter, loss))
    }
)

fit.tsir <- TSIR(X, y, sample.axis = 4L, reg.threshold = 25,
    logger = function(iter, Gammas, Sigmas, loss) {
        cat(sprintf("TSIR - iter: %d - loss: %.3f\n", iter, loss))
    }
)

plot.tsir.vs.gmlm(fit.gmlm, fit.tsir)

fit.tsir <- TSIR(X, y, sample.axis = 4L,
    logger = function(iter, Gammas, Sigmas, loss) {
        cat(sprintf("TSIR - iter: %d - loss: %.3f\n", iter, loss))
    }
)

plot.tsir.vs.gmlm(fit.gmlm, fit.tsir)



fit.gmlm <- gmlm_tensor_normal(X, y, sample.axis = 4L,
    logger = function(iter, betas, Omegas, resid, loss) {
        cat(sprintf("GMLM - iter: %d - loss: %.3f\n", iter, loss))
    }
)
beta1 <- fit.gmlm$betas[[1]]
beta1 <- (length(beta1) / norm(beta1)) * beta1

plot(beta1, type = "n")

for (i in seq_along(y)) {
    cat(sprintf("### %d\n", i))
    fit.gmlm <- gmlm_tensor_normal(X[, , , -i], y[-i], sample.axis = 4L,
        logger = function(iter, betas, Omegas, resid, loss) {
            cat(sprintf("GMLM - iter: %d - loss: %.3f\n", iter, loss))
        }
    )
    b1 <- fit.gmlm$betas[[1]]
    b1 <- (length(b1) / norm(b1)) * b1
    if (sum((beta1 - b1)^2) > sum((beta1 - b1)^2)) {
        b1 <- -b1
    }
    lines(b1, col = "#FF000080")
}



fit.tsir <- TSIR(X, y, sample.axis = 4L,
    logger = function(iter, Gammas, Sigmas, loss) {
        cat(sprintf("TSIR - iter: %d - loss: %.3f\n", iter, loss))
    }
)
beta1 <- fit.tsir$betas[[1]]
beta1 <- (length(beta1) / norm(beta1)) * beta1

plot(beta1, type = "n")

for (i in seq_along(y)) {
    cat(sprintf("### %d\n", i))
    fit.tsir <- TSIR(X[, , , -i], y[-i], sample.axis = 4L,
        reg.threshold = 25,
        logger = function(iter, Gammas, Sigmas, loss) {
            cat(sprintf("TSIR - iter: %d - loss: %.3f\n", iter, loss))
        }
    )
    b1 <- fit.tsir$betas[[1]]
    b1 <- (length(b1) / norm(b1)) * b1
    if (sum((beta1 - b1)^2) > sum((beta1 - b1)^2)) {
        b1 <- -b1
    }
    lines(b1, col = "#FF000080")
}