# we want to read in textgrids as a data frame

# first make a list of textgrid files to read in
textgrid_files <- list.files(path = "../data/text_grids", 
                         pattern = "*.TextGrid", 
                         full.names = TRUE) %>%
  purrr::set_names()

# next apply the textgrid_to_df function to our list of files to read in and store as a data frame

# if you can't read in files, see if there are line breaks somewhere in the TextGrid!!
textgrids <- map_df(textgrid_files, textgrid_to_df)

# we want to separate the file name column (source) into two columns for the participant number and condition, so we first get rid of the file extension and unnecessary information in the file names
textgrids$source <- gsub("_corrected_syllables_pauses.TextGrid", "", textgrids$source)

# now we separate the column into two
textgrids <- textgrids %>% 
  separate(source, c("participant", "condition"))

# for this analysis, we only need the bell and participant (word boundary) tiers
speech_bell <- textgrids %>%
  filter(tier_name %in% c("Bell", "Participant"))

# now we want to subset to include times only between the start and end times (when the bell sounded for each participant and condition) because this is the only data we're going to analyze

# get a separate data frame with only the speaking condition; we will append the filtered data frames for each participant/condition to this one
df_all <- speech_bell %>% filter(condition == "speakingonly")

# subset the df to only include the driving conditions, since the time filtering does not apply to the speaking only condition; we will use this in the loops below
df_driving_only <- speech_bell %>%
  filter(condition != "speakingonly")

# this nested loop first subsets the data frame that contains only the driving conditions to each participant and condition
# then it finds the onset time of the start and end bells
# then it subsets the data frame to include only times within that interval
# last it appends the subsetted data frame to the df_all we defined above and removes everything we defined (just in case)
for (p in unique(df_driving_only$participant)){
  for (cond in unique(df_driving_only$condition)){
    df_sub <- df_driving_only %>% filter(participant == p, condition == cond)
    start_time <- df_sub[df_sub$tier_name == "Bell" & df_sub$time_start > 0,]$time_start[1]
    end_time <- df_sub[df_sub$tier_name == "Bell" & df_sub$time_start > 0,]$time_start[2]
    df_sub <- df_sub %>%
      filter(time_start >= start_time & time_start <= end_time)
    df_all <- rbind(df_all, df_sub)
    rm(df_sub)
    rm(start_time)
    rm(end_time)
  }
}

# first compute duration for each interval: end_time - start_time
# then filter out NA intervals (this is when the experimenter was talking)
speech_data <- df_all %>%
  filter(content != "NA" & tier_name == "Participant") %>%
  mutate(duration = time_end - time_start)
 
# speech ratio function
speech_ratio_func <- function(speaking, total){
  x <- speaking/total
  return(x)
}
 
# calculate the total speaking time, total time, and the speech ratio (speaking/total)
speech_ratio <- speech_data %>%
  group_by(participant, condition) %>%
  summarize(speaking = sum(duration[content!="xxx"]),
            total = sum(duration),
            ratio = speech_ratio_func(speaking, total))

# read in file with scores on cognitive tests
cogtests <- read.csv("../data/kognitive_tests.csv")
colnames(cogtests) <- c("participant", "digit_span_forwards", "digit_span_backwards", "number_symbol", "vocabulary")

# read in file for pdt scores
pdt <- read.csv("../data/pdt.csv")
pdt <- pdt[,-1]

speech_ratio <- merge(merge(speech_ratio, 
                      cogtests, 
                      by="participant"),
                      pdt,
                      by="participant")

# convert condition to a factor so we can set up coding scheme
speech_ratio$condition <- as.factor(speech_ratio$condition)

# set the speakingonly condition to the reference level so the contrasts we set will apply to the correct factor levels
speech_ratio$condition <- relevel(speech_ratio$condition, ref = "speakingonly")

# look at first rows of data frame
kbl(head(speech_ratio)) %>%
  kable_styling()

# make bar plot with individual data points
speech_ratio %>%
  ggbarplot(x = "condition", y = "ratio", 
       add = c("mean_ci", "jitter"),
       color = "condition", palette = c("#440154", "#21918c", "#fde725"))

