# Clean environment and load libraries
rm(list = ls())  # Clear environment to avoid conflicts

# Source files
library(MixStable)

# Set parameters for reproducible results
set.seed(123)

# Define true mixture parameters
true_params <- list(
  component1 = list(alpha = 1.2, beta = 0.1, gamma = 1, delta = 0),
  component2 = list(alpha = 1.8, beta = -0.2, gamma = 1.5, delta = 2),
  weights = c(0.5, 0.5)
)

cat("=== Generating Mixture Data ===\n")
cat("Component 1: alpha =", true_params$component1$alpha, ", beta =", true_params$component1$beta,
    ", gamma =", true_params$component1$gamma, ", delta =", true_params$component1$delta, "\n")
cat("Component 2: alpha =", true_params$component2$alpha, ", beta =", true_params$component2$beta,
    ", gamma =", true_params$component2$gamma, ", delta =", true_params$component2$delta, "\n")
cat("Weights:", true_params$weights, "\n\n")

# Generate mixture data
n_samples <- 500
data1 <- stabledist::rstable(n_samples,
                             alpha = true_params$component1$alpha,
                             beta = true_params$component1$beta,
                             gamma = true_params$component1$gamma,
                             delta = true_params$component1$delta,
                             pm = 1)

data2 <- stabledist::rstable(n_samples,
                             alpha = true_params$component2$alpha,
                             beta = true_params$component2$beta,
                             gamma = true_params$component2$gamma,
                             delta = true_params$component2$delta,
                             pm = 1)

mixture_data <- c(data1, data2)

# Add some basic statistics about the generated data
cat("Data Statistics:\n")
cat("Total samples:", length(mixture_data), "\n")
cat("Range: [", round(min(mixture_data), 3), ",", round(max(mixture_data), 3), "]\n")
cat("Mean:", round(mean(mixture_data), 3), "\n")
cat("Std Dev:", round(sd(mixture_data), 3), "\n\n")

# Test both EM implementations with error handling
cat("=== Running EM Algorithms ===\n")

# Basic EM
cat("1. Running Basic EM Algorithm...\n")
result1 <- NULL
tryCatch({
  result1 <- em_alpha_stable(mixture_data, n_components = 2, debug = TRUE)
  cat("✓ Basic EM completed successfully\n")
}, error = function(e) {
  cat("✗ Basic EM failed with error:", e$message, "\n")
}, warning = function(w) {
  cat("⚠ Basic EM warning:", w$message, "\n")
})

cat("\n")

# EM with MLE fitting
cat("2. Running EM with MLE Fitting...\n")
result2 <- NULL
tryCatch({
  result2 <- em_fit_alpha_stable_mixture(mixture_data)
  cat("✓ EM with MLE completed successfully\n")
}, error = function(e) {
  cat("✗ EM with MLE failed with error:", e$message, "\n")
}, warning = function(w) {
  cat("⚠ EM with MLE warning:", w$message, "\n")
})

cat("\n=== Results Summary ===\n")

# Print results with better formatting
if (!is.null(result1)) {
  cat("\n--- Basic EM Results ---\n")
  if (is.list(result1)) {
    print(result1)

    # Additional analysis if structure is known
    if ("converged" %in% names(result1)) {
      cat("Convergence status:", result1$converged, "\n")
    }
    if ("iterations" %in% names(result1)) {
      cat("Number of iterations:", result1$iterations, "\n")
    }
    if ("log_likelihood" %in% names(result1)) {
      cat("Final log-likelihood:", round(result1$log_likelihood, 4), "\n")
    }
  } else {
    cat("Result structure:\n")
    str(result1)
  }
} else {
  cat("--- Basic EM Results ---\n")
  cat("Algorithm failed to produce results\n")
}

if (!is.null(result2)) {
  cat("\n--- EM with MLE Results ---\n")
  if (is.list(result2)) {
    print(result2)

    # Additional analysis if structure is known
    if ("converged" %in% names(result2)) {
      cat("Convergence status:", result2$converged, "\n")
    }
    if ("iterations" %in% names(result2)) {
      cat("Number of iterations:", result2$iterations, "\n")
    }
    if ("log_likelihood" %in% names(result2)) {
      cat("Final log-likelihood:", round(result2$log_likelihood, 4), "\n")
    }
  } else {
    cat("Result structure:\n")
    str(result2)
  }
} else {
  cat("--- EM with MLE Results ---\n")
  cat("Algorithm failed to produce results\n")
}

# Compare results if both succeeded
if (!is.null(result1) && !is.null(result2)) {
  cat("\n=== Comparison ===\n")

  # Compare log-likelihoods if available
  ll1 <- if ("log_likelihood" %in% names(result1)) result1$log_likelihood else NA
  ll2 <- if ("log_likelihood" %in% names(result2)) result2$log_likelihood else NA

  if (!is.na(ll1) && !is.na(ll2)) {
    cat("Log-likelihood comparison:\n")
    cat("Basic EM:", round(ll1, 4), "\n")
    cat("EM with MLE:", round(ll2, 4), "\n")
    cat("Difference:", round(ll2 - ll1, 4), "\n")
    if (ll2 > ll1) {
      cat("EM with MLE achieved higher likelihood\n")
    } else if (ll1 > ll2) {
      cat("Basic EM achieved higher likelihood\n")
    } else {
      cat("Both methods achieved similar likelihood\n")
    }
  }
}

# Optional: Create visualization if results are available
create_comparison_plot <- function(data, result1, result2, true_params) {
  tryCatch({
    # This is a template - adjust based on your actual result structure
    par(mfrow = c(2, 2))

    # Plot 1: Histogram of data
    hist(data, breaks = 50, freq = FALSE, main = "Mixture Data",
         xlab = "Value", col = rgb(0.7, 0.7, 0.9, 0.5))

    # Plot 2: True mixture PDF (if you have the helper functions)
    x_range <- seq(min(data), max(data), length.out = 200)
    # This would need your r_stable_pdf function
    # true_pdf <- true_params$weights[1] * r_stable_pdf(x_range, ...) + ...
    # lines(x_range, true_pdf, col = "black", lwd = 2)

    # Plots 3 & 4: Results from each method (structure dependent)
    # You would customize these based on your actual result structure

    par(mfrow = c(1, 1))

  }, error = function(e) {
    cat("Could not create comparison plot:", e$message, "\n")
  })
}

# Uncomment to create plots (if your functions support it)
# if (!is.null(result1) || !is.null(result2)) {
#   create_comparison_plot(mixture_data, result1, result2, true_params)
# }

cat("\n=== Analysis Complete ===\n")