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=""))
}

source("nt_toolbox/toolbox_wavelet_meshes/meshgrid.R")
options(repr.plot.width=3.5, repr.plot.height=3.5)

N <- c(300,200)

d <- 2

x <- array(runif(2*N[1]), c(2,N[1])) - .5

theta <- 2*pi*array(runif(1*N[2]), c(1,N[2]))
r <- .8 + .2*array(runif(1*N[2]), c(1,N[2]))
y <- rbind(cos(theta)*r, sin(theta)*r)

plotp <- function(x, col){ points(x[1,],x[2,], pch=21, cex=2, col="black", bg=col, lwd=2) }

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

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

plotp(x, 'blue')
plotp(y, 'red')

x2 <- apply(x**2, 2, sum)
y2 <- apply(y**2, 2, sum)
C <- array(rep(y2, each=N[1]), c(N[1], N[2])) + array(rep(x2, times=N[2]), c(N[1], N[2])) - 2*t(x)%*%y

p <- rep(1, N[1])/N[1]
q <- rep(1, N[2])/N[2]

gamma <- .01

xi <- exp(-C/gamma)

b <- rep(1, N[2])

a <- p/(xi%*%b)
b <- q /(t(xi)%*%a)

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

source("nt_solutions/optimaltransp_5_entropic/exo1.R")

# Insert your code here.

a_diag <- array(0, c(length(a), length(a)))
diag(a_diag) <- a

b_diag <- array(0, c(length(b), length(b)))
diag(b_diag) <- b

Pi <- (a_diag %*% xi) %*% b_diag

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

imageplot(Pi)

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

source("nt_solutions/optimaltransp_5_entropic/exo2.R")

# Insert your code here.

a_diag <- array(0, c(length(a), length(a)))
diag(a_diag) <- a

b_diag <- array(0, c(length(b), length(b)))
diag(b_diag) <- b

Pi <- (a_diag %*% xi) %*% b_diag

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))
}

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

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

plotp(x, 'blue')
plotp(y, 'red')

A <- Pi * (Pi > min(1./N)*.7)
nonzero_A <- nonzero(A)
i <- nonzero_A$I ; j <- nonzero_A$J
for (cpl in 1:length(i)){
    lines(c(x[1,i][cpl], y[1,j][cpl]), c(x[2,i][cpl], y[2,j][cpl]), lwd=2)
}


A <- Pi * (Pi > min(1./N)*.3)
nonzero_A <- nonzero(A)
i <- nonzero_A$I ; j <- nonzero_A$J
for (cpl in 1:length(i)){
    lines(c(x[1,i][cpl], y[1,j][cpl]), c(x[2,i][cpl], y[2,j][cpl]), lwd=1, lty=2)
}

N <- 200

t <- 0:(N-1)/N

Gaussian <- function(t0,sigma){ exp(-(t-t0)**2/(2*sigma**2)) }
normalize <- function(p){ p/sum(p) }

sigma <- .06
p <- Gaussian(.25,sigma)
q <- Gaussian(.8,sigma)

vmin <- .02
p <- normalize( p+max(p)*vmin)
q <- normalize( q+max(q)*vmin)

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

par(mfrow=c(2,1))

plot(t, p, col="blue", type="h", xlab="", ylab="")
plot(t, q, col="blue", type="h", xlab="", ylab="")

gamma <- (.03)**2

grid <- meshgrid_2d(t,t)
Y <- grid$X ; X <- grid$Y
xi <- exp(-(X-Y)**2/gamma)

b <- rep(1, N)

a <- p/(xi%*%b)
b <- q /(t(xi)%*%a)

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

source("nt_solutions/optimaltransp_5_entropic/exo3.R")

# Insert your code here.

a_diag <- array(0, c(length(a), length(a)))
diag(a_diag) <- a

b_diag <- array(0, c(length(b), length(b)))
diag(b_diag) <- b

Pi <- (a_diag %*% xi) %*% b_diag

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

imageplot(log(Pi+1e-5))

s <- (xi %*% (b*t))*a/p

imageplot(log(Pi+1e-5))
lines(s*N,t*N, col='red', lwd=3, xlim=c(0,N), ylim=c(0,N))

N <- 70

names <- c('disk','twodisks','letter-x','letter-z')
vmin <- .05
P <- array(0, c(N,N,length(names)))
for (i in 1:length(names)){
    p <- as.matrix( load_image(paste("nt_toolbox/data/", names[i],".png",sep=""),N) )
    p <- normalize(rescale(p)+vmin)
    P[,,i] <- p
}

K <- length(names)

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

for (i in 1:K){
    imageplot(P[,,i], '',c(2,2,i))
}

gamma <- (.04)**2

n <- 41 # width of the convolution kernel
t <- seq(-n/(2*N),n/(2*N),length=n)
g <- exp(-t**2/gamma)
g2 <- g %*% t(g)
xi <- function(x){ t( as.matrix( convolve( as.cimg( t( as.matrix( convolve(as.cimg(x), as.cimg(g2)) ) ) ), as.cimg(g2) ) ) ) }

options(repr.plot.width=7, repr.plot.height=3.5)
imageplot(P[,,1],'',c(1,2,1))
imageplot(xi(P[,,1]),'',c(1,2,2))

lambd <- rep(1, K)/K

b <- array(1, c(N,N,K))
a <- b

for (k in 1:K){
    a[,,k] <- P[,,k]/xi(b[,,k])
}

q <- rep(0, N)

for (k in 1:K){
    q <- q + lambd[k] * log(pmax(1e-19*rep(1, length(b[,,k])), b[,,k]*xi(a[,,k])))
}

q <- exp(q)

for (k in 1:K){
    b[,,k] <- q/xi(a[,,k])
}

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

source("nt_solutions/optimaltransp_5_entropic/exo4.R")

# Insert your code here.

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

imageplot(as.cimg(q))

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

source("nt_solutions/optimaltransp_5_entropic/exo5.R")

# Insert your code here.
<script>
  $(document).ready(function(){
      $('div.prompt').hide();
  });
</script>