
# Load --------------------------------------------------------------------
library(readxl)
library(dplyr)
library(tidyr)
library(ggplot2)
library(lme4)
library(lmerTest)
library(scales)
library(sjPlot)
library(nnet)
library(pscl)
library(tibble)
library(lubridate)
source("extra_functions.R")
library(gsignal)
library(patchwork)
library(mclogit)
library(ggeffects)
library(mediation)

#preprocess mappiness
mappiness_data <- read.csv('preprocessed_data/mappiness_data.csv')
mappiness_data$datetime <- as.POSIXct(mappiness_data$response_tc, format = "%d%b%Y %H:%M:%S")
mappiness_data$Year <- as.numeric(format(mappiness_data$datetime, "%Y"))
mappiness_data$Hour <- as.numeric(format(mappiness_data$datetime, "%H"))
mappiness_data[['age']]<-mappiness_data$Year-mappiness_data[['born']]
print(paste('Mean age: ',round(mean(mappiness_data[['age']]),2)))
print(paste('Std age: ',round(sd(mappiness_data[['age']]),2)))
print(paste('Female ratio: ',round(sum(mappiness_data$male==0)/length(mappiness_data$male),2)))
mappiness_data$feel_awk <- rescale(mappiness_data$feel_awk, to = c(-1, 1))

#preprocess mupsych data
music_data <- read.csv('preprocessed_data/esr_data_preprocessed.csv')
listener_data <- read.csv('preprocessed_data/listener_data_preprocessed.csv')
music_data$Lang <- NULL
music_data <- cbind(music_data,listener_data)
target_vars = c('InitArousal','MainArousal','MusicArousal','SpotifyArousal','ArousalAim','Success_Overall',
                "Success_Increase","Success_Maintain","Success_Decrease")
target <- 'MainArousal'
music_data$SpotifyArousal <- rescale(music_data$SpotifyArousal, to = c(-1, 1))

#convert timescale to a continuous time variable
mappiness_data$timeScale <- hour(mappiness_data$datetime) + 
  minute(mappiness_data$datetime)/60 + 
  second(mappiness_data$datetime)/3600

music_data$timeScale <- hour(music_data$InitStart) + 
  minute(music_data$InitStart)/60 + 
  second(music_data$InitStart)/3600

# ---- RQ1 plots ----
plot_hourly_vars(music_data,targets = c('InitArousal', 'MainArousal','FinalArousal'),
                 labels = c(InitArousal = "FeltEnergy_T1",MainArousal = "FeltEnergy_T2",
                            FinalArousal = "FeltEnergy_T3"))

plot_hourly_vars(music_data,targets = c('MusicArousal', 'SpotifyArousal'),
                 labels = c(MusicArousal = "MusicEnergy",SpotifyArousal = "SpotifyEnergy"))

p1<-plot_hourly_single_var(music_data, "ArousalAim","Energy Aim")
p2<-plot_hourly_percentage(music_data, "ArousalAim", "Energy Aims", legend_labels = c("Decrease", "Maintain", "Increase"))
combined_plot <- p1 / p2
combined_plot

freq_table <- data.frame(
  Value = sort(unique(c(
    music_data$Success_Increase,
    music_data$Success_Maintain,
    music_data$Success_Decrease
  )),decreasing = TRUE)
)
freq_table$Increase <- as.numeric(table(music_data$Success_Increase)[as.character(freq_table$Value)])
freq_table$Maintain <- as.numeric(table(music_data$Success_Maintain)[as.character(freq_table$Value)])
freq_table$Decrease <- as.numeric(table(music_data$Success_Decrease)[as.character(freq_table$Value)])
freq_table[is.na(freq_table)] <- 0
write.csv(freq_table, "success_frequencies.csv", row.names = FALSE)

