stay_probability_091724.Rmd

---
title: "stay_probability_091724"
author: "Daniel Petrie"
date: "2024-09-17"
output: html_document
editor_options: 
  chunk_output_type: console
---

# Global settings

```{r Global}
library(ggplot2)
library(tidyverse)
library(dplyr)
library(psych)
library(gridExtra)
library(ggpubr)
library(lme4)
library(lmerTest)
#Working directory
setwd("C:/Users/djpet/OneDrive/Documents/daw_resting_state")

#Full data so I can extract the correct subjects
full <- read.csv("daw_project_081624.csv", header = TRUE)
#Trial level Daw data. I do not need to run the logistic regression models. I only need to run the first part of the cleaning code to calculate the probability of first stage stay on average (i think).
luna_7t <- read.csv("cleaning/luna_7tbrainmech.csv", header = TRUE)
luna_pet <- read.csv("cleaning/luna_pet.csv", header = TRUE)
```


# Overview

This script does the following:

1) Multilevel logistic regression model with age for a table.
    - Replication of analyses from Decker et al., (2016).

2) Marginal Effects plots from logistic regression.

2) Behavioral analysis of the Daw task.

3) Simulated data depicting MB and MF learning strategies.

4) Gaussian random walks plot for second-stage choice.


# Pre-processing Data

