### R script for Sequential Importance Sampling Exercise
### Day 2, ESAC Workshop 2017, Madrid
###
### Instructions
### (1) Step through the completed code with default parameterisation;
### questions in comments are there to check that *you* understand how 
### the algorithm is designed
### (2) Try varying system parameters (sigma_v, N & kappa, and Nexp):
### observe & explain how the particle system behaves: can you think of
### a quantitative criterion we could monitor to let us known when the
### particle approximation is likely about to fail? also ,could we try a
###  defensive importance sampling proposal?
### 
### Further info: Dr Ewan Cameron <dr.ewan.cameron@gmail.com>

library(circular) # install.packages(c("circular","rgl"))
library(rgl)

### System parameters

kappa <- 500 # inverse scale parameter of the Von Mises distribution: you can plot histograms of draws from the Von Mises to get a feel for this distribution: rvonmises(number of draws,mu,kappa)
N <- 100 # pseudo-sample size for the beta distribution written as Beta(p*N,(1-p)*N): how is the Beta a conjugate prior to the binomial?
Nexp <- 30 # expected number of noise sources
sigma_v <- 0.05 # standard deviation of random velocity increments

## Function to generate mock data samples

initialise.process <- function() {

    true.object <- list()
    true.object$radialcoords <- c(-pi/2,pi/8) # arbitrary starting position
    true.object$velocity <- c(0,0)

    # observational errors on true object position
    true.object$observed.radialcoords <- c(true.object$radialcoords[1]+as.numeric(rvonmises(1,0,kappa)),-pi/2+pi*rbeta(1,(true.object$radialcoords[2]+pi/2)/pi*N,(1-(true.object$radialcoords[2]+pi/2)/pi)*N))

    # add crowded field sources to list of observations
    observed.objects <- list()
    Nextra <- rpois(1,Nexp)
    observed.objects$radialcoords <- true.object$observed.radialcoords
    if (Nextra==0) {
    } else {
        observed.objects$radialcoords <- matrix(0,nrow=(Nextra+1),ncol=2)
        ordering <- sample(1:(Nextra+1),Nextra+1,replace=F)
        observed.objects$radialcoords[ordering[1],] <- true.object$observed.radialcoords
        observed.objects$radialcoords[ordering[2:(Nextra+1)],] <- cbind(runif(Nextra,0,2*pi),runif(Nextra,-pi/2,pi/2))
    }
    
    return(list(true.object,observed.objects))
}

iterate.process <- function(true.object, add_noise = T) {

    # system equation
    velocity.increment <- rnorm(2,0,sigma_v)
    true.object$radialcoords <- true.object$radialcoords + true.object$velocity + velocity.increment/2
    true.object$velocity <- true.object$velocity/2 + velocity.increment
    # book-keeping for radial coordinates mod pi 
    if (true.object$radialcoords[2] > pi/2) {
        true.object$radialcoords[2] <- pi/2-true.object$radialcoords[2]
        true.object$velocity[2] <- -abs(true.object$velocity[2]) # why the change of sign?
    } else if (true.object$radialcoords[2] < -pi/2) {
        true.object$radialcoords[2] <- -pi/2-(pi/2+true.object$radialcoords[2])
        true.object$velocity[2] <- abs(true.object$velocity[2])
    }
    true.object$radialcoords[1] <- true.object$radialcoords[1] %% (2*pi)  

    if (!add_noise) {
        return(true.object)
    } else {

        # observational errors on true object position
        true.object$observed.radialcoords <- c(as.numeric(rvonmises(1,true.object$radialcoords[1],kappa)),-pi/2+pi*rbeta(1,(true.object$radialcoords[2]+pi/2)/pi*N,(1-(true.object$radialcoords[2]+pi/2)/pi)*N))

        # add crowded field sources to list of observations
        observed.objects <- list()
        Nextra <- rpois(1,Nexp)
        observed.objects$radialcoords <- true.object$observed.radialcoords
        if (Nextra==0) {
        } else {
            observed.objects$radialcoords <- matrix(0,nrow=(Nextra+1),ncol=2)
            ordering <- sample(1:(Nextra+1),Nextra+1,replace=F)
            observed.objects$radialcoords[ordering[1],] <- true.object$observed.radialcoords
            observed.objects$radialcoords[ordering[2:(Nextra+1)],] <- cbind(runif(Nextra,0,2*pi),runif(Nextra,-pi/2,pi/2))
        }
        return(list(true.object,observed.objects))
    }
}

twod.to.threed <- function(x,y,r=1) {
    cbind(r*cos(x)*cos(y),r*sin(x)*cos(y),r*sin(y)) # for plotting purposes: nicer to look at 3d plots than movement mod pi on the plane
}

## Illustrate generation of mock data samples and process iteration

observed.objects <- initialise.process()

clear3d()
plot3d(twod.to.threed(matrix(observed.objects[[2]]$radialcoords,ncol=2)[,1],matrix(observed.objects[[2]]$radialcoords,ncol=2)[,2]),xlab="",ylab="",zlab="",axes=F,box=F,col=hsv(0.58),xlim=c(-1,1),ylim=c(-1,1),zlim=c(-1,1))

spheres3d(0,0,0,0.97,lit=F,col="white") # make a pretty sphere/grid on which to plot our points
for (y in seq(-pi/2,pi/2,length.out=25)) {
    lines3d(twod.to.threed(seq(0,2*pi,length.out=100),rep(y,100),0.98),col="grey90")
}
for (x in seq(0,2*pi,length.out=25)) {
    lines3d(twod.to.threed(rep(x,100),seq(-pi/2,pi/2,length.out=100),0.98),col="grey90")
}

