#
# Below are the R packages needed for this routine, you can get them from CRAN
#
#  This R routine reads CERES grid data
#  From https://ceres.larc.nasa.gov/data/
#
#
# By Andy May (https://andymaypetrophysicist.com/)
#
# 3/23/2026
#
setwd("d:/Climate_Change/Clouds/CERES")
getwd()
#
#
#
# Required packages 
library(ncdf4)
library(RNetCDF)
library(terra)
library(tidyterra)
library(ggplot2)
library(dplyr)
library(readr)
library(lubridate)
library(sf)
library(rnaturalearth)
library(RColorBrewer)
library(rvest)
library(scales)
library(stringr)


# ────────────────────────────────────────────────────────────────
# Function compute latitude weights for Earth's ellipsoid
# ────────────────────────────────────────────────────────────────

lat_weights_ellipsoid <- function(r) {
  # Extract latitude centers from the raster
  lats <- yFromRow(r, 1:nrow(r))
  
  # Grid spacing (HadSST and ERSST are 5°)
  dg <- abs(lats[2] - lats[1])
  
  # WGS-84 ellipsoid parameters
  a  <- 6378.137                # semi-major axis (km)
  e2 <- 0.00669437999014        # eccentricity squared
  
  # Ellipsoidal area integrand
  integrand <- function(phi) {
    2 * pi * a^2 * (1 - e2) * cos(phi) / (1 - e2 * sin(phi)^2)^2
  }
  
  # Integrate area for each latitude band
  nlat <- length(lats)
  Area_lat <- numeric(nlat)
  
  for (j in 1:nlat) {
    phi1 <- (lats[j] - dg/2) * pi/180
    phi2 <- (lats[j] + dg/2) * pi/180
    # Divide by number of longitude cells (360/dg) to get per‑cell area
    Area_lat[j] <- integrate(integrand, phi1, phi2)$value / (360 / dg)
  }
  # Normalize by equatorial band (constant factor cancels in weighted.mean)
  eq_idx <- which.min(abs(lats))
  Wt_lat <- Area_lat / Area_lat[eq_idx]
  
  return(Wt_lat)
}

# ────────────────────────────────────────────────────────────────
# Plot themes
# ────────────────────────────────────────────────────────────────

 theme_white <- theme(
  plot.background   = element_rect(fill = "white", colour = NA),
  panel.background  = element_rect(fill = "white", colour = NA),
  panel.grid.major  = element_line(colour = "grey92"),
  panel.grid.minor  = element_line(colour = "grey96"),
  axis.text         = element_text(colour = "black"),
  axis.title        = element_text(colour = "black"),
  plot.title        = element_text(colour = "black", hjust = 0.5)
 )

# ────────────────────────────────────────────────────────────────
### Map themes
# ────────────────────────────────────────────────────────────────
 theme_map <- theme(
  text = element_text(color = "blue"),

  # Make title blue
  plot.title = element_text(
    hjust = 0.5,
    size = 14,
    color = "blue"
  ),

  # Remove axis labels ("x" and "y")
  axis.title = element_blank(),

  # Remove axis tick labels
  axis.text = element_blank(),
  axis.ticks = element_blank(),

  # Background
  plot.background = element_rect(fill = "lightgray", color = NA),
  panel.background = element_rect(fill = "lightgray", color = NA),

  # Remove default grid
  panel.grid = element_blank(),

  # Legend text blue
  legend.text = element_text(color = "blue"),
  legend.title = element_text(color = "blue"),
  legend.position = "bottom",
  legend.key.width = unit(2, "cm")
 )
# ────────────────────────────────────────────────────────────────
#
# Function: read and rotate a CERES variable
read_ceres_var <- function(file, varname) {
  r <- rast(file, subds = varname)
  rotate(r)
}
#
# Function: keep only complete years
keep_complete_years <- function(r) {
  yrs <- year(time(r))
  month_count <- table(yrs)
  complete <- as.integer(names(month_count[month_count == 12]))
  keep <- yrs %in% complete
  list(
    r = r[[keep]],
    years = yrs[keep]
  )
}
#
# Function: compute annual means
annual_mean <- function(r, years) {
  out <- tapp(r, years, fun = mean, na.rm = TRUE)
  names(out) <- sort(unique(years))
  out
}
#
# ────────────────────────────────────────────────────────────────
# Function: compute area-weighted global mean time series
# ────────────────────────────────────────────────────────────────
global_mean_ts <- function(r_yearly) {
  yrs <- as.numeric(names(r_yearly))
  Wt <- lat_weights_ellipsoid(r_yearly)

  df <- data.frame(year = yrs, mean = NA_real_)

  for (i in seq_along(yrs)) {
    mat <- as.matrix(r_yearly[[i]], wide = TRUE)
    wmat <- matrix(Wt, nrow = nrow(mat), ncol = ncol(mat), byrow = FALSE)
    df$mean[i] <- weighted.mean(mat, wmat, na.rm = TRUE)
  }
  df
}
# ────────────────────────────────────────────────────────────────
# Function: Save CSV + Plot Weighted Mean Time Series
# ────────────────────────────────────────────────────────────────

