# script that contains code for figures in the eLife manuscript "Dopamine modulates expected-value 
# signals in ventromedial prefrontal cortex that support probabilistic reward learning"

library(munsell)
library(ggplot2)
library(dplyr)
library(tidyr)
library(scales)
library(foreign)

multiplot <- function(..., plotlist=NULL, file, cols=1, layout=NULL) {
  library(grid)
  
  # Make a list from the ... arguments and plotlist
  plots <- c(list(...), plotlist)
  
  numPlots = length(plots)
  
  # If layout is NULL, then use 'cols' to determine layout
  if (is.null(layout)) {
    # Make the panel
    # ncol: Number of columns of plots
    # nrow: Number of rows needed, calculated from # of cols
    layout <- matrix(seq(1, cols * ceiling(numPlots/cols)),
                     ncol = cols, nrow = ceiling(numPlots/cols))
  }
  
  if (numPlots==1) {
    print(plots[[1]])
    
  } else {
    # Set up the page
    grid.newpage()
    pushViewport(viewport(layout = grid.layout(nrow(layout), ncol(layout))))
    
    # Make each plot, in the correct location
    for (i in 1:numPlots) {
      # Get the i,j matrix positions of the regions that contain this subplot
      matchidx <- as.data.frame(which(layout == i, arr.ind = TRUE))
      
      print(plots[[i]], vp = viewport(layout.pos.row = matchidx$row,
                                      layout.pos.col = matchidx$col))
    }
  }
}
theme_nogrid <- function (base_size = 12, base_family = "") {
  theme_bw(base_size = base_size, base_family = base_family) %+replace% 
    theme(
      panel.grid = element_blank(),
      panel.border = element_blank(),
      axis.title.y = element_text(size = rel(1.7), angle = 90, margin=margin(0,10,0,0)),
      axis.title.x = element_text(size = rel(1.7), vjust = 0.1),
      #axis.line = element_line(colour = 'black',size = rel(2)),
      axis.text.x = element_text(size = rel(1.8)),
      axis.text.y = element_text(size = rel(1.8)),
      axis.line.x = element_line(color="grey12", size = 1),
      legend.text = element_text(size = rel(1.5)),
      legend.title = element_blank(),
      plot.title=element_text(size=rel(1.5)),
      axis.ticks.x = element_blank(),
      axis.ticks.y = element_blank()
      
    )   
}

# ---------------------------------figure 1---------------------------------------------------------------

# plot for wins difference between groups 

df <- read.csv("C:\\Experiments\\DAD\\R\\beta_RQVCd_20170508.csv") # this is the name of the file that was used in the original figure creation
names(df)[names(df)=="ï..group"] <- "group"
# for reviewers, the source data is available 


#names(df)[1]="group"
df2=df %>%
  group_by(group) %>%
  summarise(m=mean(wins), se=sd(wins)/sqrt(length(wins)))

df2$group[df2$group==1]="Old"
df2$group[df2$group==2]="Young"
target <- c(2,1)
wn<-ggplot(df2, aes(y=m, x=reorder(group, target), fill=factor(group, labels=c("Young", "Old")))) +
  geom_bar(position=position_dodge(), stat="identity") +
  scale_x_discrete(name = "group") +
  scale_y_continuous(name="monetary gains (sek)", breaks=seq(110,140,10), limits=c(100, 145), oob=rescale_none) +
  #theme(axis.title.y = element_text(margin=margin(0,30,0,0))) +
  geom_errorbar(aes(ymin=m-se, ymax=m+se), size =1, width=.05, position=position_dodge(.9))+
  theme_nogrid() +
  scale_fill_brewer(palette="YlGnBu") + 
  theme(legend.position = "none") + theme(axis.title.x = element_blank()) +
  annotate("segment", x = 1, xend = 2, y = 135, yend = 135, size = 1, color="grey12") +
  annotate("segment", x = 0.5, xend = 2.5, y = 97.7, yend = 97.7, size = 1, color="grey12") +
  annotate("segment", x = 0.5, xend = 0.5, y = 97, yend = 101.8, size = 1, color="grey12") +
  annotate("segment", x = 0.5, xend = 0.5, y = 103.66, yend = 145, size = 1, color="grey12") +
  annotate("segment", x = 0.45, xend = 0.53, y = 101, yend = 103, size = 1, color="grey12") +
  annotate("segment", x = 0.45, xend = 0.53, y = 102.5, yend = 104.5, size = 1, color="grey12") +
  annotate("text", x = 1.5, y = 137, label=("*"), size = 8, color="grey12")

