#' @rdname CST_RFTemp
#' @title Temperature downscaling of a CSTools object using lapse rate correction or a reference field
#' @author Jost von Hardenberg - ISAC-CNR, \email{j.vonhardenberg@isac.cnr.it}
#' @description This function implements a simple lapse rate correction of a
#' temperature field (an object of class 's2dv_cube' as provided by `CST_Load`) as input.
#' @references Method described in H2020 ECOPOTENTIAL Deliverable No. 8.1: High resolution (1-10 km) climate, land use and ocean change scenarios \link[https://www.ecopotential-project.eu/images/ecopotential/documents/D8.1.pdf]
#' @param data An object of the class 's2dv_cube' as returned by `CST_Load`, 
#' containing the temperature fields to downscale.
#' The data object is expected to have an element named \code{$data} with at least two 
#' The number of longitudes and latitudes in the input data is expected to be even and the same. If not
#' spatial dimensions named "lon" and "lat".
#' @param oro An object of the class 's2dv_cube' as returned by `CST_Load`, 
#' containing fine scale orography (in meters).
#' @param xlim vector with longitude bounds for downscaling
#' @param ylim vector with latitude bounds for downscaling
#' @param lapse float with environmental lapse rate 
#' @param lon_dim string with name of longitude dimension
#' @param lat_dim string with name of latitude dimension
#' @param verbose logical if to print diagnostic output
#' @param nolapse logical, if true `oro` is interpreted as a fine-scale climatology and used directly for bias correction
#' @return CST_RFTemp() returns a downscaled CSTools object (i.e., of the class 's2dv_cube').
#' @export
#' @import multiApply
#' @import rainfarmr
#' @examples
#' # Generate simple synthetic data and downscale by factor 4
#' t <- rnorm(6 * 6 * 2 * 3 * 4)*10 + 273.15 + 10
#' dim(t) <- c(dataset = 1, member = 2, sdate = 3, ftime = 4, lat = 6, lon = 6)
#' lon <- seq(4, 9, 1)
#' lat <- seq(42, 47, 1)
#' exp <- list(data = t, lat = lat, lon = lon)
#' attr(exp, 'class') <- 's2dv_cube'
#' o <- runif(29*29)*3000
#' dim(o) <- c(lat = 29, lon = 29)
#' lon <- seq(3, 10, 0.25)
#' lat <- seq(41, 48, 0.25)
#' oro <- list(data = o, lat = lat, lon = lon)
#' attr(oro, 'class') <- 's2dv_cube'
#' res <- CST_RFTemp(exp, oro, xlim=c(4,8), ylim=c(43, 46), lapse=6.5)

CST_RFTemp <- function(data, oro, xlim = NULL, ylim = NULL, lapse=6.5,
                       lon_dim = "lon", lat_dim = "lat", time_dim = NULL,
                       nolapse = FALSE, verbose=FALSE) {

  if (!inherits(data, 's2dv_cube')) {
    stop("Parameter 'data' must be of the class 's2dv_cube', ",
         "as output by CSTools::CST_Load.")
  }

  res <- RFTemp(data$data, data$lon, data$lat,
                oro$data, oro$lon, oro$lat, xlim, ylim, lapse,
                lon_dim = lon_dim, lat_dim = lat_dim, time_dim = time_dim,
                nolapse = nolapse, verbose = verbose)

  data$data <- res$data
  data$lon <- res$lon
  data$lat <- res$lat

  return(data)
}

