## Script for parameter and state estimation using hidden Markov models
## Benchmark example: theta logistic population model
## Author: M.W. Pedersen, 13/7-2010
## Translated from MATLAB to R by Casper W. Berg, 13/1-2011


hmmmakekernel<-function(u,D,x){
    M = length(x);        # Number of grid cells
    dx = x[2]-x[1];       # Width of grid cell
    X = c(x, x[M]+dx)-dx/2; # Extents of envelope domain
    dX = X[2]-X[1];       # Width of grid cell in envelope domain
    P = matrix(0,nrow=M,ncol=M);    # Initialize transition matrix

    Dnum = D-(dx/2)^2/2;  # Correction term for numerical diffusion
    if(sum(Dnum<0)>0){cat("Warning: Dnum was negative! Make grid finer!",Dnum,"\n");}
    ## Fill transition probabilities into transition matrix
    npdfs=t(sapply(1:M,function(x,flaf) dnorm(flaf,u[x],sqrt(2*Dnum[x])),flaf=X))
    P=(npdfs[,1:M]+npdfs[,2:(M+1)])/2 * dX  # Trapez integration
    ## Fill transition probabilities into transition matrix
    #for( i in 1:M ){
    #    if(Dnum[i]<0){
    #    cat("Warning: Dnum was negative! Make grid finer!",Dnum,"\n");
    #}
    #    npdf = dnorm(X,u[i],sqrt(2*Dnum[i]));
    #    P[i,] = (npdf[1:M]+npdf[2:(M+1)])/2 * dX; # Trapez integration
    #}
    return(P);
}


hmmfilter<-function(par,grid,Y){
    res=list();

    N = length(Y);
    n = length(grid);

    ## Extract parameters
    theta = exp(par[1]);
    r0 = exp(par[2]);
    K = par[3];
    R = exp(2*par[4]);
    Q = exp(2*par[5]);

    zeroMat = matrix(0,nrow=N,ncol=n);
    lik=zeroMat;
    ## Compute observation likelihood for all time-steps
    for( t in 1:N){
        lik[t,] = dnorm(Y[t],grid,sd=sqrt(R));
    }

    D = Q/2*rep(1,n);                  # Variance jump rate, constant in time
    phi = zeroMat;                     #Initialize phi (state probabilities)
    pred = zeroMat;                           # Initialize vector for time-updates
    psi = rep(0,N);
    psi[1] = sum(lik[1,]);
    phi[1,] = lik[1,]/psi[1];                  # Calculate distribution in first time-step

    ## Filter
    ##for (t in 2:N){
    filt<-function(t){
        u = grid + r0*(1-(exp(grid)/K)^theta);   # Compute new mean
        P = hmmmakekernel(u,D,grid);             # Create transition matrix
        pred[t,] <<- phi[t-1,]%*%P;                # Time-update
        post = pred[t,] * lik[t,];               # Data-update, step 1,
        psi[t] <<- sum(post)+1e-20;                # Normalization constant
        phi[t,] <<- post/psi[t];                  # Data-update, step 2, normalization
    }
    lapply(2:N,filt);
    ##print(psi);
    loglikval = -sum(log(psi));

    res$loglikval=loglikval;
    res$phi=phi;
    res$pred=pred;
    return(res);
}

hmmfilter.wrap<-function(par,grid,Y){
    res=hmmfilter(par,grid,Y);
    return(res$loglikval);
}

hmmsmoother<-function(par,phi,pred,grid){
  ## Function for hidden Markov model smoothing
    N = dim(phi)[1];
    n = dim(phi)[2];
    ## Extract parameters
    theta = exp(par[1]);
    r0 = exp(par[2]);
    K = par[3];
    Q = exp(2*par[5]);

    D = Q/2*rep(1,n);      ## Variance jump rate, constant in time

    ## Smoother
    smooth = matrix(0,nrow=N,ncol=n);     ## Initialize vector for smoothed state probabilities
    smooth[N,] = phi[N,]; ## The last filter estimate is also a smoothed estimate
    for( t in N:2){
        rat = smooth[t,]/(pred[t,]+1e-16);               ## Smoothing equation, step 1
        u = grid + r0*(1-(exp(grid)/K)^theta);           ## Compute mean
        P = hmmmakekernel(u,D,grid);                     ## Create transition matrix
        smooth[t-1,] = phi[t-1,] * (rat%*%(t(P)));      ## Smoothing equation, step 2
        smooth[t-1,] = smooth[t-1,]/sum(smooth[t-1,]); ## Smoothing equation, step 3
    }
    return(smooth);
}

file = '../DATA/theta.dat';
##data = read.table(file,skip=5,nrows=1,sep='\t',header=FALSE)
data = scan(file,skip=5,sep=' ')
X=na.omit(data)
N=length(X);

##define grid
n=75; #251;
alpha=0.01;
grid=seq(min(X)+qnorm(alpha/2,0,sqrt(0.04)),max(X)+qnorm(1-alpha/2,0,sqrt(0.04)),length.out=n)

## theta r0 K R Q
guess = c(log(1), log(0.2), 800, log(sqrt(0.1)), log(sqrt(0.05)));           # Initial guess
LB = c(log(0.1), log(0.05), 100, log(sqrt(1e-3)), log(sqrt(0.005)));          # Lower bounds
UB = c(log(4), log(1), 5000, log(0.5), log(0.5));                # Upper bounds
  
time<-system.time(fit<-optim(guess,hmmfilter.wrap,grid=grid,Y=X,method="L-BFGS-B",lower=LB,upper=UB,control=list(trace=6),hessian=TRUE))
time <- unname(time["elapsed"]);

## TRUE PARAMS ARE
## Q=0.010000, R=0.040000, K=900.000000, r0=0.100000, theta=1.500000

## standard deviations
std=t(sqrt(diag(solve(fit$hessian))))
myorder=c(1,2,3,5,4)
std=std[myorder];

coef=fit$par
coef=coef[myorder]

ci.q=coef+std%o%c(-2,2)

ev <- eigen(fit$hessian)$value
eratio <- min(ev)/max(ev)

## Filtering and smoothing
filt = hmmfilter(coef,grid,X);

## Smoothing, Eq. 7
smooth = hmmsmoother(coef,filt$phi,filt$pred,grid);

## Calculate state estimates with uncertainties
GRID=grid;
for(i in 1:(N-1)) GRID=rbind(GRID,grid);
Xhat = rowSums(GRID*smooth);
XHAT=Xhat;
for(i in 1:(n-1)) XHAT=cbind(XHAT,Xhat);
Xsd = sqrt(rowSums((GRID-XHAT)^2*smooth));

## make a plot
##plot(Xhat,type="l",ylab="Log population",xlab="Time")
##lines(Xhat-2*Xsd,type="l",lty=2);
##lines(Xhat+2*Xsd,type="l",lty=2);
##lines(x,lwd=2)

results <- list(obj=fit$value,                     ## negative log L
                coef=coef,                   ## coefficients
                sd=std,                       ## standard deviations 
                confint.quad=ci.q,           ## Nx2 - matrix 95% ci 
                confint.profile=NULL,        
                confint.mcmc=NULL,           ## 0.025 0.975 quantiles
                time=c(fit=time,profile=NULL),  
                terminfo=c(maxgrad=NULL,eratio=eratio), ## MLE termination info
                convinfo=NULL       ## MCMC convergence info
                )

save("results",file="fit.RData")