```{r Cleaning}
#Creating a different sex variable for some plots/analyses
full <- full %>%
  mutate(sex = as.character(sex),
    sex_p = case_when(
    sex == "M" ~ "Male",
    sex == "F" ~ "Female",
    TRUE ~ sex 
  ))

#Making visitnum a factor. Can change accordingly
full$visitnum <- as.factor(full$visitnum)

#Making sex a factor
full$sex <- as.factor(full$sex)

#Making an ordered factor.
full$sex_p <- ordered(full$sex_p, levels = (c("Male", "Female")))

#Inverse age for linear models.
full$inv_age <- 1/full$age

#Checking for duplicate rows due to weird merging situation I put myself in.
full[which(duplicated(full$modelbased)),c(1:3, 8, 20)] #11475_2/11498_2/11589_2 are duplicates.
#These ids have duplicated Daw visits, but different rest days. Taking rows with imaging data.

full <- full %>%
  filter(!(id == "11475" & visitnum == "2") & 
           !(id == "11498" & visitnum == "2") &
           !(id == "11589" & visitnum == "2"))

#Treating visitnum as numeric?
full$visitnum_numeric <- as.numeric(full$visitnum)

#Treating id as factor for og gam fitting
full$id_fac <- as.factor(full$id)

#Age groups for plotting
full <- full %>%
  mutate(age_cat = cut(age,
                       breaks = c(10,13,17,24,34),
                       labels = c("10-13 years", "14-17 years", "18-24 years", "25-34 years"),
                       include.lowest = TRUE))

#Separating id column into "id" and "visit"
luna_7t <- luna_7t %>%
  separate(id, c("id", "vdate"))

luna_pet <- luna_pet %>%
  separate(id, c("id", "vdate"))

#Adding trial number
luna_7t <- luna_7t %>%
  group_by(id, vdate) %>%
  mutate(trial = row_number())

luna_pet <- luna_pet %>%
  group_by(id, vdate) %>%
  mutate(trial = row_number())

#Removing id with identical data that was copied across studies.
luna_pet <- luna_pet %>% group_by(id) %>% filter(id != "11565")

#Adding scanner information. Ashley slack stated that Siemens for 7T, Siemens Biograph mMR PET/MRI for PET. Seems like they only need a binary code situation here, but I will fuss with that during the neuroCombat portion of the code. This will also be useful for seeing which study came from what subje t.
luna_7t$study <- "7t"
luna_pet$study <- "PET"

#Similar id's across data sets. Check to make sure it looks correct with visit_info data set.
same_ids <- intersect(luna_7t$id, luna_pet$id)
same_ids
#"10195" ok
#"10997" ok
#"11537" ok
#"11561" ok
#"11565" Not ok. Tested on same day. Fixed above!
#"11575" ok
#"11641" ok
#"11651"
#luna_7t %>% filter(id == "11651")
#luna_pet %>% filter(id == "11651")
#visit_info %>% filter(id == "11651")

#Combining data. ADD HABIT HERE
luna_all <- bind_rows(luna_7t, luna_pet)

# compute visit number
luna_all <- luna_all %>% 
  group_by(id) %>% 
  mutate(visitnum = as.integer(factor(vdate, levels = unique(vdate))))


#Removing trials where subject did not respond during 1st or second stage
luna_all <- luna_all %>%
  filter(!(choice1 == 0 | choice2 == 0))

#Creating lagged variables. Add trial
luna_all <- luna_all %>% 
  group_by(id, visitnum) %>% 
  mutate(choice1lag = lag(choice1),
         choice2lag = lag(choice2),
         statelag = lag(state),
         moneylag = lag(money),
         visitnumlag = lag(visitnum))

# transitional vars...
luna_all <- luna_all %>% 
  mutate(commonrare = as.factor(ifelse((choice1lag == 1 & statelag == 2) |
                                         (choice1lag == 2 & statelag == 3),
                                       'Common', 
                                       'Rare')), 
         commonraredummy = ifelse(commonrare=="Common", 
                                  1, 
                                  -1), 
         moneylagdummy = ifelse(moneylag == 1, 
                                1, 
                                -1), 
         firststagestay = ifelse(choice1 == choice1lag, 
                                 1, 
                                 0), 
         stayswitchwinlose = ifelse(firststagestay==1 & moneylag==0, 'lose-stay',
                             ifelse(firststagestay==1 & moneylag==1, 'win-stay', 
                                    ifelse(firststagestay==0 & moneylag==0,
                                           'lose-switch', 
                                           'win-switch'))))
luna_all <- luna_all %>% mutate(winswitch = ifelse(stayswitchwinlose== "win-switch",
                            1,
                            0))
luna_all <- luna_all %>% mutate(
  winswitch_common = ifelse(commonrare == "Common" & 
                            stayswitchwinlose == "win-switch", 
                            1,
                            0),
  winswitch_rare = ifelse(commonrare == "Rare" &
                            stayswitchwinlose == "win-switch",
                          1,
                          0))

#Selecting columns of interest from full
full_sub <- full %>%
  dplyr::select(id, visitnum, age, study, sex, behave.date)

#Selecting columns of interest from Daw
#I think I will need firststagestay, moneylag, commonrate
luna_all_sub <- luna_all %>%
  dplyr::select(id, visitnum, vdate, trial, moneylag, moneylagdummy, commonrare, commonraredummy, firststagestay) %>%
  mutate(vdate = as.Date(vdate, format = "%Y%m%d")) %>%
  rename(behave.date = vdate,
         visitnum_daw = visitnum)
full_sub$id <- as.character(full_sub$id)

full_sub <- full_sub %>%
  mutate(behave.date = as.Date(behave.date, format = "%Y-%m-%d"))

#Remember to organize by Date. Idea is to make sure that they match up appropiately.
stay_prob <- merge(luna_all_sub, full_sub, by = c("id", "behave.date"), all.x = TRUE)

#Arrange to reorder
stay_prob <- stay_prob %>%
  arrange(id, visitnum_daw, trial)
#Adding dichotomous age
#stay_prob <- stay_prob %>%
#  mutate(age_cat = if_else(age < 18, 0, 1))
age_bins <- cut(stay_prob$age,
                breaks = c(10,13,17,24,34),
                labels = c("10-13 years", "14-17 years", "18-24 years", "25-34 years"),
                include.lowest = TRUE)
stay_prob <- stay_prob %>%
  mutate(age_cat = age_bins)

stay_prob$age_z <- scale(stay_prob$age)
```


# Multilevel logistic regression model with age

Refitting ML LR. Keep in mind that Decker appears to extract MB/MF parameters with a model without age effects regressed out. This analyses is meant to supplement that and be used soley as behavioral analyses.

```{r Logistic Regression Model}
intmodeltoplot1 <- glmer(firststagestay ~ 1 + age_z*commonraredummy*moneylagdummy + 
                          (commonraredummy*moneylagdummy|id:visitnum_daw),
                        data=stay_prob,
                        family="binomial",
                        glmerControl(optimizer = "bobyqa"))
summary(intmodeltoplot1)
```

It seems that there are age effects across the learning parameters. Let's look at marginal effects to get a sense of whether adolescents and adults are performing differently during the task.

# Marginal Effects

## Age Main Effect

