---
title: "Bayesian Inference with SVGD"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{Bayesian Inference with SVGD}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

```{r setup, include = FALSE}
knitr::opts_chunk$set(
  collapse = TRUE,
  comment = "#>",
  eval = FALSE
)
```

## Introduction

This vignette demonstrates Bayesian inference for phase-type distribution parameters using Stein Variational Gradient Descent (SVGD). SVGD is a powerful method for approximating posterior distributions without MCMC.

## The Problem

Given observed waiting times from a phase-type distribution, we want to infer the underlying parameters. For example, in the coalescent model, we observe coalescence times and want to estimate the effective population size (theta).

## Example: Coalescent Model

### Step 1: Define the Model

```{r}
library(phasic)

# Coalescent callback with parameterized rate
coalescent_callback <- function(state) {
  n <- state[1]
  if (n <= 1) {
    return(list())
  }

  # Rate = theta * n(n-1)/2
  # We parameterize as: rate = coef * theta
  coef <- n * (n - 1) / 2

  # Return: (next_state, base_weight=0, coefficients=[coef])
  list(list(c(n - 1), 0.0, c(coef)))
}

# Build parameterized graph
graph <- create_graph(
  callback = coalescent_callback,
  parameterized = TRUE,
  state_length = 1,
  nr_samples = 5  # 5 lineages
)

cat("Graph has", graph$vertices_length(), "vertices\n")
```

### Step 2: Generate Synthetic Data

```{r}
# True parameter value
true_theta <- 2.5

# Instantiate graph with true parameter
trace <- record_elimination_trace(graph, param_length = 1)
true_graph <- instantiate_from_trace(trace, theta = c(true_theta))

# Sample waiting times from the distribution
set.seed(42)
# In practice, you would use graph$sample() or similar
# For now, we'll simulate some realistic data
observed_times <- c(1.2, 2.3, 0.8, 1.9, 1.5, 2.1, 1.1, 1.7, 2.5, 1.3)

cat("Generated", length(observed_times), "observations\n")
cat("True theta:", true_theta, "\n")
```

### Step 3: Run SVGD Inference

```{r}
# Run SVGD using the high-level interface
results <- run_svgd(
  graph = graph,
  observed_data = observed_times,
  theta_dim = 1,              # One parameter (theta)
  n_particles = 100,          # Number of particles
  n_iterations = 1000,        # SVGD iterations
  learning_rate = 0.01,       # Step size
  discrete = FALSE            # Continuous distribution
)

# Extract posterior summaries
cat("\nPosterior Results:\n")
cat("Estimated theta:", results$theta_mean, "\n")
cat("Posterior std:", results$theta_std, "\n")
cat("True theta:", true_theta, "\n")
```

### Step 4: Visualize Results

```{r}
# Plot posterior particles
if (require(ggplot2)) {
  particles_df <- data.frame(theta = results$particles[, 1])

  ggplot(particles_df, aes(x = theta)) +
    geom_histogram(aes(y = after_stat(density)), bins = 30,
                   fill = "lightblue", color = "black") +
    geom_vline(xintercept = true_theta, color = "red", linetype = "dashed",
               linewidth = 1) +
    geom_vline(xintercept = results$theta_mean, color = "blue",
               linewidth = 1) +
    labs(title = "Posterior Distribution of Theta",
         x = "Theta", y = "Density",
         subtitle = sprintf("True: %.2f (red), Estimated: %.2f (blue)",
                           true_theta, results$theta_mean)) +
    theme_minimal()
}
```

## Using Trace-Based Log-Likelihood

For more control, you can create a log-likelihood function directly:

```{r}
# Record trace
trace <- record_elimination_trace(graph, param_length = 1)

# Create log-likelihood function
log_lik <- trace_to_log_likelihood(
  trace,
  observed_times,
  reward_vector = NULL,
  granularity = 100  # Uniformization granularity
)

# Create SVGD object
svgd <- SVGD$new(
  model = log_lik,
  observed_data = observed_times,
  theta_dim = 1,
  n_particles = 100,
  n_iterations = 1000,
  learning_rate = 0.01
)

# Run optimization
results <- svgd$optimize()

cat("Posterior mean:", results$theta_mean, "\n")
cat("Posterior std:", results$theta_std, "\n")
```

## Multi-Parameter Inference

For models with multiple parameters:

```{r}
# Example: Two-parameter model
# (You would define a callback with 2 parameters)

multi_param_callback <- function(state) {
  n <- state[1]
  if (n <= 1) return(list())

  # Rate = theta1 * coef1 + theta2 * coef2
  coef1 <- n * (n - 1) / 2
  coef2 <- n

  list(list(c(n - 1), 0.0, c(coef1, coef2)))
}

graph2 <- create_graph(
  callback = multi_param_callback,
  parameterized = TRUE,
  state_length = 1,
  nr_samples = 5
)

# Run SVGD with 2D parameter space
results2 <- run_svgd(
  graph = graph2,
  observed_data = observed_times,
  theta_dim = 2,  # Two parameters
  n_particles = 100,
  n_iterations = 1000
)

cat("Theta1:", results2$theta_mean[1], "±", results2$theta_std[1], "\n")
cat("Theta2:", results2$theta_mean[2], "±", results2$theta_std[2], "\n")
```

## Multivariate Models with Rewards

For multivariate phase-type distributions using reward transformations:

```{r}
# Create multivariate model
model_mv <- create_multivariate_model(
  graph,
  nr_moments = 2,
  discrete = FALSE
)

# Setup 2D rewards (n_vertices x n_features)
n_vertices <- graph$vertices_length()
rewards_2d <- matrix(
  c(rep(1, n_vertices),    # Feature 1: constant reward
    seq(1, n_vertices)),   # Feature 2: increasing reward
  ncol = 2
)

# Multivariate observations (n_obs x n_features)
observed_2d <- matrix(
  c(observed_times,
    observed_times * 2),  # Second feature scaled
  ncol = 2
)

# Run SVGD
svgd_mv <- SVGD$new(
  model = model_mv,
  observed_data = observed_2d,
  theta_dim = 1,
  n_particles = 100,
  n_iterations = 1000,
  rewards = rewards_2d
)

results_mv <- svgd_mv$optimize()
```

## Tips for SVGD

1. **Number of particles**: Start with 50-100, increase if needed
2. **Learning rate**: 0.001-0.1 typically works; use adaptive methods if available
3. **Iterations**: Monitor convergence; 500-2000 often sufficient
4. **Granularity**: Higher values (200-500) for more accurate PDF, but slower

## Monitoring Convergence

```{r}
# Check if particles have converged
if (!is.null(results$convergence_info)) {
  print(results$convergence_info)
}

# Plot particle evolution (if tracked)
# This would require modifying SVGD to track history
```

## References

- Røikjer, Hobolth & Munch (2022). Phase-type distributions in population genetics. *Statistics and Computing*, 32(5), 1-21.
- Liu & Wang (2016). Stein variational gradient descent. *NeurIPS*.