#' @rdname RFTemp
#' @title Temperature downscaling of a CSTools object using lapse rate correction (reduced version)
#' @author Jost von Hardenberg - ISAC-CNR, \email{j.vonhardenberg@isac.cnr.it}
#' @description This function implements a simple lapse rate correction of a
#' temperature field (a multidimensional array) as input
#' @references Method described in H2020 ECOPOTENTIAL Deliverable No. 8.1: High resolution (1-10 km) climate, land use and ocean change scenarios \link[https://www.ecopotential-project.eu/images/ecopotential/documents/D8.1.pdf]
#' @param data Temperature array to downscale.
#' The input array is expected to have at least two dimensions named "lon" and "lat" by default
#' (these default names can be changed with the \code{lon_dim} and \code{lat_dim} parameters)
#' @param lon Vector or array of longitudes.
#' @param lat Vector or array of latitudes.
#' @param oro Array containing fine-scale orography (in m)
#' @param xlim vector with longitude bounds for downscaling
#' @param ylim vector with latitude bounds for downscaling
#' @param lapse float with environmental lapse rate 
#' @param lon_dim string with name of longitude dimension
#' @param lat_dim string with name of latitude dimension
#' @param verbose logical if to print diagnostic output
#' @param nolapse logical, if true `oro` is interpreted as a fine-scale climatology and used directly for bias correction
#' @return RFTemp() returns a list containing the fine-scale longitudes, latitudes and the downscaled fields.
#' @export
#' @import multiApply
#' @import rainfarmr
#' @examples
#' # Generate simple synthetic data and downscale by factor 4
#' t <- rnorm(6 * 6 * 4 * 3)*10 + 273.15 + 10
#' dim(t) <- c(sdate = 3, ftime = 4, lat = 6, lon = 6)
#' lon <- seq(4, 9, 1)
#' lat <- seq(42, 47, 1)
#' o <- runif(29*29)*3000
#' dim(o) <- c(lat = 29, lon = 29)
#' lono <- seq(3, 10, 0.25)
#' lato <- seq(41, 48, 0.25)
#' res <- RFTemp(t, lon, lat, o, lono, lato, xlim=c(4,8), ylim=c(43, 46), lapse=6.5)

RFTemp <- function(data, lon, lat, oro, lonoro, latoro,
                   xlim = NULL, ylim = NULL, lapse = 6.5,
                   lon_dim = "lon", lat_dim = "lat", time_dim = NULL,
                   nolapse = FALSE, verbose=FALSE) {

  # Check/detect time_dim
  if (is.null(time_dim)) {
    time_dim_names <- c("ftime", "sdate", "time")
    time_dim_num <- which(time_dim_names %in% names(dim(data)))
    if (length(time_dim_num) > 0) {
      # Find time dimension with length > 0
      ilong <- which(dim(data)[time_dim_names[time_dim_num]] > 0)
      if (length(ilong) > 0) {
        #time_dim <- time_dim_names[time_dim_num[ilong[1]]]
        time_dim <- time_dim_names[time_dim_num[ilong]]
      } else {
        stop("No time dimension longer than zero found.")
      }
    } else {
      stop("Could not automatically detect a target time dimension ",
           "in the provided data in 'data'.")
    }
    warning(paste("Selected time dim:", time_dim,"\n"))
  }

  # Repeatedly apply .downscalet
  result <- Apply(data, target_dims=c(lon_dim, lat_dim, time_dim),
                  fun=.downscalet, lon, lat, oro, lonoro, latoro,
                  xlim=xlim, ylim=ylim, lapse=lapse,
                  nolapse=nolapse, verbose=verbose)

  names(dim(result$data))[1] <- names(dim(data))[names(dim(data))==lon_dim]
  names(dim(result$data))[2] <- names(dim(data))[names(dim(data))==lat_dim]
  result$lon=array(result$lon[1:dim(result$lon)[1]])
  result$lat=array(result$lat[1:dim(result$lat)[1]])
  names(dim(result$lon))=lon_dim
  names(dim(result$lat))=lat_dim
  return(result)
}