target_vars_lmm = c('InitArousal','MainArousal','FinalArousal','MusicArousal','SpotifyArousal','ArousalAim','Success_Overall')
output_dir <- "analysis/models"
file_path <- file.path(output_dir, "models.html")
levels = c('Morning','Afternoon','Evening', 'Late Night')
models <- list()
for(target in target_vars_lmm){
  listener <- music_data |>
    dplyr::filter(!is.na(Age)) |>
    dplyr::mutate(
      TimeofDay = factor(TimeofDay,levels=levels),
      Gender = factor(Gender),
      Target = scale(.data[[target]]),
      Age = scale(Age)
    )
  model <- lme4::lmer(
    as.formula(paste(" Target ~ TimeofDay + Age + Gender + (1|ParticipantID)")),
    data = listener
  )
  models[[target]] <- model   # store model
}
sjPlot::tab_model(
  models,
  show.re.var = TRUE,
  show.icc = TRUE,
  show.aic = TRUE,
  p.adjust = "bonferroni",
  digits = 2,
  digits.p = 3,
  dv.labels = names(models),
  title = "Linear Mixed Models",
  file = file_path)

#multinomial model for arousal aim
music_data$ArousalAimF <- factor(music_data$ArousalAim)
music_data$ArousalAimF <- relevel(music_data$ArousalAimF, ref = "0")
music_data$TimeofDay <- factor(music_data$TimeofDay)
music_data$TimeofDay <- relevel(music_data$TimeofDay, ref = "Morning")
multinomial_model <- multinom(ArousalAimF ~ TimeofDay, data = music_data)
summary(multinomial_model)
# coefficients and standard errors
coef_matrix <- summary(multinomial_model)$coefficients
se_matrix   <- summary(multinomial_model)$standard.errors
# z-values
z_matrix <- coef_matrix / se_matrix
# p-values (two-tailed)
p_matrix <- 2 * (1 - pnorm(abs(z_matrix)))
# Compute 95% CI for B (raw coefficients)
ci_lower <- coef_matrix - 1.96 * se_matrix
ci_upper <- coef_matrix + 1.96 * se_matrix

# convert to long format
matrix_to_long <- function(mat, name) {
  df <- as.data.frame(mat)
  df$Outcome <- rownames(df)
  df_long <- df %>%
    pivot_longer(-Outcome, names_to = "Predictor", values_to = name)
  return(df_long)
}

# create long dataframes
df_B    <- matrix_to_long(coef_matrix, "B")
df_SE   <- matrix_to_long(se_matrix, "SE")
df_z    <- matrix_to_long(z_matrix, "z")
df_p    <- matrix_to_long(p_matrix, "p")
df_OR   <- matrix_to_long(OR_matrix, "OR")
df_CIl  <- matrix_to_long(ci_lower, "CI_lower_B")
df_CIu  <- matrix_to_long(ci_upper, "CI_upper_B")

# combine into a single table
multi_res <- df_B %>%
  left_join(df_SE, by = c("Outcome", "Predictor")) %>%
  left_join(df_z, by = c("Outcome", "Predictor")) %>%
  left_join(df_p, by = c("Outcome", "Predictor")) %>%
  left_join(df_OR, by = c("Outcome", "Predictor")) %>%
  left_join(df_CIl, by = c("Outcome", "Predictor")) %>%
  left_join(df_CIu, by = c("Outcome", "Predictor")) %>%
  select(Outcome, Predictor, B, CI_lower_B, CI_upper_B, SE, z, p, OR)

# Round numeric columns
multi_res <- multi_res %>%
  mutate(across(c(B, CI_lower_B, CI_upper_B, SE, z, p, OR), ~round(., 3)))

write.csv(multi_res, "aims_multinomial.csv", row.names = FALSE)

# ---- RQ2. Diurnal pattern stats ----

# Mappiness Mean Hour
map_summary_df <- mappiness_data |>
  group_by(Hour) |>
  summarize(
    mean_value = mean(.data$feel_awk,na.rm=TRUE),
    ci_lower = t.test(.data$feel_awk)$conf.int[1],
    ci_upper = t.test(.data$feel_awk)$conf.int[2],
    n=n(),
    .groups = 'drop'
  )