save_weighted_mean_outputs <- function(df, varname, ylab = "W/m^2", title = NULL) {

  # Clean variable name for filenames
  base <- gsub("_mon$", "", varname)

  # Default title if none supplied
  if (is.null(title)) {
    title <- paste("Area-Weighted Global Mean:", base)
  }

  # Output filenames
  csv_file  <- paste0(base, "_area_weighted_global_mean.csv")
  png_file  <- paste0(base, "_area_weighted_global_mean.png")

  # Write CSV
  write_csv(df, csv_file)

  # Build plot
  p <- ggplot(df, aes(x = year, y = mean)) +
    geom_line(colour = "#d62728", linewidth = 1) +
    labs(
      title = title,
      x = "Year",
      y = ylab
    ) +
    theme_minimal() + theme_white +
    theme(
      panel.border = element_rect(colour = "black", fill = NA, linewidth = 0.8)
    )

  # Print to screen
  print(p)

  # Save PNG
  ggsave(
    filename = png_file,
    plot = p,
    width = 9,
    height = 5,
    dpi = 300,
    bg = "white"
  )

  message("Saved: ", csv_file, " and ", png_file)
}
#
# ────────────────────────────────────────────────────────────────
# Function to compute a trend raster
# omputes a linear trend (slope) at every grid cell of a multi-layer 
# SpatRaster, where each layer is one year of data.

# The output is a single‑layer raster where each pixel contains the 
# trend in W/m² per year.
# ────────────────────────────────────────────────────────────────
trend_raster <- function(r_yearly) {
  yrs <- as.numeric(names(r_yearly))
  app(r_yearly, fun = function(y) {
    if (all(is.na(y))) return(NA)
    coef(lm(y ~ yrs))[2]
  })
}
#
# ────────────────────────────────────────────────────────────────
# Function to plot a flux map
# ────────────────────────────────────────────────────────────────
plot_flux_map <- function(r, title_text) {
  mm <- minmax(r)

  ggplot() +
    geom_spatraster(data = r) +
    scale_fill_gradientn(
      colours = rev(RColorBrewer::brewer.pal(11, "RdYlBu")),
      limits = c(mm[1], mm[2]),
      oob = scales::squish,
      name = "(W/m^2)"
    ) +
    geom_sf(data = world_rob, fill = NA, color = "black", linewidth = 0.15) +
    geom_sf(data = grat, color = "black", linewidth = 0.2) +
    geom_sf_text(data = lon_labs, aes(label = lab), color = "blue", size = 3) +
    geom_sf_text(data = lat_labs, aes(label = lab), color = "blue", size = 3) +
    coord_sf(crs = rob, expand = FALSE) +
    theme_void() + theme_white + theme_map +
    labs(title = title_text)

}
#
# ────────────────────────────────────────────────────────────────
####### Main Program #######
# ────────────────────────────────────────────────────────────────

# Read CERES files
file_old <- "CERES_EBAF_Ed4.2.1_Subset_200003-201412.nc"
file_new <- "CERES_EBAF_Ed4.2.1_Subset_201501-202511.nc"

# nc=open.nc(file_new)
# print.nc(nc)

# Variables to process
vars <- c(
  "sfc_lw_up_all_mon",
  "sfc_lw_down_all_mon",
  "sfc_sw_down_all_mon",
  "sfc_sw_up_all_mon",
  "sfc_net_sw_all_mon",
  "sfc_net_lw_all_mon",
  "sfc_net_tot_all_mon",
  "toa_sw_all_mon",
  "toa_lw_all_mon",
  "toa_net_all_mon",          # <-- ADD THIS
  "toa_cre_sw_mon",
  "toa_cre_lw_mon",
  "cldarea_total_daynight_mon"
)
results <- list()

