options(warn=-1) # turns off warnings, to turn on: "options(warn=0)"


library(imager)
library(png)

for (f in list.files(path="nt_toolbox/toolbox_general/", pattern="*.R")) {
    source(paste("nt_toolbox/toolbox_general/", f, sep=""))
}

for (f in list.files(path="nt_toolbox/toolbox_signal/", pattern="*.R")) {
    source(paste("nt_toolbox/toolbox_signal/", f, sep=""))
}
options(repr.plot.width=3.5, repr.plot.height=3.5)

library("Matrix")

Rows <- function(n0, n1){ sparseMatrix(rep(1:n0,times=n1), 1:(n0*n1), x=rep(1, n0*n1)) }

Cols <- function(n0, n1){ sparseMatrix(rep(1:n1,each=n0), 1:(n0*n1), x=rep(1, n0*n1)) }

Sigma <- function(n0,n1){ rBind(Rows(n0,n1),Cols(n0,n1)) }

maxit <- 1e+4
tol <- 1e-9
otransp <- function(C,p0,p1){ array(perform_linprog(as.matrix(Sigma(length(p0),length(p1))),array(c(p0,p1),c(length(p0)+length(p1),1)),C,maxit,tol),c(n0,n1)) }

n0 <- 60
n1 <- 80

gauss <- function(q,a,c){ a*array(rnorm(2*q), c(2,q)) + array(rep(c, q), c(length(c),q)) }
X0 <- array(rnorm(2*n0), c(2,n0))*.3
X1 <- cbind(gauss(n1/2,.5,c(0,1.6)),
            cbind(gauss(n1/4,.3,c(-1,-1)),
                  gauss(n1/4,.3,c(1,-1))))

normalize <- function(a){ a/sum(a) }
p0 <- normalize(array(runif(n0), c(n0, 1)))
p1 <- normalize(array(runif(n1), c(n1, 1)))

myplot <- function(x,y,cex,col){ points(x,y, pch=21, cex=(cex*1.5), col="black", bg=col, lwd=2) }

options(repr.plot.width=7, repr.plot.height=5)

plot(1, type="n", axes=F, xlab="", ylab="",
     xlim=c(min(X1[1,])-.1, max(X1[1,])+.1),
     ylim=c(min(X1[2,])-.1, max(X1[2,])+.1))

for (i in 1:length(p0)){
    myplot(X0[1,i], X0[2,i], p0[i]*length(p0), 'blue')
}

for (i in 1:length(p1)){
    myplot(X1[1,i], X1[2,i], p1[i]*length(p1), 'red')
}

C <- array(rep(apply(X0**2,2,sum),n1), c(n0,n1)) + array(rep(apply(X1**2,2,sum), each=n0), c(n0,n1)) - 2*t(X0)%*%X1

gamma <- otransp(C, p0, p1)

print(paste("Number of non-zero:", length(gamma[gamma>0]), "(n0 + n1-1 =", n0 + n1-1,")"))

print(paste("Constraints deviation (should be 0):",
            norm(apply(gamma, 1, sum) - as.vector(p0)),
            norm(apply(gamma, 2, sum) - as.vector(p1)) ))

nonzero <- function(X){
    n_r <- dim(X)[1]
    idx <- which(X != 0)
    I <- mod(idx-1,n_r)+1
    J <- ceiling(idx/n_r)
    return(list(I=I, J=J))
}
nonzero_gamma <- nonzero(gamma)
I <- nonzero_gamma$I ; J <- nonzero_gamma$J

gammaij <- gamma[which(gamma != 0)]

options(repr.plot.width=7, repr.plot.height=7)

tlist = seq(0, 1, length=6)
par(mfrow=c(2,3))

for (i in 1:length(tlist)){
    t <- tlist[i]
    Xt <- (1-t)*X0[,I] + t*X1[,J]
    
    plot(1, type="n", axes=F, xlab="", ylab="", main=paste("t =", t),
     xlim=c(min(X1[1,])-.1, max(X1[1,])+.1),
     ylim=c(min(X1[2,])-.1, max(X1[2,])+.1))
    
    for (j in 1:length(gammaij)){
        myplot(Xt[1,j],Xt[2,j],gammaij[j]*length(gammaij),rgb(t,0,1-t))
    }
        
}

n0 <- 40
n1 <- n0

X0 <- array(rnorm(2*n0), c(2,n0))*.3
X1 <- cbind(gauss(n1/2,.5,c(0,1.6)),
            cbind(gauss(n1/4,.3,c(-1,-1)),
                  gauss(n1/4,.3,c(1,-1))))

p0 <- array(1, c(n0,1))/n0
p1 <- array(1, c(n1,1))/n1

C <- array(rep(apply(X0**2,2,sum),n1), c(n0,n1)) + array(rep(apply(X1**2,2,sum), each=n0), c(n0,n1)) - 2*t(X0)%*%X1

options(repr.plot.width=7, repr.plot.height=5)

plot(1, type="n", axes=F, xlab="", ylab="",
     xlim=c(min(X1[1,])-.1, max(X1[1,])+.1),
     ylim=c(min(X1[2,])-.1, max(X1[2,])+.1))

myplot(X0[1,],X0[2,],1,"blue")
myplot(X1[1,],X1[2,],1,"red")

gamma <- otransp(C, p0, p1)

options(repr.plot.width=3.5, repr.plot.height=3.5)

imageplot(gamma)

nonzero_gamma <- nonzero(gamma)
I <- nonzero_gamma$I ; J <- nonzero_gamma$J

options(repr.plot.width=7, repr.plot.height=5)

plot(1, type="n", axes=F, xlab="", ylab="",
     xlim=c(min(X1[1,])-.1, max(X1[1,])+.1),
     ylim=c(min(X1[2,])-.1, max(X1[2,])+.1))

for (k in 1:length(I)){
    lines(c(X0[1,I[k]], X1[1,J[k]]),
          c(X0[2,I[k]], X1[2,J[k]]),
          lwd=2)
}

myplot(X0[1,],X0[2,],1,"blue")
myplot(X1[1,],X1[2,],1,"red")
<script>
  $(document).ready(function(){
      $('div.prompt').hide();
  });
</script>