mus_summary_df <- music_data |>
  group_by(Hour) |>
  summarize(
    mean_value = mean(.data$ArousalAim,na.rm=TRUE),
    ci_lower = t.test(.data$ArousalAim)$conf.int[1],
    ci_upper = t.test(.data$ArousalAim)$conf.int[2],
    n=n(),
    .groups = 'drop'
  )
p1<-plot_hourly_single_var(mappiness_data,'feel_awk','Awakeness')

plot_hourly_single_var_and_cosinor(
  data = mappiness_data,
  target = "feel_awk",
  label = "Awakeness",
  nHarm = 4
)

#Compare different nHarm
fit_results <- fit_cosinor_harmonics_compare_plot(mappiness_data$timeScale, mappiness_data$feel_awk, maxHarm = 6,label='Map Raw',
                                                  criterion = 'AIC')
# Summary table
fit_results$summary
# Plot
print(fit_results$plot)

fit_results <- fit_cosinor_harmonics_compare_plot(music_data$timeScale, music_data$MusicArousal, maxHarm = 6,label='Map Raw')
# Summary table
fit_results$summary
# Plot
print(fit_results$plot)

#plot multiple cosinor and compare 
data_env <- list(
  df1 = mappiness_data,
  df2 = music_data
)

spec <- data.frame(
  df   = c("df1", "df2", "df2", "df2","df2","df2"),
  var  = c("feel_awk", "InitArousal","MainArousal", "FinalArousal","MusicArousal","ArousalAim"),
  label = c("BaselineEnergy", "FeltEnergy_T1", "FeltEnergy_T2","FeltEnergy_T3", "MusicEnergy", "EnergyAim"),
  nHarm = c(4,3,3,3,2,2) #values after trying 
)

p1<-plot_multiple_cosinor(spec, data_env)

#fit cosinors
fit_map <- fit_cosinor_harmonics(mappiness_data$timeScale, mappiness_data$feel_awk, nHarm = 4)

#aggregate timescale
map_agg <- mappiness_data %>%
  group_by(timeScale) %>%
  summarise(y = mean(feel_awk, na.rm = TRUE))

aim_agg <- music_data %>%
  group_by(timeScale) %>%
  summarise(y = mean(ArousalAim, na.rm = TRUE))

fit_map_agg <- fit_cosinor_harmonics(map_agg$timeScale, scale(map_agg$y), nHarm = 4)
fit_aim <- fit_cosinor_harmonics(aim_agg$timeScale, aim_agg$y, nHarm = 2)
# Compute phase difference
phase_df <- compute_phase_diff(
  t1 = fit_map$t,
  y1 = scale(fit_map$yhat),
  t2 = fit_aim$t,
  y2 = scale(fit_aim$yhat),
  n_points = 1000   # use high resolution!
)

target_vars = c("InitArousal", "MainArousal", "FinalArousal","MusicArousal", "ArousalAim")
harmonics = c(3,3,3,2,2)
all_phase_df <- list()
for (i in seq_along(target_vars)) {
  var_name <- target_vars[i]
  var_agg <- music_data %>%
    group_by(timeScale) %>%
    summarise(y = mean(.data[[var_name]], na.rm = TRUE), .groups = "drop")
  fit_mus <- fit_cosinor_harmonics(
    var_agg$timeScale,
    var_agg$y,
    nHarm = harmonics[i]
  )
  phase_df <- compute_phase_diff(
    t1 = fit_map$t,
    y1 = scale(fit_map$yhat),
    t2 = fit_mus$t,
    y2 = scale(fit_mus$yhat),
    n_points = 1000
  )
  phase_df$Variable <- var_name
  all_phase_df[[i]] <- phase_df
}

phase_all <- do.call(rbind, all_phase_df)

p2<-ggplot(phase_all, aes(x = Hour, y = PhaseDiff_hours, color = Variable)) +
  geom_line(size = 1) +
  theme_bw() +
  labs(
    title = "Phase Differences Between Awakeness and other Energy Variables",
    x = "Hour",
    y = "Phase Difference (hours)"
  )

