## =============================================================================
##  vanGenuchten.R
##
##  Fit or evaluate the van Genuchten water retention curve, either in its
##  original form (van Genuchten, 1980) or with the PDI extension that adds
##  adsorbed water and lets the curve reach zero water content
##  (Peters, 2013; Peters et al., 2021).
##
##  Soil Physics Lab, University of California, Merced
##  Teamrat A. Ghezzehei
##
##  How to use
##  ----------
##      source("vanGenuchten.R")
##
##      # 1. Parameters known -> predict theta at target water potentials
##      par <- c(theta_r = 0.05, theta_s = 0.42, alpha = 0.03, n = 1.6)
##      vg_predict(psi, par)                  # van Genuchten
##      vg_predict(psi, par, model = "PDI")   # van Genuchten-PDI
##
##      # 2. Data known -> fit the parameters and get fit statistics
##      fit <- vg_fit(psi, theta)                  # van Genuchten
##      fit <- vg_fit(psi, theta, model = "PDI")   # van Genuchten-PDI
##      fit            # prints parameters, standard errors, RMSE, R2, AIC
##      fit$par        # the parameters, ready to pass to vg_predict()
##
##  Units
##  -----
##  psi   : water potential expressed as suction (positive numbers; a negative
##          matric potential is accepted and its sign is dropped).
##  alpha : 1 / (unit of psi). Use the same unit for psi everywhere.
##  theta : volumetric water content (v/v).
##  h0    : PDI only, suction at zero water content, in the unit of psi.
##          The default 10^6.8 is in cm; change it if psi is not in cm.
##
##  Only base R is needed (stats::nls).
## =============================================================================


## -----------------------------------------------------------------------------
##  1. The two model equations
## -----------------------------------------------------------------------------

## van Genuchten (1980):
##     Se    = [1 + (alpha * psi)^n]^(-m),   m = 1 - 1/n
##     theta = theta_r + (theta_s - theta_r) * Se
vg_theta <- function(psi, theta_r, theta_s, alpha, n) {
  psi <- abs(psi)
  m   <- 1 - 1 / n
  Se  <- (1 + (alpha * psi)^n)^(-m)
  theta_r + (theta_s - theta_r) * Se
}

## van Genuchten-PDI (Peters, 2013):
##     theta = (theta_s - theta_r) * S_cap + theta_r * S_ad
## S_cap is the van Genuchten Se rescaled so that it is exactly 0 at psi = h0.
## S_ad describes adsorbed water: 1 for psi below ha, then decreasing linearly
## in log(psi) to 0 at h0, with a smooth corner of width b. Unless ha is
## given, ha = 1/alpha as in Peters (2013). In this model theta_r is the
## maximum adsorbed water content, not a residual water content.
vg_pdi_theta <- function(psi, theta_r, theta_s, alpha, n, h0 = 10^6.8, ha = NA) {
  psi <- pmax(abs(psi), 1e-12)          # avoid log(0)
  m   <- 1 - 1 / n
  if (is.na(ha)) ha <- 1 / alpha

  ## capillary part
  Se    <- (1 + (alpha * psi)^n)^(-m)
  Se_0  <- (1 + (alpha * h0)^n)^(-m)
  S_cap <- (Se - Se_0) / (1 - Se_0)

  ## adsorptive part
  x  <- log(psi); xa <- log(ha); x0 <- log(h0)
  b  <- 0.1 + 0.2 / n^2 * (1 - exp(-(theta_r / (theta_s - theta_r))^2))
  z  <- (xa - x) / b
  smooth <- ifelse(z > 30, z, log1p(exp(pmin(z, 30))))    # = log(1 + exp(z)), overflow-safe
  S_ad <- 1 + (x - xa + b * smooth) / (xa - x0)

  ## both parts are kept within [0, 1] (they run past it beyond h0)
  S_cap <- pmin(pmax(S_cap, 0), 1)
  S_ad  <- pmin(pmax(S_ad,  0), 1)

  (theta_s - theta_r) * S_cap + theta_r * S_ad
}


## -----------------------------------------------------------------------------
##  2. Predict theta from known parameters
## -----------------------------------------------------------------------------

## psi   : vector of target water potentials
## par   : named vector or list with theta_r, theta_s, alpha, n
##         (e.g. the $par element of a vg_fit result)
## model : "vG" (default) or "PDI"
## h0, ha: PDI only (see the header); ha = NA means 1/alpha
vg_predict <- function(psi, par, model = c("vG", "PDI"), h0 = 10^6.8, ha = NA) {
  model <- match.arg(model)
  p <- as.list(par)
  if (model == "vG") {
    vg_theta(psi, p$theta_r, p$theta_s, p$alpha, p$n)
  } else {
    vg_pdi_theta(psi, p$theta_r, p$theta_s, p$alpha, p$n, h0 = h0, ha = ha)
  }
}


