require(terra)
require(stringr)
require(cetcolor)

# ================================== Generate Raster ===========================

#' Prepare iSSF raster list
#' 
#' This function loads raster asc files and converts them into covariate data 
#' lists that can be used by `iSSF_simulate`.
#' 
#' @param files A vector of file paths to acs rasters to be read as covariate data.
#' @param layer_names An optional vector holding the names of the imported
#' @param directional An optional logical vector indicating which rasters in `files` contain directional covariate data.
#' 
#' When used by `iSSF_simulate`, rasters which are indicated to be `directional` will be read as the 
iSSF_raster_list <- function(files, layer_names, directional){
  
  if(!is(files, "character")){
    stop("`files` must be a character vector containing the paths to rasters that can be imported by `terra::rast`.")
  }

  if(!missing(layer_names) && !is(layer_names, "character")){
    stop("`layer_names` must be a character vector containing the names of the rasters imported.")
  }
  
  # If not specified, assume all layers are scalar
  if(missing(directional)){
    directional = rep.int(F, length(files))
  }
  
  if(length(directional) != length(files)){
    stop("`directional` must be a logical vector that has the same length as `files`.")
  }
  if(!missing(layer_names) && length(layer_names) != length(files)){
    stop("`layer_names` must be a character vector that has the same length as `files`.")
  }
  
  ret <- list()
  
  # loop and load raster layers
  for(i in 1:length(files)){
    file = files[i]
    
    # Load raster
    ras <- terra::rast(file)
    
    # Convert raster into a list object
    tmp <- list(
      extent = terra::ext(ras)[],
      values = matrix(values(ras), ncol = ncol(ras), nrow = nrow(ras), byrow = T),
      dims = c(nrow(ras),  ncol(ras)), 
      directional = directional[i])
    
    if(missing(layer_names)){
      tmp$name = names(ras)
    }else{
      tmp$name = layer_names[i]
    }
    
    class(tmp) <- c("iSSF_raster", "list")
        
    ret[[i]] <- tmp
  }  

  class(ret) <- c("iSSF_raster_list", "list")

  # check for duplicate names
  tmp = sapply(ret, function(x) x$name)
  if(length(tmp) != length(unique(tmp))){
    stop("Raster layers must have unique names.  Current names are set as:  ", paste(tmp, collapse = ", "))
  }
      
  return(ret)
}

print.iSSF_raster <- function(x, ...){
  
  cat(
    cat("Name:", x$name),
    sprintf(fmt = "Dims:  %i rows, %i cols", x$dims[1], x$dims[2]),
    sprintf(
      fmt = "Extent: %f, %f, %f, %f  (xmin, xmax, ymin, ymax)",
      x$extent[1], x$extent[2], x$extent[3], x$extent[4]),
    sprintf(fmt = "Directional: %s", x$directional),
    sprintf(fmt = "Value range: %f to %f", min(x$values, na.rm = T), max(x$values, na.rm = T)),
    sprintf(fmt = "Value mean %f, sd %f", mean(x$values, na.rm = T), sd(x$values, na.rm = T)), 
    sep = "\n")
}

plot.iSSF_raster <- function(x, ...){
  
  # Convert back into terra raster for plotting
  ras = terra::rast(x$values)
  ext(ras) <- x$extent
  
  if(x$directional){
    plot(ras, main = x$name, reset = T, 
         col = cetcolor::cet_pal(50, 'c2'), ...)
   }else{
    plot(ras, main = x$name, reset = T, ...)
  }
  
}

print.iSSF_raster_list <- function(x, ...){
  
  for(i in 1:length(x)){
    cat(" --- Layer", i, "\n")
    print(x[[i]])
    
  }
  
}

# =================================== Model formula ============================


# internal function for parsing iSSF formula
# Extract transform and raster name from term
iSSF_term_parse <- function(trm){
  # Possible transformations  
  lvls <- c("none", "log", "cos")
  out = list()
  
  # Check if previous step is referenced
  if(stringr::str_detect(pattern = "previous\\(.+\\)", string = trm)){
    mtch <- stringr::str_match(pattern = "previous\\((.+)\\)", string = trm)
    trm = mtch[1,2]
    out$previous = T
  }else{
    out$previous = F
  }
  
  # Extract type of transformation and raster name
  if(stringr::str_detect(pattern = "log\\(.+\\)", string = trm)){
    mtch <- stringr::str_match(pattern = "log\\((.+)\\)", string = trm)
    out$var = mtch[1,2]
    out$transform = factor("log", levels = lvls)
  }else if(stringr::str_detect(pattern = "cos\\(.+\\)", string = trm)){
    mtch <- stringr::str_match(pattern = "cos\\((.+)\\)", string = trm)
    out$var = mtch[1,2]
    out$transform = factor("cos", levels = lvls)
  }else{
    out$var = trm
    out$transform = factor("none", levels = lvls)
  }
  
  class(out) <- c("iSSF_var", "list") 
  return(out)
}

str.iSSF_var <- function(object, ...){
  
  cat("iSSF_var\n")
  invisible(NextMethod("str", ...))
}


# internal function for parsing iSSF formula
# Return character for printing in formula
iSSF_term_format <- function(trm_parsed){
  # Possible transformations  
  lvls <- c("none", "log", "cos")
  # Extract type of transformation and raster name
  if(trm_parsed$transform == "log"){
    out = sprintf("log(%s)", trm_parsed$var)
  }else if(trm_parsed$transform == "cos"){
    out = sprintf("cos(%s)", trm_parsed$var)
  }else{
    out = (trm_parsed$var)
  }
  
  if(trm_parsed$previous){
    out = sprintf("previous(%s)", trm_parsed$var)
  }
  
  return(out)
}