```{r Age Marginal Effect}
#Age
age_me <- ggpredict(intmodeltoplot1, terms = c("age_z"))

age_me_fig <- ggplot(data = age_me, 
                     aes(x = x, y = predicted)) + 
  geom_line(linewidth = 0.5) + 
  geom_ribbon(aes(ymin = conf.low, 
                  ymax = conf.high), 
              alpha = 0.2) +  # Add CI shading
  theme_modern() + 
  scale_y_continuous(labels = scales::percent_format()) +
  labs(
    title = "Main effect of age",
    x = "Age (z-scored)",
    y = "Probability of FSS"
  ) +
  theme(
    axis.title.x = element_text(size = 8, margin = margin(t = 3)),
    axis.title.y = element_text(size = 8, margin = margin(r = 3)),
    axis.text.x = element_text(size = 7),
    axis.text.y = element_text(size = 7),
    plot.title = element_text(size = 8, hjust = 0.5),
    axis.line = element_line(size = .2),
    axis.ticks = element_line(size = .2)
    )
age_me_fig


ggsave(filename = "figures/figure_S3/age_me_fig.png", 
       plot = age_me_fig,
       device = "png", 
       dpi = 500, 
       units = "mm",
       width = 60, 
       height = 60)

```

## Transition Main Effect

```{r Transition Marginal Effect}
#Transistion
trans_main1 <- ggpredict(intmodeltoplot1, terms = c("commonraredummy[-1,1]"))

trans_main_fig <- ggplot(trans_main1, 
                         aes(x = as.factor(x), 
                             y = predicted)) + 
  geom_point(size = 0.5) + 
  geom_errorbar(aes(ymin = conf.low, 
                    ymax = conf.high), 
                width = 0.2,
                size = 0.5) + 
  labs(
    title = "Main effect of tranistion type",
    x = "Transition Type",
    y = "Probability of FSS"
  ) +
    scale_y_continuous(
    #limits = c(0.7,0.9),  
    labels = scales::percent
  ) +
  theme_modern() +
  theme(
    plot.title = element_text(size = 8, hjust = 0.5),
    axis.title.x = element_text(size = 8, margin = margin(t = 3)),
    axis.title.y = element_text(size = 8, margin = margin(r = 3)),
    axis.text.x = element_text(size = 7),
    axis.text.y = element_text(size = 7),
    axis.line = element_line(size = .2),
    axis.ticks = element_line(size = .2)
    ) 

trans_main_fig

ggsave(filename = "figures/figure_S3/trans_me_fig.png", 
       plot = trans_main_fig,
       device = "png", 
       dpi = 500, 
       units = "mm",
       width = 60, 
       height = 60)
```

## Reward Main Effect

```{r Reward marginal effect}
#Money
rew_main <- ggpredict(intmodeltoplot1, terms = c("moneylagdummy"))
rew_main_fig <- ggplot(rew_main, 
                         aes(x = as.factor(x), 
                             y = predicted)) + 
  geom_point(size = 0.5) + 
  geom_errorbar(aes(ymin = conf.low, 
                    ymax = conf.high), 
                width = 0.2,
                size = 0.5) + 
  scale_y_continuous(labels = scales::percent) +
  labs(
    title = "Main effect of Reward",
    x = "Reward",
    y = "Probability of FSS"
  ) +
  theme_modern() +
  theme(
    plot.title = element_text(size = 8, hjust = 0.5),
    axis.title.x = element_text(size = 8, margin = margin(t = 3)),
    axis.title.y = element_text(size = 8, margin = margin(r = 3)),
    axis.text.x = element_text(size = 7),
    axis.text.y = element_text(size = 7),
    axis.line = element_line(size = .2),
    axis.ticks = element_line(size = .2)
    ) 
rew_main_fig

ggsave(filename = "figures/figure_S3/rew_me_fig.png", 
       plot = rew_main_fig,
       device = "png", 
       dpi = 500, 
       units = "mm",
       width = 60, 
       height = 60)
```

## Age x transition interaction

