library(knitr)
knitr::opts_chunk$set(cache=TRUE, 
                      autodep=TRUE,
                      dev=c('png', 'pdf', 'svg'),
                      fig.retina=2,
                      warning=FALSE,
                      message=FALSE)

library(devtools)
if (! require('supunsup', quietly=TRUE)) {
  devtools::install_github('kleinschmidt/phonetic-sup-unsup')
  require('supunsup')
}

# pre-parsed + excluded data from package
dat <- supunsup::supunsup_clean

library(lme4)
library(dplyr)

dat_mod <- dat %>%
  filter(trialSupCond == 'unsupervised' | supCond == 'unsupervised') %>%
  mutate_for_lmer()

dat_fit <- glmer(respP ~ trial.s * vot_rel.s * supCond * bvotCond +
                   (trial.s * vot_rel.s | subject),
                 data = dat_mod,
                 family = 'binomial',
                 control = glmerControl(optimizer = 'bobyqa'))

library(ggplot2)
theme_set(theme_bw())

four_colors <-
  c("#255984",     # blue
    "#CC8A2E",     # yellow
    "#CC5B2E",     # red-orange
    "#208C5E")     # green

four_colors_saturated <-
  c("#1461A1",
    "#F99710",
    "#F95210",
    "#0BAB66")


scale_color_discrete <- function(...) {
  scale_color_manual(values = four_colors_saturated, ...)
}

scale_fill_discrete <- function(...) {
  scale_fill_manual(values = four_colors_saturated, ...)
}



# size for three-panel results figures
res_w <- 11
res_h <- 5

# formatting boilerplate:
format_results_plot <- function(p) {
  p + 
    scale_color_discrete('Shift (ms)', drop=FALSE) + 
    scale_fill_discrete('Shift (ms)', drop=FALSE) +
    scale_x_continuous('VOT (ms)', breaks=seq(-20, 80, by=20)) + 
    scale_y_continuous('Proportion /p/ response') +
    scale_linetype_discrete('Condition')
}

plot_category_bounds <- function(cat_bounds, dodge_w = 0.75) {
  ggplot(cat_bounds,
         aes(x=factor(bvotCond, levels=rev(levels(bvotCond))),
             y=boundary_vot,
             ymin=boundary_vot - 1.96*boundary_vot_se,
             ymax=boundary_vot + 1.96*boundary_vot_se,
             color=bvotCond,
             linetype=factor(supCond,
               levels=c('unsupervised', 'supervised', 'mixed')),
             group=paste(shift, supCond))) +
    geom_pointrange(size=1.5, position=position_dodge(w=dodge_w)) +
    geom_point(aes(y=boundary_vot_true), shape=1, size=6) + 
    scale_color_discrete(drop=FALSE) +
    scale_linetype_manual(drop=FALSE,
                          values = c(1, 2, 3)) + 
    theme(legend.position='none') +
    scale_x_discrete('Shift (ms VOT)') +
    scale_y_continuous('/b/-/p/ boundary (ms VOT)',
                       breaks = seq(10, 60, by=5)) +
    coord_flip()
}

add_experiment <- function(data_) {
  data_ %>%
    mutate(experiment = 
             ifelse(supCond == 'mixed', 'Experiment 4',
                    ifelse(bvotCond == 20, 'Experiment 2',
                           ifelse(bvotCond == 30, 'Experiment 3',
                                  'Experiment 1'))))
}

dat_pred <- make_prediction_data(dat, dat_mod) %>% add_experiment

# raw average respond-P probability
respP_by_thirds <- dat %>%
  mutate(thirds=ntile(trial, 3)) %>%
  select(-trial) %>%
  left_join(bin_trials(dat)) %>%
  group_by(supCond, trialSupCond, trial_range, bvotCond, vot) %>%
  summarise(respP = mean(respP)) %>%
  add_experiment

respP_by_thirds_unlab <- respP_by_thirds %>%
  mutate(type='data') %>%               # for plotting along w/ glmer fits
  filter(trialSupCond == 'unsupervised')

cat_bounds <- category_boundaries(dat_mod, dat_fit) %>% add_experiment

n_subj <- supunsup::supunsup %>%
  group_by(subject) %>%
  summarise %>%
  nrow

n_subj_excluded <- supunsup::supunsup_excluded %>%
  group_by(subject) %>%
  summarise %>%
  nrow

supunsup::supunsup_clean %>%
  group_by(supCond, bvotCond, subject) %>%
  summarise %>%
  tally %>%
  kable(caption = 'Subjects in each condition')