for (v in vars) {

  # Read both files
  r1 <- read_ceres_var(file_old, v)
  r2 <- read_ceres_var(file_new, v)

  # Merge time series
  r <- c(r1, r2)

  # Keep complete years
  k <- keep_complete_years(r)
  r_complete <- k$r
  yrs <- k$years

  # Annual means
  r_yearly <- annual_mean(r_complete, yrs)

  # Global mean time series
  ts <- global_mean_ts(r_yearly)

  # Save weighted means
  save_weighted_mean_outputs(
    df = ts,
    varname = v,
    ylab = "W/m^2",
    title = paste("Area-Weighted Global Mean:", v)
  )

  # Trend raster
  tr <- trend_raster(r_yearly)

  # Store results
  results[[v]] <- list(
    yearly = r_yearly,
    ts = ts,
    trend = tr
  )
} # End for (v
#
# ────────────────────────────────────────────────────────────────
# Surface EEI
# Pull out the entire surface net total flux
# ────────────────────────────────────────────────────────────────

EEI <- results[["sfc_net_tot_all_mon"]]

# Extract the annual mean maps.
EEI_yearly <- EEI$yearly

# Pull out the global annual time series
EEI_ts <- EEI$ts

# Extract the spatial trend map
EEI_trend <- EEI$trend

# Set up maps
# Robinson projection
rob <- "+proj=robin +lon_0=0 +x_0=0 +y_0=0 +datum=WGS84 +units=m +no_defs"
world <- ne_countries(scale = "medium", returnclass = "sf")

# Transform world map
world_rob <- st_transform(world, crs = rob)

# Build graticule in geographic coords, then transform
grat <- st_graticule(
  lat = seq(-60, 60, 15),
  lon = seq(-180, 180, 30)
) |> st_transform(rob)

# Custom longitude labels (bottom only)
lon_labs <- data.frame(
  lon = seq(-180, 180, 30),
  lat = -80,                     # place labels below map
  lab = paste0(seq(-180, 180, 30), "°")
) |>
  st_as_sf(coords = c("lon", "lat"), crs = 4326) |>
  st_transform(rob)

lat_labs <- data.frame(
  lon = -170,                     # place labels left of the map
  lat = seq(-60, 60, 15),
  lab = paste0(seq(-60, 60, 15), "°")
) |>
  st_as_sf(coords = c("lon", "lat"), crs = 4326) |>
  st_transform(rob)

colourPalette <- c('blue', 'lightblue', 'lightgreen', brewer.pal(9, 'YlOrRd'))

# Strictly speaking this is not "EEI," it is sort of a surface EEI, EEI is at the TOA
#
# EEI_yearly (2001): Surface net total radiation = (SW↓ – SW↑) + (LW↓ – LW↑)  
# i.e., net shortwave + net longwave at the surface
# Positive downward (energy gained by the surface)

 p1 <- plot_flux_map(EEI_yearly[[1]], "2001 Surface Net Radiation (Net SW+LW) W/m^2")
 p2 <- plot_flux_map(EEI_trend, "2001-2024 Surface Net Radiation Trend (W/m^2 per year)")

 
 ggsave("Net_SW_LW_surface_2001.png", p1, width = 9, height = 5, dpi = 300, bg = "white")
 ggsave("Net_Radiation_surface_trend_2001_2024.png", p2, width = 9, height = 5, dpi = 300, bg = "white")

  
write_csv(EEI_ts, "Surface_net radiation_mean_2001_2024.csv")

# ────────────────────────────────────────────────────────────────
# Compute TOA EEI (true EEI) and compare to Surface "EEI"
# ────────────────────────────────────────────────────────────────

# --- 1. Pull TOA net flux (downward positive)
# 
TOA <- results[["toa_net_all_mon"]]

TOA_yearly <- TOA$yearly
TOA_ts     <- TOA$ts
# Extract the spatial trend map
TOA_trend <- TOA$trend


 p1 <- plot_flux_map(TOA_yearly[[1]], "2001 TOA Net Radiation (Net SW+LW) W/m^2")
 p2 <- plot_flux_map(TOA_trend, "2001-2024 TOA Net Radiation Trend (W/m^2 per year)")
 ggsave("Net_SW_LW_TOA_2001.png", p1, width = 9, height = 5, dpi = 300, bg = "white")
 ggsave("Net_Radiation_TOA_trend_2001_2024.png", p2, width = 9, height = 5, dpi = 300, bg = "white")

# ============================================================
# Compute anomaly rasters for TOA and Surface EEI
# ============================================================

# Multi-year mean maps
TOA_mean_map  <- mean(TOA_yearly,  na.rm = TRUE)
EEI_mean_map  <- mean(EEI_yearly,  na.rm = TRUE)

# Anomaly rasters (each layer = year minus multi-year mean)
TOA_anom_yearly <- TOA_yearly - TOA_mean_map
EEI_anom_yearly <- EEI_yearly - EEI_mean_map

