# This function is for decadal-daily data loading.
# The time input (ftimemin & ftimemax) units in the recipe is month. For example,
# ftimemin = 1 means the first month, and ftimemax = 3 means the 3rd month.
# For daily data, if initial month is November, ftimemin = 1 equals to c(1:30), 
# and ftimemax = 3, which means January next year, equals to c(31+31:31+31+30).
#
#NOTE: For the datasets that having one file per sdate, we can use array [syear, time] to define time since
# the Start() call doesn't need time_across = 'chunk' (there is only one chunk). BUT to 
# keep the script simple, we use indices always and ignore leap year difference.
# Only distinguish 360-day and non 360-day calendar.

get_daily_time_ind <- function(ftimemin, ftimemax, initial_month, sdates, calendar) {
#NOTE: "sdates" is not needed if leap year is not considered

  if (!calendar %in% c('360_day', '365_day', 'noleap', 'standard', 'proleptic_gregorian',
                       'gregorian')) 
    stop("The calendar is not recognized. Please contact maintainers.")

  total_months <- (initial_month + ftimemin - 1):(initial_month + ftimemax - 1)
  total_months <- ifelse(total_months > 12, total_months %% 12, total_months)
  total_months <- ifelse(total_months == 0, 12, total_months)

  if (calendar == '360-day') {
    daily_time_ind <- ((ftimemin - 1) * 30 + 1): (((ftimemin - 1) * 30) + (ftimemax - ftimemin + 1) * 30)

  } else {

    total_days_sum <- sum(days_in_month(total_months))
    if (ftimemin == 1) {
      daily_time_ind <- 1:total_days_sum
    } else {
      before_start_months <- initial_month:(initial_month + ftimemin - 2)
      before_start_months <- ifelse(before_start_months > 12, before_start_months %% 12, before_start_months)
      before_start_months <- ifelse(before_start_months == 0, 12, before_start_months)
      before_start_total_days_sum <- sum(days_in_month(before_start_months))
      daily_time_ind <- (before_start_total_days_sum + 1):(before_start_total_days_sum + total_days_sum)
    }

#    if (calendar %in% c('365_day', 'noleap')) {
#      # daily_time_ind
#    } else if (calendar %in% c('standard', 'proleptic_gregorian', 'gregorian')) {
#
##      leap <- leap_year(sdates)
#      sdates_full <- ymd(paste0(sdates, sprintf('%02d', initial_month), '01'))
#      idx_min <- sdates_full + months(ftimemin)
#      idx_max <- sdates_full + months(ftimemax + 1) - days(1)
#
#      day0229 <- rep(NA, length(sdates))
#      for (i_sdate in 1:length(sdates_full)) {
#        days <- seq(idx_min[i_sdate],idx_max[i_sdate], by = 'days')
#        day0229[i_sdate] <- sum(format(days, "%m%d") == "0229")
#      }
#
#      if (all(day0229 == 0)) {  # no 0229
#        # daily_time_ind
#      } else {
#        # get extra indices and remove 0229 & extra indices after retrieval
#        #TODO: Make Start() can handle time as a list or array for this case
#        daily_time_ind <- c(daily_time_ind, 
#                            (tail(daily_time_ind, 1) + 1):((tail(daily_time_ind, 1) + 1) + (max(day0229) - 1)))
#      }
#  
#    }
  }

  return(daily_time_ind)
}

#==========================================

# This function is used to pick out 0229 after loading data. But it is not used in the 
# script now.
correct_daily_for_leap <- function(data = NULL, time_attr, return_time = TRUE) {
  # data & time_attr: [syear, time]
  #NOTE: return_time = TRUE, only return time_attr; FALSE, only return data

  new_tmp_time_attr <- array(NA, dim = dim(time_attr))
  if (!is.null(data)) {
    new_data <- array(NA, dim = dim(time_attr))
  } else {
    new_data <- NULL
  }
  for (i_sdate in 1:dim(time_attr)[1]) {
    ind_not_0229 <- !(format(time_attr[i_sdate, ], "%m%d") == "0229")
    new_tmp_time_attr[i_sdate, 1:sum(ind_not_0229)] <- time_attr[i_sdate, ][ind_not_0229]
    if (!is.null(data)) {
      new_data[i_sdate, 1:sum(ind_not_0229)] <- data[i_sdate, ][ind_not_0229]
    }
  }
  real_time_length <- min(which(is.na(new_tmp_time_attr),arr.ind = T)[, 2]) - 1
  new_tmp_time_attr <- new_tmp_time_attr[, 1:real_time_length]
  new_tmp_time_attr <- as.POSIXct(new_tmp_time_attr, tz = 'UTC', origin = '1970-01-01')
  if (!is.null(data)) {
    new_data <- new_data[, 1:real_time_length]
  }

  if (any(!new_tmp_time_attr %in% time_attr | any(is.na(new_tmp_time_attr)))) {
    stop("Wrong time attribute modification.")
  }
  if (return_time) {
    return(list(time_attr = new_tmp_time_attr))
  } else {
    return(list(data = new_data))
  }
}

#==========================================
# This function generates the path for Start() call. It shouldn't be needed when Start() is improved.
# table, grid, version: A list with variables as name. E.g., list(tas = 'Amon')
get_dcpp_path <- function(archive, exp.name, table, grid, version, sdates) {

  if (length(table) == 1) {
    # only one variable, or the variables share a common path pattern
    fcst.path <- file.path(archive$src, archive$System[[exp.name]]$src$startR,
                           '$ensemble$', table, '$var$', grid, 'v*')
    fcst.files <- paste0('$var$_', table, '_*_$dcpp$_s$syear$-$ensemble$_', grid, '_$chunk$.nc')
  } else {
    # path pattern depends on the variable
    fcst.path <- file.path(archive$src, archive$System[[exp.name]]$src$startR,
                           '$ensemble$', '$table$', '$var$', '$grid$', '$version$')
    fcst.files <- paste0('$var$_', '$table$', '_*_$dcpp$_s$syear$-$ensemble$_', '$grid$', '_$chunk$.nc')
  }
  path_list <- file.path(fcst.path, fcst.files)
  dcppB_string <- "dcppB-forecast"
  ## NOTE: CanESM5 has a 'special' path pattern... Adding the global expression '*'
  ## only works if Start() is loading both dcppA and dcppB. For dcppB only,
  ## it crashes.
  if (exp.name == "CanESM5") {
    if (all(sdates >= archive$System[[exp.name]]$src$first_dcppB_syear)) {
      # Replace first instance of $dcpp$ with the actual subdirectory name
      path_list <- str_replace(path_list, fixed("$dcpp$"), "dcppB-forecast_i1p2")
    } else {
      # Add global expression to the string so that it will find 'dcppB-forecast_i1p2'
      dcppB_string <- paste0(dcppB_string, "*")
    }
  }
  dcpp_list <- vector('list', length = length(sdates))
  names(dcpp_list) <- as.character(sdates)
  for (i_sdate in 1:length(sdates)) {
    if ((!is.null(archive$System[[exp.name]]$src$first_dcppB_syear)) && 
        (sdates[i_sdate] >= archive$System[[exp.name]]$src$first_dcppB_syear)) {
      dcpp_list[[i_sdate]] <- dcppB_string
    } else {
      dcpp_list[[i_sdate]] <- "dcppA-hindcast"
    }
  }
  return(list(path_list = path_list, dcpp_list = dcpp_list))
}

