# =============================================================================
# 03_student_vs_gaussian.R — Student-t NBC vs Gaussian NBC comparison
#
# Demonstrates the classification-collapse failure mode of Gaussian NBC on
# fat-tailed financial data and shows how Student-t NBC avoids it.
#
# Key quantities compared:
#   - Overall accuracy and Stress-class recall
#   - Per-observation log-density at extreme returns
#   - Decision boundary shift driven by tail behaviour
# =============================================================================

library(kronxNBC)
library(naivebayes)   # for naive_bayes() — Gaussian NBC

# ---- 1. Shared training data -------------------------------------------------

set.seed(42L)
n  <- 300L
mk <- n / 3L

X_train <- rbind(
  # Calm
  cbind(
    log_return         = rnorm(mk, 0.0002, 0.003),
    rolling_volatility = abs(rnorm(mk, 0.004,  0.001)),
    drawdown           = rnorm(mk, -0.002, 0.002),
    transition_stress  = abs(rnorm(mk, 0.001, 0.0005)),
    residence_pressure = as.numeric(rpois(mk, 1L)),
    ruin_proxy         = rbeta(mk, 1, 20)
  ),
  # Steady
  cbind(
    log_return         = rnorm(mk, 0.0005, 0.005),
    rolling_volatility = abs(rnorm(mk, 0.008,  0.002)),
    drawdown           = rnorm(mk, -0.008, 0.004),
    transition_stress  = abs(rnorm(mk, 0.003, 0.001)),
    residence_pressure = as.numeric(rpois(mk, 3L)),
    ruin_proxy         = rbeta(mk, 2, 10)
  ),
  # Stress: fat-tailed t(3)
  cbind(
    log_return         = rt(mk, df = 3L) * 0.012,
    rolling_volatility = abs(rnorm(mk, 0.022,  0.005)),
    drawdown           = rnorm(mk, -0.030, 0.010),
    transition_stress  = abs(rnorm(mk, 0.015, 0.005)),
    residence_pressure = as.numeric(rpois(mk, 12L)),
    ruin_proxy         = rbeta(mk, 5, 3)
  )
)

y_train <- factor(
  rep(c("Calm", "Steady", "Stress"), each = mk),
  levels = c("Calm", "Steady", "Stress")
)

# ---- 2. Stress-concentrated test set -----------------------------------------
#
# 60 Stress observations drawn from t(3) + 20 Calm + 20 Steady.
# This mimics a crisis window where the model must identify Stress correctly.

set.seed(99L)
n_stress <- 60L
n_calm   <- 20L
n_steady <- 20L

X_test <- rbind(
  # Stress: same generative model as training Stress
  cbind(
    log_return         = rt(n_stress, df = 3L) * 0.012,
    rolling_volatility = abs(rnorm(n_stress, 0.022, 0.005)),
    drawdown           = rnorm(n_stress, -0.030, 0.010),
    transition_stress  = abs(rnorm(n_stress, 0.015, 0.005)),
    residence_pressure = as.numeric(rpois(n_stress, 12L)),
    ruin_proxy         = rbeta(n_stress, 5, 3)
  ),
  # Calm
  cbind(
    log_return         = rnorm(n_calm, 0.0002, 0.003),
    rolling_volatility = abs(rnorm(n_calm, 0.004, 0.001)),
    drawdown           = rnorm(n_calm, -0.002, 0.002),
    transition_stress  = abs(rnorm(n_calm, 0.001, 0.0005)),
    residence_pressure = as.numeric(rpois(n_calm, 1L)),
    ruin_proxy         = rbeta(n_calm, 1, 20)
  ),
  # Steady
  cbind(
    log_return         = rnorm(n_steady, 0.0005, 0.005),
    rolling_volatility = abs(rnorm(n_steady, 0.008, 0.002)),
    drawdown           = rnorm(n_steady, -0.008, 0.004),
    transition_stress  = abs(rnorm(n_steady, 0.003, 0.001)),
    residence_pressure = as.numeric(rpois(n_steady, 3L)),
    ruin_proxy         = rbeta(n_steady, 2, 10)
  )
)

y_test <- factor(
  c(rep("Stress", n_stress), rep("Calm", n_calm), rep("Steady", n_steady)),
  levels = c("Calm", "Steady", "Stress")
)

