library(RTMB)
library(Matrix)
library(fmesher)

load("slogo.RData")
pol$x<-c(pol$x,pol$x[1]) # following methods are fuzzy about having a closed polygon 
pol$y<-c(pol$y,pol$y[1])

# Observations
dat <- as.list(read.table("plastictime.dat", header=TRUE))

# setup mesh
mesh <- fm_mesh_2d(boundary = sf::st_polygon(list(cbind(pol$x, pol$y))), max.edge=c(.05))
spde <- fm_fem(mesh)
plot(mesh, edge.col="gray")
points(dat$x, dat$y, cex=dat$count/max(dat$count)*5)
projObs <- fm_basis(mesh, loc = cbind(dat$x, dat$y))

# add objects to data set 
dat$mesh <- mesh
dat$spde <- spde
dat$timeidx <- dat$year+1
dat$projObs <- projObs

# define parameters
par <- list()
par$mu <- 0
par$logTau <- 0  
par$logKappa <- 0  
par$logLambda <- array(0, dim=c(dat$mesh$n, max(dat$timeidx)))
par$transPhi<-0

# setup likelihood (here gaussian field for log-intensity of pois observations) 
jnll <- function(par){
  getAll(par, dat)
  tau <- exp(logTau)
  kappa <- exp (logKappa)
  phi <- 2*plogis(transPhi)-1
  Q <- tau*(kappa^4*spde$c0+2*kappa^2*spde$g1+spde$g2)
  d1 <- function(x)dgmrf(x, Q=Q, log=TRUE)
  d2 <- function(x)dautoreg(x, phi=phi, log=TRUE)
  ret <- -dseparable(d1,d2)(logLambda-mu)
  logPredObs <- (projObs%*%logLambda)[cbind(1:length(count), timeidx)]
  ret <- ret -sum(dpois(count, exp(logPredObs), log=TRUE))
  sdField <- sqrt(1/(4*pi*tau^2*kappa^2))
  range <- sqrt(8)/kappa
  ADREPORT(sdField)
  ADREPORT(range)
  ret
}

obj <- MakeADFun(jnll, par, random="logLambda", silent=FALSE)
fit <- nlminb(obj$par, obj$fn, obj$gr)

# Extracting estimates 
sdr <- sdreport(obj)
pl <- as.list(sdr,"Est")
plsd <- as.list(sdr,"Std")
plr <- as.list(sdr,"Est", report=TRUE)
plrsd <- as.list(sdr,"Std", report=TRUE)

pdf("spacetime.pdf", width=15, height=10)
  lam<-exp(pl$logLambda)
  br<-seq(min(lam)-.1, max(lam)+.1, length=501)
  cc<-viridis::mako(500)
  tc<-apply(col2rgb(cc)/255, 2, function(x)rgb(x[1],x[2],x[3],.8))
  nt<-nrow(mesh$graph$tv)
  for(ii in 1:max(dat$timeidx)){
    lambda<-lam[,ii]
    lamCol<-tc[as.numeric(cut(lambda,breaks = br))]
    plot(mesh, col="gray", lwd=.1)
    title(ii)
    lamv<-apply(mesh$graph$tv,1,function(idx)mean(lambda[idx]))
    lamc<-tc[as.numeric(cut(lamv,breaks = br))]
    dummy<-sapply(1:nrow(mesh$graph$tv),function(i){idx<-mesh$graph$tv[i,];polygon(mesh$loc[idx,1], mesh$loc[idx,2], col=lamc[i], border=NA)})
    fields::image.plot(as.matrix(lam), col=tc, type="n", legend.only=TRUE)
  }
  #points(dat$x, dat$y, cex=ifelse(dat$count==0, 1, dat$count/max(dat$count)*5),
  #pch=ifelse(dat$count==0,1,16), col=ifelse(dat$count==0,"red","blue"))
dev.off()