# ST 453 Final Project - Bootstrap File
# Trace Brown
# 04/21/26

#
# This script fits a logistic regression model to the real heart failure data,
# then draws one pairs/case bootstrap sample (resample (X, y) rows with
# replacement) and refits the model on the bootstrap sample.
#
# Usage: Rscript bootstrap.r <seed>
#   seed  - integer seed for the bootstrap resample
#
# Example: Rscript bootstrap.r 42

# No R packages are permitted for use in this assignment.



# -----------------------------------------------------------------------------
# Command line arguments
# -----------------------------------------------------------------------------

args = commandArgs(trailingOnly=TRUE)

if(length(args) < 1){
  stop("Usage: Rscript bootstrap.r <seed>")
}

seed = as.integer(args[1])

print(paste("Running bootstrap with seed =", seed))



# -----------------------------------------------------------------------------
# Helper functions from previous homework assignments
# -----------------------------------------------------------------------------

# Gaussian elimination function (from HW2)
gaussian_elimination = function( M, b){
  # Function to return row echelon form via Gaussian elimination
  
  p = dim(M)[1]
  aug = cbind( M, b)
  
  for(k in 1:(p-1)){
    
    # Partial pivoting: find row with largest value in column k
    max_idx = k
    for(i in (k+1):p){
      if(abs(aug[i,k]) > abs(aug[max_idx,k]))  max_idx = i
    }
    
    # Swap rows if needed
    if(max_idx != k){
      temp = aug[k,]
      aug[k,] = aug[max_idx,]
      aug[max_idx,] = temp
    }
    
    # Skip if pivot is zero
    if(abs(aug[k,k]) < 1e-12)  next
    
    # Eliminate entries below pivot
    for(i in (k+1):p){
      factor = aug[i,k] / aug[k,k]
      aug[i,] = aug[i,] - factor * aug[k,]
    }
  }
  
  M_tilde = aug[,1:p]
  b_tilde = aug[,p+1]
  
  return(list( M=M_tilde, b=b_tilde))
}


# Back substitution function (from HW2)
back_substitution = function( M, b){
  # Function to solve Mx = b where M is upper triangular
  
  p = dim(M)[1]
  x = rep( NA, p)
  
  x[p] = b[p] / M[p,p]
  for(i in (p-1):1){
    
    x[i] = b[i] / M[i,i]
    for(j in (i+1):p)  x[i] = x[i] - M[i,j] * x[j] / M[i,i]
  }
  
  return(x)
}


# Function to solve linear system Ax = b using Gaussian elimination
solve_system = function( A, b){
  row_echelon = gaussian_elimination( A, b)
  x = back_substitution( row_echelon$M, row_echelon$b)
  return(x)
}



# -----------------------------------------------------------------------------
# Newton-Raphson for Logistic Regression (same as midterm)
# -----------------------------------------------------------------------------

# Logistic function
logistic = function( z){
  return( 1 / (1 + exp(-z)))
}