# subset main textgrids data frame to only include syllable tier
# filter out NA intervals (this is when the experimenter was talking)
# compute duration for each interval: end_time - start_time
syllable_data <- textgrids %>%
  filter(tier_name == "syllables", content != "NA") %>%
  mutate(duration = time_end - time_start)

# calculate the total number of syllables, total time, and speaking rate (syllables/second)

speaking_rate <- syllable_data %>%
  group_by(participant, condition) %>%
  filter(content!="xxx" & content!="" & content != "PRAG" & content != "NA") %>%
  summarize(syllables = n(),
            total = sum(duration),
            rate = syllables/total)

speaking_rate <- merge(merge(speaking_rate, 
                             cogtests, 
                             by="participant"),
                       pdt,
                       by="participant")

# convert condition to a factor so we can set up coding scheme
speaking_rate$condition <- as.factor(speaking_rate$condition)

# set the speakingonly condition to the reference level so the contrasts will apply to the correct factor levels
speaking_rate$condition <- relevel(speaking_rate$condition, ref = "speakingonly")


# look at first rows of data frame
kbl(head(speaking_rate)) %>%
  kable_styling()

# make bar plot with individual data points
speaking_rate %>%
  ggbarplot(x = "condition", y = "rate",
       add = c("mean_ci", "jitter"),
       color = "condition", palette = c("#440154", "#21918c", "#fde725"))

# import libraries we will need
import spacy
from spacy.matcher import Matcher
import glob
from praatio import textgrid
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import re

# load nlp trained model
nlp = spacy.load("de_core_news_sm")

# empty list for textgrids
tg_files = []

# read in textgrid files and append them to the list tg_files
for file in glob.glob("../data/text_grids/*.TextGrid"):
    tg_files.append(file)

# make sure we have the correct number of files    
len(tg_files) 

# define a function to compute mean length utterance
def compute_mlu(file):
    """
    This function computes the mean length utterance for each file by means of the following steps:
    1. Splits the string of the file name and saves the first two elements, which are the participant number and condition
    2. Reads the textgrid file
    3. Creates an object that consists of the participant tier of the textgrid (the text the participant uttered)
    4. Creates a list of all text from the participant's speech task
    5. Removes NA and xxx (NA is when the experimenter was talking; xxx are pauses in the speech stream)
    6. Parses the text using the nlp model "de_cores_news_sm"
    7. Counts the number of morphemes in the continuous speech
    8. Counts the number of utterances (speech between pauses)
    9. Computes the mean length utterance = morphemes/utterances
    10. Returns strings of the participant number and condition and the mean length utterance in a list
    """
    info = re.split('_|/', file)
    tg = textgrid.openTextgrid(file, includeEmptyIntervals=True)
    participant_tier = tg.getTier('Participant')
    labelList = [entry.label for entry in participant_tier.entries]
    text = ' '.join(labelList)
    remove = ['xxx', 'na']
    text_new = [word for word in text.split() if word.lower() not in remove]
    new_text = " ".join(text_new)
    doc = nlp(new_text)
    morphs = len([token.morph for token in doc])
    matcher = Matcher(nlp.vocab)
    pattern = [{"LOWER": "na", "OP": "*"}, {"LOWER": "xxx", "OP": "*"}]
    matcher.add("UTTERANCE_BOUNDARIES", [pattern])
    doc2 = nlp(text)
    matches = matcher(doc2)
    utterances = len(matches)-1
    MLU = round(morphs/utterances,2)
    return [info[4], info[5], morphs, utterances, MLU]
  

# map our function over the list of TextGrid files
MLU_all = map(compute_mlu, tg_files)

# converts map object to a list and converts list to a pandas data frame and includes participant number, condition, and mean length utterance (MLU)
data = list(MLU_all)
df_mlu = pd.DataFrame(data, columns = ['participant', 'condition', 'morphemes', 'utterances', 'MLU']) 

# look at the first few rows of the data frame to make sure everything worked like it should
df_mlu.head()

# Create an array with the colors you want to use
colors = ["#440154", "#21918c", "#fde725"]
# Set your custom color palette
sns.set_palette(sns.color_palette(colors))

