## corr-fit.r
##
## Copyright (C) 2004 Havard Rue
##
## This program is free software; you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation; either version 2 of the License, or (at
## your option) any later version.
##
## This program is distributed in the hope that it will be useful, but
## WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
## General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program; if not, write to the Free Software
## Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
##
## The author's contact information:
##
##      H{\aa}vard Rue
##      Department of Mathematical Sciences
##      The Norwegian University of Science and Technology
##      N-7491 Trondheim, Norway
##      Voice: +47-7359-3533    URL  : http://www.math.ntnu.no/~hrue
##      Fax  : +47-7359-3524    Email: havard.rue@math.ntnu.no
##


require(nlme,quietly=T)
options(digits=22)

NH     = 2
NNH    = (NH+1)^2
NPARAM = (NH*(NH+1)/2)+NH

# type = 1 ...... exponential
#        2 ...... gaussian
#        3 ...... matern, give nu as well
#                                        
G = list(corr.ref=NA, weight=NA, ndim=128, range=10,min.eigval=0,
  use.fftw=FALSE, restart=0,
  type=1, matern.nu=1, matern.scale=1)

G$use.fftw = require(rimage,quietly=T)    ## use fftw if present
G$use.fftw = FALSE                        ## didn't gave any speedup, must be the interface

type.tag = function()
{
    if (G$type == 1) return ("exp")
    if (G$type == 2) return ("gauss")
    if (G$type == 3) return (paste("matern-", G$matern.nu, sep=""))
    stop("unknown G$type")
}

matern.scale = function(nu)
{
    ## find the scaling such that d=1 is the range

    ## first find an approximate region
    high=sqrt(nu)
    while(matern(1,nu=nu,scale=high) > 0.05) high=high/0.9
    low=high
    while(matern(1,nu=nu,scale=low) < 0.05) low=low*0.9
    
    ## then solve
    optimise(f = function(s) (matern(1,nu=nu,scale=s)-0.05)^2,  tol = .Machine$double.eps^0.5,interval=c(low,high))$minimum
}

matern <- function(d,nu=1,scale=1)
{
    d = pmax(d, 0.0000001)
    return(1/gamma(nu)*2^(1-nu)*(scale*d)^nu*besselK(scale*d,nu=nu))
}

# build small coofs matrix, elm is the diagonal.
coofs.matrix = function(coofs, elm=1)
{
    coofs.m      = matrix(nrow=NH+1,ncol=NH+1)
    coofs.m[1,1] = elm
    k = 1
    for(i in 2:(NH+1))
    {
        for(j in 1:i)
        {
            coofs.m[i,j] = coofs.m[j,i] = coofs[k]
            k = k + 1
        }
    }
    return (coofs.m)
}

# build matrix based on the coefficients, elm is the diagonal.
build.matrix = function(coofs,elm=1)
{
    coofs.m  = coofs.matrix(coofs,elm=elm)
    x        = matrix(data=0, nrow=G$ndim, ncol=G$ndim)
    i        = abs(-NH:NH)+1
    ii       = ((-NH:NH) + G$ndim)%%G$ndim + 1
    x[ii,ii] = coofs.m[i,i]
    return(x)
}

is.legal.coof = function(coofs, ndim=4096,elm=1)
{
    # return TRUE/FALSE if this is legal on a larger torus
    nd = G$ndim
    G$ndim <<- ndim
    min.eig = min(Re(fft(build.matrix(coofs, elm=elm))))
    G$ndim <<- nd
    return (paste("is.legal.coof", ndim, min.eig, min.eig > 0))
}

# compute the cov.matrix
cov.matrix = function(coofs, elm=1)
{
    nn = floor(G$ndim/2)
    x = build.matrix(coofs, elm=elm)
    #print(x)
    if (G$use.fftw)
    {
        eig         = Re(fftw(x,dir=-1))
        G$min.eigval  <<-  min(eig)
        cov.matrix = Re(fftw(1/eig,dir=1))/length(x)
    }
    else
    {
        eig         = Re(fft(x))
        G$min.eigval  <<-  min(eig)
        cov.matrix = Re(fft(1/eig,inverse=T)/length(x))
    }
    cov.matrix = cov.matrix[1:nn,1:nn]
    return(cov.matrix)
}