# Newton-Raphson algorithm for logistic regression MLE
newton_raphson_logistic = function( y, X, max_iter=200, tol=1e-8){
  # Function to estimate logistic regression coefficients via Newton-Raphson.
  # Returns: beta (coefficients), cov_mat (inverse Fisher information),
  #          se (standard errors), iter, converged
  
  n = dim(X)[1]
  p = dim(X)[2]
  
  # Initialize beta at zero
  beta = rep( 0, p)
  
  for(iter in 1:max_iter){
    
    # Compute probabilities
    eta = X %*% beta
    prob = logistic(eta)
    
    # Avoid numerical issues at boundaries
    prob[prob < 1e-10] = 1e-10
    prob[prob > 1 - 1e-10] = 1 - 1e-10
    
    # Gradient: X'(y - p)
    gradient = t(X) %*% (y - prob)
    
    # Hessian: -X' W X, where W = diag(p * (1-p))
    W = as.vector(prob * (1 - prob))
    XtWX = t(X) %*% (W * X)
    
    # Newton-Raphson update: beta = beta + (X'WX)^{-1} * gradient
    delta = solve_system( XtWX, as.vector(gradient))
    beta_new = beta + delta
    
    # Check convergence
    if(sqrt(sum((beta_new - beta)^2)) < tol){
      
      # Covariance matrix (inverse of Fisher information at MLE)
      prob_final = logistic(X %*% beta_new)
      prob_final[prob_final < 1e-10] = 1e-10
      prob_final[prob_final > 1 - 1e-10] = 1 - 1e-10
      W_final = as.vector(prob_final * (1 - prob_final))
      XtWX_final = t(X) %*% (W_final * X)
      
      cov_mat = matrix( 0, nrow=p, ncol=p)
      for(j in 1:p){
        e_j = rep( 0, p)
        e_j[j] = 1
        cov_mat[,j] = solve_system( XtWX_final, e_j)
      }
      
      se = sqrt(diag(cov_mat))
      
      return(list( beta=beta_new, cov_mat=cov_mat, se=se,
                   iter=iter, converged=TRUE))
    }
    
    beta = beta_new
  }
  
  # Did not converge; still return covariance estimate at final iterate
  prob_final = logistic(X %*% beta)
  prob_final[prob_final < 1e-10] = 1e-10
  prob_final[prob_final > 1 - 1e-10] = 1 - 1e-10
  W_final = as.vector(prob_final * (1 - prob_final))
  XtWX_final = t(X) %*% (W_final * X)
  
  cov_mat = matrix( 0, nrow=p, ncol=p)
  for(j in 1:p){
    e_j = rep( 0, p)
    e_j[j] = 1
    cov_mat[,j] = solve_system( XtWX_final, e_j)
  }
  
  se = sqrt(diag(cov_mat))
  
  return(list( beta=beta, cov_mat=cov_mat, se=se,
               iter=max_iter, converged=FALSE))
}



# -----------------------------------------------------------------------------
# Load and prepare real data
# -----------------------------------------------------------------------------

print("Loading real data...")

dat = read.csv("heart_failure_clinical_records_dataset.csv")

# Construct design matrix with intercept and selected predictors
X_real = cbind( 1, 
                dat$age, 
                dat$ejection_fraction, 
                dat$serum_creatinine, 
                dat$time)
y_real = dat$DEATH_EVENT

n = dim(X_real)[1]
p = dim(X_real)[2]

print(paste("Sample size: n =", n))
print(paste("Number of parameters: p =", p))

# Standardize using ORIGINAL data's means and SDs.
# This is held fixed across bootstrap samples so estimated coefficients
# are on a consistent scale (and directly comparable to the midterm fit).
X_means = colMeans(X_real[, 2:p])
X_sds = apply(X_real[, 2:p], 2, sd)

X_scaled = X_real
for(j in 2:p){
  X_scaled[,j] = (X_real[,j] - X_means[j-1]) / X_sds[j-1]
}



# -----------------------------------------------------------------------------
# Fit model to real data: this is the point estimate to bootstrap around
# -----------------------------------------------------------------------------

print("Fitting model to real data...")

fit_real = newton_raphson_logistic( y_real, X_scaled)
beta_hat_real = fit_real$beta
se_real       = fit_real$se

print(paste("Real-data fit converged:", fit_real$converged, 
            "in", fit_real$iter, "iterations"))
print("Real-data beta estimates:")
print(round(beta_hat_real, 6))



# -----------------------------------------------------------------------------
# Pairs/case bootstrap: resample (X, y) rows and refit
# -----------------------------------------------------------------------------

set.seed(seed)

# Resample row indices with replacement (size = n, matches original sample)
boot_idx = sample( 1:n, size=n, replace=TRUE)
X_boot   = X_scaled[boot_idx, , drop=FALSE]
y_boot   = y_real[boot_idx]

# Refit the logistic regression on the bootstrap sample
fit_boot = newton_raphson_logistic( y_boot, X_boot)

print(paste("Seed:", seed,
            "| Converged:", fit_boot$converged,
            "| Iterations:", fit_boot$iter,
            "| Boot events:", sum(y_boot)))



# -----------------------------------------------------------------------------
# Save results for this bootstrap iteration
# -----------------------------------------------------------------------------

result = list(
  seed          = seed,
  beta_hat_real = beta_hat_real,
  se_real       = se_real,
  beta_boot     = fit_boot$beta,
  se_boot       = fit_boot$se,
  converged     = fit_boot$converged,
  iter          = fit_boot$iter,
  boot_events   = sum(y_boot),
  n             = n,
  p             = p,
  X_means       = X_means,
  X_sds         = X_sds
)

out_file = paste0("bootstrap_results/boot_seed_", seed, ".rds")
saveRDS(result, file=out_file)

print(paste("Results saved to", out_file))