#'
#' Samples a pair of correlated \eqn{G(X)} random dot product graphs
#'
#' @param X \eqn{n \times d} random matrix
#' @param rho Numeric scalar in the unit interval, the target Pearson correlation between
#' the adjacency matrieces of the original and the generated graph.
#'
#' @return A list of two igraph objects, named \code{A} and \code{B}, which are
#' two graphs whose adjacency matrix entries are correlated with \code{rho}.
#'
#' @author Minh Tang <scent.of.time@gmail.com>
#' @export
#

suppressMessages(library(igraph))

## rdpg.sample <- function(X) {
##     P <- X %*% t(X)
##     n <-  nrow(P)
##     U <- matrix(0, nrow = n, ncol = n)
##     U[col(U) > row(U)] <- runif(n*(n-1)/2)
##     U <- (U + t(U))
##     diag(U) <- runif(n)
##     A <- (U < P) + 0 ;
##     diag(A) <- 0
##     return(graph.adjacency(A,"undirected"))
## }

## rdpg.sample2 <- function(X1, X2) {
##     P <- X %*% t(X2)
##     n1 <- nrow(P)
##     n2 <- ncol(P)
##     U <- matrix(runif(n1*n2), n1, n2)
##     A <- ifelse(P > U, 1, 0)

    
##     return(graph.adjacency(A,"undirected"))
## }


## sample.correlated <- function(X1, X2, rho) {
##     P <- X1 %*% t(X2)
##     n1 <- nrow(P)
##     n2 <- ncol(P)
    
##     avec <- A[col(A) > row(A)]
##     pvec <- P[col(P) > row(P)]
##     bvec <- numeric(n*d)
##     uvec <- runif(n*d)

##     idx1 <- which(avec == 1)
##     idx0 <- which(avec == 0)

##     bvec[idx1] <- (uvec[idx1] < (rho + (1 - rho)*pvec[idx1])) + 0
##     bvec[idx0] <- (uvec[idx0] < (1 - rho)*pvec[idx0]) + 0

##     B <- matrix(0, nrow = n, ncol = d)
##     B[col(B) > row(B)] <- bvec
## #    B <- B + t(B)
## #    diag(B) <- 0

## }

## Correlated RDPG
rdpg.sample.correlated <- function(X, rho) {
    P <- X %*% t(X)
    n <-  nrow(P)
    U <- matrix(0, nrow = n, ncol = n)
    U[col(U) > row(U)] <- runif(n*(n-1)/2)
    U <- (U + t(U))
    diag(U) <- runif(n)
    A <- (U < P) + 0 ;
    diag(A) <- 0

    avec <- A[col(A) > row(A)]
    pvec <- P[col(P) > row(P)]
    bvec <- numeric(n*(n-1)/2)

    uvec <- runif(n*(n-1)/2)

    idx1 <- which(avec == 1)
    idx0 <- which(avec == 0)

    bvec[idx1] <- (uvec[idx1] < (rho + (1 - rho)*pvec[idx1])) + 0
    bvec[idx0] <- (uvec[idx0] < (1 - rho)*pvec[idx0]) + 0

    B <- matrix(0, nrow = n, ncol = n)
    B[col(B) > row(B)] <- bvec
    B <- B + t(B)
    diag(B) <- 0

    return(list(A = graph.adjacency(A,"undirected"), B = graph.adjacency(B,"undirected")))
}

## Correlated ER
rg.sample.correlated.gnp <- function(P,sigma){
    n <-  nrow(P)
    U <- matrix(0, nrow = n, ncol = n)
    U[col(U) > row(U)] <- runif(n*(n-1)/2)
    U <- (U + t(U))
    diag(U) <- runif(n)
    A <- (U < P) + 0 ;
    diag(A) <- 0

    avec <- A[col(A) > row(A)]
    pvec <- P[col(P) > row(P)]
    bvec <- numeric(n*(n-1)/2)

    uvec <- runif(n*(n-1)/2)

    idx1 <- which(avec == 1)
    idx0 <- which(avec == 0)

    bvec[idx1] <- (uvec[idx1] < (sigma + (1 - sigma)*pvec[idx1])) + 0
    bvec[idx0] <- (uvec[idx0] < (1 - sigma)*pvec[idx0]) + 0

    B <- matrix(0, nrow = n, ncol = n)
    B[col(B) > row(B)] <- bvec
    B <- B + t(B)
    diag(B) <- 0

    return(list(A = graph.adjacency(A,"undirected"), B = graph.adjacency(B,"undirected")))
}

# non-igraph version of correlated SBM
rg.sample.SBM.correlated <- function(n, B, rho, sigma, conditional = FALSE){
    if(!conditional){
        tau <- sample(c(1:length(rho)), n, replace = TRUE, prob = rho)
    }
    else{
        tau <- unlist(lapply(1:2,function(k) rep(k, rho[k]*n)))
    }
    P <- B[tau,tau]
    return(list(adjacency=rg.sample.correlated.gnp(P, sigma),tau=tau))
}


## igraph version of correlated ER pair
sample_correlated_gnp_pair <- function (n, p, rho, directed = FALSE, permutation = NULL)
{
  igraph::sample_correlated_gnp_pair(n,rho,p,directed,permutation)
}

## generate a SBM using igraph, then generate a correlated pair
sample_correlated_sbm_pair <- function(n, pref.matrix, block.sizes, rho){

  K <- length(block.sizes)
  # Make the first graph
  graph1 <- sample_sbm(n,pref.matrix,block.sizes)

  # Make two graphs which will be used to make the
  # second graph
  corr.matrix <- (1-rho)*pref.matrix
  Z0 <- sample_sbm(n,corr.matrix,block.sizes)
  Z1 <- sample_sbm(n,corr.matrix+rho,block.sizes)

  graph2 <- Z1 %s% graph1 %u% (Z0-graph1)

  list(graph1=graph1,graph2=graph2)
}