#----------------------------------------------------------------------------------
# Use gribr package to load GRIB files
# Assume that all the messages have the same metadata; one message is one time step
# If exp, has.memb is a number; if obs, has.memb = NULL
# syear_time_dim is the time attr dim of exp as the input for obs
#----------------------------------------------------------------------------------
GrbLoad <- function (dat, time_step = 1, has.memb = NULL, syear_time_dim = NULL,
                     regrid = NULL) {
  library(gribr)

  result <- vector('list', length = length(dat))
  times <- vector('list', length = length(dat))
  times <- lapply(times, '[<-', rep(NA, length(time_step))) #NOTE: length is 0 (slower in loop?)
#  times <- lapply(times, '[<-', .POSIXct(rep(NA, length(time_step)), tz = 'UTC'))

  for (dat_i in 1:length(dat)) {

    file_to_load <- grib_open(dat[[dat_i]])

    #----------------------------------------
    # HOW TO FIND THE VALUE OF EACH FTIME STEP?
    #----------------------------------------
    #NOTE: ValidityTime is not considered now. So if the time frequency is less than daily, it has problem.

#    # METHOD 1: Get first message to figure out the validityDate/Time of each message
#    #NOTE: gm1$validityDate should be "s", "m", "h", etc. according to document. But our files have "1".
    gm1 <- grib_get_message(file_to_load, 1)
#    first_ftime <- as.character(gm1$validityDate)
#    first_ftime_hour <- gm1$validityTime
#    # For monthly data
#    #NOTE: may not be correct because it is calculated by the first message
#    cdo_time_attr <- clock::add_months(as.POSIXct(first_ftime, 
#                                       format = "%Y%m%d", tz = 'UTC'), time_step - 1)
#    cdo_time <- format(cdo_time_attr, "%Y%m%d") 

    # METHOD 2: Use cdo showtimestamp (DEPENDENCY!)
    #TODO: Change to method 1 because can't predict what cdo will produce
    cdo_time <- system(paste0("cdo showtimestamp ", dat[[dat_i]]), intern = T)
    cdo_time <- strsplit(cdo_time, " ")[[length(cdo_time)]]
    cdo_time <- cdo_time[which(cdo_time != "")]
    # Check if there is member dim or not
    has.memb <- ifelse((length(unique(cdo_time)) == length(cdo_time)), FALSE, TRUE)
    if (has.memb) memb_dim_length <- length(cdo_time)/length(unique(cdo_time))
    cdo_time <- unique(cdo_time)[time_step]  #"2000-12-01T00:00:00"
    cdo_time_attr <- as.POSIXct(gsub('T', ' ', cdo_time), tz = 'UTC')
    cdo_time <- sapply(sapply(cdo_time, strsplit, "T"), '[[', 1)
    cdo_time <- gsub('-', '', cdo_time)
    
    #----------------------------------------

    # all members + ftimes: length should be memb*ftime (e.g., 51*7)
    ## Method 1: use grib_select and real values to filter
    memb_ftime <- grib_select(file_to_load, list(validityDate = cdo_time))
    if (inherits(memb_ftime, 'gribMessage')) memb_ftime <- list(memb_ftime)

#    ## Method 2: Calculate which messages are the desired ones
#    gm <- grib_get_message(file_to_load, time_step)
#    if (length(time_step) == 1) {
#      gm <- list(gm)
#    }

  ##################################################################
  # Get data as an array [longitude, latitude, (memb*)time]
  ##################################################################
  if (grepl("reduced", gm1$gridType)) {
    #NOTE: Need to call gribr::grib_expand_grids because I don't know how to make .Call("gribr_redtoreg") work outside that function
    #      https://github.com/nawendt/gribr/blob/main/src/redtoreg.c
    values_l <- vector('list', length = length(memb_ftime))
    for (gm_i in 1:length(memb_ftime)) {
      values_l[[gm_i]] <- grib_expand_grids(memb_ftime[[gm_i]])
    }
    result[[dat_i]] <- array(unlist(values_l), dim = c(longitude = gm1$Nj * 2, latitude = gm1$Nj, time = length(values_l)))
    # Save memory
    rm(values_l); gc()

  } else {
    result[[dat_i]] <- .grib_expand_grids(memb_ftime)
  }

  ##################################################################
  # Get metadata
  ##################################################################
  ## (1-1) Everything from the first message of first file
  if (dat_i == 1) {
    ## (1-1) Everything from the first message of first file
#    dims <- dim(result[[dat_i]])
#    attributes(result) <- gm1
#    # turn result into array again
#    dim(result[[dat_i]]) <- dims

    ## (1-2) Only save the necessary attributes
    attr(result, 'edition') <- gm1$edition
    attr(result, 'shortName') <- gm1$shortName
    #NOTE: Tune varaible name!!
      attr(result, 'shortName') <- tune_var_name(gm1$shortName)
    attr(result, 'name') <- gm1$name
    attr(result, 'units') <- gm1$units
#    attr(result, 'validityDate') <- gm1$validityDate
#    attr(result, 'validityTime') <- gm1$validityTime

    ## (2) Lat and lon
    latlon <- grib_latlons(gm1, expand = TRUE)
    attr(result, 'latitude') <- unique(as.vector(c(latlon$lats)))
    attr(result, 'longitude') <- unique(as.vector(c(latlon$lons)))
    # Save memory (though it's small)
    rm(latlon); gc()

    #NOTE: Find another way to check regular grid; Ni/Nj not always exist
#    if (has.key(gm1, "Nx") && has.key(gm1, "Ny")) {
#      nx <- gm1$Nx
#      ny <- gm1$Ny
#    } else {
#      nx <- gm1$Ni
#      ny <- gm1$Nj
#    }
#    if (length(lats) != ny | length(lons) != nx) {
#      stop("Latitude and Longitude seem to be non-regular grid.")
#    }

  }

#--------------------------------
#NOTE: Just use cdo_time
#  ## (3) Date and time: Need to get from each massage
#  for (time_i in 1:length(time_step)) {
#    gm1 <- gm[[time_i]]
#    #NOTE: What's the correct time?
##    dates <- gm1$validityDate #Date of validity of the forecast
##    times <- gm1$validityTime
##    dates <- gm1$dataDate # Reference date
#    times[[dat_i]][time_i] <- as.POSIXct(
#                       lubridate::ymd_hms(paste0(paste(gm1$year,gm1$month,gm1$day, '-'), ' ',
#                                          paste(gm1$hour, gm1$minute, gm1$second, ':')))
#                     )
#  }
   times[[dat_i]] <- cdo_time_attr
#--------------------------------

  ##################################################################
  # regrid
  ##################################################################
  if (!is.null(regrid)) {
    # result[[dat_i]]: [longitude, latitude, time]
    res_data <- s2dv::CDORemap(result[[dat_i]], lons = attr(result, 'longitude'), lats = attr(result, 'latitude'),
                               grid = regrid$type, method = regrid$method, force_remap = TRUE)
    if (dat_i == length(dat)) {
      attr(result, 'longitude') <- res_data$lons
      attr(result, 'latitude') <- res_data$lats
    }
    result[[dat_i]] <- res_data$data_array
  }


  ##################################################################
  # Save memory
  rm(memb_ftime); rm(gm1); gc()
  grib_close(file_to_load)  # Doesn't impact memory
  ##################################################################
}  #for loop dat

  # Turn result list into array
  attr <- attributes(result)
  res_dim <- c(dim(result[[1]]), syear = length(result)) #[longitude, latitude, (memb*)time, syear]
  result <- unlist(result)
  dim(result) <- res_dim

  # Generate date/time attributes
  times <- array(unlist(times), dim = c(time = length(time_step), syear = length(dat),
                                        sday = 1, sweek = 1))
  times <- s2dv::Reorder(times, c('sday', 'sweek', 'syear', 'time'))
  if (!is.null(syear_time_dim)) dim(times) <- syear_time_dim
  times <- as.POSIXct(times, origin = '1970-01-01', tz = 'UTC')

  # Reshape and reorder array
  if (is.null(has.memb) | has.memb == FALSE) {  # obs doesn't have memb; reshape syear/time dim
    result <- s2dv::Reorder(result, c("syear", "time", "latitude", "longitude"))
    result <- array(result, dim = c(dat = 1, var = 1, 
                                    syear_time_dim, dim(result)[3:4],
				    ensemble = 1))
  } else {
    nmemb <- as.numeric(dim(result)[3])/length(time_step)
    result <- array(result, dim = c(dim(result)[1:2], ensemble = nmemb,
                                    time = length(time_step), dim(result)[4]))
    result <- s2dv::Reorder(result, c("syear", "time", "latitude", "longitude", "ensemble"))
    dim(result) <- c(dat = 1, var = 1, sday = 1, sweek = 1, dim(result))
  }

  # Add attributes back
  attr$dim <- dim(result)
  attributes(result) <- attr
  attr(result, 'time') <- times

  # Save memory
  rm(times); rm(attr); gc()

  return(result)
}

