# R version: 4.2.0
library(tidyverse)
library(data.table)

data = read.table("V1C2_trackdata_wb_sg.txt") 

# include Bavarian speakers only
data_bav <- data %>%
  subset(variety == "Bavarian")

## by-subject scaling of f0 (z-scores)
data_scaled <- data_bav %>%
  group_by(session) %>%
  mutate(scaledF0 = scale(Hz) %>% as.numeric()) %>%
  ungroup()

###################
### FLMM 2: Age
## select, rename and convert relevant columns for FLMM
curve_info_age = data_scaled %>%
  select(c("bundle", "session", "word", "repetition", "scaledF0", "t.norm", "age")) %>%
  rename(n_long = bundle, subject_long = session, word_long = word, combi_long = repetition,
         y_vec = scaledF0, t = t.norm,
         covariate.1 = age) %>%
  mutate_at(c("n_long", "subject_long", "word_long"), as.factor)

curve_info_age = curve_info_age %>%
  mutate_at(c("n_long", "subject_long", "word_long"), as.integer)

## set younger Bavarians as reference mean
curve_info_age$covariate.1 <- if_else(curve_info_age$covariate.1 == "younger", 0, 1)

## get mean and range of data points per f0 trajectory included in the analysis
n_age = count(curve_info_age,n_long)
mean(n_age$n)
range(n_age$n)

## convert data.frame to data.table
curve_info_age <- as.data.table(curve_info_age)
curve_info_age = setDT(curve_info_age)

## unload add-on packages (necessary for all FLMM functions to work properly)
invisible(lapply(paste0("package:", names(sessionInfo()$otherPkgs)),
                 detach,
                 character.only = TRUE, unload = TRUE))

## model data
library(sparseFLMM) # v.0.4.2

covariate = TRUE
num_covariates = 1
covariate_form = "by"
interaction = FALSE

# model computation takes approx. 6 minutes
results_age <- sparseFLMM(curve_info = curve_info_age,
                      covariate = covariate,
                      num_covariates = num_covariates, covariate_form = covariate_form,
                      interaction = interaction,
                      bf_covs = c(5,5,5), m_covs = list(c(2,3),c(2,3),c(2,3)),
                      use_famm = TRUE)

## Figure 3
# define reference mean and plot grid
intercept <- results_age$fpc_famm_hat_tri_constr$intercept
y_mean <- results_age$fpc_famm_hat_tri_constr$famm_cb_mean$value
my_grid <- results_age$my_grid

# save this plot with jpeg() as JPEG, height = 380, width = 960
# set plot device to two rows, two columns
par(mfrow = c(1,2))

## plot effect of covariate.1
y_cov1 <- results_age$fpc_famm_hat_tri_constr$famm_cb_covariate.1$value
se_cov1 <- results_age$fpc_famm_hat_tri_constr$famm_cb_covariate.1$se

plot(my_grid, y_cov1, t = "l", lwd = 3, xlab = "Normalized time",
     ylab = "", ylim = c(-0.75, 0.75), cex.axis = 1.8,
     main = expression(paste("a. ", italic("Age "), f[1](t))), cex.main = 2.2, cex.lab = 2.0)
title(ylab = expression(paste(Delta, " Scaled f0")), line = 2.5, cex.lab = 2.0)

# add confidence bands
lines(x = my_grid, y = (y_cov1 + 2*se_cov1), lty = 2, lwd = 3)
lines(x = my_grid, y = (y_cov1 - 2*se_cov1), lty = 2, lwd = 3)

# add zero reference line
abline(h = 0, lty = 3)

## plot summed effects curves: effect of Age on scaled f0
# define summed effects
none <- intercept + y_mean
Age <- intercept + y_mean + y_cov1

# define scaling of the y-axis
miny <- -0.75
maxy <- 0.75

plot(my_grid, none, col = "black", lty = "solid", lwd = 4,t = "l",
     ylim = c(miny, maxy), ylab = "", xlab = "", xaxt="n", yaxt = "n")
par(new = TRUE)
plot(my_grid, Age, col = "black", lty = "dashed", lwd = 4,
     ylim = c(miny, maxy),
     ylab = "",
     xlab = "Normalized time",
     main = expression(paste("b. Vowel+sonorant-sequences")), t = "l", cex.axis=1.8,
     cex.main = 2.2, cex.lab = 2.0)
title(ylab = "Summed effects: scaled f0", line = 2.5, cex.lab = 2.0)

# add zero reference line
abline(h = 0, lty = 3)

# add legend text
leg.text <- c(bquote(atop("Dialect, younger",
                          "(baseline,"~f[0](t)~")")),
              expression(paste("Dialect, older")))
legend("bottomright", legend = leg.text, cex = 1.8, bty = "n",
       col = c("black","black"),
       lty = c("solid","dashed"), 
       lwd = c(3, 3))

# set plot device back to default: one row, one column
par(mfrow = c(1,1))
dev.off()