# visualize group and individual data
# bar plot shows the mean of the mean length utterance per condition with errorbars showing 95% confidence intervals
# points represent individual data points to show the spread of the data
sns.barplot(data=df_mlu, x="condition", y="MLU", estimator="mean", errorbar=("ci", 95), hue="condition")
sns.stripplot(data=df_mlu, x="condition", y="MLU", color="gray", legend=False)
plt.show()

def count_filled_pauses(file):
    """
    This function counts the number of filled pauses (äh/ähm) in the speech task for each participant and condition with the following steps:
    1. Splits the string of the file name and saves the first two elements, which are the participant number and condition
    2. Reads the textgrid file
    3. Creates an object that consists of the participant tier of the textgrid (the text the participant uttered)
    4. Creates a list of all text from the participant's speech task
    5. Finds instances of 'äh' or 'ähm' and counts them
    6. Returns strings of the participant number and condition and the number of filled pauses saved in a list
    """
    info = re.split('_|/', file)
    tg = textgrid.openTextgrid(file, includeEmptyIntervals=True)
    if info[5] != 'speakingonly':
        bell_tier = tg.getTier('Bell')
        start = bell_tier.entries[1][0]
        end = bell_tier.entries[1][1]
        croppedTG = tg.crop(start, end, mode='truncated', rebaseToZero=False)
    participant_tier = tg.getTier('Participant')
    labelList = [entry.label for entry in participant_tier.entries]
    text = ' '.join(labelList)
    matcher = Matcher(nlp.vocab)
    pattern = [{"LOWER": "äh", "OP": "*"}, 
               {"LOWER": "ähm", "OP": "*"}, 
               {"LOWER": "hm", "OP": "*"}, 
               {"LOWER": "mm", "OP": "*"}]
             # {"LOWER": "uff", "OP": "*"}]
    matcher.add("FILLED_PAUSES", [pattern])
    doc = nlp(text)
    matches = matcher(doc)
    FP = len(matches)
    remove = ['xxx', 'na', 'prag', '']
    duration_list = []
    for start, stop, label in participant_tier.entries:
        if label.lower() not in remove:
            duration_list.append(stop - start)
    total_time = sum(duration_list)
    FP_rate = FP/total_time
    return [info[4], info[5], FP, total_time, FP_rate]

FP_all = map(count_filled_pauses, tg_files)

data_fp = list(FP_all)
df_fp = pd.DataFrame(data_fp, columns = ['participant', 'condition', 'number_filled_pauses', 'time_speaking', 'filled_pauses_rate']) 
df_fp.sort_values(by=['participant']).head()

sns.barplot(data=df_fp, x="condition", y="filled_pauses_rate", estimator="mean", errorbar=("ci", 95), hue="condition")
sns.stripplot(data=df_fp, x="condition", y="filled_pauses_rate", color="gray", legend=False)
plt.show()

# write a function to read in the data how we want it (no header, comma separator, and only the first three columns)
read <- function(x){
  fread(x, header = FALSE, sep = ",", select = 1:3)
}

# make a list of file names in the directory where we have the ascii data files
file_names <- list.files(path = "../data/ascii_data_files",
                         pattern = "*.dat",
                         full.names = TRUE) %>%
  purrr::set_names()

# map over file names to read them in and assign an ID column as the file name
driving <- map_dfr(file_names, read, .id = "ID")

# rename columns
colnames(driving) <- c("ID", "time", "velocity", "latpos")

# remove file path so we just have participant number and task condition
driving$ID <- gsub("../data/ascii_data_files/p_", "", driving$ID)
driving$ID <- gsub(".dat", "", driving$ID)

# separate the ID column into two columns of participant and condition
driving <- driving %>%
  separate(ID, c("participant", "condition")) %>%
  filter(condition %in% c("easy", "hard"))

# take a look at the first few rows of the data frame
kbl(head(driving)) %>%
  kable_styling()

# write a function to read in the event files like we want them: no header, separator is a space, skipping the first row
read_evt <- function(x){
  fread(x, header = FALSE, sep = " ", skip = 1)
}

# make a list of the event files in the evt_files directory
evt_files <- list.files(path = "../data/evt_files",
                         pattern = "*.evt",
                         full.names = TRUE) %>%
  purrr::set_names()

