```{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)
library(pROC)
```


```{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),
         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),
         sl = ifelse(is.na(sl), 0, sl),
         tl = ifelse(is.na(tl), NA, tl),
         ll = ifelse(is.na(ll), NA, ll),
         al = ifelse(is.na(al), NA, al)
) %>%
  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) %>%
  unnest(local_stats)

#! ===============================================
#! visualize correlations
#DEBUG inspect_cor(seqs, show_plot = T)
```

```{r models}
#! ===============================================
#! prepare data for modeling (remove na, etc)
seqs <- seqs %>% 
  filter(!is.na(correct), !is.na(rt)) %>%
  #mutate(correct=factor(as.numeric(correct),labels=c("I","C")))
  mutate(correct=as.numeric(correct))

#FIXME remove outcomes before dummy out the data and imputing missing values

# replace factors with dummy data
#seqs.dummy <- predict(dummyVars(~.,data=seqs[,-1]),seqs[,-1])

# impute missing values
#seqs.imputed <- predict(preProcess(seqs.dummy, "bagImpute"), seqs.dummy)
#DEBUG View(seqs.imputed)

#! ===============================================
#! split into train/test (consider class-inbalence for the correct)
set.seed(42)
train_indexes <- createDataPartition(seqs$correct,
                                     times = 1,
                                     p = 0.7,
                                     list = F)

seqs.train <- seqs[train_indexes,]
seqs.test <- seqs[-train_indexes,]

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

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

# PLS variable importance
plot(varImp(pls.fit.accuracy), main="Accuracy - Variable Importance")


#==================================================#
# Train PLS model (rt)
train_data_x <- train_data %>% select(-rt,-a,-correct) # all except rt
train_data_y <- (train_data %>% select(rt))$rt # only rt

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

# PLS variable importance
plot(varImp(pls.fit.rt), main="RT - Variable Importance")

pls.predicted.rt <- predict(pls.fit.rt, test_data)

#FIXME
confusionMatrix(pls.predicted.rt,test_data$rt)
colAUC(pls.predicted.rt,test_data$rt, plotROC=T)

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

glm.fit.correct
varImp(glm.fit.correct)


glm.predicted.correct <- predict(glm.fit.correct, test_data, type="prob")

#FIXME
confusionMatrix(glm.predicted.correct, test_data$correct)
colAUC(glm.predicted.correct, test_data$correct, plotROC=T)

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

glm.fit.correct.old

glm.fit.predicted.old <- predict(glm.fit.correct.old, test_data, type="prob")

#FIXME
confusionMatrix(glm.fit.predicted.old, test_data$correct, )
colAUC(glm.fit.predicted.old, test_data$correct, plotROC=T)
#TODO use pROC to viz AUX roc(test_data$correct,glm.fit.predicted.old)
```