transform_units_precipitation <- function(data, original_units, new_units,
                                          var_name, freq, flux = FALSE, ncores = NULL,
                                          var_index = 1) {

## TODO consider higher frequencies (e.g. 6hourly) 
## could create a constant hours <- 24 or hours <- 6 and use the same code

  data_arr <- data[[1]]$data
  data_list <- asplit(data_arr, which(names(dim(data_arr)) == "var"))


  if (original_units == "ms-1") {
    if (new_units == "mm") {
      data_list[[var_index]] <- data_list[[var_index]] * 3600 * 24 * 1000
      data[[1]]$attrs$Variable$metadata[[var_name]]$units <- "mm"
    } else if (new_units == "m") {
      data_list[[var_index]] <- data_list[[var_index]] * 3600 * 24
      data[[1]]$attrs$Variable$metadata[[var_name]]$units <- 'm'
    } else if (new_units == "kgm-2s-1") {
      data_list[[var_index]] <- data_list[[var_index]] * 1000
      data[[1]]$attrs$Variable$metadata[[var_name]]$units <- 'kg m-2 s-1'
    } else {
      stop(paste("Unknown transformation from", original_units, "to", new_units))
    }
  } else if (original_units == "mm") {
    if (new_units == "ms-1") {
      data_list[[var_index]] <- data_list[[var_index]] / (3600 * 24 * 1000)
      data[[1]]$attrs$Variable$metadata[[var_name]]$units <- 'm s-1'
    } else if (new_units == "m") {
      data_list[[var_index]] <- data_list[[var_index]] / 1000
      data[[1]]$attrs$Variable$metadata[[var_name]]$units <- 'm'
    } else if (new_units == "kgm-2s-1") {
      data_list[[var_index]] <- data_list[[var_index]] / (3600 * 24)
      data[[1]]$attrs$Variable$metadata[[var_name]]$units <- 'kg m-2 s-1'
    } else {
      stop(paste("Unknown transformation from", original_units, "to", new_units))
    }
  } else if (original_units == "m") {
    if (new_units == "ms-1") {
      data_list[[var_index]] <- data_list[[var_index]] / (3600 * 24)
      data[[1]]$attrs$Variable$metadata[[var_name]]$units <- 'm s-1'
    } else if (new_units == "mm") {
      data_list[[var_index]] <- data_list[[var_index]] * 1000
      data[[1]]$attrs$Variable$metadata[[var_name]]$units <- 'mm'
    } else if (new_units == "kgm-2s-1") {
      data_list[[var_index]] <- data_list[[var_index]] * 1000 / (3600 * 24)
      data[[1]]$attrs$Variable$metadata[[var_name]]$units <- 'kg m-2 s-1'
    } else {
      stop(paste("Unknown transformation from", original_units, "to", new_units))
    }
  } else if (original_units == "kgm-2s-1") {
    if (new_units == "ms-1") {
      data_list[[var_index]] <- data_list[[var_index]] / 1000
      data[[1]]$attrs$Variable$metadata[[var_name]]$units <- 'm s-1'
    } else if (new_units == "mm") {
      data_list[[var_index]] <- data_list[[var_index]] * 3600 * 24
      data[[1]]$attrs$Variable$metadata[[var_name]]$units <- 'mm'
    } else if (new_units == "m") {
      data_list[[var_index]] <- data_list[[var_index]] * 3600 * 24 / 1000
      data[[1]]$attrs$Variable$metadata[[var_name]]$units <- 'm'
    } else {
      stop(paste("Unknown transformation from", original_units, "to", new_units))
    }
  } else {
    stop("Unknown precipitation units transformation")
  }

  if (flux) {
    data[[1]]$attrs$Variable$metadata[[var_name]]$units <- paste0(
      data[[1]]$attrs$Variable$metadata[[var_name]]$units, "/day")
  } else {
    ## TODO: Shouldn't use time dimension, need it for Compute()
    if (freq == "monthly_mean") { # could it not be mean?
      time_pos <- which(lapply(data[[1]]$attrs$Variable$metadata[[var_name]]$dim,
                               function(x) {x$name}) == 'time')
      cal <- tolower(data[[1]]$attrs$Variable$metadata[[var_name]]$dim[[time_pos]]$calendar)
      data_list[[var_index]] <-
        Apply(list(data_list[[var_index]], data[[1]]$attrs$Dates),
              target_dim = list(c('syear'), c('syear')),
              extra_info = list(cal = cal, days_in_month = .days_in_month),
              fun = function(x, y) {
                date <- as.Date(y, "%Y-%m-%d")
                num_days <- .days_in_month(date, cal = .cal)
                res <- x * num_days
              }, ncores = ncores)$output1
    }
  }

  # Combine list back to array
  data_arr <- array(unlist(data_list),
                    dim = c(dim(data_list[[1]]), var = length(data_list)))
  data[[1]]$data <- aperm(data_arr, match(names(dim(data[[1]]$data)),
                          names(dim(data_arr))))

  return(data)
}


.days_in_month <- function(dates, cal) {
  n_days <- array()
  for (x in 1:length(dates)) {
    if (cal %in% c('gregorian', 'standard', 'proleptic_gregorian')) {
      N_DAYS_IN_MONTHS <- lubridate:::N_DAYS_IN_MONTHS
      if (leap_year(year(dates[x]))) {
        N_DAYS_IN_MONTHS[2] <- N_DAYS_IN_MONTHS[2] + 1
      }
    } else if (cal %in% c('360', '360_day')) {
      N_DAYS_IN_MONTHS <- rep(30, 12)
      names(N_DAYS_IN_MONTHS) <- month.abb
    } else if (cal %in% c('365', '365_day')) {
      N_DAYS_IN_MONTHS <- lubridate:::N_DAYS_IN_MONTHS
    } else {
      stop("Unknown calendar")
    } 
    month_x <- month(dates[x], label = TRUE, locale = "C")
    n_days[x] <- N_DAYS_IN_MONTHS[month_x]
    # n_days[month_x == "Feb" & leap_year(x)] <- 29L
  }
  return(n_days)
}