# map over files to read them in and assign an ID column that includes the file name
df_evt <- map_dfr(evt_files, read_evt, .id = "ID")

# rename columns
colnames(df_evt) <- c("ID", "signal", "time")

# remove file path from ID to only have participant and condition
df_evt$ID <- gsub("../data/evt_files/p_", "", df_evt$ID)
df_evt$ID <- gsub(".evt", "", df_evt$ID)

# separate ID column into two columns: one for participant and one for condition
# reshape into wide format so that we have the min time (starting bell) and max time (ending bell) in a row for each participant and condition
events <- df_evt %>%
  separate(ID, c("participant", "condition")) %>%
  pivot_wider(names_from = signal, values_from = time)

# rename the new columns
colnames(events)[3:4] <- c("time_min", "time_max")

# merge the driving and events data frames so that we have the min and max time in a column for each participant and condition
driving_evt <- merge(driving, events, by = c("participant", "condition"))

# filter to only include times within the bell sounds for each participant and condition and then remove time_min and time_max columns
driving_filtered <- driving_evt %>%
  group_by(participant, condition) %>%
  filter(time > time_min & time < time_max) %>%
  select(!c(time_min, time_max))

# take a look at the first few rows
kbl(head(driving_filtered)) %>%
  kable_styling()

# compute adjusted time
adjusted_time <- driving_filtered %>%
  group_by(participant, condition) %>%
  mutate(time_adj = time - time[1]) %>%
  select(-time)

# move columns around
adjusted_time <- adjusted_time[c(1:2,5,3:4)]

# make new columns to match speech data frame
adjusted_time$content <- NA
adjusted_time$tier_name <- NA

speech <- speech_bell %>%
  filter(condition!="speakingonly") %>%
  group_by(participant, condition) %>%
  filter(time_start > 0) %>%
  mutate(time_adj = time_start - time_start[1]) %>%
  #mutate(duration = time_end - time_start) %>%
  select(!c(id, time_start, time_end, tier))

# make new columns to match driving data frame
speech$velocity <- NA
speech$latpos <- NA

# move around columns
speech <- speech[c(3:7,1:2)]

# rbind columns and sort by adjusted time so we have everything in chronological order
data_all <- rbind(adjusted_time, speech)

data_all <- data_all %>%
  group_by(participant, condition) %>%
  arrange(time_adj, .by_group = TRUE)

# fill in NAs in tier name column with "Participant"
data_all$tier_name <- ifelse(is.na(data_all$tier_name), "Participant", data_all$tier_name)

# filter to only include Participant tier
data_all <- data_all %>%
  filter(tier_name == "Participant") %>%
  select(!tier_name)

# change first row of the content column for each participant to be NA in text so we can fill later
data_all <- data_all %>%
  group_by(participant, condition) %>%
  mutate(content = ifelse(row_number() == 1, "NA", content))


# fill in content column with last non-NA value
data_all <- data_all %>%
  group_by(participant, condition) %>%
  mutate(content = na.locf(content))

# interpolate missing values for velocity and lateral position (linear interpolation)
# and calculate duration for all intervals
# change last NA to 0
data_all_filled <- data_all %>%
  group_by(participant, condition) %>%
  mutate(velocity = na.approx(velocity, na.rm = FALSE)) %>%
  mutate(latpos = na.approx(latpos, na.rm = FALSE)) %>%
  mutate(duration = dplyr::lead(time_adj)-time_adj) %>%
  mutate(duration = ifelse(is.na(duration), 0, duration)) 

data_all_filled_test <- data_all_filled %>%
  group_by(participant, condition) %>%
  mutate(group = findInterval(time_adj, seq(from = min(data_all_filled$time_adj), to = max(data_all_filled$time_adj), by=5)))

# subset to parts where participant is speaking
data_all_sub <- data_all_filled_test %>%
  filter(content != "NA", content != "")

check <- data_all_sub %>%
  group_by(participant, condition, group) %>%
  summarize(n = n())

# summarize for sd of velocity and latpos over 5 sec intervals
data_all_gc_prep <- data_all_sub %>%
  group_by(participant, condition, group) %>%
  summarize_at(c("velocity", "latpos"), sd, na.rm=TRUE)

