# =============================================================================
# FPScausal - Full pipeline test: SCALAR outcome
# =============================================================================

library(FPScausal)

cat("============================================================\n")
cat("FPScausal - Scalar Outcome Pipeline Test\n")
cat("============================================================\n\n")

# ---- 1. Simulate data ----
cat("Step 1: Simulate data (n=2000, setting='LL')\n")

dat <- simulate_fps_data(
  n                      = 2000,
  setting                = "LL",
  outcome_type           = "scalar",
  include_functional_cov = TRUE,
  seed                   = 42
)

cat("  Treatment X :", nrow(dat$X), "x", ncol(dat$X), "\n")
cat("  Outcome Y   : length", length(dat$Y), "\n")
cat("  Confounders :", nrow(dat$C), "x", ncol(dat$C), "\n")
cat("  Functional D:", nrow(dat$D), "x", ncol(dat$D), "\n\n")

# ---- 2. Estimate weights ----
cat("Step 2: Estimate FPS weights\n")

w_obj <- fps_weighting(
  treatment   = dat$X,
  t_grid      = dat$t_grid,
  domain      = c(0, 1),
  domain_name = "t",
  pve         = 0.95,
  covariates  = list(
    scalar     = dat$C,
    functional = list(dat$D)
  ),
  cov_t_grids = list(dat$t_grid),
  cov_domains = list(c(0, 1))
)

print(w_obj)

cat("\n--- Diagnostic plots ---\n")
print(plot(w_obj, type = "weights"))
print(plot(w_obj, type = "balance"))
print(plot(w_obj, type = "fpca_treatment"))
cov_plots <- plot(w_obj, type = "fpca_covariates")
if (!is.null(cov_plots)) print(cov_plots[[1]])

# ---- 3. Estimate effect (no bootstrap) ----
cat("\nStep 3: Estimate causal effect (no bootstrap)\n")

eff <- fps_effect_estimation(
  outcome    = dat$Y,
  fps_object = w_obj,
  true_beta  = dat$true_beta
)
print(eff)
summary(eff)
# Comparison plot: uses analytical CI since bootstrap was not run
print(plot(eff, type = "comparison"))

# ---- 4. Estimate effect with bootstrap ----
cat("\nStep 4: Estimate causal effect (bootstrap B=300)\n")

eff_boot <- fps_effect_estimation(
  outcome    = dat$Y,
  fps_object = w_obj,
  bootstrap  = TRUE,
  B          = 300,
  alpha      = 0.05,
  true_beta  = dat$true_beta,
  seed       = 123
)

print(plot(eff_boot, type = "effect"))
print(plot(eff_boot, type = "comparison"))
print(plot(eff_boot, type = "significance"))
summary(eff_boot)

# ---- 5. Binary outcome ----
cat("\nStep 5: Binary outcome (linear probability model)\n")

Y_bin   <- as.integer(dat$Y > median(dat$Y))
eff_bin <- fps_effect_estimation(Y_bin, w_obj)
print(eff_bin)
print(plot(eff_bin, type = "effect"))
print(plot(eff_bin, type = "comparison"))

# ---- 6. All four settings with combined plot ----
cat("\nStep 6: All four simulation settings\n")

settings <- c("LL", "LN", "NL", "NN")

beta_results <- lapply(settings, function(s) {
  cat(sprintf("  Setting %s ... ", s))
  d   <- simulate_fps_data(2000, setting = s, outcome_type = "scalar",
                            include_functional_cov = TRUE, seed = 1)
  w   <- fps_weighting(d$X, d$t_grid, c(0, 1),
                        covariates  = list(scalar = d$C,
                                           functional = list(d$D)),
                        cov_t_grids = list(d$t_grid),
                        cov_domains = list(c(0, 1)))
  eff <- fps_effect_estimation(d$Y, w, true_beta = d$true_beta)
  m_w <- FPScausal:::.compute_error_metrics(eff$beta,             d$true_beta)
  m_u <- FPScausal:::.compute_error_metrics(eff$beta_unweighted,  d$true_beta)
  cat(sprintf("ISE_w=%.4f  ISE_u=%.4f\n", m_w["ISE"], m_u["ISE"]))
  list(
    t_grid    = d$t_grid,
    weighted  = eff$beta,
    unweighted= eff$beta_unweighted,
    true_beta = d$true_beta,
    setting   = s
  )
})

# Combined plot: one panel per setting
df_all <- do.call(rbind, lapply(beta_results, function(r) {
  tg <- r$t_grid
  rbind(
    data.frame(t = tg, value = r$true_beta,   curve = "True",       setting = r$setting),
    data.frame(t = tg, value = r$weighted,     curve = "Weighted",   setting = r$setting),
    data.frame(t = tg, value = r$unweighted,   curve = "Unweighted", setting = r$setting)
  )
}))
df_all$curve   <- factor(df_all$curve,   levels = c("True", "Weighted", "Unweighted"))
df_all$setting <- factor(df_all$setting, levels = settings)

p_settings <- ggplot(df_all, aes(x = t, y = value,
                                  colour   = curve,
                                  linetype = curve)) +
  geom_line(linewidth = 0.85) +
  geom_hline(yintercept = 0, linetype = "dotted", colour = "grey60") +
  facet_wrap(~ setting, ncol = 2,
             labeller = labeller(setting = c(
               LL = "LL: linear-linear",
               LN = "LN: linear-nonlinear",
               NL = "NL: nonlinear-linear",
               NN = "NN: nonlinear-nonlinear"
             ))) +
  scale_colour_manual(
    values   = c("True" = "black", "Weighted" = "#D62728", "Unweighted" = "#4393C3"),
    name     = NULL
  ) +
  scale_linetype_manual(
    values   = c("True" = "dashed", "Weighted" = "solid", "Unweighted" = "solid"),
    name     = NULL
  ) +
  labs(x = "t", y = expression(beta(t)),
       title = "Causal effect across simulation settings") +
  theme_bw() +
  theme(
    legend.position = "top",
    plot.title      = element_text(hjust = 0.5, face = "bold"),
    strip.text      = element_text(face = "bold")
  )
print(p_settings)

cat("\n============================================================\n")
cat("Scalar outcome test PASSED.\n")
cat("============================================================\n")