```{r Age x transition}
#Age x transition
agetrans_int <- ggpredict(intmodeltoplot1, terms = c("age_z", "commonraredummy"))


agetrans_fig <- ggplot(agetrans_int,
                       aes(x = x,
                           y = predicted,
                           group = group,
                           color = group)) +
  geom_line(linewidth = 0.5, 
            show.legend = FALSE) +
  geom_ribbon(aes(ymin = conf.low, 
                  ymax = conf.high,
                  fill = group), 
              alpha = 0.2,
              color = NA) +  
  scale_y_continuous(labels = scales::percent) +
  scale_fill_discrete(name = "Transition") +
  guides(color = "none") +
  theme_modern() + 
  labs(
    title = "Age x transition type interaction",
    x = "Age (z-scored)",
    y = "Probability of FSS") +
  theme(
    plot.title = element_text(size = 8, hjust = 0.5),
    axis.title.x = element_text(size = 8, margin = margin(t = 3)),
    axis.title.y = element_text(size = 8, margin = margin(r = 3)),
    axis.text.x = element_text(size = 7),
    axis.text.y = element_text(size = 7),
    axis.line = element_line(size = .2),
    axis.ticks = element_line(size = .2),
    legend.text = element_text(size = 7),
    legend.title = element_text(size = 8),
    legend.key.size = unit(0.3, "cm"),
    ) 
agetrans_fig

ggsave(filename = "figures/figure_S3/age_trans_int_fig.png", 
       plot = agetrans_fig,
       device = "png", 
       dpi = 500, 
       units = "mm",
       width = 60, 
       height = 60)
```

## Age x reward interaction

```{r Age x reward}
# Age x reward
agerew_int <- ggpredict(intmodeltoplot1, terms = c("age_z", "moneylagdummy"))


agerew_fig <- ggplot(agerew_int,
                       aes(x = x,
                           y = predicted,
                           group = group,
                           color = group)) +
  geom_line(linewidth = 0.5, 
            show.legend = FALSE) +
  geom_ribbon(aes(ymin = conf.low, 
                  ymax = conf.high,
                  fill = group), 
              alpha = 0.2,
              color = NA) +  
  scale_y_continuous(labels = scales::percent) +
  scale_fill_discrete(name = "Reward") +
  guides(color = "none") +
  theme_modern() + 
  labs(
    title = "Age x reward interaction",
    x = "Age (z-scored)",
    y = "Probability of FSS") +
  theme(
    plot.title = element_text(size = 8, hjust = 0.5),
    axis.title.x = element_text(size = 8, margin = margin(t = 3)),
    axis.title.y = element_text(size = 8, margin = margin(r = 3)),
    axis.text.x = element_text(size = 7),
    axis.text.y = element_text(size = 7),
    axis.line = element_line(size = .2),
    axis.ticks = element_line(size = .2),
    legend.text = element_text(size = 7),
    legend.title = element_text(size = 8),
    legend.key.size = unit(0.3, "cm"),
    ) 
agerew_fig

ggsave(filename = "figures/figure_S3/age_rew_int_fig.png", 
       plot = agerew_fig,
       device = "png", 
       dpi = 500, 
       units = "mm",
       width = 60, 
       height = 60)
```

## Transition x reward interaction

```{r transition x reward}
#transition x reward
trans_rew <- ggpredict(intmodeltoplot1, terms = c("commonraredummy", "moneylagdummy"))

# Plot with pre-calculated jittered positions
transrew_fig <- ggplot(trans_rew, 
       aes(x = as.factor(x), 
           y = predicted,
           group = group,
           color = group)) + 
  geom_point(size = 0.5) + 
  geom_errorbar(aes(ymin = conf.low, 
                    ymax = conf.high), 
                width = 0.2,
                size = 0.5) + 
  scale_y_continuous(labels = scales::percent) +
  scale_color_discrete(name = "Reward") +
  theme_modern() +
  labs(
    title = "Transition type x reward interaction",
    x = "Transition type",
    y = "Probability of FSS"
  ) +
  theme(
    plot.title = element_text(size = 8, hjust = 0.5),
    axis.title.x = element_text(size = 8, margin = margin(t = 3)),
    axis.title.y = element_text(size = 8, margin = margin(r = 3)),
    axis.text.x = element_text(size = 7),
    axis.text.y = element_text(size = 7),
    axis.line = element_line(size = .2),
    axis.ticks = element_line(size = .2),
    legend.text = element_text(size = 7),
    legend.title = element_text(size = 8),
    legend.key.size = unit(0.3, "cm"),
    ) 

transrew_fig

ggsave(filename = "figures/figure_S3/trans_rew_int_fig.png", 
       plot = transrew_fig,
       device = "png", 
       dpi = 500, 
       units = "mm",
       width = 60, 
       height = 60)
```

