---
title: "A Simple Race Model"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{A Simple Race Model}
  %\VignetteEncoding{UTF-8}
  %\VignetteEngine{knitr::rmarkdown}
editor_options: 
  markdown: 
    wrap: 72
---

This vignette shows the basic simulate-and-fit workflow for a race model. The
example has two accumulators feeding response `R1` and one accumulator feeding
response `R2`.

```{r}
library(AccumulatR)
```

## Define the model

We use three accumulators. `R1_A` and `R1_B` both feed response `R1`, while
`R2` feeds response `R2`. Parameters are grouped by default; the
`set_parameters()` call gives the two drift parameters separate values and
shares the threshold gaps, starting-point ranges, and numerical values of
the LBA drift spread and RDM diffusion scale.

```{r}
model <- race_spec() |>
  add_accumulator("R1_A", "LBA") |>
  add_accumulator("R1_B", "RDM") |>
  add_accumulator("R2", "lognormal") |>
  add_pool("R1", c("R1_A", "R1_B")) |>
  add_outcome("R1", "R1") |>
  add_outcome("R2", "R2") |>
  set_parameters(
    separate = list(v = c("R1_A", "R1_B"), m = TRUE, lognormal.s = TRUE),
    share = list(
      B_shared = c("R1_A.B", "R1_B.B"),
      A_shared = c("R1_A.A", "R1_B.A"),
      noise_shared = c("R1_A.sv", "R1_B.s")
    ),
    rename = c(t0 = "t0_shared")
  ) |>
  finalize_model()

true_params <- c(
  R1_A.v = 2,
  R1_B.v = 3,
  B_shared = 1,
  A_shared = 0.3,
  noise_shared = 1,
  R2.m = log(0.4),
  R2.s = 0.18,
  t0_shared = 0.05
)
```

## Simulate data

We generate response outcomes and response times for each trial.

```{r}
set.seed(123456)

n_trials <- 2000
params_df <- build_param_matrix(model, true_params, n_trials = n_trials)

sim <- simulate(model, params_df)

data_df <- data.frame(
  trials = sim$trials,
  R = factor(sim$R),
  rt = sim$rt,
  stringsAsFactors = FALSE
)

table(data_df$R)
```
## Evaluate the likelihood

We prepare the data, build a model context, and evaluate the log-likelihood at
the true parameter values.

```{r}
prepared <- prepare_data(model, data_df)
ctx <- make_context(model)

params_df_true <- build_param_matrix(
  model,
  true_params,
  n_trials = n_trials
)

ll_true <- as.numeric(log_likelihood(ctx, prepared, params_df_true))
ll_true
```

For comparison, we can evaluate a clearly misspecified parameter set.

```{r}
wrong_params <- true_params
wrong_params["t0_shared"] <- 0.08

params_df_wrong <- build_param_matrix(
  model,
  wrong_params,
  n_trials = n_trials
)

ll_wrong <- as.numeric(log_likelihood(ctx, prepared, params_df_wrong))
ll_wrong
```

## Estimate parameters with `optim()`

We estimate six parameters: `R1_A.v`, `R1_B.v`, shared `B`, `R2.m`, `R2.s`,
and shared `t0`. The RDM drift, `B`, `R2.s`, and `t0` are estimated on the
log scale to keep them positive. The shared `A` and variability parameters
are held fixed in this example.

```{r eval=FALSE}
neg_loglik <- function(theta) {
  est <- true_params
  est["R1_A.v"] <- theta[["R1_A.v"]]
  est["R1_B.v"] <- exp(theta[["log_R1_B.v"]])
  est["B_shared"] <- exp(theta[["log_B_shared"]])
  est["R2.m"] <- theta[["R2.m"]]
  est["R2.s"] <- exp(theta[["log_R2.s"]])
  est["t0_shared"] <- exp(theta[["log_t0_shared"]])
  params_df <- build_param_matrix(
    model,
    est,
    n_trials = n_trials
  )
  ll <- log_likelihood(ctx, prepared, params_df)
  -as.numeric(ll)
}

start <- c(
  R1_A.v = 1.5,
  log_R1_B.v = log(1.5),
  log_B_shared = log(1.0),
  R2.m = log(0.32),
  log_R2.s = log(0.15),
  log_t0_shared = log(0.03)
)

fit <- optim(start, neg_loglik, method = "Nelder-Mead", control = list(maxit = 4000, reltol = 1e-9))

fit_params <- c(
  R1_A.v = fit$par[["R1_A.v"]],
  R1_B.v = exp(fit$par[["log_R1_B.v"]]),
  B_shared = exp(fit$par[["log_B_shared"]]),
  R2.m = fit$par[["R2.m"]],
  R2.s = exp(fit$par[["log_R2.s"]]),
  t0_shared = exp(fit$par[["log_t0_shared"]])
)
target <- c(
  R1_A.v = true_params[["R1_A.v"]],
  R1_B.v = true_params[["R1_B.v"]],
  B_shared = true_params[["B_shared"]],
  R2.m = true_params[["R2.m"]],
  R2.s = true_params[["R2.s"]],
  t0_shared = true_params[["t0_shared"]]
)
data.frame(true = target, recovered = fit_params, miss = abs(target - fit_params))
```