##==============================================================================
## Creation of a one-dimensional finite difference grid
##==============================================================================

setup.grid.1D <- function(x.up=0,	x.down=NULL, L=NULL,
  N=NULL, dx.1 =NULL, p.dx.1 = rep(1,length(L)),
  max.dx.1 = L, dx.N =NULL, p.dx.N = rep(1,length(L)),
  max.dx.N = L) {

## Check on the input

  for (i in 1:length(L)) {
  if (is.null(x.down[i]) && is.null(L[i]))
    stop (paste("Error in setup.grid.1D! Grid zone",i,": either the length (L) or the end point (x.down) should be specified"))
  }

## If the interval lengths are given, create the end points of each zone 

  if (is.null(x.down[1])) {
    if (any(L<0)) stop (paste("Error in setup.grid.1D! L[",which(L<0),"] < 0",sep=""))
    x.down <- cumsum(L)
  }
  
## If the interval lengths are given, create the end points of each zone 

  if (is.null(L[1])) {
    L <- diff(c(x.up,x.down))
    if (any(L<0)) stop (paste("Error in setup.grid.1D! L[",which(L<0),"] < 0",sep=""))
  }

## Check wether all info is available in each zone 

  for (i in 1:length(L)) {
    if (is.null(N[i]) && is.null(dx.1[i]) && is.null(dx.N[i]))
    stop (paste("Error in setup.grid.1D! Both N[",i,"], dx.1[",i,"] and dx.N[",i,"] are NULL"))
  }

## Calculation of the grid cell sizes

  dx <- vector()

## Power law function that controls the increase in grid size. The 
## root of this function determines the power factor that corresponds to a 
## desired number of cells N 

  f.root <- function(p,dx,N,L) dx*(p^(N)-1)/(p-1) - L

## Loop over all grid zones
  for (i in 1:length(L)) {

## Option 1: only N[i] is specified, all grid cells are equal
    if (!(is.null(N[i])) && (is.null(dx.1[i])) && (is.null(dx.N[i])))  { 
      if (N[i] < 1) stop (paste("Error in setup.grid.1D! N[",i,"] < 1",sep=""))
      A.dx  <- rep(L[i]/N[i],N[i])
    }
    
## Option 2: dx.1[i] specified, N[i] = NULL and dx.N[i] = NULL 

    if ((is.null(N[i])) && !(is.null(dx.1[i])) && (is.null(dx.N[i])))  { 
      if (dx.1[i]<=0)
        stop (paste("Error in setup.grid.1D! dx.1[",i,"] <= 0",sep=""))
      if (dx.1[i] > L[i])
        stop (paste("Error in setup.grid.1D! dx.1[",i,"] > L[",i,"]",sep=""))
  
     # use gradual increase of grid cell size at upstream interface 
      A.dx <- vector()
      pos <- vector()
      A.dx[1] <- dx.1[i]
      pos[1] <- dx.1[i]
      j <- 1
      while (pos[j] < L[i]) {
        j <- j + 1
        A.dx[j] <- min(max.dx.1[i],A.dx[j-1]*p.dx.1[i])
        pos[j] <- pos[j-1] + A.dx[j]
      }
     # rescaling to fit the interval length 
      A.dx <- A.dx*(L[i]/pos[length(pos)])
    }

## Option 3: dx.N[i] specified, N[i] = NULL and dx.1[i] = NULL 

    if ((is.null(N[i])) && (is.null(dx.1[i])) && !(is.null(dx.N[i])))  { 
      if (dx.N[i]<=0)
        stop (paste("Error in setup.grid.1D! dx.N[",i,"] <= 0",sep=""))
      if (dx.N[i] > L[i])
        stop (paste("Error in setup.grid.1D! dx.N[",i,"] > L[",i,"]",sep=""))
  
     # use gradual increase of grid cell size at downstream interface 
      A.dx <- vector()
      pos <- vector()
      A.dx[1] <- dx.N[i]
      pos[1] <- dx.N[i]
      j <- 1
      while (pos[j] < L[i]) {
        j <- j + 1
        A.dx[j] <- min(max.dx.N[i],A.dx[j-1]*p.dx.N[i])
        pos[j] <- pos[j-1] + A.dx[j]
      }
     # rescaling to fit the interval length 
      A.dx <- A.dx*(L[i]/pos[length(pos)])
     # reversing the A.dx vector
      A.dx <- A.dx[length(A.dx):1]
    }

## Option 4: N[i] and dx.1[i] are specified, dx.N[i] = NULL 

    if (!(is.null(N[i])) && !(is.null(dx.1[i])) && (is.null(dx.N[i])))  { 

      if (N[i] < 1) stop (paste("Error in setup.grid.1D! N[",i,"] < 1",sep=""))
      if (dx.1[i]<=0)
        stop (paste("Error in setup.grid.1D! dx.1[",i,"] <= 0",sep=""))
      if (dx.1[i] > L[i])
        stop (paste("Error in setup.grid.1D! dx.1[",i,"] > L[",i,"]",sep=""))
  
      # estimate power in power law 
      p.estim <- uniroot(f=f.root,dx=dx.1[i],N=N[i],L=L[i],lower=1.001,upper=2)$root
      # use gradual increase of grid cell size at upper interface 
      A.dx <- vector()
      pos <- vector()
      A.dx[1] <- dx.1[i]
      pos[1] <- dx.1[i]
      j <- 1
      while (j < N[i]) {
        j <- j + 1
        A.dx[j] <- min(max.dx.1[i],A.dx[j-1]*p.estim)
        pos[j] <- pos[j-1] + A.dx[j]
      }
     # rescaling to fit the interval length 
      A.dx <- A.dx*(L[i]/pos[length(pos)])
    }

## Option 5: N[i] and dx.N[i] are specified, dx.1[i] = NULL 

    if (!(is.null(N[i])) && (is.null(dx.1[i])) && !(is.null(dx.N[i])))  { 

      if (N[i] < 1) stop (paste("Error in setup.grid.1D! N[",i,"] < 1",sep=""))
      if (dx.N[i]<=0)
        stop (paste("Error in setup.grid.1D! dx.N[",i,"] <= 0",sep=""))
      if (dx.N[i] > L[i])
        stop (paste("Error in setup.grid.1D! dx.N[",i,"] > L[",i,"]",sep=""))
  
      # estimate power in power law 
      p.estim <- uniroot(f=f.root,dx=dx.N[i],N=N[i],L=L[i],lower=1.001,upper=2)$root
      # use gradual increase of grid cell size at upper interface 
      A.dx <- vector()
      pos <- vector()
      A.dx[1] <- dx.N[i]
      pos[1] <- dx.N[i]
      j <- 1
      while (j < N[i]) {
        j <- j + 1
        A.dx[j] <- min(max.dx.N[i],A.dx[j-1]*p.estim)
        pos[j] <- pos[j-1] + A.dx[j]
      }
     # rescaling to fit the interval length 
      A.dx <- A.dx*(L[i]/pos[length(pos)])
     # reversing the A.dx vector
      A.dx <- A.dx[length(A.dx):1]
    }

## Option 6: dx.1[i] and dx.N[i] are specified, N[i] = NULL 

    if (!(is.null(dx.1[i])) && !(is.null(dx.N[i])))  { 

      if (dx.1[i]<=0)
        stop (paste("Error in setup.grid.1D! dx.1[",i,"] <= 0",sep=""))
      if (dx.1[i] > L[i]/2)
        stop (paste("Error in setup.grid.1D! dx.1[",i,"] > L[",i,"]",sep=""))
      if (dx.N[i]<=0)
        stop (paste("Error in setup.grid.1D! dx.N[",i,"] <= 0",sep=""))
      if (dx.N[i] > L[i]/2)
        stop (paste("Error in setup.grid.1D! dx.N[",i,"] > L[",i,"]",sep=""))

      # estimate power in power law 
      if (!is.null(N[i])) {
      N.A <- as.integer(N[i]/2) 
      N.B <- N[i] - N.A
      p.dx.1[i] <- uniroot(f=f.root,dx=dx.1[i],N=N.A,L=0.5*L[i],lower=1.001,upper=2)$root
      p.dx.N[i] <- uniroot(f=f.root,dx=dx.N[i],N=N.B,L=0.5*L[i],lower=1.001,upper=2)$root
      }
      # use gradual increase of grid cell size at upsteam interface
      A.dx  <- vector()
      pos <- vector()
      A.dx[1] <- dx.1[i]
      pos[1] <- dx.1[i]
      j <- 1
      while ((is.null(N[i])&&(pos[j] < 0.5*L[i]))||(!is.null(N[i])&&(j < N.A))) {
         j <- j + 1
         A.dx[j] <- min(max.dx.1[i],A.dx[j-1]*p.dx.1[i])
         pos[j] <- pos[j-1] + A.dx[j]
      }
      A.dx <- A.dx*(0.5*L[i]/pos[length(pos)])

      # use gradual increase of grid cell size at downsteam interface
      B.dx  <- vector()
      pos <- vector()
      B.dx[1] <- dx.N[i]
      pos[1] <- dx.N[i]
      j <- 1
      while ((is.null(N[i])&&(pos[j] < 0.5*L[i]))||(!is.null(N[i])&&(j < N.B))) {
        j <- j + 1
        B.dx[j] <- min(max.dx.N[i],B.dx[j-1]*p.dx.N[i])
        pos[j] <- pos[j-1] + B.dx[j]
      }
      B.dx <- B.dx*(0.5*L[i]/pos[length(pos)])
      # assemble the distance vector 
      A.dx <- c(A.dx,B.dx[length(B.dx):1])
    }

  ## calculation of present zone has finished, add distances to existing array 
  
    dx <- c(dx,A.dx)
  }

  ## positions and distances

  x.int  <- x.up + diffinv(dx)
  x.mid  <- colMeans(rbind(x.int[-1],x.int[-(length(dx)+1)]))
  dx.aux <- c(dx[1]/2,diff(x.mid),dx[length(dx)]/2)

  ## Packaging of results

  Res <- list(x.up = x.up,
	  			    x.down = x.down[length(x.down)],
		  		    x.mid = x.mid,   # position of centre of the grid cells, vector of length N
              x.int  = x.int,  # position of the grid cell interfaces , vector of length N+1
              dx = dx  ,       # thickness of the grid cells , vector length N
              dx.aux = dx.aux, # auxiliary vector with distances between centre of adjacent cells, first and last: half of cell size, vector of length N+1
              N = length(dx))  # total number of grid cells
  class(Res) <- "grid.1D"
  return(Res)
}

##==============================================================================
## S3 method: Plotting of a one-dimensional finite difference grid
##==============================================================================

plot.grid.1D <- function(x,...) {
  mf <- par(mfrow=c(2,1))
  on.exit(par(mf))
    
  plot(x$x.int,main="position of cells",ylab="x",xlab="index",...)
  points(0.5+(1:x$N),x$x.mid,pch=16)
  legend("topleft",c("x.int","x.mid"),pch=c(1,16))
  plot(x$dx.aux,main="box thickness",ylab="dx",xlab="index",...)
  points(0.5+(1:x$N),x$dx,pch=16)
  legend("topleft",c("dx.aux","dx.mid"),pch=c(1,16))
}