#########################################################################
#########################################################################

.grib_expand_grids <- function(gribMessages, vector = FALSE) {
  # gribMessages is a list of multiple messages
  gribMessage <- gribMessages[[1]]

  if (gribr::has.key(gribMessage, "Nx") && gribr::has.key(gribMessage, "Ny")) {
    nx <- gribMessage$Nx
    ny <- gribMessage$Ny
  } else {
    nx <- gribMessage$Ni
    ny <- gribMessage$Nj
  }

  if (is.null(nx) || is.null(ny)) {
    stop("Unsupported grid type: ", gribMessage$gridType)
  }

  if (grepl("reduced", gribMessage$gridType)) {
    #TODO: This part is not used now. 
    nx <- ny * 2
    values <- .Call("gribr_redtoreg", nx, gribMessage$pl,
                    gribMessage$values)
    values <- matrix(values, nx, ny,
                     byrow = gribMessage$jPointsAreConsecutive)

#    values_l <- vector('list', length = length(gribMessages))
#    for (gm_i in 1:length(gribMessages)) {
#      values <- .Call("gribr_redtoreg", nx, gribMessages[[gm_i]]$pl, 
#                      gribMessages[[gm_i]]$values)
#      values <- matrix(values, nx, ny,
#                       byrow = gribMessage$jPointsAreConsecutive)
#      values_l[[gm_i]] <- values
#    }

  } else {
#    values <- matrix(gribMessage$values, nx, ny,
#                     byrow = gribMessage$jPointsAreConsecutive)
    values_l <- lapply(gribMessages, '[[', 'values')
    values_l <- lapply(values_l, matrix, nx, ny, byrow = gribMessage$jPointsAreConsecutive)
    values <- array(unlist(values_l), dim = c(longitude = nx, latitude = ny, time = length(values_l)))
  }

  if (vector) {
    values <- as.numeric(values)
  }

  values
}


tune_var_name <- function(x) {
  if (x == '2t') {
    res <- "tas"
  } else if (x == 'tprate') {
    res <- "prlr"
  } else if (x == "10si") {
    res <- "sfcWind"
  } else if (x == "2d") {
    res <- "tdps"
  } else if (x == "msl") {
    res <- "psl"
  } else if (x == "sst") {
    res <- "tos"
  } else if (x == "tp") {
    res <- "prlr"
  } else if (x == "mx2t24") {
    res <- "tasmax"
  } else if (x == "mn2t24") {
    res <- "tasmin"
  } else if (x == "Maximum temperature at 2 metres since previous post-processing") {
    res <- "tasmax"
  } else if (x == "Minimum temperature at 2 metres since previous post-processing") {
    res <- "tasmin"
  } else {
    res <- x
    warning("Conversion name", x, "needed?")
  }
  return(res)
}