## Three-way interaction

```{r Three-way interaction}
#Three-way interaction
three <- ggpredict(intmodeltoplot1, terms = c("age_z", "moneylagdummy", "commonraredummy")) 


three_fig <- ggplot(three, aes(x = x,
                               y = predicted,
                               group = group,
                               color = group)) + 
  geom_line(linewidth = 0.5, 
            show.legend = FALSE) +
  geom_ribbon(aes(ymin = conf.low, 
                  ymax = conf.high,
                  fill = group), 
              alpha = 0.2,
              color = NA) + 
  theme_modern() + 
  labs(
    title = "Age x transition x reward interaction",
    x = "Age (z-scored)",
    y = "Probability of FSS") +
  theme(
    plot.title = element_text(size = 8, hjust = 0.5),
    axis.title.x = element_text(size = 8, margin = margin(t = 3)),
    axis.title.y = element_text(size = 8, margin = margin(r = 3)),
    axis.text.x = element_text(size = 7),
    axis.text.y = element_text(size = 7),
    axis.line = element_line(size = .2),
    axis.ticks = element_line(size = .2),
    legend.text = element_text(size = 7),
    legend.title = element_text(size = 8),
    legend.key.size = unit(0.3, "cm"),
    strip.text = element_text(size = 7)
    ) +
  scale_fill_discrete(name = "Reward") +
  scale_y_continuous(
    labels = scales::percent
  ) +
  facet_wrap(~facet, 
             labeller = labeller(facet = c("-1" = "Rare transition",
                                           "1" = "Common transition")))

three_fig

ggsave(filename = "figures/figure_S4/three_fig.png", 
       plot = three_fig,
       device = "png", 
       dpi = 500, 
       units = "mm",
       width = 140, 
       height = 60)
```


# Behavioral figure of the Daw task.

Similar to Figure 2A in Decker et al., (2016)

```{r Behavioral Fig}
# Summarize the data and calculate proportions. This is also useful to get a sense of missing responses.
summary_data <- stay_prob %>%
  group_by(age_cat, commonrare, moneylag) %>%
  summarise(
    count_total = n(),
    count_stay = sum(firststagestay)
  ) %>%
  mutate(
    probability_stay = count_stay / count_total,
    se = sqrt(probability_stay * (1 - probability_stay) / count_total)
  ) %>%
  ungroup()
# Create the bar graph with error bars, faceting by age category
#This one is similar to Decker.
raw <- ggplot(summary_data %>% filter(complete.cases(.)), 
              aes(x = factor(moneylag), 
                  y = probability_stay, 
                  fill = factor(commonrare))) +
  geom_bar(stat = "identity", 
           position = position_dodge(width = 0.9), 
           width = 0.5) +
  geom_errorbar(aes(ymin = probability_stay - se, 
                    ymax = probability_stay + se), 
                position = position_dodge(width = 0.9), 
                width = 0.2,
                size = 0.3) +
  scale_x_discrete(labels = c("0" = "No Reward", 
                              "1" = "Reward"),
                   guide = guide_axis(n.dodge = 1)) +
  scale_fill_brewer(palette = "Dark2") +
   scale_y_continuous(labels = scales::percent) +
  labs(
    x = "Outcome of Previous Trial",
    y = "Probability of FSS",
    fill = "Transition type"
  ) +
  facet_grid(~age_cat) +
  theme_modern() +
   theme(
    axis.title.x = element_text(size = 8, margin = margin(t = 3)),
    axis.title.y = element_text(size = 8, margin = margin(r = 3)),
    axis.text.x = element_text(size = 7),
    axis.text.y = element_text(size = 7),
    axis.line = element_line(size = .2),
    axis.ticks = element_line(size = .2),
    legend.text = element_text(size = 7),
    legend.title = element_text(size = 8),
    legend.key.size = unit(0.3, "cm"),
    strip.text = element_text(size = 7, angle = 0)) +
  coord_cartesian(ylim = (c(0.5,1)))


raw
ggsave(filename = "figures/figure_S2/raw_beh.png", 
       plot = raw,
       device = "png", 
       dpi = 500, 
       units = "mm",
       width = 180, 
       height = 60)
```

# Simulated MB/MF learning strategies.

## Model-Free

