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("plastic.dat", header=TRUE))

# setup mesh
mesh <- fm_mesh_2d(boundary = sf::st_polygon(list(cbind(pol$x, pol$y))), max.edge=c(.02))
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$projObs <- projObs

# define parameters
par <- list()
par$mu <- 0
par$logTau <- 0  
par$logKappa <- 0  
par$logLambda <- numeric(dat$mesh$n)

# setup likelihood (here gaussian field for log-intensity of pois observations) 
jnll <- function(par){
  getAll(par, dat)
  tau <- exp(logTau)
  kappa <- exp (logKappa)
  Q <- tau*(kappa^4*spde$c0+2*kappa^2*spde$g1+spde$g2)
  ret <- -dgmrf(logLambda-mu, Q=Q, log = TRUE)
  logPredObs <- projObs%*%logLambda
  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)

lambda<-exp(pl$logLambda)
cc<-viridis::mako(500)
tc<-apply(col2rgb(cc)/255, 2, function(x)rgb(x[1],x[2],x[3],.8))
plot(mesh, col="gray", lwd=.1)
nt<-nrow(mesh$graph$tv)
lamv<-apply(mesh$graph$tv,1,function(idx)mean(lambda[idx]))
lamc<-tc[as.numeric(cut(lamv,breaks = 501))]
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(lambda), 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"))