library(CSTools)

## Function to obtain month with respect to forecast time
get_month_number <- function(start_month, add_months) {
  start_index <- as.numeric(start_month)
  if (is.na(start_index) || start_index < 1 || start_index > 12) {
    stop("Invalid start month. Use a two-digit number between '01' and '12'.")
  }
  new_month_number <- (((start_index - 2 + add_months) %% 12) +1) 
  return(sprintf("%02d", new_month_number))
}



## Function to merge s2dv_cubes
merge_s2dv_cubes <- function(cubes, new_dim = "ftime") {
  # Ensure all cubes are valid s2dv_cube objects
  stopifnot(all(sapply(cubes, function(x) inherits(x, "s2dv_cube"))))
  
  # Start with the first cube in the list as the base for merging
  merged_cube <- cubes[[1]]
  
  dim_names <- c(names(dim(merged_cube$data)), 'ftime')
  time_dim <- dim(merged_cube$data)['time']
  
  # Loop through the remaining cubes and combine
  for (i in 2:length(cubes)) {
    cube <- cubes[[i]]
    
    # Check if dimensions match (except the new dimension)
    matching_dims <- all(
      merged_cube$dims[-which(names(merged_cube$dims) == new_dim)] == 
        cube$dims[-which(names(cube$dims) == new_dim)]
    )
    if (!matching_dims) stop("Cubes have mismatched dimensions (except for the merging dimension)")

    merged_cube$data <- abind::abind(merged_cube$data, cube$data,
                                     along = length(dim(cubes[[1]]$data))+1,
                                     new.names = "new_dim")
    
    # Update the dimension size for the new dimension
    merged_cube$dims[new_dim] <- merged_cube$dims[new_dim] + 1
    merged_cube$dims[new_dim] <- length(cubes)

    # Update the coordinates for the new dimension
    new_dim_value <- paste0(new_dim, merged_cube$dims[new_dim])  # Generate a default name
    merged_cube$coords[[new_dim]] <- c(merged_cube$coords[[new_dim]], new_dim_value)
  }
  
  names(dim(merged_cube$data)) <- dim_names
  merged_cube$dims['time'] <- time_dim
  
  # Return the combined s2dv_cube
  return(merged_cube)
}



##### Function to separate forecast time and days from the dimension time for daily data
process_monthly_data <- function(data, recipe) {
  sdate_month <- substr(recipe$Analysis$Time$sdate, 1, 2)
  fmonths <- recipe$Analysis$Time$ftime_min : recipe$Analysis$Time$ftime_max
  ncores <- recipe$Analysis$ncores

  hcst_aux <- sapply(length(fmonths), function(x) NULL)
  hcst_period_aux <- NULL
  hcst_dims <- dim(data$hcst$data)
  hcst_dims["time"] <- 31  
  
  obs_aux <-  sapply(length(fmonths), function(x) NULL)
  obs_period_aux <- NULL
  obs_dims <- dim(data$obs$data)
  obs_dims["time"] <- 31   
  
  if (!is.null(recipe$Analysis$Time$fcst_year)) {
    fcst_aux <-  sapply(length(fmonths), function(x) NULL)
    fcst_period_aux <- NULL
    fcst_dims <- dim(data$fcst$data)
    fcst_dims["time"] <- 31  
  }

  
  for (ft in fmonths) {
    eval_month <- get_month_number(sdate_month, ft)
    
    ########## Hindcast ##########
    ## Pad the data with NA in time dimension to 31
    hcst_period <- array(NA, dim = hcst_dims)
    # Select period corresponding to forecast month
    hcst_period_aux <- CST_SelectPeriodOnData(data$hcst, 
                                              start = list(1, as.numeric(eval_month)), 
                                              end = list(31, as.numeric(eval_month)),
                                              ncores = ncores)
    ## Obtain the number of days in the month
    month_days <- as.numeric(dim(hcst_period_aux$data)["time"][[1]])
    ## Include the data in the NA padded array to fill months with less days than 31
    hcst_period[,,,,,1:month_days,,,] <- hcst_period_aux$data
    hcst_period_aux$data <- hcst_period
    hcst_aux[[ft]] <- hcst_period_aux
    
    ######### Observations ##########
    obs_period <- array(NA, dim = obs_dims)
    obs_period_aux <- CST_SelectPeriodOnData(data$obs, 
                                             start = list(1, as.numeric(eval_month)), 
                                             end = list(31, as.numeric(eval_month)),
                                             ncores = ncores)
    obs_period[,,,,,1:month_days,,,] <- obs_period_aux$data
    obs_period_aux$data <- obs_period
    obs_aux[[ft]] <- obs_period_aux
    
    ########## Forecast #########
    if (!is.null(recipe$Analysis$Time$fcst_year)) {
      fcst_period <- array(NA, dim = fcst_dims)
      fcst_period_aux <- CST_SelectPeriodOnData(data$fcst, 
                                                start = list(1, as.numeric(eval_month)), 
                                                end = list(31, as.numeric(eval_month)),
                                                ncores = ncores)
      fcst_period[,,,,,1:month_days,,,] <- fcst_period_aux$data
      fcst_period_aux$data <- fcst_period
      fcst_aux[[ft]] <- fcst_period_aux
    }
  }
  
  ## Merge all forecast months into single s2dv_cube
  merged_hcst <- merge_s2dv_cubes(hcst_aux, new_dim = "ftime")
  merged_hcst$coords$time <- as.integer(1:31)
  merged_hcst$coords$ftime <- as.integer(fmonths)
  
  merged_obs <- merge_s2dv_cubes(obs_aux, new_dim = "ftime")
  merged_obs$coords$time <- as.integer(1:31)
  merged_obs$coords$ftime <- as.integer(fmonths)

  if (!is.null(recipe$Analysis$Time$fcst_year)) {
    merged_fcst <- merge_s2dv_cubes(fcst_aux, new_dim = "ftime")
    merged_fcst$coords$time <- as.integer(1:31)
    merged_fcst$coords$ftime <- as.integer(fmonths)
  } else
    merged_fcst <- NULL
  
  ## Adjust attributes ??
  # merged_hcst$attrs$Dates <-
  
  return(list("hcst" = merged_hcst, 
              "fcst" = merged_fcst,
              "obs" = merged_obs))
}