grad = function(coof, elm=1)
{
    # return the derivative of the covariance matrix wrt each parameter as a list
    nn     = floor(G$ndim/2)
    grad  = list()
    coof.d = vector("double", NPARAM)
    aa     = Re(fft(build.matrix(coof, elm=elm)))^(-2)
    for(i in 1:NPARAM)
    {
        coof.d[]  = 0
        coof.d[i] = 1
        grad[[i]] = Re(fft( - aa* Re(fft(build.matrix(coof.d,elm=0))), inverse=TRUE)[1:nn,1:nn])/G$ndim^2
    }
    return(grad)
}

grad.corr = function(coof, elm=1)
{
    #derivative of the correlation function
    cov.der = grad(coof, elm=elm)
    cov.mat = cov.matrix(coof,elm=elm)
    rho.der = list()
    for(i in 1:NPARAM)
    {
        rho.der[[i]] = cov.der[[i]]/cov.mat[1,1] - cov.mat/cov.mat[1,1]^2*cov.der[[i]][1,1]
    }
    return(rho.der)
}

test.grad = function(coof)
{
    coofs = coof
    for(idx in 1:NPARAM)
    {
        d = 0.000001
        coof.new = coofs
        coof.new[idx] = coof.new[idx] + d
        coof.nnew = coofs
        coof.nnew[idx] = coof.nnew[idx] - d
        d.1 = (cov.matrix(coof.new) - cov.matrix(coof.nnew))/2/d
        d.2 = grad(coofs)[[idx]]
        print(rms(d.1-d.2))
    }
}

test.grad.corr = function(coof)
{
    coofs = coof
    for(idx in 1:NPARAM)
    {
        d = 0.0001
        coof.new      = coofs
        coof.new[idx] = coof.new[idx] + d
        coof.nnew      = coofs
        coof.nnew[idx] = coof.nnew[idx] - d
        a = cov.matrix(coof.new)
        a = a/a[1,1]
        b = cov.matrix(coof.nnew)
        b = b/b[1,1]
        
        d.1 = (a-b)/2/d
        d.2 = grad.corr(coofs)[[idx]]
        print(rms(d.1-d.2))
    }
}

# compute the reference correlation function
corr.matrix.ref = function(range)
{
    nn = floor(G$ndim/2)
    corr.ref = matrix(0,nrow=nn,ncol=nn)
    if (range > 0)
    {
        if (G$type == 1)
        { 
            for(j in 1:nn)
            {
                i = 1:nn
                corr.ref[i,j] = exp(-3*sqrt((i-1)^2+(j-1)^2)/range)
            }
        }
        if (G$type == 2)
        {
            for(j in 1:nn)
            {
                i = 1:nn
                corr.ref[i,j] = exp(-3*(((i-1)^2+(j-1)^2)/range^2))
            }
        }
        if (G$type == 3)
        {
            for(j in 1:nn)
            {
                i = 1:nn
                corr.ref[i,j] = matern(sqrt((i-1)^2+(j-1)^2)/range, nu=G$matern.nu, scale=G$matern.scale)
            }
        }
    }
    else
    {
        corr.ref[1,1] = 1
    }
    return(corr.ref)
}

# compute the weight matrix
weight.matrix = function(range)
{
    nn = floor(G$ndim/2)
    weight = matrix(data=0,nrow=nn,ncol=nn)
    for(j in 1:nn)
    {
        i = 1:nn
        d = sqrt((j-1)^2 + (i-1)^2)
        d = pmax(1,d)
        weight[i,j] = (1+ range/d)/d
    }
    return(weight/mean(weight))
}

cost.function = function(coofs)
{
    params = as.vector(coofs)
    
    nn     = floor(G$ndim/2)
    cov.m  = cov.matrix(params)
    c.m    = cov.m/cov.m[1,1]

    if (G$min.eigval > 0)
    {
        cost <- 0
        for(j in 1:nn)
        {
            i = 1:j
            cost = cost + sum((c.m[i,j]-G$corr.ref[i,j])^2*G$weight[i,j])
        }
        return(cost)
    }
    else
    {
        return(NA)
    }
}

cost.function.grad = function(coofs)
{
    grad = vector("double", NPARAM)

    if (is.na(cost.function(coofs))) 
    {
        grad[]  = NA
        return (grad)
    }
    
    params = as.vector(coofs)

    grad[] = 0
    nn     = floor(G$ndim/2)
    cov.m  = cov.matrix(params)
    c.m    = cov.m/cov.m[1,1]

    grad.corr = grad.corr(params)
    
    ## this version should be faster, we only need to sum (i,j), i>j, i < N/2
    for(k in 1:NPARAM)
    {
        for(j in 1:nn)
        {
            i = 1:j
            grad[k] = grad[k] + sum(2*(c.m[i,j]-G$corr.ref[i,j])*G$weight[i,j]*grad.corr[[k]][i,j])
        }
    }
    
    return(grad)
}

    