# ---- 3. Fit both classifiers -------------------------------------------------

st_model  <- student_t_naive_bayes(X_train, y_train)
gnb_model <- naive_bayes(X_train, y_train, usekernel = FALSE)

# ---- 4. Overall accuracy and Stress recall -----------------------------------

st_pred  <- predict(st_model,  newdata = X_test, type = "class")
gnb_pred <- predict(gnb_model, newdata = X_test, type = "class")

report <- function(pred, label) {
  acc    <- mean(pred == y_test)
  recall <- mean(pred[y_test == "Stress"] == "Stress")
  fpr    <- mean(pred[y_test != "Stress"] == "Stress")
  cat(sprintf("\n%s\n", label))
  cat(sprintf("  Overall accuracy : %.1f%%\n",   100 * acc))
  cat(sprintf("  Stress recall    : %.1f%%\n",   100 * recall))
  cat(sprintf("  False alarm rate : %.1f%%\n\n", 100 * fpr))
  cat("  Confusion matrix:\n")
  print(table(Actual = y_test, Predicted = pred))
}

report(st_pred,  "=== Student-t NBC (kronxNBC) ===")
report(gnb_pred, "=== Gaussian NBC  (naivebayes) ===")

# ---- 5. Tail log-density comparison -----------------------------------------
#
# For a single extreme log_return observation we show that the Gaussian model
# assigns a much more negative log-density, which drives the MAP decision away
# from Stress and toward whatever class is least penalised.

cat("\n=== Tail log-density at extreme log_return ===\n\n")

# 95th percentile of |log_return| in Stress training data
x_extreme <- quantile(
  abs(X_train[y_train == "Stress", "log_return"]),
  probs = 0.95
)
cat(sprintf("Extreme |log_return| value: %.5f\n\n", x_extreme))

for (regime in c("Calm", "Steady", "Stress")) {
  mu_st <- st_model$params$mu[regime, "log_return"]
  sd_st <- st_model$params$sd[regime, "log_return"]
  nu_st <- st_model$params$nu[regime, "log_return"]

  ld_st  <- stats::dt((x_extreme - mu_st) / sd_st, df = nu_st, log = TRUE) - log(sd_st)
  ld_gn  <- stats::dnorm(x_extreme, mean = mu_st, sd = sd_st, log = TRUE)

  cat(sprintf(
    "  %-7s  nu=%3g  | Student-t log-f: %+8.2f  |  Gaussian log-f: %+8.2f  |  diff: %+6.2f\n",
    regime, nu_st, ld_st, ld_gn, ld_st - ld_gn
  ))
}

cat("\nPositive diff = Student-t assigns HIGHER log-density.\n")
cat("Summed over 6 features, this compounds to a large posterior shift\n")
cat("that keeps Stress observations correctly classified under Student-t NBC.\n")

# ---- 6. Posterior probability for extreme Stress observations ---------------

cat("\n=== Posterior P(Stress) for 10 most extreme log_return observations ===\n\n")

stress_rows    <- which(y_test == "Stress")
extreme_stress <- stress_rows[order(abs(X_test[stress_rows, "log_return"]),
                                    decreasing = TRUE)[1:10]]

st_prob  <- predict(st_model,  newdata = X_test, type = "prob")
gnb_prob <- predict(gnb_model, newdata = X_test, type = "prob")

out_df <- data.frame(
  log_return    = round(X_test[extreme_stress, "log_return"], 5),
  St_P_Stress   = round(st_prob[extreme_stress,  "Stress"], 4),
  GNB_P_Stress  = round(gnb_prob[extreme_stress, "Stress"], 4),
  St_pred       = as.character(st_pred[extreme_stress]),
  GNB_pred      = as.character(gnb_pred[extreme_stress])
)
rownames(out_df) <- NULL
print(out_df)

# ---- 7. nu values selected by Student-t NBC ----------------------------------

cat("\n=== Degrees-of-freedom selected per (class, feature) ===\n\n")
print(st_model$params$nu)

cat("\nNote: low nu on Stress features reflects the t(3) generative model.\n")
cat("Calm/Steady features receive higher nu (closer to Gaussian).\n")
cat("This asymmetric tail-thickness is the core classification advantage\n")
cat("of Student-t NBC over Gaussian NBC.\n")
