```{r setup, message=FALSE, include=FALSE, paged.print=FALSE}
#! ===============================================
#! load required packages

library(ggplot2)
library(tidyverse)
library(stringi)
library(pls)
library(caret)
library(here)
library(tsibble)
library(broom)
library(rsample)
library(inspectdf)
library(caTools)
```


```{r preprocessing}

#! ===============================================
#! load data set and set running window size
load(here('notebooks/data/CL2015.RData'))
window_size <- 8

#! ===============================================
#! A function to mark lures in a sequence
with_lures <- function(stimulus, stimulus_type, n) {
  sapply(1:length(stimulus), function(i) {
    lures <- c(as.character(stimulus[i-n-1]), as.character(stimulus[i-n+1]))
    are_valid_trials <- i>n && all(!is.na(c(lures,stimulus[i])))
    ifelse(are_valid_trials && stimulus[i] %in% lures,
           "lure", 
           as.character(stimulus_type[i]))
  })
}

#! ===============================================
#! Preprocess data set to add t,tl,l,ll,u,ul,s,sl,a,al
#! a and al are respectively accuracy and recent accuracy
seqs <- NB %>%
  group_by(participant, block, condition) %>%
  mutate(n = ifelse(condition=='2-back',2,3)) %>%
  mutate(stimulus_type = with_lures(stimulus, stimulus_type, n)) %>%
  mutate(tl = slide2_dbl(stimulus_type, rt, ~length(which(.x=='target'))/length(which(!is.na(.y))), .partial=T,.size=window_size),
         ll = slide2_dbl(stimulus_type, rt, ~length(which(.x=='lure'))/length(which(!is.na(.y))), .partial=T, .size=window_size),
         sl = slide_dbl(stimulus_type, ~sum(sort(table(.x), decreasing = T)[1:2]) - 1, .partial=T, .size=window_size),
         sl = ifelse(is.na(sl), 0, sl),
         tl = ifelse(is.na(tl), NA, tl),
         ll = ifelse(is.na(ll), NA, ll),
         ul = slide_dbl(stimulus, ~max(table(.))-1, .partial=T, .size=window_size),
         vl = slide_dbl(stimulus, ~length(unique(.)), .partial=T, .size=window_size),
         al = slide2_dbl(correct, rt, ~length(which(.x))/length(which(!is.na(.y))), .partial=T, .size=window_size)) %>%
  nest(.key='local_stats') %>%
  #mutate(stimuli = map(local_stats, ~paste0(.x$stimulus,collapse = ''))) %>%
  mutate(a  = map_dbl(local_stats, ~length(which(.x$correct)))) %>%
  mutate(t  = map_dbl(local_stats, ~length(which(.x$stimulus_type=='target')))) %>%
  mutate(l  = map_dbl(local_stats, ~length(which(.x$stimulus_type=='lure')))) %>%
  mutate(s  = map_dbl(local_stats, ~sum(sort(table(.x$stimulus), decreasing = T)[1:2]) - 1)) %>%
  mutate(v  = map_dbl(local_stats, ~length(unique(.x$stimulus)))) %>%
  mutate(local_stats = map(local_stats, ~.x %>% select(-trial,-stimulus,-stimulus_type,-choice))) %>%
  ungroup() %>%
  select(-participant,-block,-condition)

#! ===============================================
#! visualize correlations
inspect_cor(seqs %>% unnest(local_stats), show_plot = T)
```

```{r models}
#! ===============================================
#! prepare data for modeling (remove na, etc)
#! it also restructures "correct" column to avoid caret errors. C stands for "CORRECT", and I is "INCORRECT"
data <- seqs %>% 
  unnest(local_stats) %>%
  mutate(correct=factor(as.numeric(correct),labels=c("C","I"))) %>%
  filter(!is.na(correct), !is.na(rt))

#! ===============================================
#! Prepare train and test partials
shuff <- sample(nrow(data))
split <- nrow(data) * 0.8

train_data <- data[1:split,]
test_data <- data[(split+1):nrow(data),]

#! ===============================================
#! training parameters for the PLS models
plsTrControl <- trainControl(
    method = "cv",
    number = 5
  )

#==================================================#
# Train PLS model (accuracy)
model_pls_accuracy <- train(
  a ~ .-rt-al-correct,
  data = train_data,
  method = "pls",
  tuneLength = 20,
  trControl = plsTrControl,
  preProc = c("zv","center","scale"))
 
# Check CV profile
plot(model_pls_accuracy)

# PLS variable importance
plot(varImp(model_pls_accuracy), main="Accuracy - Variable Importance")


#==================================================#
# Train PLS model (rt)
train_data_x <- data %>% select(-rt,-a,-correct)
train_data_y <- (data %>% select(rt))$rt

model_pls_rt <- train(
  train_data_x,
  train_data_y,
  method = "pls",
  tuneLength = 20,
  trControl = plsTrControl,
  preProc = c("center","scale"))
 
# Check CV profile
plot(model_pls_rt)

# PLS variable importance
plot(varImp(model_pls_rt), main="RT - Variable Importance")

predicted_rt_data <- predict(model_pls_rt, test_data)

#FIXME
confusionMatrix(predicted_rt_data,test_data$rt)
colAUC(predicted_rt_data,test_data$rt, plotROC=T)

#==================================================#
# training control params for "correct" column
trControl <- trainControl(
    method = "cv",
    number = 5,
    classProbs = T,
    summaryFunction = twoClassSummary
  )
  
model_glm_correct <- train(
  correct ~ .-rt-a-al,
  data = train_data,
  method = "glm",
  family = "binomial",
  trControl = trControl
)

model_glm_correct
varImp(model_glm_correct)


predicted_correct_data <- predict(model_glm_correct, test_data, type="prob")

#FIXME
confusionMatrix(predicted_correct_data, test_data$correct)
colAUC(predicted_correct_data, test_data$correct, plotROC=T)

#==================================================#
## OLD MODEL (only global features)
model_glm_correct_old <- train(
  correct ~ n+t+v,
  data = train_data,
  method = "glm",
  family = "binomial",
  trControl = trControl
)

model_glm_correct_old

predicted_old_correct_data <- predict(model_glm_correct_old, test_data, type="prob")

#FIXME
confusionMatrix(test_data$correct, predicted_old_correct_data)
colAUC(predicted_old_correct_data,test_data$correct, plotROC=T)
```