################################################################################
################################################################################
################################################################################


system("test -d ./figs || mkdir ./figs", intern=FALSE)


tt = as.integer(Sys.getenv("GMRFAPPROX_TYPE"))
if (is.na(tt))  tt = 3 # Matern
G$type = tt

nu = seq(from=0.5,to=0.95,by=0.05)
#nu <- c(9.5, 1.05, 1.1, 1.15, 1.2, 1.25) 

for (count in 1:length(nu))
{
  coofs        = matrix(data=rep(0,NPARAM),nrow=NPARAM,ncol=1)
  coofs.prev   = matrix(data=rep(0,NPARAM),nrow=NPARAM,ncol=1)
  ranges        = seq(from=0.0,to=100,by=0.05)
  all.coofs     = matrix(NA,nrow = length(ranges),ncol = NPARAM+1)
  all.coofs[,1] = ranges

  
  G$matern.nu = nu[count]
  print(paste("*******************************Matern, nu =", G$matern.nu))

  for (counter in 1:length(ranges))
  {
    range = ranges[counter]

    G$matern.scale = matern.scale(G$matern.nu)
    G$ndim = max(64, 2^(floor(log(5*max(0.000001,range))/log(2))+1))
    
    print(paste("************************************* counter", counter, " range =", range))
    
    G$range       = range
    G$weight      = weight.matrix(range)
    G$corr.ref    = corr.matrix.ref(G$range)

    # keep track of old solutions to improve the initial value
    coofs.prev   = coofs

    coofs.initial = coofs = coofs.prev  

    # ensure that the inital value is ok
    if (is.na(cost.function(coofs.initial))) coofs = coofs.prev

    #test.grad.corr(coofs)
    
    for (met  in c("BFGS"))
    {
        solution = optim(coofs,
          cost.function,gr=cost.function.grad,method=met,hessian=FALSE,
          control=list(maxit=500,trace=6,type=3,reltol=.Machine$double.eps^(1./2.)/1000, REPORT=1))
        
        if (solution$convergence != 0) 
        {
            print(paste(sprintf("%-15s",met), ": convergence ", solution$convergence, " message",
                        solution$message))
        }
        else
        {
            print(paste(sprintf("%-15s",met), ": value",solution$value, " nfunc_evals",
                         solution$counts[1]))
        }

        coofs = solution$par

        all.coofs[counter,2:(NPARAM+1)] = coofs

        print(paste("rms in initialisation", sqrt(mean((coofs-coofs.initial)^2))))

        xx = c(G$range, solution$value, as.vector(coofs), met, solution$counts)
        write(xx,file=paste("trace-optim-", type.tag(),".dat",sep=""),ncol=length(xx),append=T)
    }

    c.m      = cov.matrix(as.vector(coofs))
    variance = c.m[1,1]
    print(paste("variance ",variance))
    c.m      = c.m/variance
    nn       = floor(G$ndim/2)

    fnm = paste("./figs/match-range-",G$range,"-", type.tag(),".ps",sep="")
    postscript(file=fnm,paper="a4",)
    par(mfrow=c(2,2))
    plot(0:(nn-1),c.m[1:nn,1],type="l",main=paste("Exact and matched correlation for range",range))
    lines(0:(nn-1),G$corr.ref[1:nn,1])
    plot(0:(nn-1),G$corr.ref[1:nn,1] - c.m[1:nn,1],type="l",
         main=paste("Exact - matched correlation for range",range))
    plot(0:(nn-1),diag(c.m),type="l",main=paste("Exact and matched correlation for range(diag)",range))
    lines(0:(nn-1),diag(G$corr.ref))
    plot(0:(nn-1),diag(G$corr.ref) - diag(c.m),type="l",
         main=paste("Exact - matched correlation (diag) for range",range))
    dev.off()

    max.err = max(abs(c.m[1:nn,1:nn]- G$corr.ref[1:nn,1:nn]))
    stdev   = sqrt(mean((c.m[1:nn,1:nn]- G$corr.ref[1:nn,1:nn])^2))
    xx = c(1, nu[count], G$range, as.vector(coofs.matrix(as.vector(coofs))), variance, 
      is.legal.coof(coofs,ndim=1024), max.err, stdev)
     write(xx, file=paste("COOFS-", type.tag(), ".dat", sep=""), ncol = length(xx), append=T)

  }

}