library(dplyr)
library(purrr)
library(tidyr)
library(knitr)
library(kableExtra)
library(ggplot2)

new_guids <- readRDS("data/autistic-guis-added-20250507.rds")
d_meta <- readRDS("data/cdi-metadata.rds")
d_cdi <- readRDS("data/asd_na-osg-2025-05-20.rds") |>
    mutate(
        dataset = if_else(subjectkey %in% new_guids, 2, 1),
        dataset = factor(dataset, 1:2, c("old", "new"))
    ) |>
    left_join(
        d_meta |> select(num_item_id, word)
    )

d_vsoa <- readRDS("data/vsoa-autistic-nonautistic-ndar-id-fix-remodel-v2.rds")
d_vsoa_diff <- readRDS("data/vsoa-autistic-nonautistic-diff-ndar-id-fix-remodel-v2.rds")
d_vsoa_diff_old <- readRDS("data/vsoa-autistic-nonautistic-diff-old.rds")
    #rename(vsoa_NA = na, vsoa_ASD = asd, vsoa_diff = diff,
    #       ci_l_diff = ci_l, ci_u_diff = ci_u)

trim_vsoa <- function(.data) {
    .data |>
        mutate(
            vsoa_NA = if_else(vsoa_NA <   1,  1, vsoa_NA),
            vsoa_NA = if_else(vsoa_NA > 680, NA, vsoa_NA),
            vsoa_ASD = if_else(vsoa_ASD <   1,  1, vsoa_ASD),
            vsoa_ASD = if_else(vsoa_ASD > 680, NA, vsoa_ASD),
            vsoa_diff = vsoa_ASD - vsoa_NA
        ) |>
        drop_na()
}
d_vsoa_diff_trimmed <- trim_vsoa(d_vsoa_diff)
d_vsoa_diff_old_trimmed <- trim_vsoa(d_vsoa_diff_old)

d_ppt <- d_cdi |>
    select(subjectkey, nproduced, sex, group, interview_age) |>
    distinct()

d_aoa <- readRDS("data/CDI_English (American)_WS_14May2025_AoA.rds")
d_vsoa_aoa <- left_join(
    d_vsoa |> filter(group != "ASD-NA"),
    d_aoa |> select(num_item_id, aoa = aoa_glm), by = "num_item_id"
) |>
    mutate(
        group = factor(group, c("ASD", "NA"), c("autistic", "not aut."))
    ) |>
    relocate(aoa, .before = vsoa)

d_vsoa_aoa_trimmed <- d_vsoa_aoa |>
    mutate(
        vsoa = if_else(vsoa <   1,  1, vsoa),
        vsoa = if_else(vsoa > 680, NA, vsoa),
        aoa  = if_else( aoa <  12, 12,  aoa),
        aoa  = if_else( aoa >  30, NA,  aoa)
    ) |>
    drop_na()

vsoa_corr_tbl <- function(.data, .x, .y, ..., digits = 3) {
    cor_methods <- c("pearson", "spearman", "kendall")
    map(cor_methods,
        function(df, method) {
            df |>
                ungroup() |>
                group_by(...) |>
                summarize(r = cor({{ .x }}, {{ .y }}, method = method,
                                  use = "complete.obs"))
        },
        df = .data
    ) |>
        set_names(cor_methods) |>
        list_rbind(names_to = "method") |>
        pivot_wider(names_from = method, values_from = r) |>
        kable(digits = digits)
}

vsoa_corr_fig <- function(.data, .x, .y, ...) {
    p <- ggplot(.data, aes(x = {{ .x }}, y = {{ .y }})) +
        geom_point() +
        geom_smooth(method = lm) +
        theme_bw(base_size = 18) +
        facet_wrap(vars(...))
    
    return(p)
}

d_vsoa |>
    filter(group != "ASD-NA") |>
    mutate(
        extreme = if_else(vsoa <= 1, 1,
                          if_else(vsoa >= max(d_ppt$nproduced), 2, NA)),
        extreme = factor(extreme, 1:2, c("low", "high"))
    ) |>
    drop_na(extreme) |>
    count(group, extreme) |>
    pivot_wider(id_cols = group, names_from = extreme,
                values_from = n, values_fill = 0) |>
    kable()