#' Lapse-rate temperature correction downscaling
#'
#' @description Downscales a temperature field using a lapse-rate
#' correction based on a reference orography. Time-averaging is done on all 
#' dimensions after the first two.
#' @author Jost von Hardenberg, \email{j.vonhardenberg@isac.cnr.it}
#' @param lon vector of input longitudes
#' @param lat vector of input latitudes
#' @param t matrix of input temperature data
#' @param lono vector of orography longitudes
#' @param lato vector of orography latitudes
#' @param oro matrix of topographical elevations (in meters) 
#' @param xlim vector of longitude bounds
#' @param ylim vector of latitude bounds
#' @param radius smoothing radius expressed in longitude units (default is half a large-scale pixel)
#' @param lapse environmental lapse rate (in K/Km)
#' @param nolapse logical, if true `oro` is interpreted as a fine-scale climatology and used directly for bias correction
#' @param verbose logical if to print diagnostic output
#' @return A downscaled temperature matrix
#' @export
#' @examples
#' lon=5:20
#' lat=35:40
#' t=runif(16*6); dim(t)=c(16,6)
#' lono=seq(5,20,0.1)
#' lato=seq(35,40,0.1)
#' o=runif(151*51)*2000; dim(o)=c(151,51)
#' td=.downscalet(t,lon, lat,o,lono,lato,c(8,12),c(36,38))

.downscalet <- function(t, lon, lat, oro, lono, lato,
                        xlim = NULL, ylim = NULL,
                        radius = 0, lapse = 6.5, nolapse = FALSE,
                        verbose = FALSE) {
    tdim <- dim(t)
    ldim <- FALSE
    if (length(tdim)==3) {
        nt = tdim[3]
    } else if (length(tdim)==2) {
        nt = 1
        dim(t) <- c(tdim, 1)
    } else if  (length(tdim)<2) {
        stop("Input array must have at least two dimensions")
    } else {
        ldim <- TRUE
        dim(t) <- c(tdim[1:2], time = prod(tdim[-(1:2)])) 
        nt = dim(t)[3]
    }

    dxl <- (lon[2] - lon[1]) / 2
    dyl <- (lat[2] - lat[1]) / 2
    if(radius==0) {
        radius <- dxl
    }
    if(is.null(xlim)) {
        xlim <- c(lon[1] + dxl, lon[length(lon)] - dxl)
    }
    if(is.null(ylim)) {
        ylim <- c(lat[1] + dyl, lat[length(lat)] - dyl)
        if (ylim[1] > ylim[2]) {
           ylim <- ylim[2:1] 
        }
    }
    #Add buffer
    lon1 <- xlim[1] - radius
    lon2 <- xlim[2] + radius
    lat1 <- ylim[1] - radius
    lat2 <- ylim[2] + radius
   
    orocut <- oro[(lono <= lon2) & (lono >= lon1),
                  (lato <= lat2) & (lato >= lat1)]
    lonocut <- lono[(lono <= lon2) & (lono >= lon1)]
    latocut <- lato[(lato <= lat2) & (lato >= lat1)]

    if(length(lonocut)==0 | length(latocut)==0) {
        stop("Orography not available for selected area")
    }
   
    if(verbose) { 
        print("Nearest neighbours interpolation")
    }
    tcut <- .interpnn(lon, lat, t, lonocut, latocut)
    # Fine-scale smooth interpolated input field
    dxol <- lonocut[2] - lonocut[1]
    nrad <- as.integer(radius / abs(dxol) + 0.5)
    nx <- length(lonocut)
    ny <- length(latocut)
    tout <- array(0., c(nx, ny, nt))
    for (i in seq_len(nt)) {
        if(verbose) { 
            print(paste("Smooth ",i))
        }
        tout[,, i] <- .smooth(tcut[,,i], nrad)
    }

    if (nolapse) {
        # oro is a reference fine-scale field: use that directly
        # for bias correcting the interpolated input field's temporal mean
        t1 <- orocut - apply(tout, c(1,2), mean)
    } else {
        # Lapse-rate correction
        ns <- length(lonocut)
        orocuts <- .smooth(orocut, nrad)
        t1 <- -(orocut-orocuts)*lapse/1000.
    }
    if(verbose) { 
        print("Adding")
    }
    for (i in seq_len(nt)) {
        tout[,, i] <- tout[,, i] + t1
    }
    tdown <- tout[(lonocut <= xlim[2]) & (lonocut >= xlim[1]),
                  (latocut <= ylim[2]) & (latocut >= ylim[1]),]
    if(ldim) {
        dim(tdown) <- c(dim(tdown)[1:2], tdim[-(1:2)])
    }
    lonocut <- lonocut[(lonocut <= xlim[2]) & (lonocut >= xlim[1])]
    latocut <- latocut[(latocut <= ylim[2]) & (latocut >= ylim[1])]

    return(list(data=tdown, lon=lonocut, lat=latocut))
}

