h <- horizons(nc_spc)

# Columns we will use
cols <- c("pedon_key", "hzn_desgn", "hzn_top", "hzn_bot",
          "sand", "silt", "clay", "db_od", "oc",
          "w6clod",   # θ at  6 kPa  =   61.2 cm head  (pF 1.79)
          "w3cld",    # θ at 33 kPa  =  336.7 cm head  (pF 2.53)  ← field capacity
          "w15l2")    # θ at 1500 kPa= 15300   cm head  (pF 4.18) ← wilting point

head(h[, intersect(cols, names(h))], 6)

library(aqp)
set.seed(1)
sub <- nc_spc[sample(1:length(nc_spc), 12), ]
plotSPC(sub, color = "clay",
        col.label = "Clay (%)",
        name = "hzn_desgn",
        cex.names = 0.7)
title("Raw KSSL profiles — North Carolina (n = 12)")

library(mpspline2)

h <- horizons(nc_spc) %>%
  rename(PEDON_ID = pedon_key,
         top      = hzn_top,
         bot      = hzn_bot,
         bd       = db_od,
         soc      = oc) %>%
  filter(!is.na(top), !is.na(bot), bot > top)

std_d     <- c(0, 5, 15, 30, 60, 100)
soil_vars <- c("sand", "silt", "clay", "bd", "soc")
ret_vars  <- c("w6clod", "w3cld", "w15l2")
all_vars  <- c(soil_vars, ret_vars)

spline_var <- function(df, var, std_d) {
  # Drop rows with NA in this variable
  dat <- df %>%
    select(PEDON_ID, top, bot, value = all_of(var)) %>%
    filter(!is.na(value)) %>%
    as.data.frame()

  if (nrow(dat) == 0) return(NULL)

  # mpspline2 expects columns: id, top, bot, value
  sp <- tryCatch(
    mpspline2::mpspline(dat, var_name = "value", d = std_d, vlow = 0),
    error = function(e) NULL
  )
  if (is.null(sp)) return(NULL)

  # Extract standard-depth estimates for each profile
  map_dfr(names(sp), function(pid) {
    est <- sp[[pid]]$est_dcm   # mean value per standard-depth interval
    n   <- length(std_d) - 1
    tibble(
      PEDON_ID  = pid,
      depth_top = std_d[seq_len(n)],
      depth_bot = std_d[seq_len(n) + 1],
      !!var    := est[seq_len(n)]
    )
  })
}

# Run for all variables (~1 min for 487 profiles × 8 variables)
harmonised_list <- map(all_vars, ~ spline_var(h, .x, std_d))
names(harmonised_list) <- all_vars

# Join all variables by PEDON_ID + depth interval
harmonised <- reduce(harmonised_list, full_join,
                     by = c("PEDON_ID", "depth_top", "depth_bot")) %>%
  mutate(depth = paste0(depth_top, "-", depth_bot))

cat("Harmonised rows:", nrow(harmonised), "\n")

pid   <- "10NC001A"
std_d <- c(0, 5, 15, 30, 60, 100)

# Simulated raw horizons (realistic NC Ultisol — clay increases with depth)
raw_prof <- data.frame(
  top = c(  0,  8, 22, 45,  78, 105),
  bot = c(  8, 22, 45, 78, 105, 145),
  clay = c(11, 16, 27, 36,  41,  43)
) |> dplyr::mutate(depth_mid = (top + bot) / 2)

# Simulated spline estimates at standard depth midpoints
sp_prof <- data.frame(
  depth_top = c( 0,  5, 15, 30,  60),
  depth_bot = c( 5, 15, 30, 60, 100),
  clay      = c(10.8, 14.6, 24.1, 34.3, 40.6)
) |> dplyr::mutate(depth_mid = (depth_top + depth_bot) / 2)

ggplot() +
  geom_rect(data = raw_prof,
            aes(xmin = 0, xmax = clay, ymin = top, ymax = bot),
            fill = "sienna3", alpha = 0.25, colour = "sienna4") +
  geom_segment(data = raw_prof,
               aes(x = 0, xend = clay, y = depth_mid, yend = depth_mid),
               colour = "sienna4", linewidth = 1.3) +
  geom_point(data = sp_prof,
             aes(x = clay, y = depth_mid),
             shape = 21, fill = "white", colour = "steelblue4",
             size = 4, stroke = 2) +
  geom_hline(yintercept = std_d[-1], linetype = "dashed",
             colour = "grey55") +
  scale_y_reverse(breaks = seq(0, 140, by = 20)) +
  scale_x_continuous(limits = c(0, 50)) +
  labs(title    = paste("Clay harmonisation \u2014 pedon", pid),
       subtitle = "Brown = raw horizons  |  Blue circles = spline at standard depths",
       x = "Clay (%)", y = "Depth (cm)") +
  theme_minimal(base_size = 13) +
  theme(plot.subtitle = element_text(colour = "grey40", size = 11))