# plot for efficient choice difference between groups 

df3=df %>%
  group_by(group) %>%
  summarise(m=mean(obj_optimal), se=sd(obj_optimal)/sqrt(length(obj_optimal)))

df3$group[df3$group==1]="Old"
df3$group[df3$group==2]="Young"
target <- c(2,1)
ec<-ggplot(df3, aes(y=m, x=reorder(group, target), fill=factor(group, labels=c("Young", "Old")))) +
  geom_bar(position=position_dodge(), stat="identity") +
  scale_x_discrete(name = "group") +
  scale_y_continuous(name="efficient choices", breaks=seq(0.5,0.7,0.1), limits=c(0.5, 0.7), oob=rescale_none, labels=percent) +
  geom_errorbar(aes(ymin=m-se, ymax=m+se), size =1, width=.05, position=position_dodge(.9))+
  theme_nogrid() +
  scale_fill_brewer(palette="YlGnBu") + 
  #theme(axis.title.y = element_text(margin=margin(0,30,0,0))) +
  theme(legend.position = "none") + theme(axis.title.x = element_blank()) +
  annotate("segment", x = 1, xend = 2, y = 0.65, yend = 0.65, size = 1, color="grey12") +
  annotate("segment", x = 0.5, xend = 2.5, y = 0.49, yend = 0.49, size = 1, color="grey12") +
  annotate("segment", x = 0.5, xend = 0.5, y = 0.48, yend = 0.7, size = 1, color="grey12") +
  annotate("text", x = 1.5, y = 0.665, label=("*"), size = 8, color="grey12")

#plot for adaptive switches difference between groups

df4=df %>%
  group_by(group) %>%
  summarise(m=mean(propadswitch), se=sd(propadswitch)/sqrt(length(propadswitch)))

df4$group[df4$group==1]="Old"
df4$group[df4$group==2]="Young"
target <- c(2,1)
ads<-ggplot(df4, aes(y=m, x=reorder(group, target), fill=factor(group, labels=c("Young", "Old")))) +
  geom_bar(position=position_dodge(), stat="identity") +
  scale_x_discrete(name = "group") +
  scale_y_continuous(name="adaptive switches", breaks=seq(0.5,0.8,0.1), limits=c(0.5, 0.8), oob=rescale_none, labels=percent) +
  geom_errorbar(aes(ymin=m-se, ymax=m+se), size =1, width=.05, position=position_dodge(.9))+
  theme_nogrid() +
  scale_fill_brewer(palette="YlGnBu") + 
  theme(legend.position = "none") + theme(axis.title.x = element_blank()) +
  annotate("segment", x = 1, xend = 2, y = 0.74, yend = 0.74, size = 1, color="grey12") +
  annotate("segment", x = 0.5, xend = 2.5, y = 0.485, yend = 0.485, size = 1, color="grey12") +
  annotate("segment", x = 0.5, xend = 0.5, y = 0.48, yend = 0.8, size = 1, color="grey12") +
  annotate("text", x = 1.5, y = 0.76, label=("*"), size = 8, color="grey12")

# plot for no switches difference between groups

df5=df %>%
  group_by(group) %>%
  summarise(m=mean(switches), se=sd(switches)/sqrt(length(switches)))

