# MOSAIC_CREATION_UTILS.R
# ======================
# Utility functions for creating weekly mosaics from daily satellite imagery.
# These functions support cloud cover assessment, date handling, and mosaic creation.

#' Safe logging function
#' @param message The message to log
#' @param level The log level (default: "INFO")
#' @return NULL (used for side effects)
#'
safe_log <- function(message, level = "INFO") {
  if (exists("log_message")) {
    log_message(message, level)
  } else {
    if (level %in% c("ERROR", "WARNING")) {
      warning(message)
    } else {
      message(message)
    }
  }
}

#' Generate a sequence of dates for processing
#'
#' @param end_date The end date for the sequence (Date object)
#' @param offset Number of days to look back from end_date
#' @return A list containing week number, year, and a sequence of dates for filtering
#' 
date_list <- function(end_date, offset) {
  # Input validation
  if (!lubridate::is.Date(end_date)) {
    end_date <- as.Date(end_date)
    if (is.na(end_date)) {
      stop("Invalid end_date provided. Expected a Date object or a string convertible to Date.")
    }
  }
  
  offset <- as.numeric(offset)
  if (is.na(offset) || offset < 1) {
    stop("Invalid offset provided. Expected a positive number.")
  }
  
  # Calculate date range
  offset <- offset - 1  # Adjust offset to include end_date
  start_date <- end_date - lubridate::days(offset)
  
  # Extract week and year information
  week <- lubridate::isoweek(end_date)
  year <- lubridate::isoyear(end_date)
  
  # Generate sequence of dates
  days_filter <- seq(from = start_date, to = end_date, by = "day")
  days_filter <- format(days_filter, "%Y-%m-%d")  # Format for consistent filtering
  
  # Log the date range
  safe_log(paste("Date range generated from", start_date, "to", end_date))
  
  return(list(
    "week" = week,
    "year" = year,
    "days_filter" = days_filter,
    "start_date" = start_date,
    "end_date" = end_date
  ))
}

#' Create a weekly mosaic from available VRT files
#'
#' @param dates List from date_list() with date range info
#' @param field_boundaries Field boundaries for image cropping
#' @param daily_vrt_dir Directory containing VRT files
#' @param merged_final_dir Directory with merged final rasters
#' @param output_dir Output directory for weekly mosaics
#' @param file_name_tif Output filename for the mosaic
#' @param create_plots Whether to create visualization plots (default: TRUE)
#' @return The file path of the saved mosaic
#'
create_weekly_mosaic <- function(dates, field_boundaries, daily_vrt_dir, 
                                merged_final_dir, output_dir, file_name_tif,
                                create_plots = TRUE) {
  # Find VRT files for the specified date range
  vrt_list <- find_vrt_files(daily_vrt_dir, dates)
  
  # Find final raster files for fallback
  raster_files_final <- list.files(merged_final_dir, full.names = TRUE, pattern = "\\.tif$")
  
  # Process the mosaic if VRT files are available
  if (length(vrt_list) > 0) {
    safe_log("VRT list created, assessing cloud cover for mosaic creation")
    
    # Calculate cloud cover statistics
    missing_pixels_count <- count_cloud_coverage(vrt_list, field_boundaries)
    
    # Create mosaic based on cloud cover assessment
    mosaic <- create_mosaic(vrt_list, missing_pixels_count, field_boundaries, raster_files_final)
    
  } else {
    safe_log("No VRT files available for the date range, creating empty mosaic with NA values", "WARNING")
    
    # Create empty mosaic if no files are available
    if (length(raster_files_final) == 0) {
      stop("No VRT files or final raster files available to create mosaic")
    }
    
    mosaic <- terra::rast(raster_files_final[1]) %>% 
      terra::setValues(NA) %>%
      terra::crop(field_boundaries, mask = TRUE)
    
    names(mosaic) <- c("Red", "Green", "Blue", "NIR", "CI")
  }
  
  # Save the mosaic (without mask files to avoid breaking other scripts)
  file_path <- save_mosaic(mosaic, output_dir, file_name_tif, create_plots, save_mask = FALSE)
  
  safe_log(paste("Weekly mosaic processing completed for week", dates$week))
  
  return(file_path)
}