```{r Model-free simulation}
set.seed(1738)  # for reproducibility

# Create a dataframe with simulated data
simulated_data_modelfree <- expand.grid(
  age_cat = c(0, 1),  # Adolescents and Adults
  commonrare = c("Common", "Rare"),  # Common and Rare
  moneylag = c("No Reward", "Reward")  # Unrewarded and Rewarded trials
)

# Function to generate simulated probabilities
generate_probabilities <- function(age_cat, commonrare, moneylag) {
  # Probability of first stage stay
  if (moneylag == "Reward") {
    # Rewarded trials: higher probability of staying
    probability_stay <- ifelse(commonrare == "Common", 0.8, 0.8 - 0.05)
  } else {
    # Unrewarded trials: lower probability of staying
    probability_stay <- ifelse(commonrare == "Common", 0.3, 0.3 - 0.05)
  }
  
  return(probability_stay)
}

# Apply the function to generate probabilities
simulated_data_modelfree$probability_stay <- mapply(generate_probabilities, simulated_data_modelfree$age_cat, simulated_data_modelfree$commonrare, simulated_data_modelfree$moneylag)

# Add SEM (standard error of the mean) column with a constant value
simulated_data_modelfree$sem <- 0.03  # Adjust this value as needed

# Print the simulated data
print(simulated_data_modelfree)

mf <- ggplot(simulated_data_modelfree %>% 
               filter(complete.cases(.)), 
             aes(x = factor(moneylag), 
                 y = probability_stay, 
                 fill = factor(commonrare))) +
  geom_bar(stat = "identity", 
           position = position_dodge(width = 0.9), 
           width = 0.7) +
  scale_fill_brewer(palette = "Dark2") +
  labs(
    x = "Outcome of Previous Trial",
    y = "Probability of First-Stage Stay",
    fill = "Transition type ",
    title = "Model-Free Learner"
  ) +
  theme_modern() +
     theme(
    plot.title = element_text(size = 8, hjust = 0.5),
    axis.title.x = element_text(size = 8, margin = margin(t=3)),
    axis.title.y = element_text(size = 8, margin = margin(r=3)),
    axis.text.x = element_text(size = 7),
    axis.text.y = element_blank(),
        axis.line = element_line(size = .2),
    axis.ticks = element_line(size = .2),
legend.text = element_text(size = 6),
    legend.title = element_text(size = 7)
)
mf

ggsave(filename = "figures/daw_task_image/mf_sim.png", 
       plot = mf,
       device = "png", 
       dpi = 500, 
       units = "mm",
       width = 90, 
       height = 60)
```

## Model-Based

```{r Model-based simulation}
set.seed(1738)  # for reproducibility

# Create a dataframe with simulated data
simulated_data_modelbased <- expand.grid(
  age_cat = c(0, 1),  # Adolescents and Adults
  commonrare = c("Common", "Rare"),  # Common and Rare
  moneylag = c("No Reward", "Reward")  # Unrewarded and Rewarded trials
)

# Function to generate simulated probabilities
generate_probabilities <- function(age_cat, commonrare, moneylag) {
  # Probability of first stage stay
  if (moneylag == "Reward") {
    # Rewarded trials
    if (commonrare == "Common") {
      probability_stay <- 0.8
    } else {
      probability_stay <- 0.35
    }
  } else {
    # Unrewarded trials
    if (commonrare == "Common") {
      probability_stay <- 0.2
    } else {
      probability_stay <- 0.7
    }
  }
  
  return(probability_stay)
}

# Apply the function to generate probabilities
simulated_data_modelbased$probability_stay <- mapply(generate_probabilities, simulated_data_modelbased$age_cat, simulated_data_modelbased$commonrare, simulated_data_modelbased$moneylag)

# Add SEM (standard error of the mean) column with a constant value
simulated_data_modelbased$sem <- 0.03  # Adjust this value as needed

# Print the simulated data
print(simulated_data_modelbased)

mb <- ggplot(simulated_data_modelbased %>% 
               filter(complete.cases(.)), 
             aes(x = factor(moneylag), 
                 y = probability_stay, 
                 fill = factor(commonrare))) +
  geom_bar(stat = "identity", 
           position = position_dodge(width = 0.9), 
           width = 0.7) +
  scale_fill_brewer(palette = "Dark2") +
  labs(
    x = "Outcome of Previous Trial",
    y = "Probability of First-Stage Stay",
    fill = "Transition type ",
    title = "Model-Based Learner"
  ) +
  theme_modern() +
     theme(
    plot.title = element_text(size = 8, hjust = 0.5),
    axis.title.x = element_text(size = 8, margin = margin(t=3)),
    axis.title.y = element_text(size = 8, margin = margin(r=3)),
    axis.text.x = element_text(size = 7),
    axis.text.y = element_blank(),
    axis.line = element_line(size = .2),
    axis.ticks = element_line(size = .2),
legend.text = element_text(size = 6),
    legend.title = element_text(size = 7)
)
mb

ggsave(filename = "figures/daw_task_image/mb_sim.png", 
       plot = mb,
       device = "png", 
       dpi = 500, 
       units = "mm",
       width = 90, 
       height = 60)
```