# ============================================================
# Plot anomaly maps
# ============================================================

# Example: first year (2001)
p_toa_anom_2001 <- plot_flux_map(TOA_anom_yearly[[1]],
                                 "2001 TOA Net Radiation Anomaly (W/m^2)")

p_sfc_anom_2001 <- plot_flux_map(EEI_anom_yearly[[1]],
                                 "2001 Surface Net Radiation Anomaly (W/m^2)")

ggsave("TOA_anomaly_2001.png",  p_toa_anom_2001, width=9, height=5, dpi=300)
ggsave("Surface_anomaly_2001.png", p_sfc_anom_2001, width=9, height=5, dpi=300)

# ============================================================
# Compute TOA minus Surface Net Radiation and anomaly difference
# How much energy enters the climate system at the TOA vs. 
# how much reaches the surface? In W/m^2
# ============================================================
# ============================================================
# Plot TOA - Surface difference map
# ============================================================

EEI_difference_yearly <- TOA_yearly - EEI_yearly
p_diff_2001 <- plot_flux_map(EEI_difference_yearly[[1]],
                             "2001 TOA minus Surface Net Radiation (W/m^2)")

ggsave("TOA_minus_Surface_2001.png",
       p_diff_2001, width=9, height=5, dpi=300)

# ============================================================
# Plot TOA - Anomaly Surface difference map
# ============================================================

EEI_difference_anom_yearly <- TOA_anom_yearly - EEI_anom_yearly
p_diff_anom_2001 <- plot_flux_map(EEI_difference_anom_yearly[[1]],
                                  "2001 TOA - Surface Radiation Anomaly (W/m^2)")

ggsave("TOA_minus_Surface_anomaly_2001.png",
       p_diff_anom_2001, width=9, height=5, dpi=300)


# ============================================================
# --- Convert both TOA and Surface to anomalies
# --- Here these are 'ts' dataframes of the weighted avg means
# ============================================================

surf_df <- EEI_ts
toa_df  <- TOA_ts

surf_df$anom <- surf_df$mean - mean(surf_df$mean, na.rm = TRUE)
toa_df$anom  <- toa_df$mean  - mean(toa_df$mean,  na.rm = TRUE)

# --- 3. Merge into a single comparison dataframe
cmp <- data.frame(
  year        = surf_df$year,
  surf_anom   = surf_df$anom,
  toa_anom    = toa_df$anom
)

write_csv(cmp, "Surface_vs_TOA_EEI_anomalies.csv")

# --- Compute linear trends (W/m^2 per decade)
surf_trend <- coef(lm(surf_anom ~ year, data = cmp))[2] * 10
toa_trend  <- coef(lm(toa_anom  ~ year, data = cmp))[2] * 10

cat("Surface net radiation trend (W/m^2 per decade): ", surf_trend, "\n")
cat("TOA EEI trend (W/m^2 per decade): ", toa_trend,  "\n")

# ────────────────────────────────────────────────────────────────
# Time series plot of both
# ────────────────────────────────────────────────────────────────

# --- Plot comparison of anomalies
p_cmp <- ggplot(cmp, aes(x = year)) +
  geom_line(aes(y = surf_anom, color = "Surface Net Radiation"), linewidth = 1.1) +
  geom_line(aes(y = toa_anom,  color = "TOA Net Flux (EEI)"), linewidth = 1.1) +
  labs(
    title = "Surface vs TOA Energy Imbalance (Anomalies)",
    x = "Year",
    y = "Anomaly (W/m^2)",
    color = ""
  ) +
  theme_minimal() + theme_white +
  theme(
    panel.border = element_rect(colour = "black", fill = NA, linewidth = 0.8),
    legend.position = "bottom"
  )

print(p_cmp)

ggsave(
  "Surface_vs_TOA_EEI_anomalies.png",
  p_cmp,
  width = 9,
  height = 5,
  dpi = 300,
  bg = "white"
)

# ────────────────────────────────────────────────────────────────
# Time series plot
# ────────────────────────────────────────────────────────────────

p_ts <- ggplot(EEI_ts, aes(x = year, y = mean)) +
  geom_line(color = "#d62728", linewidth = 1) +
  labs(
    title = "Global Mean Surface net radiation (Net SW+LW)",
    x = "Year",
    y = "W/m^2"
  ) +
  theme_minimal() + theme_white +
  theme(
    panel.border = element_rect(colour = "black", fill = NA, linewidth = 0.8)
  )

print(p_ts)

ggsave(
  "Surface_net_radiation_mean_2001_2024.png",
  p_ts,
  width = 9,
  height = 5,
  dpi = 300,
  bg = "white"
)