# learn VERSO phylogenetic tree from data
learn.VERSO.phylogenetic.tree <- function( D, 
                                           alpha = 10^-3, 
                                           beta = 10^-3, 
                                           initialization = NULL, 
                                           keep_equivalent = FALSE, 
                                           marginalize = FALSE,
                                           num_rs = 10, 
                                           num_iter = 10000, 
                                           n_try_bs = 1000, 
                                           num_processes = Inf, 
                                           verbose = TRUE,
                                           log_file = "" ) {
    
    # we remove any NA from data
    if(any(is.na(D))) {
        D[which(is.na(D),arr.ind=TRUE)] <- -3
    }
    storage.mode(D) <- "integer"
    
    equivalent_solutions <- list()
    
    # perform a number of num_rs restarts
    if(num_processes > 1) {
        
        parallel_cl <- makeCluster(num_processes,outfile=log_file)
        
        clusterExport(parallel_cl,varlist=c("D","alpha","beta","initialization","keep_equivalent","marginalize","num_rs","num_iter","n_try_bs","verbose","log_file"),envir=environment())
        clusterExport(parallel_cl,c("MCMC","initialize.B","move.B", "relabeling", "prune.and.reattach", "compute.C"),envir=environment())
        clusterSetRNGStream(parallel_cl,iseed=round(runif(1)*100000))
        
        mcmc_res <- parLapply(parallel_cl,1:num_rs,function(x) {
            
            MCMC( D = D, 
                  alpha = alpha, 
                  beta = beta, 
                  initialization = initialization, 
                  keep_equivalent = keep_equivalent, 
                  marginalize = marginalize, 
                  num_rs = num_rs, 
                  num_iter = num_iter, 
                  n_try_bs = n_try_bs, 
                  rs_i = x, 
                  verbose = verbose, 
                  log_file = log_file )
            
        })
        
        stopCluster(parallel_cl)
        
    }
    else {
        mcmc_res <- list()
        for(i in seq_len(num_rs)) {
            mcmc_res[[i]] <- MCMC( D = D, 
                                   alpha = alpha, 
                                   beta = beta, 
                                   initialization = initialization, 
                                   keep_equivalent = keep_equivalent, 
                                   marginalize = marginalize, 
                                   num_rs = num_rs, 
                                   num_iter = num_iter, 
                                   n_try_bs = n_try_bs, 
                                   rs_i = i, 
                                   verbose = verbose, 
                                   log_file = log_file )
        }
    }

    if(num_rs > 1) {
        best_lik <- c()
        for(i in  1:num_rs) {
            best_lik <- c(best_lik, mcmc_res[[i]]$log_likelihood)
        }
        idx_bjl <- which(best_lik == max(best_lik))
        best_mcmc <- mcmc_res[[idx_bjl[1]]]
        for(i in 2:length(idx_bjl)) {
            best_mcmc$equivalent_solutions <- c(best_mcmc$equivalent_solutions,mcmc_res[[idx_bjl[i]]]$equivalent_solutions)
        }
    }
    else {
        best_mcmc <-  mcmc_res[[1]]
    }
    
    # renaming
    rownames(best_mcmc$B) <- paste0("G",seq_len(nrow(best_mcmc$B)))
    colnames(best_mcmc$B) <- c("Reference",colnames(D)[as.numeric(colnames(best_mcmc$B)[2:ncol(best_mcmc$B)])])
    best_mcmc$C <- matrix(paste0("G",best_mcmc$C),ncol=1)
    rownames(best_mcmc$C) <- rownames(D)
    colnames(best_mcmc$C) <- "Genotype"
    
    return(best_mcmc)
    
}

