# install_git("https://www.gitlab.ci.inf.usi.ch/SQUIC/SQUIC_R.git")
library(SQUIC)

#install.packages("Matrix",repos = "http://cran.us.r-project.org")
library(Matrix)

#install.packages("devtools",repos = "http://cran.us.r-project.org")
library(devtools)

#devtools::install_github('ramhiser/datamicroarray')
library(datamicroarray)

# install.packages("emstreeR",repos = "http://cran.us.r-project.org")
library(emstreeR)

# install.packages("caret",repos = "http://cran.us.r-project.org")
library(caret)

# ======================================== #

# ======================================== #
#             Linear DA Function           #
# ======================================== #
# data_train   = input training data in the form p (dimensions) x n (samples)
# labels_train = vector with the labels of the training data
# data_test    = input testing data in the form p (dimensions) x n (samples)
# labels_test  = vector with the labels of the training data
# theta_est    = estimated precision matrix
Linear_DA <- function(data_train, labels_train, data_test, labels_test, theta_est){
        
        # Input validation could be added here
        n_train <- length(labels_train)
        n_test <- length(labels_test)
        K <- max(labels_train)

        rho <- matrix(0, ncol = K, nrow = n_test)

        for (k in 1:K) {
                # prior probability for class k
                prior <- length(which(labels_train == k)) / n_train
                # sample mean for class k
                mu <- rowMeans(data_train[, which(labels_train == k)])
                mu_matrix <- matrix(mu, nrow = nrow(data_test), ncol = ncol(data_test), byrow = FALSE)
                # x - 0.5*\mu
                adjusted_data <- data_test - 0.5 * mu_matrix
                # Linear Discriminant function
                lda_score <- diag(t(adjusted_data) %*% theta_est %*% mu_matrix + log(prior))
                rho[, k] <- lda_score
        }
        predicted <- apply(rho, 1, which.max)

        err = (abs(predicted-labels_test)>0.0)*1
        acc = 1-sum(err)/length(err)

        print(paste("acc=",acc));

        return(list(acc = acc, predicted = predicted))
}

# datasets under consideration
datasets      <- c('burczynski','yeoh','shipp','alon')
lambda_set    <- c(0.7, 0.8, 0.8, 0.8)
eta           <- 0.1


ACC        <- c(); 
time       <- c();
time_opt   <- c();
n_iter     <- c();
nnz_X      <- c();
nnz_W      <- c();

ACC_M      <- c(); 
time_M     <- c(); 
time_opt_M <- c();
n_iter_M   <- c();
nnz_X_M    <- c();
nnz_W_M    <- c();

ACC_M_c      <- c(); 
time_M_c     <- c(); 
time_opt_M_c <- c();
n_iter_M_c   <- c();
nnz_X_M_c    <- c();
nnz_W_M_c    <- c();

# Load data (n x p)
data(list=datasets, package = 'datamicroarray')

for (i in 1:length(datasets)){
        set.seed(1991)
        lambda      <- lambda_set[i];
        
        ############################
        # Load and split data in test/train
        ############################
        curr_dataset  <- datasets[i];
        data          <- eval(as.symbol(curr_dataset))$x;
        data          <- data.matrix(data, rownames.force = NA);
        labels        <- as.numeric(eval(as.symbol(curr_dataset))$y);

        train_index   <- createDataPartition(labels, p = .7, list = FALSE)
        train_index   <- as.numeric(train_index)
        data_train    <- t(data[ train_index,])
        data_test     <- t(data[-train_index,])
        labels_train  <- labels[train_index]
        labels_test   <- labels[-train_index]

        # Normalizing data
        n_train <- dim(data_train)[2];
        a       <- (apply(data_train,1,var))*(n_train-1)/n_train;
        a       <- 1/sqrt(a);
        D       <- Diagonal(x=a);
        data_train_scale <- D%*%data_train;

        ############################
        # Scalar Tunning Param.
        ############################
        squic_out   <- SQUIC(Y = data_train_scale , lambda = lambda, tol=1e-3, tol_inv=1e-3, M = NULL, verbose = 1)
        time[i]     <- squic_out$info_time_total;
        time_opt[i] <- squic_out$info_time_optimize;
        n_iter[i]   <- length(squic_out$info_objective);
        theta_est   <- D %*% squic_out$X %*% D;

        ACC[i]   <- Linear_DA(data_train, labels_train,data_test, labels_test, theta_est)$acc;
        nnz_X[i] <- nnzero(squic_out$X)/dim(squic_out$X)[1]
        nnz_W[i] <- nnzero(squic_out$W)/dim(squic_out$W)[1]

        ############################
        # Matrix Tunning Param.
        ############################
        p   <- nrow(data_train_scale);
        mst <- ComputeMST(as.matrix(data_train_scale));
        G   <- sparseMatrix(dims = c(p,p), i=mst$to, j=mst$from , x=rep(1, p));
        G   <- forceSymmetric(G,uplo="L") + Diagonal(nrow(G)); 

        squic_out     <- SQUIC(Y = data_train_scale, lambda = 0.95,  tol=1e-3,tol_inv=1e-3, M = eta*G, verbose = 1)
        time_M[i]     <- squic_out$info_time_total;
        time_opt_M[i] <- squic_out$info_time_optimize;
        n_iter_M[i]   <- length(squic_out$info_objective);
        theta_est     <- D %*% squic_out$X %*% D;

        ACC_M[i]   <- Linear_DA(data_train, labels_train,data_test, labels_test, theta_est)$acc;
        nnz_X_M[i] <- nnzero(squic_out$X)/dim(squic_out$X)[1]
        nnz_W_M[i] <- nnzero(squic_out$W)/dim(squic_out$W)[1]

        ############################
        # Control
        ############################
        squic_out_M_c     <- SQUIC(Y = data_train_scale, lambda = lambda,  tol=1e-3,tol_inv=1e-3, M = lambda*G, verbose = 1)
        time_M_c[i]       <- squic_out_M_c$info_time_total;
        time_opt_M_c[i]   <- squic_out_M_c$info_time_optimize;
        n_iter_M_c[i]     <- length(squic_out_M_c$info_objective);
        theta_est         <- D %*% squic_out_M_c$X %*% D;

        ACC_M_c[i]   <- Linear_DA(data_train, labels_train,data_test, labels_test, theta_est)$acc;
        nnz_X_M_c[i] <- nnzero(squic_out$X)/dim(squic_out$X)[1]
        nnz_W_M_c[i] <- nnzero(squic_out$W)/dim(squic_out$W)[1]

}
# Plot time and ACC for all datasets
times_all    <- cbind(time,time_M,time_M_c);
ACC_all      <- cbind(ACC,ACC_M,ACC_M_c);
n_iter_all   <- cbind(n_iter,n_iter_M,n_iter_M_c);
nnz_X_all    <- cbind(nnz_X,nnz_X_M,nnz_X_M_c);
nnz_W_all    <- cbind(nnz_W,nnz_W_M,nnz_W_M_c);

options(width=150)

print(data.frame(datasets,times_all, ACC_all,n_iter_all,nnz_X_all,nnz_W_all,lambda_set), digits = 3)
