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

library(png)
library(pracma)
library(imager)

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=5, repr.plot.height=5)

n <- 128
V <- perform_blurring( array( rnorm(n*n*2), c(n,n,2) ), c(40), "per" )

myplot <- function(V){ plot_vf(V[seq(1,n,6),seq(1,n,6),]) }

myplot(V)

normalize <- function(V){ V/ array(rep( pmax( array(1e-9, dim(V)[1:2]), sqrt(apply(V**2, c(1,2), sum)) ), 2 ), dim(V)) }

myplot(normalize(V))

grid <- meshgrid_2d(0:(n-1), 0:(n-1))
Y <- grid$X ; X <- grid$Y
mu <- sin(X*pi/n)**2
mu <- -4*(mu + t(mu))
mu[1,1] <- 1

A <- function(V){ FFT <- fft( div(V[,,1], V[,,2], bound="per") )
                  return( Re( fft( FFT/mu, inverse=TRUE)/length(FFT) ) ) }

ProjI <- function(V){ V - grad(A(V), bound="per") }

U <- ProjI(V)

myplot(U)

myplot(V-U)

f <- as.matrix(load_image("nt_toolbox/data/lena.png", 2*n))
f <- f[(n-round(n/2)+1):(n+round(n/2)),(n-round(n/2)+1):(n+round(n/2))]

U <- normalize(ProjI(V))

periodic <- function(P){ mod(P,n) }

extend1 <- function(f){ ext <- array(0, c(dim(f)[1], dim(f)[2]+1))
                        ext[,1:dim(f)[2]] <- f
                        ext[,dim(ext)[2]] <- f[,1]
                        return(ext) }

extend <- function(f){ t(extend1(t(extend1(f)))) }

myinterp <- function(f1, Pi){
    dim_f1 <- dim(f1)
    f1 <- as.cimg(t(f1))
    locations <- data.frame(x=as.vector(Pi[,,2]), y=as.vector(Pi[,,1]))
    return( as.matrix(as.cimg(imager::interp(f1, locations))) )
}

grid <- meshgrid_2d(1:n, 1:n)
Y <- grid$X ; X <- grid$Y
P <- array(0, c(dim(X),2)) ; P[,,1] <- X ; P[,,2] <- Y

W <- function(f, U){ myinterp(extend(f), periodic(P - U)) }

rho <- 2
options(repr.plot.width=5, repr.plot.height=5)
imageplot(W(f, rho*U))

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

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

## Insert your code here.

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

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

## Insert your code here.

nu <- 1/10

mu <- 2*nu

Wt <- function(V, U){ x <- W(V[,,1],U)
                      y <- W(V[,,2],U)
                      out <- array(0, dim(V))
                      out[,,1] <- x ; out[,,2] <- y
                      return(out)}

tau <- .5

V <- normalize(ProjI(V))

g <- f

g <- W(g, tau*U)
V <- Wt(V, tau*U)

s1 <- c(2:n,1)
s2 <- c(n,1:(n-1))

Delta <- function(g){
    if (length(dim(g))==2) { 1/4.*(g[s1,] + g[s2,] + g[,s1] + g[,s2]) - g }
    else if (length(dim(g))==3) { 1/4.*(g[s1,,] + g[s2,,] + g[,s1,] + g[,s2,]) - g } }

V <- V + tau*nu*Delta(V)
g <- g + tau*mu*Delta(g)

V <- ProjI(V)

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

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

## Insert your code here.