# perform MCMC
MCMC <- function( D, 
                  alpha = 10^-3, 
                  beta = 10^-3, 
                  initialization = NULL, 
                  keep_equivalent = FALSE, 
                  marginalize = FALSE, 
                  num_rs = 10, 
                  num_iter = 10000, 
                  n_try_bs = 1000, 
                  rs_i = NULL, 
                  verbose = TRUE, 
                  log_file = "" ) {
    
    # initialize B
    if(is.null(initialization)) {
        B <- initialize.B(D)
        
    }
    else {
        B <- initialization
        storage.mode(B) <- "integer"
        if(rs_i > 1) {
            # Restars after the first one start from a different point if 
            # a initial tree was passed (it could be generated by TRAIT or by user)
            for(i in seq(1,ncol(B))) {
                B <- move.B(B)   
            }
        }
    }

    # compute C given B
    res <- compute.C(B,D,alpha,beta,marginalize)
    C <- res$C
    lik <- res$lik
    
    # initialize result variables (best solution for the current restart)
    B_best <- B
    C_best <- C
    lik_best <- lik
    count_lik_best_cons <- 0
    equivalent_solutions <- list()
    if(keep_equivalent) {
        equivalent_solutions[[1]] <- list(B=B_best,C=C_best,log_likelihood=lik_best)
    }
    
    # repeat MCMC moves until num_iter number of iterations is performed
    j = 1
    while(j <= num_iter) {
       
        if(verbose&&(j%%500)==0) {
            cat("\r", 
                "Current best lik. = ",format(lik_best, digit = 2, nsmall = 2), 
                " | Restart # ",rs_i,"/",num_rs," | Iter # ",j, " | Likelihood not improved for ", count_lik_best_cons,"/",n_try_bs," iterations", 
                "     ", 
                sep='', 
                file = log_file, 
                append = TRUE)
        }
        
        # perform a move on B
        B_tmp <- move.B(B)
        
        # compute C at maximun likelihood given B_tmp and returns its likelihood
        res <- compute.C(B_tmp,D,alpha,beta,marginalize)
        C_tmp <- res$C
        lik_tmp <- res$lik
        
        # if likelihood at current step is better than best likelihood, replace best model with current
        if(lik_tmp>lik_best) {
            B_best <- B_tmp
            C_best <- C_tmp
            lik_best <- lik_tmp
            count_lik_best_cons <- 0
            B <- B_best
            C <- C_best
            lik <- lik_best
            equivalent_solutions <- list()
            if(keep_equivalent) {
                equivalent_solutions[[1]] <- list(B=B_best,C=C_best,log_likelihood=lik_best)
            }
        }
        else {

            if(lik_tmp==lik_best) {
                if(keep_equivalent) {
                    equivalent_solutions[[(length(equivalent_solutions)+1)]] <- list(B=B_tmp,C=C_tmp,log_likelihood=lik_tmp)
                }
            }

            # update count
            count_lik_best_cons <- count_lik_best_cons + 1
            if(count_lik_best_cons>n_try_bs) {
                # print a message
                if(verbose) {
                    cat("\r",
                        "Current best lik. = ",format(lik_best, digit = 2, nsmall = 2), 
                        " | Restart # ",rs_i,"/",num_rs," | Iter # ",j, " | Likelihood not improved for ", (count_lik_best_cons-1),"/",n_try_bs," iterations", 
                        "     ", 
                        "\n", 
                        sep='')
                }
                break;
            }
            # take the current state with a probability proportional to the ratio of the two likelihoods
            rho <- min(exp((lik_tmp-lik)),1)
            if(runif(1)<=rho) {
                B <- B_tmp
                C <- C_tmp
                lik <- lik_tmp
            }
        }
        j = j + 1
    }

    return(list(B=B_best,C=C_best,log_likelihood=lik_best,equivalent_solutions=equivalent_solutions))
    
}

# randomly initialize B
initialize.B <- function( D ) {
    
    # data structure where to save B
    B <- array(0L,c((ncol(D)+1),(ncol(D)+1)))
    rownames(B) <- c("r",seq_len(ncol(D)))
    colnames(B) <- c("r",sample(seq_len(ncol(D))))
    diag(B) <- 1L
    B[,1] <- 1L
    
    # add arcs with probability 50% to obtain a random tree topology
    p <- 0.50
    for(i in 2:(nrow(B)-1)) {
        if(runif(1)<p) {
            B[(i+1),] <- B[i,] + B[(i+1),]
        }
    }
    B[which(B>1)] <- 1L
    
    # return B
    return(B)
    
}

# performing either relabeling or edge changing moves on B
move.B <- function( B ) {
    
    # sample a random probability of choosing a move
    p <- runif(1)
    
    # perform pairwise relabeling with 55% probability
    if(p<0.55) {
        
        # nodes relabeling
        B <- relabeling(B=B)
        
    }
    # perform structural moves with 40% probability
    else if(p>=0.55&&p<0.95) {
        
        # prune and reattach
        B <- prune.and.reattach(B=B)
        
    }
    # perform full relabeling with 5% probability
    else if(p>=0.95) {
        
        # random relabeling of all clones
        colnames(B) <- c("r",sample(seq_len((ncol(B)-1))))
        
    }
    
    # return B
    return(B)
    
}