combined_plot <- p1 / p2
combined_plot

results_table <- run_cosinor_first_harmonic(
  ref_df = mappiness_data,
  ref_var = "feel_awk",
  music_df = music_data,
  target_vars = c("InitArousal", "MainArousal", "FinalArousal","MusicArousal", "SpotifyArousal","ArousalAim"),
  harmonics = c(3, 3, 3,2, 1, 2))

results_table[,2:10]<-round(results_table[,2:10],2)
results_table[,5]<-round(results_table[,5])+round((results_table[,5]%%1)*60)/100
write.csv(results_table, "cosinor_tables.csv", row.names = FALSE)

# ---- Regressions ----
cos_1440 <- generate_cosinor_1440(fit_map)

# 4. Attach to music data
music_data <- attach_cosinor_energy(music_data, cos_1440)
#scale
music_data<- music_data %>%
  mutate(
    UserID = factor(ParticipantID),
    Energy = scale(CosinorEnergy),
    Velocity = scale(CosinorVelocity),
    Acceleration = scale(CosinorAcceleration),
  )

output_dir <- "analysis/models/rq2"
file_path <- file.path(output_dir, "models.html")
models <- list()

for(target in target_vars_lmm){
  music_data<- music_data %>%
    mutate(Target = scale(.data[[target]]))
  model <- lme4::lmer(
    as.formula(paste("Target ~ Energy + Velocity + (1|ParticipantID)")),
    data = music_data
  )
  models[[target]] <- model   # store model
}
sjPlot::tab_model(
  models,
  show.re.var = TRUE,
  show.icc = TRUE,
  show.aic = TRUE,
  p.adjust = "bonferroni",
  digits = 2,
  digits.p = 3,
  dv.labels = names(models),
  title = "Linear Mixed Models",
  file = file_path
)

# ----Congruence ----
#VELOCITY congruence variable with all 3 outcomes
aim_data <- music_data %>%
  dplyr::filter(!is.na(ArousalAim))
for (i in 1:nrow(aim_data)) {
  if (aim_data$CosinorVelocity[i]>0 & aim_data$ArousalAim[i] == 1) {
    aim_data$CongruenceVel[i] <- aim_data$CosinorVelocity[i]
  } else if ((aim_data$CosinorVelocity[i]<0 & aim_data$ArousalAim[i] == -1)) {
    aim_data$CongruenceVel[i] <- abs(aim_data$CosinorVelocity[i])
  } else if ((aim_data$CosinorVelocity[i]>0 & aim_data$ArousalAim[i] == -1)) {
    aim_data$CongruenceVel[i] <- -1*aim_data$CosinorVelocity[i]
  } else if ((aim_data$CosinorVelocity[i]<0 & aim_data$ArousalAim[i] == 1)) {
    aim_data$CongruenceVel[i] <- -1*abs(aim_data$CosinorVelocity[i])
  } else if ((aim_data$CosinorVelocity[i]>0 & aim_data$ArousalAim[i] == 0)) {
    aim_data$CongruenceVel[i] <- -1*aim_data$CosinorVelocity[i]
  } else if ((aim_data$CosinorVelocity[i]<0 & aim_data$ArousalAim[i] == 0)) {
    aim_data$CongruenceVel[i] <- aim_data$CosinorVelocity[i]
  }}
#congruence table
aim_data <- aim_data %>%
  mutate(
    CongruenceLabel = case_when(
      ArousalAim %in% c(-1, 1) & sign(CosinorVelocity) == ArousalAim ~ "Congruent",
      ArousalAim %in% c(-1, 1) & sign(CosinorVelocity) != ArousalAim ~ "Incongruent",
      TRUE ~ NA_character_
    )
  )
aim_subset <- aim_data %>%
  dplyr::filter(ArousalAim %in% c(-1, 1))
table(aim_subset$ArousalAim, aim_subset$CongruenceLabel)