df5$group[df5$group==1]="Old"
df5$group[df5$group==2]="Young"
target <- c(2,1)
sw<-ggplot(df5, aes(y=m, x=reorder(group, target), fill=factor(group, labels=c("Young", "Old")))) +
  geom_bar(position=position_dodge(), stat="identity") +
  scale_x_discrete(name = "group") +
  scale_y_continuous(name="switches", breaks=seq(50,80,10), limits=c(40, 85), oob=rescale_none) +
  geom_errorbar(aes(ymin=m-se, ymax=m+se), size =1, width=.05, position=position_dodge(.9))+
  theme_nogrid() +
  theme(axis.title.y = element_text(margin=margin(0,17,0,0))) +
  scale_fill_brewer(palette="YlGnBu") + 
  theme(legend.position = "none") + theme(axis.title.x = element_blank()) +
  annotate("segment", x = 0.5, xend = 0.5, y = 37, yend = 41.8, size = 1, color="grey12") +
  annotate("segment", x = 0.5, xend = 0.5, y = 43.66, yend = 85, size = 1, color="grey12") +
  annotate("segment", x = 0.45, xend = 0.53, y = 41, yend = 43, size = 1, color="grey12") +
  annotate("segment", x = 0.45, xend = 0.53, y = 42.5, yend = 44.5, size = 1, color="grey12") +
  annotate("segment", x = 1, xend = 2, y =77, yend = 77, size = 1, color="grey12") +
  annotate("segment", x = 0.5, xend = 2.5, y = 37.75, yend = 37.75, size = 1, color="grey12") +
  annotate("text", x = 1.5, y = 81, label=("n.s."), size = 7, color="grey12")

multiplot(wn, sw,ec,ads,  cols=2)


# ------------------------------figure 3-b, d, e---------------------------------------------------------

#3b

df3=df %>%
  group_by(group) %>%
  summarise(m=mean(vmPFC_Qch_Qonly_Bayes), se=sd(vmPFC_Qch_Qonly_Bayes)/sqrt(length(vmPFC_Qch_Qonly_Bayes)))

df3$group[df3$group==1]="old"
df3$group[df3$group==2]="young"
target <- c(2,1)
ggplot(df3, aes(y=m, x=reorder(group, target), fill=factor(group, labels=c("young", "old")))) +
  geom_bar(position=position_dodge(), stat="identity") +
  scale_x_discrete(name = "group") +
  scale_y_continuous(name="parameter estimate", breaks=seq(0,10,2), limits=c(0, 10), oob=rescale_none) +
  geom_errorbar(aes(ymin=m-se, ymax=m+se), size =1, width=.05, position=position_dodge(.9))+
  theme_nogrid() +
  scale_fill_brewer(palette="YlGnBu") + 
  theme(axis.title.y = element_text(margin=margin(0,20,0,0))) +
  theme(legend.position = "none") + theme(axis.title.x = element_blank()) +
  annotate("segment", x = 0.5, xend = 2.5, y = 0, yend = 0, size = 1) +
  annotate("segment", x = 0.5, xend = 0.5, y = 0, yend =10, size = 1) +
  annotate("segment", x = 1, xend = 2, y = 9, yend = 9, size = 1) +
  annotate("segment", x = 1, xend = 2, y = 9, yend = 9, size = 1) +
  annotate("text", x = 1.5, y = 9.2, label=("*"), size = 8)

# 3d

res<-as.data.frame(residuals(lm(df$vmPFC_Qch_Qonly_Bayes~df$age+df$model_fit, na.action=na.exclude)))
colnames(res)<-"res_QvmPFC_agemf"

df5 = df %>%
  group_by(group)
df5 <- df[, c("group", "age_vmpfc_residuals", "wins")]
df5$group[df5$group==1]="old"
df5$group[df5$group==2]="young"
df5<-cbind(df5,res)