# perform relabeling
relabeling <- function( B ) {
    
    # relabeling
    chosen <- sample(2:ncol(B),2,replace=FALSE)
    tmp <- colnames(B)[chosen[1]]
    colnames(B)[chosen[1]] <- colnames(B)[chosen[2]]
    colnames(B)[chosen[2]] <- tmp
    return(B)
    
}

# perform prune and reattach
prune.and.reattach <- function( B ) {
    
    # change one arch
    is_not_valid <- TRUE
    while(is_not_valid) {
        
        # select source node
        ch_1 <- sample(x=2:nrow(B),size=1)
        ch_1_gen <- B[ch_1,1:ch_1]
        remaing_node <- as.numeric(which(apply(B[,1:ch_1], c(1), FUN = function(x){!all(x == ch_1_gen)})))
        
        # chose the target node from the nodes not included in the subtree where ch_1 is the root
        if(length(remaing_node) > 1) {
            
            ch_2 <- sample(x=remaing_node,size=1)
            
        } else if(length(remaing_node)==1) {
            
            ch_2 <- remaing_node
            
        } else {
            # if there aren't any nodes, select a different source
            next
        }

        # a pair of two nodes is valid if the nodes are not already directly connected
        if(!(all(B[ch_1,1:ch_2]==B[ch_2,1:ch_2])&sum(B[ch_1,])==(sum(B[ch_2,])+1))) {
            is_not_valid <- FALSE
        }
        
    }
    
    descendent_nodes <- setdiff(1:nrow(B), remaing_node)
    
    # extract descendent node submatrix
    rem_B <- B[remaing_node,remaing_node,drop=FALSE]
    new_B <- matrix(data = 0L, nrow = nrow(rem_B), ncol = ncol(B))
    colnames(new_B) <- c(colnames(B)[remaing_node], colnames(B)[descendent_nodes])
    new_B[1:length(remaing_node),1:length(remaing_node)] <- rem_B
    gen_ch_2 <- new_B[which(colnames(new_B)==colnames(B)[ch_2]),1:length(remaing_node)]
    desc_B <- cbind(matrix(rep(gen_ch_2,each=length(descendent_nodes)),  nrow = length(descendent_nodes)), 
                    B[descendent_nodes,descendent_nodes,drop=FALSE])
    new_B <- rbind(new_B,desc_B)
    
    return(new_B)
    
}

# compute attachments matrix C at maximum likelihood given B and D
compute.C <- function( B, D, alpha = 10^-3, beta = 10^-3, marginalize = FALSE) {
    
    # determine indeces to order D such that it matches B
    idx_srt <- as.integer(colnames(B)[2:ncol(B)])
    
    # go through each patient and compute likelihood for all possible attachments
    lik_matrix <- array(0L,c(nrow(D),ncol(B)))
    curr_D <- cbind(rep(1,nrow(D)),D[,idx_srt,drop=FALSE])
    for(k in seq_len(nrow(B))) {
        curr_C = matrix(rep(0L,nrow(B)),nrow=1)
        curr_C[1,k] <- 1L
        r_D_tilde <- (curr_C%*%B)*2
        sum_cell <- as.matrix(sweep(curr_D,MARGIN=2,r_D_tilde,"+"))
        lik_matrix[,k] <- log(beta^Rfast::rowsums(sum_cell==2)) +
            log((1-beta)^Rfast::rowsums(sum_cell==0)) + 
            log((alpha)^Rfast::rowsums(sum_cell==1)) + 
            log((1-alpha)^Rfast::rowsums(sum_cell==3))
    }
    
    # compute maximum likelihood attachments
    C <- Rfast::rowMaxs(lik_matrix,value=FALSE)
    storage.mode(C) <- "integer"
    
    lik <- sum(Rfast::rowMaxs(lik_matrix,value=TRUE))
    
    if(marginalize==TRUE) {
        lik_time <- sum(Rfast::rowsums(lik_matrix))
    }
    else {
        lik_time <- sum(Rfast::rowMaxs(lik_matrix,value=TRUE))
    }
    
    # return maximum likelihood attachments
    return(list(C=C,lik=lik))
    
}