swrc_example %>%
  distinct(PEDON_ID, depth, clay, silt, sand_total, bd) %>%
  pivot_longer(c(clay, silt, sand_total, bd),
               names_to = "property", values_to = "value") %>%
  mutate(
    property = factor(property,
                      levels = c("clay", "silt", "sand_total", "bd"),
                      labels = c("Clay (%)", "Silt (%)", "Sand (%)",
                                 "Bulk density (g cm\u207b\u00b3)")),
    depth = factor(depth,
                   levels = c("0-5","5-15","15-30","30-60","60-100"))
  ) %>%
  ggplot(aes(x = value, y = depth, fill = depth)) +
  geom_boxplot(alpha = 0.75, outlier.size = 0.6) +
  facet_wrap(~ property, scales = "free_x", ncol = 2) +
  scale_fill_manual(values = c("0-5"    = "#8B2500",
                                "5-15"   = "#CC5500",
                                "15-30"  = "#E8A020",
                                "30-60"  = "#3985C8",
                                "60-100" = "#1A4580")) +
  labs(x = NULL, y = "Depth interval") +
  theme_minimal(base_size = 12) +
  theme(legend.position = "none",
        strip.text = element_text(face = "bold"))

head_to_pF <- function(h) log10(h)

depth_pal <- c("0-5"    = "#8B2500",
               "5-15"   = "#CC5500",
               "15-30"  = "#E8A020",
               "30-60"  = "#3985C8",
               "60-100" = "#1A4580")

swrc_example %>%
  mutate(pF = head_to_pF(matric_head),
         depth = factor(depth,
                        levels = c("0-5","5-15","15-30","30-60","60-100"))) %>%
  ggplot(aes(x = pF, y = water_content,
             group = PEDON_ID, colour = depth)) +
  geom_line(alpha = 0.35, linewidth = 0.5) +
  geom_point(alpha = 0.5, size = 1) +
  facet_wrap(~ depth, ncol = 3) +
  scale_colour_manual(values = depth_pal) +
  labs(
    title   = "Observed water retention by depth interval",
    x       = expression(pF == log[10](h~"[cm]")),
    y       = expression(theta~(m^3~m^{-3})),
    caption = "Each line = one soil profile; points = measured tensions (6, 33, 1500 kPa)"
  ) +
  theme_minimal(base_size = 12) +
  theme(legend.position = "none",
        strip.text = element_text(face = "bold"))

knitr::include_graphics("figures/texture-triangle.png")

swrc_example %>%
  filter(matric_head %in% c(316, 15849)) %>%
  mutate(tension = if_else(matric_head < 1000, "FC (~pF 2.5)", "WP (~pF 4.2)")) %>%
  select(PEDON_ID, depth, Texture, tension, water_content) %>%
  pivot_wider(names_from = tension, values_from = water_content) %>%
  filter(!is.na(`FC (~pF 2.5)`), !is.na(`WP (~pF 4.2)`)) %>%
  ggplot(aes(x = `WP (~pF 4.2)`, y = `FC (~pF 2.5)`, colour = Texture)) +
  geom_abline(slope = 1, intercept = 0, linetype = "dashed", colour = "grey60") +
  geom_point(alpha = 0.6, size = 1.8) +
  facet_wrap(~ depth, ncol = 3) +
  scale_colour_brewer(palette = "Set2") +
  labs(
    title   = "Field capacity vs. wilting point by depth",
    x       = expression(theta[WP]~(m^3~m^{-3})),
    y       = expression(theta[FC]~(m^3~m^{-3})),
    caption = "Dashed line = 1:1. Points above line = plant-available water"
  ) +
  theme_minimal(base_size = 12) +
  theme(legend.position = "bottom",
        strip.text = element_text(face = "bold"))

# Remove rows where water content was not measured (prevents NaN gradients)
swrc_df <- swrc_df %>% filter(!is.na(theta_n))

set.seed(2025)
ids     <- unique(swrc_df$PEDON_ID)
tr_ids  <- sample(ids, floor(0.70 * length(ids)))
val_ids <- sample(setdiff(ids, tr_ids), floor(0.15 * length(ids)))
te_ids  <- setdiff(ids, c(tr_ids, val_ids))

train_df <- swrc_df[swrc_df$PEDON_ID %in% tr_ids,  ]
val_df   <- swrc_df[swrc_df$PEDON_ID %in% val_ids, ]
test_df  <- swrc_df[swrc_df$PEDON_ID %in% te_ids,  ]

cat(sprintf("Train: %d profiles | Val: %d | Test: %d\n",
            length(tr_ids), length(val_ids), length(te_ids)))
# Train: 280 profiles | Val: 60 | Test: 60

fit <- fit_swrc(
  train_df   = train_df,
  x_inputs   = c("clay", "silt", "bd_gcm3", "soc", "Depth_num"),
  val_df     = val_df,
  epochs     = 80,
  batch_size = 256,
  lambdas    = norouzi_lambdas("norouzi"),
  verbose    = TRUE
)