#' Find VRT files within a date range
#'
#' @param vrt_directory Directory containing VRT files
#' @param dates List from date_list() function containing days_filter
#' @return Character vector of VRT file paths
#' 
find_vrt_files <- function(vrt_directory, dates) {
  # Get all VRT files in directory
  vrt_files <- list.files(here::here(vrt_directory), full.names = TRUE)
  
  if (length(vrt_files) == 0) {
    warning("No VRT files found in directory: ", vrt_directory)
    return(character(0))
  }
  
  # Filter files by dates
  vrt_list <- purrr::map(dates$days_filter, ~ vrt_files[grepl(pattern = .x, x = vrt_files)]) %>%
    purrr::compact() %>%
    purrr::flatten_chr()
  
  # Log results
  safe_log(paste("Found", length(vrt_list), "VRT files for the date range"))
  
  return(vrt_list)
}

#' Count missing pixels (clouds) in rasters
#'
#' @param vrt_list List of VRT files to analyze
#' @param field_boundaries Field boundaries vector for masking
#' @return Data frame with cloud coverage statistics
#'
count_cloud_coverage <- function(vrt_list, field_boundaries) {
  if (length(vrt_list) == 0) {
    warning("No VRT files provided for cloud coverage calculation")
    return(NULL)
  }
  
  tryCatch({
    # Calculate total pixel area using the first VRT file
    total_pix_area <- terra::rast(vrt_list[1]) |>
      terra::subset(1) |>
      terra::setValues(1) |>
      terra::crop(field_boundaries, mask = TRUE) |>
      terra::global(fun = "notNA")
    
    # Process each raster to detect clouds and shadows
    processed_rasters <- list()
    cloud_masks <- list()
      
    # Create data frame for missing pixels count
    missing_pixels_df <- data.frame(
      filename = vrt_list,
      notNA = numeric(length(vrt_list)),
      total_pixels = numeric(length(vrt_list)),
      missing_pixels_percentage = numeric(length(vrt_list)),
      thres_5perc = numeric(length(vrt_list)),
      thres_40perc = numeric(length(vrt_list))
    )
    
    # Process each VRT file to calculate cloud coverage
    for (i in seq_along(vrt_list)) {
      tryCatch({
        # Load the raster
        current_raster <- terra::rast(vrt_list[i]) |>
          terra::crop(field_boundaries, mask = TRUE)
        
        # Calculate valid pixels (assuming CI band is the 5th layer)
        if (terra::nlyr(current_raster) >= 5) {
          ci_layer <- current_raster[[5]]  # CI layer
        } else {
          ci_layer <- current_raster[[1]]  # Fallback to first layer
        }
        
        notna_count <- terra::global(ci_layer, fun = "notNA")$notNA
        missing_pixels_df$notNA[i] <- notna_count
        missing_pixels_df$total_pixels[i] <- total_pix_area$notNA
        missing_pixels_df$missing_pixels_percentage[i] <- round(100 - ((notna_count / total_pix_area$notNA) * 100))
        missing_pixels_df$thres_5perc[i] <- as.integer(missing_pixels_df$missing_pixels_percentage[i] < 5)
        missing_pixels_df$thres_40perc[i] <- as.integer(missing_pixels_df$missing_pixels_percentage[i] < 45)
        
        # Store processed raster for potential use in mosaic creation
        processed_rasters[[i]] <- current_raster
        
        # Create a simple cloud mask (1 = valid data, 0 = missing/cloud)
        cloud_mask <- terra::setValues(ci_layer, 1)
        cloud_mask[is.na(ci_layer)] <- 0
        cloud_masks[[i]] <- cloud_mask
        
      }, error = function(e) {
        safe_log(paste("Error processing", basename(vrt_list[i]), ":", e$message), "WARNING")
        # Set default values for failed processing
        missing_pixels_df$notNA[i] <- 0
        missing_pixels_df$total_pixels[i] <- total_pix_area$notNA
        missing_pixels_df$missing_pixels_percentage[i] <- 100
        missing_pixels_df$thres_5perc[i] <- 0
        missing_pixels_df$thres_40perc[i] <- 0
      })
    }
    
    # Store processed rasters and cloud masks as attributes
    attr(missing_pixels_df, "cloud_masks") <- cloud_masks
    attr(missing_pixels_df, "processed_rasters") <- processed_rasters
    
    # Log results
    safe_log(paste(
      "Cloud cover assessment completed for", length(vrt_list), "files.",
      sum(missing_pixels_df$thres_5perc), "files with <5% cloud cover,",
      sum(missing_pixels_df$thres_40perc), "files with <45% cloud cover"
    ))
      return(missing_pixels_df)
  }, error = function(e) {
    warning("Error in cloud coverage calculation: ", e$message)
    return(NULL)
  })
}