#' Nearest neighbour interpolation
#'
#' @description The input field is interpolated onto the output coordinate grid using nearest neighbours
#' @author Jost von Hardenberg, \email{j.vonhardenberg@isac.cnr.it}
#' @param z matrix with the input field to interpolate (assumed to include also a third time dimension)
#' @param lon vector of input longitudes
#' @param lat vector of input latitudes
#' @param lonp vector of output longitudes
#' @param latp vector of output latitudes
#' @return The interpolated field.
#' @examples
#' lon=5:11
#' lat=35:40
#' z=runif(7*6*2); dim(z)=c(7,6,2)
#' lonp=seq(5,10,0.2)
#' latp=seq(35,40,0.2)
#' zo <- .interpnn(lon, lat, z, lonp, latp) 

.interpnn <- function(lon, lat, z, lonp, latp) {

    nx <- length(lonp)
    ny <- length(latp)
    nt <- dim(z)[3]

    # Interpolate nn assuming a regular grid 
    zo <- array(0., c(nx, ny, nt))
    dy <- lat[2]-lat[1]
    dx <- lon[2]-lon[1]
    jj <- ((latp-(lat[1]-dy/2)) %/% dy + 1)
    ii <- ((lonp-(lon[1]-dx/2)) %/% dx + 1)

    if((ii[1]<=0)|(ii[length(ii)]>length(lon))|
       (jj[1]<=0)|(jj[length(jj)]>length(lat))) {
           stop("Downscaling area not contained in input data")
    }
    for (k in seq_len(nt)) {
        for (j in seq_along(jj)) {
            zo[,j,k] <- z[ii,jj[j],k]
        }
    }
    names(dim(zo)) <- names(dim(z))
    return(zo)
}

#' Smoothening using convolution with a circular kernel
#'
#' @description The input field is convolved with a circular kernel with equal
#' weights. Takes into account missing values.
#' @author Jost von Hardenberg, \email{j.vonhardenberg@isac.cnr.it}
#' @param z matrix with the input field to smoothen, with dimensions `c(ns, ns)`
#' @param sdim the smoothing kernel radius in pixel
#' @return The smoothened field.
#' @export
#' @examples
#' z <- rnorm(64 * 64)
#' dim(z) <- c(64, 64)
#' zs <- smooth(z, 8)
#' sd(zs)
#' # [1] 0.1334648

.smooth <- function(z, sdim) {
  imask <- !is.finite(z)
  z[imask] <- 0.
  nsx <- dim(z)[1]
  nsy <- dim(z)[2]
  mask <- matrix(0., nsx, nsy)
  for (i in 1:nsx) {
    for (j in 1:nsy) {
      kx <- i - 1
      ky <- j - 1
      if (i > nsx / 2 + 1) {
        kx <- i - nsx - 1
      }
      if (j > nsy / 2 + 1) {
        ky <- j - nsy - 1
      }
      r2 <- kx * kx + ky * ky
      #mask[i, j] <- exp( - (r2 / (sdim * sdim)) / 2)
      mask[i, j] <- (r2 <= (sdim * sdim))
    }
  }
  fm <- fft(mask)
  zf <- Re(fft(fm * fft(z), inverse = TRUE)) / sum(mask) / length(fm)
  if (sum(imask) > 0) {
           z[!imask] <- 1.0
           zf <- zf / (Re(fft(fm * fft(z), inverse = TRUE)) /
                       sum(mask) / length(fm))
  }
  zf[imask] <- NA
  return(zf)
}