#' Create an iSSF formula
#'
#' An internal function for creating a model formula for `iSSF_simulation`. The
#' resulting list of integer vectors indicates which covaraite layers should be
#' multiplied (interactions) and added (linear effects).
#'
#' @param beta A named vector giving the beta values for each model term.
#' @param raster_names An ordered character vector containing the names of the
#'   raster layers contained in the `iSSF_raster_list`.
#'
#'   To be read by the c++ code, the formula will be made of a list of integer
#'   vectors. Each list element is a linear term (to be added added with each
#'   other), while multiple terms within the same element will be interactions
#'   (to be multiplied with each other). The terms will be stored as raster
#'   layer indices (converted from layer names into layer indices)
iSSF_formula <- function(beta){
  
  # Get beta names
  # All these names must have a matching raster layer name in iSSF_raster_list 
  trms <- names(beta)

  out <- list()
  frm_string <- list()

  # --- Loop model terms and build formula object
  for(trm in trms){

    #  split interacting terms
    tmp <- strsplit(trm, split = ":")[[1]]
    tmp_parsed <- lapply(tmp, FUN = iSSF_term_parse)
    class(tmp_parsed) <- c("iSSF_term", "list")
    out <- append(out, list(tmp_parsed))
  }
  
  # --- Make sure terms are valid
  chk <- any(sapply(out, function(term){
    any(sapply(term, function(var){
      (var$var %in% c("ta", "sl")) & var$prev
    }))
  }))
  if(chk) stop("Cannot use previous step length or turnign angles in simulations.")
  
  ret <- list(
    coef = beta,
    terms = out)
  class(ret) <- "iSSF_formula"
  return(ret)
}

print.iSSF_term <- function(x){
  cat("iSSF_term", paste(lapply(x, FUN = iSSF_term_format), collapse = ":"), "\n")
}

str.iSSF_term <- function(object, ...){
  print(object)
  invisible(NextMethod("str", ...))
}


# Return a list of indicies listing term names to raster covariate layers
iSSF_ras_idx <- function(issf_frm, raster_names){

  # trm = issf_frm$terms[[1]]

  idx <- list()
  trns <- list()
  prev <- list()

  # Loop linear terms
  for(trm_i in 1:length(issf_frm$terms)){

    # Extract indexes
    tmp <- sapply(issf_frm$terms[[trm_i]], function(x){
        # turning angle and step length are indexed by negative values
        if(x$var == "ta"){
          return(-2)
        }else if(x$var == "sl"){
          return(-1)
        }else if(x$var %in% raster_names){
          # Raster variables are indexed by their list index, minus 1 (cpp indexing)
          return(which(raster_names == x$var)[1] - 1)
        }else{
          stop("Variable `", x$var, "` is not a raster layer name or either of the reserved values: `ta` or `sl`.")
        }
      })
    idx[[trm_i]] <- tmp

    # Extract transforms
    tmp <- sapply(issf_frm$terms[[trm_i]], `[[`, "transform")
    trns[[trm_i]] <- tmp

    # Extract transforms
    tmp <- sapply(issf_frm$terms[[trm_i]], `[[`, "previous")
    prev[[trm_i]] <- as.integer(tmp)
  }
  
  out = list(idx = idx, transform = trns, previous = prev)
  
  return(out)
}


print.iSSF_formula <- function(x, ...){

  cat("iSSF formula terms and coefficients.\n")
  print(x$coef, row.names = F)
  
}


# ============================== Plot movement parameters ======================


plot_sl_dist <- function(p = c(0, 0.99), scale, shape){
  
  q <- qgamma(p, shape = shape, scale = scale)
  
  x = seq(q[1], q[2], length.out = 1000)
  
  y = dgamma(x, shape = shape, scale = scale)
  
  plot(x,y, ylab = "Density", xlab = "Step Length", type = 'l')
  mtext(sprintf("scale = %.3f; shape = %.3f", scale, shape), adj = 0)
  
}

plot_ta_dist <- function(kappa){
  
  x_vm <- seq(-pi, pi, length.out = 1000)
  plot(x_vm, circular::dvonmises(
    x = circular::circular(x_vm),
    mu =  circular::circular(0),
    kappa = kappa), t = 'l', 
       xlab = "Turning Angle (rads)", ylab = "Density")
  mtext(sprintf("kappa = %.3f", kappa), adj = 0)
  
}


# =================================== Simulation ===============================

iSSF_simulation <- function(
    issf_frm, ras_lst, position_start, bearing_start,
    n_steps, len_track, shape, scale, kappa){
  
  # Extract raster names
  raster_names <- sapply(ras_lst, `[[`, "name")
  
  # ------ Build model formula
  # Get raster layer indicies for formula
  idxs <- iSSF_ras_idx(issf_frm, raster_names)
  
  agent <- cpp_iSSF_simulation(
    position_start = position_start, 
    bearing_start = bearing_start, 
    n_steps = n_steps, 
    len_track = len_track, iSSF_raster_list = ras_lst, 
    coef = issf_frm$coef, 
    ras_idx =  idxs$idx, ras_trns = idxs$transform, ras_prev = idxs$previous,
    scale = scale, shape, kappa)
  
}