#' Create a mosaic from VRT files based on cloud coverage
#'
#' @param vrt_list List of VRT files to create mosaic from
#' @param missing_pixels_count Cloud coverage statistics from count_cloud_coverage()
#' @param field_boundaries Field boundaries vector for masking (optional)
#' @param raster_files_final List of processed raster files to use as fallback
#' @return A SpatRaster object with the mosaic
#'
create_mosaic <- function(vrt_list, missing_pixels_count, field_boundaries = NULL, raster_files_final = NULL) {
  # If no VRT files, create an empty mosaic
  if (length(vrt_list) == 0) {
    if (length(raster_files_final) == 0 || is.null(field_boundaries)) {
      stop("No VRT files available and no fallback raster files or field boundaries provided")
    }
    
    safe_log("No images available for this period, creating empty mosaic with NA values", "WARNING")
    
    x <- terra::rast(raster_files_final[1]) |>
      terra::setValues(NA) |>
      terra::crop(field_boundaries, mask = TRUE)
    
    names(x) <- c("Red", "Green", "Blue", "NIR", "CI")
    return(x)
  }
  
  # If missing pixel count was not calculated, use all files
  if (is.null(missing_pixels_count)) {
    safe_log("No cloud coverage data available, using all images", "WARNING")
    
    if (length(vrt_list) == 1) {
      # Handle single file case
      x <- terra::rast(vrt_list[1])
    } else {
      # Handle multiple files case
      rsrc <- terra::sprc(vrt_list)
      x <- terra::mosaic(rsrc, fun = "max")
    }
    names(x) <- c("Red", "Green", "Blue", "NIR", "CI")
    return(x)
  }
  
  # Check if we have processed rasters from cloud detection
  processed_rasters <- attr(missing_pixels_count, "processed_rasters")
  cloud_masks <- attr(missing_pixels_count, "cloud_masks")
  
  if (!is.null(processed_rasters) && length(processed_rasters) > 0) {
    safe_log("Using cloud-masked rasters for mosaic creation")
    
    # Determine best rasters to use based on cloud coverage
    index_5perc <- which(missing_pixels_count$thres_5perc == max(missing_pixels_count$thres_5perc))
    index_40perc <- which(missing_pixels_count$thres_40perc == max(missing_pixels_count$thres_40perc))
    
    # Create mosaic based on available cloud-free images
    if (sum(missing_pixels_count$thres_5perc) > 1) {
      safe_log("Creating max composite from multiple cloud-free images (<5% clouds)")
      
      # Use the cloud-masked rasters instead of original files
      cloudy_rasters_list <- processed_rasters[index_5perc]
      if (length(cloudy_rasters_list) == 1) {
        x <- cloudy_rasters_list[[1]]
      } else {
        rsrc <- terra::sprc(cloudy_rasters_list)
        x <- terra::mosaic(rsrc, fun = "max")
      }
      
      # Also create a composite mask showing where data is valid
      mask_list <- cloud_masks[index_5perc]
      if (length(mask_list) == 1) {
        mask_composite <- mask_list[[1]]
      } else {
        mask_rsrc <- terra::sprc(mask_list)
        mask_composite <- terra::mosaic(mask_rsrc, fun = "max")
      }
      attr(x, "cloud_mask") <- mask_composite
      
    } else if (sum(missing_pixels_count$thres_5perc) == 1) {
      safe_log("Using single cloud-free image (<5% clouds)")
      
      # Use the cloud-masked raster
      x <- processed_rasters[[index_5perc[1]]]
      attr(x, "cloud_mask") <- cloud_masks[[index_5perc[1]]]
      
    } else if (sum(missing_pixels_count$thres_40perc) > 1) {
      safe_log("Creating max composite from partially cloudy images (<40% clouds)", "WARNING")
      
      # Use the cloud-masked rasters
      cloudy_rasters_list <- processed_rasters[index_40perc]
      if (length(cloudy_rasters_list) == 1) {
        x <- cloudy_rasters_list[[1]]
      } else {
        rsrc <- terra::sprc(cloudy_rasters_list)
        x <- terra::mosaic(rsrc, fun = "max")
      }
      
      # Also create a composite mask
      mask_list <- cloud_masks[index_40perc]
      if (length(mask_list) == 1) {
        mask_composite <- mask_list[[1]]
      } else {
        mask_rsrc <- terra::sprc(mask_list)
        mask_composite <- terra::mosaic(mask_rsrc, fun = "max")
      }
      attr(x, "cloud_mask") <- mask_composite
      
    } else if (sum(missing_pixels_count$thres_40perc) == 1) {
      safe_log("Using single partially cloudy image (<40% clouds)", "WARNING")
      
      # Use the cloud-masked raster
      x <- processed_rasters[[index_40perc[1]]]
      attr(x, "cloud_mask") <- cloud_masks[[index_40perc[1]]]
      
    } else {
      safe_log("No cloud-free images available, using all cloud-masked images", "WARNING")
      
      # Use all cloud-masked rasters
      if (length(processed_rasters) == 1) {
        x <- processed_rasters[[1]]
      } else {
        rsrc <- terra::sprc(processed_rasters)
        x <- terra::mosaic(rsrc, fun = "max")
      }
      
      # Also create a composite mask
      if (length(cloud_masks) == 1) {
        mask_composite <- cloud_masks[[1]]
      } else {
        mask_rsrc <- terra::sprc(cloud_masks)
        mask_composite <- terra::mosaic(mask_rsrc, fun = "max")
      }
      attr(x, "cloud_mask") <- mask_composite
    }
  } else {
    # Fall back to original behavior if no cloud-masked rasters available
    safe_log("No cloud-masked rasters available, using original images", "WARNING")
    
    # Determine best rasters to use based on cloud coverage  
    index_5perc <- which(missing_pixels_count$thres_5perc == max(missing_pixels_count$thres_5perc))
    index_40perc <- which(missing_pixels_count$thres_40perc == max(missing_pixels_count$thres_40perc))
    
    # Create mosaic based on available cloud-free images
    if (sum(missing_pixels_count$thres_5perc) > 1) {
      safe_log("Creating max composite from multiple cloud-free images (<5% clouds)")
      
      cloudy_rasters_list <- vrt_list[index_5perc]
      if (length(cloudy_rasters_list) == 1) {
        x <- terra::rast(cloudy_rasters_list[1])
      } else {
        rsrc <- terra::sprc(cloudy_rasters_list)
        x <- terra::mosaic(rsrc, fun = "max")
      }
      
    } else if (sum(missing_pixels_count$thres_5perc) == 1) {
      safe_log("Using single cloud-free image (<5% clouds)")
      
      x <- terra::rast(vrt_list[index_5perc[1]])
      
    } else if (sum(missing_pixels_count$thres_40perc) > 1) {
      safe_log("Creating max composite from partially cloudy images (<40% clouds)", "WARNING")
      
      cloudy_rasters_list <- vrt_list[index_40perc]
      if (length(cloudy_rasters_list) == 1) {
        x <- terra::rast(cloudy_rasters_list[1])
      } else {
        rsrc <- terra::sprc(cloudy_rasters_list)
        x <- terra::mosaic(rsrc, fun = "max")
      }
      
    } else if (sum(missing_pixels_count$thres_40perc) == 1) {
      safe_log("Using single partially cloudy image (<40% clouds)", "WARNING")
      
      x <- terra::rast(vrt_list[index_40perc[1]])
      
    } else {
      safe_log("No cloud-free images available, using all images", "WARNING")
      
      if (length(vrt_list) == 1) {
        x <- terra::rast(vrt_list[1])
      } else {
        rsrc <- terra::sprc(vrt_list)
        x <- terra::mosaic(rsrc, fun = "max")
      }
    }
  }
  
  # Set consistent layer names
  names(x) <- c("Red", "Green", "Blue", "NIR", "CI")
  return(x)
}

