---
title: "Writing Stan code"
output: rmarkdown::html_vignette
vignette: >
  %\VignetteIndexEntry{Writing Stan code}
  %\VignetteEngine{knitr::rmarkdown}
  %\VignetteEncoding{UTF-8}
---

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

The package allows for only limited models as, e.g., neither random slopes, nor interaction effects are allowed. Imposing this restriction was a design decision, as it would require duplicating functionality of general purposes packages. Instead, the package itself provides some basic fitting that should be sufficient for most simple cases. However, below you will find example of how to incorporate cumulative history into a model written in Stan. This way, you can achieve maximal flexibility but still save time by reusing the code.

## Stan model

This is a complete Stan code for a model with log-normal distribution for multiple runs from a single experimental session of a single participant. The history time-constant `tau` is fitted, whereas constants are used for other cumulative history parameters.

```{stan output.var="example_model", eval=FALSE}
data{
    // --- Complete time-series ---
    int<lower=1> rowsN;     // Number of rows in the COMPLETE multi-timeseries table including mixed phase.
    real duration[rowsN];   // Duration of a dominance/transition phase
    int istate[rowsN];      // Index of a dominance istate, 1 and 2 code for two competing clear states, 3 - transition/mixed.
    int is_used[rowsN];     // Whether history value must used to predict duration or ignored
                            // (mixed phases, warm-up period, last, etc.)
    int run_start[rowsN];   // 1 marks a beginning of the new time-series (run/block/etc.)
    real session_tmean[rowsN]; // Mean dominance phase duration for both CLEAR percepts. Used to scale time-constant.
    
    // --- A shorter clear-states only time-series ---
    int clearN;                  // Number of rows in the clear-states only time-series
    real clear_duration[clearN]; // Duration for clear percepts only.
    
    // --- Cumulative history parameters
    real<lower=0, upper=1> history_starting_values[2]; // Starting values for cumulative history at the beginning of the run
    real<lower=0, upper=1> mixed_state;                // Mixed state signal strength
}
parameters {
    real<lower=0> tau; // history time-constant
    
    // linear model for mu
    real a;
    real bH;
    
    // variance
    real<lower=0> sigma;
}
transformed parameters{
    vector[clearN] mu; // vector of computed mu for each clear percept
  
    {
        // temporary variables
        real current_history[2]; // current computed history
        real tau_H;              // tau in the units of time
        real dH;                 // computed history difference
        int iC = 1;              // Index of clear percepts used for fitting

        // matrix with signal levels
        matrix[2, 3] level = [[1, 0, mixed_state], 
                              [0, 1, mixed_state]];

        for(iT in 1:rowsN){
            // new time-series, recompute absolute tau and reset history state
            if (run_start[iT]){
                // reset history
                current_history = history_starting_values;

                // Recompute tau in units of time. 
                // This is relevant only for multiple sessions / participants.
                // However, we left this code for generality.
                tau_H = session_tmean[iT] * tau;
            }

            // for valid percepts, we use history to compute mu
            if (is_used[iT] == 1){
                // history difference
                dH = current_history[3-istate[iT]] - current_history[istate[iT]];

                // linear model for mu
                mu[iC] = a + bH * dH;
                iC += 1;
            }

            // computing history for the NEXT episode
            // see vignette on cumulative history
            for(iState in 1:2){
                current_history[iState] = level[iState, istate[iT]] + 
                  (current_history[iState] - level[iState, istate[iT]]) * exp(-duration[iT] / tau_H);
            }
        }
    }
}
model{
  // sampling individual parameters
  tau ~ lognormal(log(1), 0.75);
  a ~ normal(log(3), 5);
  bH ~ normal(0, 1);
  sigma ~ exponential(1);
  
  // sampling data using computed mu and sampled sigma
  clear_duration ~ lognormal(exp(mu), sigma);
}
```


## Data preparation
The `data` section defines model inputs. Hopefully, the comments make understanding it fairly straightforward. However, it has several features that although are not needed for the limited single session / single session make it easier to generalized the code for more complicated cases. 

For example, not all dominance phases are used for fitting. Specifically, all mixed perception phases, first dominance phase for each percept (not enough time to form reliably history) and last dominance phase (curtailed by the end of the block) are excluded. Valid dominance phases are marked in `is_used` vector. Their total number is stored in `clearN` variable and the actual dominance durations in `clear_duration`. The latter is not strictly necessary but allows us to avoid a loop and vectorize the sampling statement `clear_duration ~ lognormal(exp(mu), sigma);`.

In addition, `session_tmean` is a vector rather than a scalar. This is not necessary for a single session example here but we opted to use as it will better generalize for more complicated cases.

bistability package provides a service function `preprocess_data()` that simplifies the process of preparing the data. However, you need to perform the last step, forming a list of inputs for Stan sampling, yourself.
```{r eval=FALSE}
# function that checks data for internal consistency and returns a preprocessed table
df <- bistablehistory::preprocess_data(br_single_subject, 
                                       state="State",
                                       duration="Duration",
                                       run="Block")

# data for Stan model
stan_data <- list(
  # complete time-series
  rowsN = nrow(df),
  duration = df$duration,
  istate = df$istate,
  is_used = df$is_used,
  run_start = df$run_start,
  session_tmean = df$session_tmean,
  
  # only valid clear percepts
  clearN = sum(df$is_used),
  clear_duration = df$duration[df$is_used == 1],
  
  # history parameters, all fixed to default values
  history_starting_values = c(0, 0),
  mixed_state = 0.5
)
```

## Using the model
You can use this model either with `rstan` or `cmdstanr` packages. Below is in an example using `cmdstanr`, assuming that model file is called `example.stan`. 
```{r eval=FALSE}
# compile the model
model <- cmdstanr::cmdstan_model("example.stan")

# sample model
fit <- model$sample(data=stan_data, chains=1)

# extract posterior samples for tau parameter
tau <- fit$draws(variables = "tau")
```