data_all_gc_prep[is.na(data_all_gc_prep)] <- 0

# create an empty data frame for the for loop
sp_rat_temp <- data.frame(group = numeric(), speech_ratio = numeric())

sp_rat_temp <- data_all_sub %>%
  group_by(participant, condition, group) %>%
  summarize(speaking = sum(duration[content!="xxx"]),
            total = sum(duration),
            speech_ratio = speech_ratio_func(speaking, total)) %>%
  select(-c(speaking, total))

# merge data frames
data_all_gc_prep <- merge(data_all_gc_prep, sp_rat_temp, by = c("participant", "condition", "group"))

data_all_gc_visualize <- data_all_gc_prep %>%
  select(participant, condition, group, velocity, latpos, speech_ratio) %>%
  gather(key = "variable", value = "value", -c(participant, condition, group))

# add time variable to line up start time for all participants
data_all_gc_prep2 <- data_all_gc_prep %>%
  group_by(participant, condition) %>%
  arrange(group, .by_group = TRUE) %>%
  mutate(time = rep(1:n()))

se <- function(x) sd(x)/sqrt(length(x))

time_series_plot_data <- data_all_gc_visualize %>%
  group_by(group, condition, variable) %>%
  summarize(mean = mean(value),
            ci_lower = mean(value) - (1.96*se(value)),
            ci_upper = mean(value) + (1.96*se(value)))

time_series_plot_data %>%
  filter(variable != "latpos") %>%
  ggplot(aes(x = group, y = mean)) + 
  facet_wrap(~condition) +
  geom_line(aes(color = variable), size = 1) +
  geom_ribbon(aes(ymin = ci_lower, ymax = ci_upper, fill=variable), alpha = 0.1) +
  scale_color_manual(values = c("#00AFBB", "#E7B800"), labels = c("Speech ratio", "Velocity (sd)")) +
  scale_fill_manual(values = c("#00AFBB", "#E7B800"), labels = c("Speech ratio", "Velocity (sd)")) +
  labs(x="Time(s)", y="Value", color="Variable", fill="Variable", title="Plot of speech ratio and variability in velocity over time") + 
  theme_minimal() +
  theme(legend.position = "bottom")

time_series_plot_data %>%
  filter(variable != "velocity") %>%
  ggplot(aes(x = group, y = mean)) + 
  facet_wrap(~condition) +
  geom_line(aes(color = variable), size = 1) +
  geom_ribbon(aes(ymin = ci_lower, ymax = ci_upper, fill=variable), alpha = 0.1) +
  scale_color_manual(values = c("#E7B800", "#00AFBB"), labels = c("Lateral position (sd)", "Speech ratio")) +
  scale_fill_manual(values = c("#E7B800", "#00AFBB"), labels = c("Lateral position (sd)", "Speech ratio")) +
  labs(x="Time(s)", y="Value", color="Variable", fill="Variable", title="Plot of speech ratio and variability in lateral position over time") + 
  theme_minimal() + 
  theme(legend.position = "bottom")

# get a list of the conditions (should be easy and hard)
conditions <- unique(data_all_gc_prep2$condition)

# create a list of the different combinations of IVs and DVs we want to test
combos <- c("speech_ratio ~ velocity", "velocity ~ speech_ratio", "speech_ratio ~ latpos", "latpos ~ speech_ratio")

# this function maps over each condition and IV/DV pairs we want to test, performs a Granger causality test, and saves the test statistic and p value in a data frame
granger_tests <- map_dfr(conditions, function(condition){
  map_dfr(combos, function(combo){
    df_condition <- pdata.frame(data_all_gc_prep2[data_all_gc_prep2$condition==condition,], index = c("participant", "time"), drop.index=FALSE)
    test <- pgrangertest(eval(parse(text=combo)), 
                         data = df_condition, 
                         # test = "Wbar", 
                         order = 1L)
    data_pgtests <- tibble(condition = condition,
                           Z_tilde = test$statistic,
                           p_value = test$p.value,
                           test = combo)
  })
})

kbl(granger_tests) %>%
    kable_styling()