knitr::include_graphics("figures/training-history.png")

sel_tex <- c("Sandy Loam", "Loam", "Clay Loam", "Clay")
test_sub <- test_df[test_df$Texture %in% sel_tex, ]

# Step 1 – predict at observed tensions for every test profile
test_sub$theta_pred <- predict_swrc(fit, test_sub)

# Step 2 – RMSE per profile; pick the best one per texture × depth cell
best_pids <- test_sub |>
  dplyr::group_by(PEDON_ID, Texture, Depth_label, Depth_num) |>
  dplyr::summarise(
    rmse = sqrt(mean((theta_pred - theta_n)^2, na.rm = TRUE)),
    .groups = "drop"
  ) |>
  dplyr::group_by(Texture, Depth_label, Depth_num) |>
  dplyr::slice_min(rmse, n = 1, with_ties = FALSE) |>
  dplyr::ungroup()

# Step 3 – dense prediction for each best profile
rep_profiles <- test_sub |>
  dplyr::inner_join(
    best_pids |> dplyr::select(PEDON_ID, Texture, Depth_label, Depth_num),
    by = c("PEDON_ID", "Texture", "Depth_label", "Depth_num")
  ) |>
  dplyr::group_by(PEDON_ID, Texture, Depth_label, Depth_num) |>
  dplyr::slice(1) |>
  dplyr::ungroup()

dense_best <- predict_swrc_dense(fit, newdata = rep_profiles, n_points = 500)

# Step 4 – observed points for those same profiles only
obs_pts <- test_sub |>
  dplyr::inner_join(
    best_pids |> dplyr::select(PEDON_ID, Texture, Depth_label, Depth_num),
    by = c("PEDON_ID", "Texture", "Depth_label", "Depth_num")
  ) |>
  dplyr::mutate(pF = log10(matric_head))

knitr::include_graphics("figures/swrc-curves.png")

pred_vals <- predict_swrc(fit, test_df)

data.frame(
  theta_obs  = test_df$theta_n,
  theta_pred = pred_vals,
  Depth      = test_df$Depth_label
) %>%
  ggplot(aes(x = theta_obs, y = theta_pred, colour = Depth)) +
  geom_abline(slope = 1, intercept = 0, linetype = "dashed", colour = "grey50") +
  geom_point(alpha = 0.5, size = 1.5) +
  coord_equal() +
  scale_colour_brewer(palette = "YlOrBr", direction = -1) +
  labs(x = expression(theta[obs]~(m^3~m^{-3})),
       y = expression(theta[pred]~(m^3~m^{-3})),
       colour = "Depth") +
  theme_minimal(base_size = 13)

knitr::include_graphics("figures/pred-obs.png")

pid      <- te_ids[1]
pedon_df <- distinct(test_df[test_df$PEDON_ID == pid, ],
                     Depth_num, Depth_label, .keep_all = TRUE)

fc_pred <- predict_swrc(fit, pedon_df %>% mutate(matric_head = 316))
wp_pred <- predict_swrc(fit, pedon_df %>% mutate(matric_head = 15849))

profile_df <- data.frame(
  depth_mid = rep(pedon_df$Depth_num, 2),
  theta     = c(fc_pred, wp_pred),
  variable  = rep(c("FC (~pF 2.5)", "WP (~pF 4.2)"), each = nrow(pedon_df))
) %>% pivot_wider(names_from = variable, values_from = theta)

ggplot(profile_df, aes(y = depth_mid)) +
  geom_ribbon(aes(xmin = `WP (~pF 4.2)`, xmax = `FC (~pF 2.5)`),
              fill = "steelblue", alpha = 0.2) +
  geom_path(aes(x = `FC (~pF 2.5)`, colour = "FC (~pF 2.5)"), linewidth = 1.3) +
  geom_path(aes(x = `WP (~pF 4.2)`, colour = "WP (~pF 4.2)"), linewidth = 1.3) +
  geom_point(aes(x = `FC (~pF 2.5)`, colour = "FC (~pF 2.5)"), size = 3.5) +
  geom_point(aes(x = `WP (~pF 4.2)`, colour = "WP (~pF 4.2)"), size = 3.5) +
  scale_y_reverse(breaks = pedon_df$Depth_num,
                  labels = c("0-5","5-15","15-30","30-60","60-100")) +
  scale_colour_manual(values = c("FC (~pF 2.5)" = "steelblue4",
                                 "WP (~pF 4.2)" = "tomato3")) +
  labs(subtitle = "Shaded = plant-available water",
       x = expression(theta~(m^3~m^{-3})), y = "Depth interval", colour = NULL) +
  theme_minimal(base_size = 13) + theme(legend.position = "bottom")

knitr::include_graphics("figures/vertical-profile.png")