spheres3d(twod.to.threed(observed.objects[[1]]$radialcoords[1],observed.objects[[1]]$radialcoords[2]),radius=0.025,col="magenta",lit=F) # true object position
spheres3d(twod.to.threed(observed.objects[[1]]$observed.radialcoords[1],observed.objects[[1]]$observed.radialcoords[2]),radius=0.015,col="green",lit=F) # observed object position

# plot the true object path for 100 steps to illustrate its random motion
true.object <- observed.objects[[1]]
old.object <- true.object
for (i in 1:100) {
old.object <- true.object
true.object <- iterate.process(true.object,add_noise=F)
points3d(twod.to.threed(true.object$radialcoords[1],true.object$radialcoords[2]),col="magenta")
segments3d(twod.to.threed(c(true.object$radialcoords[1],old.object$radialcoords[1]),c(true.object$radialcoords[2],old.object$radialcoords[2])),col="magenta")
}

## Track object with sequential importance sampler

current.observation <- initialise.process()
Nparticles <- 1000 # size of the particle approximation to our current state of knowledge 
particles <- list() # initialise supposing we initially know the exact location and (zero) velocity of the true object
particles$radialcoords <- cbind(rep(current.observation[[1]]$radialcoords[1],Nparticles),rep(current.observation[[1]]$radialcoords[2],Nparticles))
particles$velocities <- cbind(rep(0,Nparticles),rep(0,Nparticles))

# illustrate initial system state
clear3d()
plot3d(twod.to.threed(matrix(current.observation[[2]]$radialcoords,ncol=2)[,1],matrix(current.observation[[2]]$radialcoords,ncol=2)[,2]),xlab="",ylab="",zlab="",axes=F,box=F,col=hsv(0.58),xlim=c(-1,1),ylim=c(-1,1),zlim=c(-1,1))
spheres3d(0,0,0,0.97,lit=F,col="white")
    for (y in seq(-pi/2,pi/2,length.out=25)) {
        lines3d(twod.to.threed(seq(0,2*pi,length.out=100),rep(y,100),0.98),col="grey90")
    }
    for (x in seq(0,2*pi,length.out=25)) {
        lines3d(twod.to.threed(rep(x,100),seq(-pi/2,pi/2,length.out=100),0.98),col="grey90")
    }
spheres3d(twod.to.threed(current.observation[[1]]$radialcoords[1],current.observation[[1]]$radialcoords[2]),radius=0.025,col="magenta",lit=F)
spheres3d(twod.to.threed(current.observation[[1]]$observed.radialcoords[1],current.observation[[1]]$observed.radialcoords[2]),radius=0.015,col="green",lit=F)

for (iter in 1:200) { # main algorithm loop: sequential importance sampling

    # noisy observation of new system state
    current.observation <- iterate.process(current.observation[[1]])

    # importance sampling proposal given current particle approximation: this is just evolving our particles each by the random system equation: why?
    for (i in 1:Nparticles) {
        update <- iterate.process(list('radialcoords'=particles$radialcoords[i,],'velocity'=particles$velocities[i,]),add_noise=F)
        particles$velocities[i,] <- update$velocity
        particles$radialcoords[i,] <- update$radialcoords
    }

    # importance sampling weights
    weights <- numeric(Nparticles)
    for (i in 1:Nparticles) { # what are we summing over here?
        weights[i] <- sum(as.numeric(dvonmises(matrix(current.observation[[2]]$radialcoords,ncol=2)[,1],particles$radialcoords[i,1],kappa))*dbeta((matrix(current.observation[[2]]$radialcoords,ncol=2)[,2]+pi/2)/pi,(particles$radialcoords[i,2]+pi/2)/pi*N,(1-(particles$radialcoords[i,2]+pi/2)/pi)*N))
    }

    # resampling: why do we bother with this? how would the weighting step above need to be adjusted if we didn't want to resample?
    resample.ids <- sample(1:Nparticles,Nparticles,prob=weights,replace=T)
    particles$velocities <- particles$velocities[resample.ids,]
    particles$radialcoords <- particles$radialcoords[resample.ids,]

    # illustrate current particle approximation
    clear3d() # note: the rgl window can be reoriented with the mouse if you lose sight of your particles on the dark side of the moon
    plot3d(twod.to.threed(matrix(current.observation[[2]]$radialcoords,ncol=2)[,1],matrix(current.observation[[2]]$radialcoords,ncol=2)[,2]),xlab="",ylab="",zlab="",axes=F,box=F,col=hsv(0.58),xlim=c(-1,1),ylim=c(-1,1),zlim=c(-1,1))
    spheres3d(0,0,0,0.97,lit=F,col="white")
    for (y in seq(-pi/2,pi/2,length.out=25)) {
        lines3d(twod.to.threed(seq(0,2*pi,length.out=100),rep(y,100),0.98),col="grey90")
    }
    for (x in seq(0,2*pi,length.out=25)) {
        lines3d(twod.to.threed(rep(x,100),seq(-pi/2,pi/2,length.out=100),0.98),col="grey90")
    }
    spheres3d(twod.to.threed(current.observation[[1]]$radialcoords[1],current.observation[[1]]$radialcoords[2]),radius=0.025,col="magenta",lit=F)
    spheres3d(twod.to.threed(current.observation[[1]]$observed.radialcoords[1],current.observation[[1]]$observed.radialcoords[2]),radius=0.015,col="green",lit=F)
    points3d(twod.to.threed(particles$radialcoords[,1],particles$radialcoords[,2],r=0.99),col="grey90")
    
    readline(prompt="Press [enter] to continue; Control-C to stop")
    eval(parse(text=paste("snapshot3d(\"tracking",sprintf("%03i",iter),".png","\")",sep=""))) # write out the current plot to file so we can make an animated gif later on
}

system("convert -delay 10 *.png tracking.gif") # supposes you have the imageMagick package default installed on many mac and unix systems