#' Save a mosaic raster to disk
#'
#' @param mosaic_raster A SpatRaster object to save
#' @param output_dir Directory to save the output
#' @param file_name Filename for the output raster
#' @param plot_result Whether to create visualizations (default: FALSE)
#' @param save_mask Whether to save cloud masks separately (default: FALSE)
#' @return The file path of the saved raster
#'
save_mosaic <- function(mosaic_raster, output_dir, file_name, plot_result = FALSE, save_mask = FALSE) {
  # Validate input
  if (is.null(mosaic_raster)) {
    stop("No mosaic raster provided to save")
  }
  
  # Create output directory if it doesn't exist
  dir.create(output_dir, recursive = TRUE, showWarnings = FALSE)
  
  # Create full file path
  file_path <- here::here(output_dir, file_name)
  
  # Get cloud mask if it exists
  cloud_mask <- attr(mosaic_raster, "cloud_mask")
  
  # Save raster
  terra::writeRaster(mosaic_raster, file_path, overwrite = TRUE)
  
  # Save cloud mask if available and requested
  if (!is.null(cloud_mask) && save_mask) {
    # Create mask filename by adding _mask before extension
    mask_file_name <- gsub("\\.(tif|TIF)$", "_mask.\\1", file_name)
    mask_file_path <- here::here(output_dir, mask_file_name)
    
    # Save the mask
    terra::writeRaster(cloud_mask, mask_file_path, overwrite = TRUE)
    safe_log(paste("Cloud/shadow mask saved to:", mask_file_path))
  } else if (!is.null(cloud_mask)) {
    safe_log("Cloud mask available but not saved (save_mask = FALSE)")
  }
  
  # Create plots if requested
  if (plot_result) {
    # Plot the CI band
    if ("CI" %in% names(mosaic_raster)) {
      terra::plot(mosaic_raster$CI, main = paste("CI map", file_name))
    }
    
    # Plot RGB image
    if (all(c("Red", "Green", "Blue") %in% names(mosaic_raster))) {
      terra::plotRGB(mosaic_raster, main = paste("RGB map", file_name))
    }
    
    # Plot cloud mask if available
    if (!is.null(cloud_mask)) {
      terra::plot(cloud_mask, main = paste("Cloud/shadow mask", file_name), 
                 col = c("red", "green"))
    }
    
    # If we have both RGB and cloud mask, create a side-by-side comparison
    if (all(c("Red", "Green", "Blue") %in% names(mosaic_raster)) && !is.null(cloud_mask)) {
      old_par <- par(mfrow = c(1, 2))
      terra::plotRGB(mosaic_raster, main = "RGB Image")
      
      # Create a colored mask for visualization (red = cloud/shadow, green = clear)
      mask_plot <- cloud_mask
      terra::plot(mask_plot, main = "Cloud/Shadow Mask", col = c("red", "green"))
      par(old_par)
    }
  }
  
  # Log save completion
  safe_log(paste("Mosaic saved to:", file_path))
  
  return(file_path)
}