# correct success to binary
for (i in 1:nrow(aim_data)) {
  if (aim_data$Success_Overall[i]==-1)
    aim_data$Success_Overall[i] <- 0
}

#GLLM for predicting success
aim_data <- aim_data %>%
  mutate(
    Success = as.numeric(Success_Overall),   # must be 0/1
    ArousalAim = factor(ArousalAim),
    UserID = factor(ParticipantID),
    
    Energy = scale(CosinorEnergy),
    Velocity = scale(CosinorVelocity),
    Congruence = scale(CongruenceVel),
    MusicEnergy = scale(MusicArousal),
    InitEnergy = scale(InitArousal),
    SpotifyEnergy = scale(SpotifyArousal)
  )

#Main effects
model <- glmer(
  Success ~ Energy + Congruence + ArousalAim + (1 | UserID),
  data = aim_data,
  family = binomial
)

summary(model)
pred <- ggpredict(model, terms = c("Congruence", "ArousalAim"))
plot(pred)

betas <- fixef(model)
coefs <- summary(model)$coefficients

# Profile likelihood CI
ci <- confint(model, method = "profile")

# Keep only fixed effects
ci_fixed <- ci[rownames(ci) %in% names(betas), ]

# Ensure correct order
ci_fixed <- ci_fixed[names(betas), ]

# --- Build results table ---
results_table <- data.frame(
  Term = names(betas),
  Estimate = as.numeric(betas),
  SE = coefs[, "Std. Error"],
  CI_lower = ci_fixed[, 1],
  CI_upper = ci_fixed[, 2],
  p_value = coefs[, "Pr(>|z|)"],
  row.names = NULL
)

# round values
results_table <- results_table %>%
  mutate(across(where(is.numeric), ~ round(., 3)))

# --- Export to Excel ---
write.csv(results_table, "success_model_results.csv", row.names = TRUE)

#interaction model
model_int <- glmer(
  Success ~ Energy + Congruence * ArousalAim + (1 | UserID),
  data = aim_data,
  family = binomial
)

# --- Mean energy data every 5 minutes for publication---
music_agg <- music_data %>%
  mutate(
    minute_of_day = hour(InitStart) * 60 + minute(InitStart),
    minute_bin = floor(minute_of_day / 5)
  ) %>%
  pivot_longer(
    cols = c(InitArousal, MainArousal, FinalArousal, 
             MusicArousal, SpotifyArousal, ArousalAim, Success_Overall), 
    names_to = "variable",
    values_to = "value"
  ) %>%
  group_by(minute_bin, variable) %>%
  summarise(
    y = mean(value, na.rm = TRUE),
    .groups = "drop"
  ) %>%
  complete(minute_bin = 0:(24*60/5 - 1), variable) %>%
  pivot_wider(
    names_from = variable,
    values_from = y)

mappinessc_agg <- mappiness_data %>%
  mutate(
    minute_of_day = hour(datetime) * 60 + minute(datetime),
    minute_bin = floor(minute_of_day / 5)
  ) %>%
  pivot_longer(
    cols = c(feel_awk), 
    names_to = "variable",
    values_to = "value"
  ) %>%
  group_by(minute_bin, variable) %>%
  summarise(
    y = mean(value, na.rm = TRUE),
    .groups = "drop"
  ) %>%
  complete(minute_bin = 0:(24*60/5 - 1), variable) %>%
  pivot_wider(
    names_from = variable,
    values_from = y)

data_agg <- mappinessc_agg %>%
  left_join(music_agg, by = "minute_bin")


data_agg <- data_agg %>%
  rename(
    FeltEnergy_T1 = InitArousal,
    FeltEnergy_T2 = MainArousal,
    FeltEnergy_T3 = FinalArousal,
    Awakeness = feel_awk,
    EnergyAim = ArousalAim,
    MusicEnergy = MusicArousal,
    SpotifyEnergy = SpotifyArousal,
    Success = Success_Overall)

write.csv(data_agg, "aggregate_energy_data.csv", row.names = FALSE)