## -----------------------------------------------------------------------------
##  3. Fit the parameters to measured data
## -----------------------------------------------------------------------------

## psi, theta : measured water potentials and water contents (same length)
## model      : "vG" (default) or "PDI"
## h0, ha     : PDI only, held fixed during the fit (ha = NA means 1/alpha)
## start      : optional named list of starting values; otherwise a few
##              combinations of alpha and n are tried and the best is kept
##
## Returns a list (class "vg_fit") with
##   $model, $par, $se, $fitted, $residuals, $RMSE, $R2, $AIC, $nls
vg_fit <- function(psi, theta, model = c("vG", "PDI"), h0 = 10^6.8, ha = NA,
                   start = NULL) {
  model <- match.arg(model)

  ## keep complete cases; psi as suction
  ok  <- is.finite(psi) & is.finite(theta)
  dat <- data.frame(psi = abs(psi[ok]), theta = theta[ok])
  if (nrow(dat) < 5) stop("vg_fit: need at least 5 valid (psi, theta) pairs.")

  ## the formula nls will fit
  form <- if (model == "vG") {
    theta ~ vg_theta(psi, theta_r, theta_s, alpha, n)
  } else {
    theta ~ vg_pdi_theta(psi, theta_r, theta_s, alpha, n, h0 = h0, ha = ha)
  }

  ## parameter bounds (theta_r >= 0, theta_s <= 1, n > 1)
  lower <- c(theta_r = 0,              theta_s = 0.5 * max(dat$theta), alpha = 1e-6, n = 1.01)
  upper <- c(theta_r = max(dat$theta), theta_s = 1,                    alpha = 1e3,  n = 20)

  ## starting values: alpha near 1/psi at mid-range theta, and a few n values
  mid     <- (max(dat$theta) + min(dat$theta)) / 2
  psi_mid <- dat$psi[which.min(abs(dat$theta - mid))]
  if (psi_mid <= 0) psi_mid <- median(dat$psi[dat$psi > 0])

  starts <- if (!is.null(start)) list(start) else {
    grid <- expand.grid(a = c(1, 0.2, 5) / psi_mid, n = c(1.2, 1.5, 2, 3))
    lapply(seq_len(nrow(grid)), function(i)
      list(theta_r = 0.5 * min(dat$theta), theta_s = max(dat$theta),
           alpha = grid$a[i], n = grid$n[i]))
  }

  ## try each start, keep the converged fit with the smallest residual sum of squares
  best <- NULL
  for (s in starts) {
    f <- try(nls(form, data = dat, start = s, lower = lower, upper = upper,
                 algorithm = "port"), silent = TRUE)
    if (!inherits(f, "try-error") && (is.null(best) || deviance(f) < deviance(best))) best <- f
  }
  if (is.null(best)) stop("vg_fit: the fit did not converge. Check the units of psi or give 'start'.")

  ## statistics
  par    <- coef(best)
  se     <- summary(best)$coefficients[, "Std. Error"]
  fitted <- fitted(best)
  res    <- dat$theta - fitted
  RMSE   <- sqrt(mean(res^2))
  R2     <- 1 - sum(res^2) / sum((dat$theta - mean(dat$theta))^2)

  structure(list(model = model, par = par, se = se,
                 h0 = if (model == "PDI") h0 else NA, ha = if (model == "PDI") ha else NA,
                 fitted = as.numeric(fitted), residuals = as.numeric(res),
                 RMSE = RMSE, R2 = R2, AIC = AIC(best), n_obs = nrow(dat), nls = best),
            class = "vg_fit")
}

## Readable printout of a fit
print.vg_fit <- function(x, digits = 4, ...) {
  cat("Model:", if (x$model == "vG") "van Genuchten (1980)" else "van Genuchten-PDI", "\n")
  if (x$model == "PDI") cat("h0 (fixed):", x$h0, "  ha (fixed):",
                            if (is.na(x$ha)) "1/alpha" else x$ha, "\n")
  tab <- data.frame(estimate = signif(x$par, digits), std_error = signif(x$se, digits))
  tab <- rbind(tab, m = c(signif(1 - 1 / x$par[["n"]], digits), NA))
  print(tab, na.print = "")
  cat(sprintf("\nObservations: %d   RMSE: %.4g   R2: %.4f   AIC: %.2f\n",
              x$n_obs, x$RMSE, x$R2, x$AIC))
  invisible(x)
}