ggplot(df5, aes(x=res_QvmPFC_agemf, y=wins)) +
  geom_point(aes(color=group), shape = 19, size=2) +
  geom_smooth(method=lm, fullrange=TRUE, color="gray30", fill="gray65", size=0.8)+
  theme_nogrid() + 
  scale_y_continuous(name="monetary gains (sek)", breaks=seq(100,150,10), limits=c(100, 150), oob=rescale_none) + 
  scale_x_discrete(name="residuals age - Q in vmPFC", labels=element_blank()) + 
  theme(axis.title.y = element_text(margin=margin(0,20,0,0))) +
  theme(axis.title.x = element_text(margin=margin(20,0,0,0))) +
  scale_color_manual(values=c("#aad400", "#3fb5c2")) + 
  annotate("segment", x = -14, xend = 15, y = 100, yend = 100, size = 1) +
  annotate("segment", x = -14, xend = -14, y = 100, yend =150, size = 1)

# 3e

res<-as.data.frame(residuals(lm(df$BP_NAcc~df$age, na.action=na.exclude)))
colnames(res)<-"res_BPNAc_QvmPFC"

df7 = df %>%
  group_by(group)
df7 <- df[, c("subject", "group", "BP_NAcc", "vmPFC_Qch_Qonly_Bayes")]
df7$group[df7$group==1]="old"
df7$group[df7$group==2]="young"
df7<-cbind(df7,res)

ggplot(df7, aes(x=res_BPNAc_QvmPFC, y=vmPFC_Qch_Qonly_Bayes)) +
  geom_point(aes(color=group), shape = 19, size=2) +
  geom_smooth(method=lm, fullrange=TRUE, color="gray30", fill="gray65", size=0.8)+
  theme_nogrid() + 
  scale_y_continuous(name="Q in vmPFC", 
                     breaks=seq(-15,20,5), 
                     limits=c(-15, 20)) +
  #oob=rescale_none) + 
  scale_x_continuous(name="NAcc binding potential (residuals)",
                     breaks=seq(-0.5,0.5,0.25), 
                     limits=c(-0.5, 0.5)) +
  theme(axis.title.y = element_text(margin=margin(0,20,0,0))) +
  theme(axis.title.x = element_text(margin=margin(20,0,0,0))) +
  scale_color_manual(values=c("#aad400", "#3fb5c2")) + 
  annotate("segment", x = -0.5, xend = 0.5, y = -15, yend = -15, size = 1) +
  annotate("segment", x = -0.5, xend = -0.5, y = -15, yend =20, size = 1)


# ------------------------------figure 4-----------------------------------------------------------------

df <- as.data.frame(read.spss("C:\\Experiments\\DAD\\R\\beta_RQVCd_20170508.sav")) # you can use the source data instead (titled source_data_supp_fig2)

names(df)[1]="group"

df_long=gather(df, "var", "val", 5:41)
df4=df_long %>%
  filter((var=="left_nacc_R_fwe05")|(var=="left_nacc_Q_fwe05")|(var=="right_nacc_R_fwe05")|(var=="right_nacc_Q_fwe05")) %>%
  group_by(group, var) %>%
  
  summarise(m=mean(val), se=sd(val)/sqrt(length(val)))
df4$group[df4$group==1]="bOld"
df4$group[df4$group==2]="aYoung"
target<-c(4,2,8,6,3,1,7,5)

ggplot(df4, aes(y=m, x=reorder(var,target), fill=factor(group, labels=c("Young", "Old")))) +
  geom_bar(position=position_dodge(), stat="identity") +
  scale_x_discrete(name = "", labels=element_blank()) +
  scale_y_continuous(name="Parameter estimate", breaks=seq(-0.2,1.2,0.2), limits=c(-0.2,1), oob=rescale_none) +
  geom_errorbar(aes(ymin=m-se, ymax=m+se), size=1, width=.09, position=position_dodge(.9)) +
  theme_nogrid() +
  scale_fill_brewer(palette="YlGnBu") +
  annotate("text", x = 1.5, y = -0.17, label="left NAcc", size = 6) +
  annotate("text", x = 3.5, y = -0.17, label="right NAcc", size = 6) +
  annotate("text", x = 1, y = -0.05, label="R", size = 6) +
  annotate("text", x = 2, y = -0.05, label="Q", size = 6) +
  annotate("text", x = 3, y = -0.05, label="R", size = 6) +
  annotate("text", x = 4, y = -0.05, label="Q", size = 6) +
  theme(axis.ticks.x = element_blank()) + theme(legend.position=c(0.85, 0.8)) 