# Simulating Gaussian random walks.

Gaussian random walks plot for second-stage choice.

```{r Gaussian Random Walks Second Stage}
# Set parameters
n_trials <- 200
n_series <- 4  
mean_noise <- 0
sd_noise <- 0.025
lower_bound <- 0.25
upper_bound <- 0.75

# Function to simulate one time series
simulate_series <- function(n_trials, mean_noise, sd_noise, lower_bound, upper_bound) {
  prob <- numeric(n_trials)
  prob[1] <- runif(1, lower_bound, upper_bound)  # Initialize first trial

  for (t in 2:n_trials) {
    prob[t] <- prob[t-1] + rnorm(1, mean_noise, sd_noise)
    # Reflecting boundaries
    if (prob[t] > upper_bound) prob[t] <- 2 * upper_bound - prob[t]
    if (prob[t] < lower_bound) prob[t] <- 2 * lower_bound - prob[t]
  }
  
  return(prob)
}

# Simulate four time series
set.seed(1738) #Fetty Wap
data <- data.frame(
  trial = rep(1:n_trials, n_series),
  prob = c(simulate_series(n_trials, mean_noise, sd_noise, lower_bound, upper_bound),
           simulate_series(n_trials, mean_noise, sd_noise, lower_bound, upper_bound),
           simulate_series(n_trials, mean_noise, sd_noise, lower_bound, upper_bound),
           simulate_series(n_trials, mean_noise, sd_noise, lower_bound, upper_bound)),
  group = factor(rep(c("Group 1", "Group 1", "Group 2", "Group 2"), each = n_trials)),
  choice = factor(rep(c("Choice 1", "Choice 2", "Choice 1", "Choice 2"), each = n_trials))
)

# Plot the simulated time series
reward_probabilities <- ggplot(data, aes(x = trial, y = prob, color = interaction(group,choice))) +
  geom_line(aes(linetype = group), size = 0.5) +
  scale_y_continuous(limits = c(0.2, 0.8)) +
   scale_color_manual(
    values = c("Group 1.Choice 1" = "#FF0000", 
               "Group 1.Choice 2" = "#800080", 
               "Group 2.Choice 1" = "#FF0000", 
               "Group 2.Choice 2" = "#800080"), # Custom colors
    labels = c("Red Alien 1", 
               "Red Alien 2", 
               "Purple Alien 1", 
               "Purple Alien 2"),          # Custom labels
    name = "Group & Choice"                  # Legend title
  ) +
  labs(x = "Trial", y = "Probability of Reward", title = "Simulated Reward Probabilities") +
  theme_modern() +
    theme(
    axis.title.x = element_text(size = 8, margin = margin(t = 3)),
    axis.title.y = element_text(size = 8, margin = margin(r = 3)),
    axis.text.x = element_text(size = 7),
    axis.text.y = element_text(size = 7),
    plot.title = element_text(size = 8, hjust = 0.5),
    axis.line = element_line(size = .2),
    axis.ticks = element_line(size = .2),
    legend.text = element_text(size = 6),
    legend.title = element_text(size = 7),
    legend.key.size = unit(0.3, "cm")
    ) +
  guides(
    color = guide_legend(override.aes = list(linetype = c("solid", "dashed", "solid", "dashed"))),
    linetype = "none")
reward_probabilities

ggsave(filename = "figures/daw_task_image/reward_probabilities.png", 
       plot = reward_probabilities,
       device = "png", 
       dpi = 500, 
       units = "mm",
       width = 90, 
       height = 60)
```