dat %>%
  group_by(bvotCond) %>%
  filter(subject == first(subject)) %>%
  group_by(bvotCond, vot) %>%
  tally() %>%
  ggplot(aes(x=vot, y=n, fill=factor(bvotCond))) +
  geom_bar(stat='identity') +
  facet_grid(.~bvotCond) +
  scale_x_continuous('VOT (ms)', breaks=seq(-20, 80, by=20))

dat %>%
  filter(bvotCond == 0) %>%
  group_by(supCond) %>%
  filter(subject == first(subject)) %>%
  group_by(supCond, labeled, vot) %>%
  tally %>%
  ggplot(aes(x=vot, y=n, fill=labeled)) +
  geom_bar(stat='identity') +
  scale_fill_manual(values = c('black', 'gray')) + 
  facet_grid(.~supCond) +
  scale_x_continuous('VOT (ms)', breaks=seq(-20, 80, by=20))

se <- function(x) {sd(x)/length(x)}

# responses and accuracy on labeled trials across all experiments:
labeled_summary <- dat %>%
  filter(labeled == 'labeled') %>%
  add_experiment %>%
  mutate(labelCat = respCategory) %>%
  group_by(labelCat, subject) %>%
  summarise(acc = mean(labelCat == respCat),
            respP = mean(respP)) %>%
  summarise_each(funs(mean, se))

ggplot(labeled_summary, aes(x=labelCat, y=respP_mean)) +
  geom_bar(stat='identity') +
  scale_x_discrete('Labeled as') +
  scale_y_continuous('Proportion /p/ responses', breaks=c(0, 0.5, 1))

dat_ex1 <- dat %>%
  add_experiment %>%
  filter(experiment == 'Experiment 1')

expt1_respP <- respP_by_thirds_unlab %>%
  filter(experiment == 'Experiment 1')

predict_and_plot(filter(dat_pred, experiment == 'Experiment 1'),
                 dat_fit,
                 show_se=TRUE) %>%
  format_results_plot + 
  geom_point(data = expt1_respP, aes(y=respP)) +
  geom_line(data = expt1_respP, aes(y=respP))

cat_bounds %>%
  filter(experiment == 'Experiment 1') %>%
  plot_category_bounds() +
  scale_y_continuous('/b/-/p/ boundary (ms VOT)',
                     breaks = seq(10, 60, by=5),
                     limits = c(18, 32))

expt23_respP <- respP_by_thirds_unlab %>%
  filter(experiment %in% c('Experiment 2', 'Experiment 3'))

predict_and_plot(filter(dat_pred, experiment %in% c('Experiment 2', 'Experiment 3')),
                 dat_fit,
                 show_se=TRUE) %>%
  format_results_plot + 
  geom_point(data = expt23_respP, aes(y=respP)) + 
  geom_line(data = expt23_respP, aes(y=respP))

cat_bounds %>%
  filter(experiment %in% c('Experiment 2', 'Experiment 3')) %>%
  plot_category_bounds()

expt4_respP <- respP_by_thirds_unlab %>%
  filter(experiment == 'Experiment 4' | supCond == 'unsupervised')

predict_and_plot(filter(dat_pred, supCond %in% c('mixed', 'unsupervised')),
                 dat_fit,
                 show_se=TRUE) %>%
  format_results_plot + 
  geom_point(data = expt4_respP, aes(y=respP)) + 
  geom_line(data = expt4_respP, aes(y=respP))

cat_bounds %>%
  filter(experiment == 'Experiment 4' | supCond == 'unsupervised') %>%
  plot_category_bounds(dodge_w=1)

library(stargazer)

var_name_subs <- list(
  c(':', ' : '),
  c('vot_rel.s', 'VOT'),
  c('bvotCond', 'Shift'),
  c('supCond', 'unsup-vs-'),
  c('trial.s', 'Trial'))

stargazer(dat_fit, float=FALSE, single.row=TRUE,
          covariate.labels = str_replace_multi(names(fixef(dat_fit)),
                                               var_name_subs, TRUE),
          digits = 2, star.cutoffs = c(0.05, 0.01, 0.001),
          column.labels=c('Experiment 1', 'Experiment 2'), align=TRUE,
          intercept.bottom=FALSE, model.numbers=FALSE, 
          dep.var.labels.include=FALSE, dep.var.caption='', 
          keep.stat = c('n'), type='html')