# --------------------------supplementary figure 1-------------------------------------------------------------


#names(df)[1]="group" # you can use the source data instead (titled source_data_supp_fig1)
df2<-df
df2$group<-3-df2$group
df_long=gather(df2, "var", "val", 5:34)
df4=df_long %>%
  filter((var=="BP_NAcc")|(var=="BP_Cau")|(var=="BP_Put")|(var=="BP_OFC")|(var=="BP_vmPFC")|(var=="BP_vlPFC")|(var=="BP_dlPFC")) %>%
  group_by(group, var) %>%
  summarise(m=mean(na.omit(val)), se=sd(na.omit(val))/sqrt(length(na.omit(val))))

positions<-c("BP_Put","BP_Cau","BP_NAcc","BP_vmPFC","BP_dlPFC","BP_OFC","BP_vlPFC")
lbels <- c("Put", "Cau", "NAcc", "vmPFC", "dlPFC", "OFC", "vlPFC")

ggplot(df4, aes(y=m, x=var, fill=factor(group, labels=c("Old", "Young")))) +
  geom_bar(position=position_dodge(), stat="identity") +
  scale_x_discrete(limits=positions, name = "", labels=lbels) + 
  geom_errorbar(aes(ymin=m-se, ymax=m+se), size=1, width=.09, position=position_dodge(.9)) + 
  theme_nogrid() +
  scale_fill_manual(values=c("#7fcdbb", "#edf8b1")) +
  scale_y_continuous(name="binding potential", breaks=seq(0,2,0.2), limits=c(0, 2.2), oob=rescale_none)


# -------------------------supplementary figure 2-------------------------------------------------------------

df <- as.data.frame(read.spss("C:\\Experiments\\DAD\\R\\beta_RQVCd_20170508.sav")) # you can use the source data instead (titled source_data_supp_fig2)

names(df)[1]="group"

df_long=gather(df, "var", "val", 5:41)
df4=df_long %>%
  filter((var=="left_NAcc_R_RW05")|(var=="left_NAcc_Q_RW05")|(var=="right_NAcc_R_RW05")|(var=="right_NAcc_Q_RW05")) %>%
  group_by(group, var) %>%
  
  summarise(m=mean(val), se=sd(val)/sqrt(length(val)))
df4$group[df4$group==1]="Old"
df4$group[df4$group==2]="Young"
target<-c(4,2,8,6,3,1,7,5)

ggplot(df4, aes(y=m, x=reorder(var,target), fill=factor(group, labels=c("Young", "Old")))) +
  geom_bar(position=position_dodge(), stat="identity") +
  scale_x_discrete(name = "", labels=element_blank()) +
  scale_y_continuous(name="Parameter estimate", breaks=seq(-0.2,1.2,0.2), limits=c(-0.2,1), oob=rescale_none) +
  geom_errorbar(aes(ymin=m-se, ymax=m+se), size=1, width=.09, position=position_dodge(.9)) +
  theme_nogrid() +
  scale_fill_brewer(palette="YlGnBu") +
  annotate("text", x = 1.5, y = -0.17, label="left NAcc", size = 6) +
  annotate("text", x = 3.5, y = -0.17, label="right NAcc", size = 6) +
  annotate("text", x = 1, y = -0.05, label="R", size = 6) +
  annotate("text", x = 2, y = -0.05, label="Q", size = 6) +
  annotate("text", x = 3, y = -0.05, label="R", size = 6) +
  annotate("text", x = 4, y = -0.05, label="Q", size = 6) +
  theme(axis.ticks.x = element_blank()) + theme(legend.position=c(0.85, 0.8)) 
