#' Set function parameters to an environment.
#'
#' This function is designed to help debug functions. It will attempt to set all
#' the default parameter values to the specified environment (global environment
#' by default). This is useful for when you want to execute code within the 
#' function definition interactively but need the parameters set in the current 
#' environment.
#'
#' **Warning:** This function will modify the global environment and therefore 
#' violates CRAN policy
#' ["Packages should not modify the global environment (user’s workspace)"]
#' (https://cran.r-project.org/web/packages/policies.html#Source-packages).
#'
#' @param FUN the function to assign parameters to an environment.
#' @param envir the environment to assign the variables to. Defaults to the 
#'        global environment.
#' @param verbose whether to return the data frame invisibly or to print the results.
#' @return a data frame where row names correspond to the parameter name with 
#'        two columns: `set` which is logical indicating if the variable was set 
#'        and `value` with a character representation of the variable value.
set_function_params <- function(FUN, envir = globalenv(), verbose = interactive()) {
    params <- formals(FUN)
    params_set <- data.frame(row.names = names(params),
                             set = rep(FALSE, length(params)),
                             value = rep(NA_character_, length(params)))
    for(param in names(params)) {
        value <- params[[param]]
        if(!missing(value)) {
            if(is.character(value)) {
                assign(param, value, envir = envir)
                params_set[param,]$value <- value
            } else {
                assign(param, eval(value), envir = envir)
                params_set[param,]$value <- eval(value)
            }
            params_set[param,]$set <- TRUE
        }
    }
    if(verbose) {
        return(params_set)
    } else {
        invisible(params_set)
    }
}

# List objects in the current environment
ls()

# Call the function
param_set_result <- set_function_params(clav::profile_plot)

# Check to see if the parameters are actually set
ls()

param_set_result

set.seed(2112)
x <- rnorm(100, mean = 0, sd = 1)
(mean1 <- mean(x))

(sd1 <- sd(x))

x[sample(length(x), length(x) * 0.2, replace = FALSE)] <- NA
(mean2 <- mean(x, na.rm = TRUE))

(sd2 <- sd(x, na.rm = TRUE))

x[is.na(x)] <- mean(x, na.rm = TRUE)
(mean3 <- mean(x))

(sd3 <- sd(x))

n_samples <- 1000
percent_missing <- 0.10
sd_diffs <- data.frame(sample = 1:n_samples,
                       sd_drop_miss = numeric(n_samples),
                       sd_impute_miss = numeric(n_samples))
for(i in seq_len(n_samples)) {
    x2 <- x
    x2[sample(length(x), length(x) * percent_missing, replace = FALSE)] <- NA
    sd_diffs[i,]$sd_drop_miss <- sd(x2, na.rm = TRUE)
    x2[is.na(x2)] <- mean(x2, na.rm = TRUE)
    sd_diffs[i,]$sd_impute_miss <- sd(x2)
}

sd_diffs |> 
    reshape2::melt(id.vars = 'sample', variable.name = 'calculation_type', value.name = 'sd') |>
    ggplot(aes(x = sd, color = calculation_type)) +
        geom_vline(xintercept = sd(x)) +
        geom_density() +
        xlab('Standard Deviation') +
        theme_minimal()

t.test(sd_diffs$sd_drop_miss, sd_diffs$sd_impute_miss)

n <- 100
mean_x <- 0
mean_y <- 0
sd_x <- 1
sd_y <- 1
rho <- 0.18

set.seed(2112)
df <- mvtnorm::rmvnorm(
    n = 100,
    mean = c(mean_x, mean_y),
    sigma = matrix(c(sd_x^2, rho * (sd_x * sd_y),
                     rho * (sd_x * sd_y), sd_y^2), 2, 2)) |>
    as.data.frame() |>
    dplyr::rename(x = V1, y = V2)

cor.test(df$x, df$y)

df_miss <- df
df_miss[sample(n, size = 0.2 * n, replace = FALSE),]$x <- NA
cor.test(df_miss$x, df_miss$y)

df_miss[is.na(df_miss$x),] <- mean(df$x, na.rm = TRUE)
cor.test(df_miss$x, df_miss$y)

remotes::install_github('jbryer/ipeds')

data(surveys, package = 'ipeds')

levels <- c('1' = 'Required',
            '5' = 'Test Optional',
            '3' = 'Test Blind')

admissions <- data.frame()
years <- 2014:2023
for(i in years) {
    adm <- getIPEDSSurvey('ADM', year = i)
    adm$admcon7 <- factor(adm$admcon7,
                          levels = names(levels),
                          labels = levels)
    adm$year <- i
    admissions <- rbind(admissions, adm[,c('year', 'admcon7')])
}

tab <- table(admissions$year, admissions$admcon7) |>
    prop.table(1) |>
    reshape2::melt() |>
    dplyr::rename(Year = Var1, Test_Use = Var2)

p <- ggplot(tab, aes(x = Year, y = value, group = Test_Use, color = Test_Use)) +
    geom_path(linewidth = 1.5) +
    geom_point(shape = 16, size = 4, color = 'white') +
    geom_point(shape = 1, size = 4) +
    scale_color_manual('Test usage: ', 
                       values = c('Required' = '#7209B7', 
                                  'Test Optional' = '#84DB8F', 
                                  'Test Blind' = '#E2BD6B')) +
    scale_y_continuous(labels = scales::percent) +
    scale_x_continuous(breaks = c(years, max(years) + 1),
                       labels = c(years, max(years) + 1),
                       limits = c(min(years), max(years) + 1)) +
    theme_minimal() +
    # xlim(c(2014, 2024)) +
    theme(legend.position = 'bottom') +
    ggtitle('Use of high stakes assessments (SAT/ACT) for college admissions') +
    ylab('Percentage of Institutions')

directlabels::direct.label(p, "last.qp")

monty_hall <- function(pick, switch) {
    interactive <- FALSE
    if(missing(pick)) {
        interactive <- TRUE
        cat('Pick your door:')
        pick <- LETTERS[menu(c('A', 'B', 'C'))]
    } else {
        if(!pick %in% LETTERS[1:3]) {
            stop('pick must be either A, B, or C')
        }
    }
    doors <- c('win', 'lose', 'lose')
    doors <- sample(doors) # Shuffle the doors
    names(doors) <- LETTERS[1:3]
    if(doors[pick] == 'win') {
        show <- sample(names(doors[!names(doors) %in% pick]), size = 1)
    } else {
        show <- doors[!names(doors) %in% pick] == 'lose'
        show <- names(which(show == TRUE))
    }
    if(missing(switch)) {
        interactive <- TRUE
        cat(paste0('Showing door ', show, '. Do you want to switch your choice?'))
        switch <- menu(c('yes', 'no')) == 1
    }
    if(switch) {
        pick <- names(doors)[!names(doors) %in% c(show, pick)]
    }
    win <- unname(doors[pick] == 'win')
    if(interactive) {
        if(win) {
            cat('You win!')
        } else {
            cat('Sorry, you lost.')
        }
        invisible(win)
    } else {
        return(win)
    }
}

monty_hall()

n_games <- 1000
mh_switch <- logical(n_games)
mh_no_switch <- logical(n_games)
for(i in 1:n_games) {
    pick <- sample(LETTERS[1:3], size = 1)
    mh_switch[i] <- monty_hall(pick = pick, switch = TRUE)
    mh_no_switch[i] <- monty_hall(pick = pick, switch = FALSE)
}

mean(mh_switch)

mean(mh_no_switch)

remotes::install_github('jbryer/VisualStats')

library(VisualStats)
data('distributions', package = 'VisualStats')
plot_distributions(dist = 'norm',
                   xvals = c(-1, 0, 0.5),
                   xmin = -4,
                   xmax = 4)

plot_distributions(dist = 'binom',
                   xvals = c(1, 3),
                   xmin = 0,
                   xmax = 10,
                   args = list(size = 10, prob = 0.35))

remotes::install_github('jbryer/medley')

library(dplyr)
library(ggplot2)

#' Simulate random samples to estimate confidence intervals and bootstrap
#' estimates.
#'
#' @param pop a numeric vector representing the population.
#' @param n sample size for each random sample from the population.
#' @param n_samples the number of random samples.
#' @param n_boot number of bootstrap samples to take for each sample.
#' @param seed a seed to use for the random process.
#' @param cv critical value to use for calculating confidence intervals.
#' @return a data.frame with the sample and bootstrap mean and confidence
#'        intervals along with a logical variable indicating whether a Type I
#'        error would have occurred with that sample.
bootstrap_clt_simulation <- function(
        pop,
        n = 30,
        n_samples = 500,
        n_boot = 500,
        cv = abs(qt(0.025, df = n - 1)),
        seed,
        verbose = interactive()
) {
    if(missing(seed)) {
        seed <- sample(100000)
    }
    results <- data.frame(
        seed = 1:n_samples,
        samp_mean = numeric(n_samples),
        samp_se = numeric(n_samples),
        samp_ci_low = numeric(n_samples),
        samp_ci_high = numeric(n_samples),
        samp_type1 = logical(n_samples),
        boot_mean = numeric(n_samples),
        boot_ci_low = numeric(n_samples),
        boot_ci_high = numeric(n_samples),
        boot_type1 = logical(n_samples)
    )
    if(verbose) {
        pb <- txtProgressBar(min = 0, max = n_samples, style = 3)
    }
    for(i in 1:n_samples) {
        if(verbose) {
            setTxtProgressBar(pb, i)
        }
        set.seed(seed + i)
        samp <- sample(pop, size = n)
        boot_samp <- numeric(n_boot)
        for(j in 1:n_boot) {
            boot_samp[j] <- sample(samp, size = length(samp), replace = TRUE) |>
                mean()
        }
        results[i,]$seed <- seed + i
        results[i,]$samp_mean <- mean(samp)
        results[i,]$samp_se <- sd(samp) / sqrt(length(samp))
        results[i,]$samp_ci_low <- mean(samp) - cv * results[i,]$samp_se
        results[i,]$samp_ci_high <- mean(samp) + cv * results[i,]$samp_se
        results[i,]$samp_type1 <- results[i,]$samp_ci_low > mean(pop) |
            mean(pop) > results[i,]$samp_ci_high
        results[i,]$boot_mean <- mean(boot_samp)
        results[i,]$boot_ci_low <- mean(boot_samp) - cv * sd(boot_samp)
        results[i,]$boot_ci_high <- mean(boot_samp) + cv * sd(boot_samp)
        results[i,]$boot_type1 <- results[i,]$boot_ci_low > mean(pop) |
            mean(pop) > results[i,]$boot_ci_high
    }
    if(verbose) {
        close(pb)
    }
    return(results)
}

pop_unif <- runif(1e5, 0, 1)
ggplot(data.frame(x = pop_unif), aes(x = x)) + geom_density()

results_unif <- bootstrap_clt_simulation(pop = pop_unif, seed = 42, verbose = FALSE)

tab <- table(results_unif$samp_type1, results_unif$boot_type1, useNA = 'ifany')
tab

pop_skewed <- rnbinom(1e5, 3, .5)
ggplot(data.frame(x = pop_skewed), aes(x = x)) + geom_density(bw = 0.75)

results_skewed <- bootstrap_clt_simulation(pop = pop_skewed, seed = 42, verbose = FALSE)
mean(results_skewed$samp_type1) # Percent of samples with Type I error

mean(results_skewed$boot_type1) # Percent of bootstrap estimates with Type I error

# CLT vs Bootstrap Type I error rate
table(results_skewed$samp_type1, results_skewed$boot_type1, useNA = 'ifany')

results_differ <- results_skewed |>
    dplyr::filter(!samp_type1 & boot_type1)
results_differ

set.seed(results_differ[1,]$seed)
samp <- sample(pop_skewed, size = 30)
boot_samp <- numeric(500)
for(j in 1:500) {
    boot_samp[j] <- sample(samp, size = length(samp), replace = TRUE) |>
        mean()
}
cv = abs(qt(0.025, df = 30 - 1))
mean(pop_skewed)

ci <- c(mean(samp) - cv * sd(samp) / sqrt(30), mean(samp) + cv * sd(samp) / sqrt(30))
ci

mean(pop_skewed) < ci[1] | mean(pop_skewed) > ci[2]

ci_boot <- c(mean(boot_samp) - cv * sd(boot_samp), mean(boot_samp) + cv * sd(boot_samp))
ci_boot

mean(pop_skewed) < ci_boot[1] | mean(pop_skewed) > ci_boot[2]

# Define some parameters
sig_level <- .95  # Significance level, 95% here
es <- 0.15        # Effect size in standard units
null_val <- 0     # null value

#' Calculate the standard error
#' 
#' This function will calculate the standard error from a vector of observations or with a given
#' sample standard deviation and sample size.
#' 
#' @param x numeric vector of observations.
#' @param sigma the sample standard deviation.
#' @param n sample size.
standard_error <- function(x, sigma = sd(x), n = length(x)) {
    if(!missing(x)) { # Some basic error checking
        if(sigma != sd(x)) { warning('The sample standard deviation (sigma) is not equal to sd(x)')}
        if(n != length(x)) { warning('The sample size (n) is not equal to length(x).' )}
    }
    return(sigma / sqrt(n))
}
# Create a data.frame with varying sample sizes and the corresponding standard error
df <- data.frame(
    n = 10:400,
    se = standard_error(sigma = 1, n = 10:400)
)
cv <- abs(qnorm((1 - sig_level) / 2)) # Critical value (z test statistic)
df$ci_low <- es - cv * df$se
df$ci_high <- es + cv * df$se
df$sig <- null_val < df$ci_low | null_val > df$ci_high
min_n <- df$n[df$sig] |> min()
ggplot(df, aes(x = n, y = se, color = sig)) + 
    geom_path() +
    geom_point() +
    scale_color_brewer(paste0('p < ', (1 - sig_level)), type = 'qual', palette = 6) +
    ggtitle(paste0('Minumum n for p < ', (1 - sig_level), ': ', min_n),
            subtitle = paste0('effect size: ', es, '; null value: ', null_val))

#' Calculate p-value from a chi-squared test with varying sample sizes
#'
#' This algorithm will start with an initial sample size (`n_start`) and perform a chi-squared test
#' with a vector of counts equal to `n * probs`. This will repeat increasing the sample size by
#' `n_step` until the p-value from the chi-squared test is less than `p_stop`.
#'
#' @param vector of cell probabilities. The sum of the values must equal 1.
#' @param sig_level significance level.
#' @param p_stop the p-value to stop estimating chi-squared tests.
#' @param max_n maximum n to attempt if `p_value` is never less than `p_stop`.
#' @param min_cell_size minimum size per cell to perform the chi-square test.
#' @param n_start the starting sample size.
#' @param n_step the increment for each iteration.
#' @return a data.frame with three columns: n (sample size), p_value, and sig (TRUE if
#'         p_value < sig_level).
#' @importFrom DescTools power.chisq.test CramerV
chi_squared_power <- function(
        probs,
        sig_level = 0.05,
        p_stop = 0.01,
        power = 0.80,
        power_stop = 0.90,
        max_n = 100000,
        min_cell_size = 10,
        n_start = 10,
        n_step = 10
) {
    if(sum(probs) != 1) { # Make sure the sum is equal to 1
        stop('The sum of the probabilities must equal 1.')
    } else if(length(unique(probs)) == 1) {
        stop('All the probabilities are equal.')
    }

    n <- n_start
    p_values <- numeric()
    power_values <- numeric()
    df <- ifelse(is.vector(probs),
                 length(probs) - 1,
                 min(dim(probs)) - 1) # Degrees of freedom
    repeat {
        x <- (probs * n) |> round()
        if(all(x > min_cell_size)) {
            cs <- chisq.test(x, rescale.p = TRUE, simulate.p.value = FALSE)
            p_values[n / n_step] <- cs$p.value
            pow <- DescTools::power.chisq.test(
                n = n,
                w = DescTools::CramerV(as.table(x)),
                df = df,
                sig.level = sig_level
            )
            power_values[n / n_step] <- pow$power
            if((cs$p.value < p_stop & pow$power > power_stop) | n > max_n) {
                break;
            }
        } else {
            p_values[n / n_step] <- NA
            power_values[n / n_step] <- NA
        }
        n <- n + n_step
    }
    result <- data.frame(n = seq(10, length(p_values) * n_step, n_step),
                         p_value = p_values,
                         sig = p_values < sig_level,
                         power = power_values)
    class(result) <- c('chisqpower', 'data.frame')
    attr(result, 'probs') <- probs
    attr(result, 'sig_level') <- sig_level
    attr(result, 'p_stop') <- p_stop
    attr(result, 'power') <- power
    attr(result, 'power_stop') <- power_stop
    attr(result, 'max_n') <- max_n
    attr(result, 'n_step') <- n_step
    return(result)
}

#' Plot the results of chi-squared power estimation
#'
#' @param x result of [chi_squared_power()].
#' @param plot_power whether to plot the power curve.
#' @param plot_p whether to plot p-values.
#' @param digits number of digits to round to.
#' @param segement_color color of the lines marking where power and p values exceed threshold.
#' @param sgement_linetype linetype of the lines marking where power and p values exceed threshold.
#' @param p_linetype linetype for the p-values.
#' @param power_linetype linetype for the power values.
#' @param title plot title. If missing a title will be automatically generated.
#' @parma ... currently not used.
#' @return a ggplot2 expression.
plot.chisqpower <- function(
        x,
        plot_power = TRUE,
        plot_p = TRUE,
        digits = 4,
        segment_color = 'grey60',
        segment_linetype = 1,
        p_linetype = 1,
        power_linetype = 2,
        title,
        ...
) {
    pow <- attr(x, 'power')

    p <- ggplot(x[!is.na(x$p_value),], aes(x = n, y = p_value))

    if(plot_power) {
        if(any(x$power > pow, na.rm = TRUE)) {
            min_n_power <- min(x[x$power > pow,]$n, na.rm = TRUE)
            p <- p +
                geom_segment(
                    x = 0,
                    xend = min_n_power,
                    y = pow,
                    yend = pow,
                    color = segment_color,
                    linetype = segment_linetype) +
                ggplot2::annotate(
                    geom = 'text',
                    x = 0,
                    y =  pow,
                    label = paste0('Power = ',  pow),
                    vjust = -1,
                    hjust = 0) +
                geom_segment(
                    x = min_n_power,
                    xend = min_n_power,
                    y = pow,
                    yend = 0,
                    color = segment_color,
                    linetype = segment_linetype) +
                ggplot2::annotate(
                    geom = 'text',
                    x = min_n_power, y = 0,
                    label = paste0('n = ', prettyNum(min_n_power, big.mark = ',')),
                    vjust = 1,
                    hjust = -0.1)
        }
        p <- p +
            geom_path(
                aes(y = power),
                color = '#7570b3',
                linetype = power_linetype)
    }
    if(plot_p) {
        if(any(x$sig, na.rm = TRUE)) {
            p <- p +
                geom_segment(
                    x = 0,
                    xend = min(x[x$sig,]$n, na.rm = TRUE),
                    y = attr(x, 'sig_level'),
                    yend = attr(x, 'sig_level'),
                    color = segment_color,
                    linetype = segment_linetype) +
                ggplot2::annotate(
                    geom = 'text',
                    x = 0,
                    y =  attr(x, 'sig_level'),
                    label = paste0('p = ',  attr(x, 'sig_level')),
                    vjust = -1,
                    hjust = 0) +
                geom_segment(
                    x = min(x[x$sig,]$n, na.rm = TRUE),
                    xend = min(x[x$sig,]$n, na.rm = TRUE),
                    y = attr(x, 'sig_level'),
                    yend = 0,
                    color = segment_color,
                    linetype = segment_linetype) +
                ggplot2::annotate(
                    geom = 'text',
                    x = min(x[x$sig,]$n, na.rm = TRUE),
                    y = 0,
                    label = paste0('n = ', prettyNum(min(x[x$sig,]$n, na.rm = TRUE), big.mark = ',')),
                    vjust = 1,
                    hjust = -0.1)
        }
        p <- p +
            geom_path(
                alpha = 0.7,
                linetype = p_linetype)
            # geom_point(aes(color = sig), size = 1) +
            # scale_color_brewer(paste0('p < ', attr(x, 'sig_level')), type = 'qual', palette = 6)
    }

    if(missing(title)) {
        if(any(x$power > pow, na.rm = TRUE) & any(x$sig, na.rm = TRUE)) {
            min_n <- min(x[x$sig & x$power > pow,]$n, na.rm = TRUE)
            title <- paste0('Smallest n where p < ', attr(x, 'sig_level'), ' and power > ', pow, ': ',
                            prettyNum(min_n, big.mark = ','))
        } else {
            title <- paste0('No n found where p < ', attr(x, 'sig_level'), ' and power > ', pow)
        }
    }

    p <- p +
        ylim(c(0, 1)) +
        ylab('') +
        xlab('Sample Size') +
        ggtitle(title,
                subtitle = paste0('Probabilities: ', paste0(round(attr(x, 'probs'), digits = digits), collapse = ', ')))

    return(p)
}

csp1 <- chi_squared_power(probs =  c(.33, .25, .42))
csp1[csp1$sig,]$n |> min(na.rm = TRUE) # Smallest n that results in p < 0.05

plot(csp1)

csp3 <- chi_squared_power(probs = c(.25, .25, .24, .26), max_n = 20000)
csp3[csp3$sig,]$n |> min(na.rm = TRUE) # Smallest n that results in p < 0.05

plot(csp3)

M <- as.table(rbind(c(762, 327, 468), c(484, 239, 477)))
dimnames(M) <- list(gender = c("Femal", "Male"),
                    party = c("Democrat", "Independent", "Republican"))
M

sum(M)

chisq.test(M)

DescTools::CramerV(M) # Effect size

DescTools::power.chisq.test(n = sum(M),
                            w = DescTools::CramerV(M),
                            df = min(dim(M)) - 1,
                            sig.level = 1 - sig_level)

M_prob <- M / sum(M) # Convert the counts to percentages
csp4 <- chi_squared_power(probs = M_prob)
plot(csp4)

size <- 5 # Test size (i.e. number of questions)
p <- 1/4 # Probability of randomly getting correct answer

test <- sample(c(TRUE, FALSE), size = size, prob = c(p, 1 - p), replace = TRUE)
test

#' Simulate how long until a specified number of responses are correct
#' @param size test size.
#' @param prob probability of randomly getting correct answer
#' @param stop_score the score on the test we wish to achieve. Value of 1
#'        indicates a perfect score.
simulate_test <- function(size, p, stop_score = 1) {
    n <- 0
    repeat{
        n <- n + 1
        test <- sample(c(TRUE, FALSE),
                       size = size,
                       prob = c(p, 1 - p),
                       replace = TRUE)
        if(mean(test) >= stop_score) {
            break
        }
    }
    return(n)
}

(num_tests <- simulate_test(size = size, p = p))

simulations <- integer(1000)
for(i in 1:length(simulations)) {
    simulations[i] <- simulate_test(size = size, p = p)
}
mean(simulations)

median(simulations)

ggplot(data.frame(x = simulations), aes(x = x)) +
    geom_histogram(aes(y = ..density..), bins = 50, fill = 'grey70') +
    geom_density(color = 'blue') +
    ggtitle('Distribution of simulation results')

dist <- dbinom(x = 0:size, size = size, prob = p)
ggplot(data.frame(x = 0:size,                          
                  prob = dist,
                  label = paste0(round(100 * dist, digits = 2), '%')),
       aes(x = x, y = prob, label = label)) +
    geom_bar(stat = 'identity', fill = 'grey50') +
    geom_text(vjust = -0.5)

(p_all_correct <- dbinom(x = size, size = size, prob = p))

1 / p_all_correct

geom_dist <- data.frame(x = 0:5000,
                        y = dgeom(0:5000, prob = dbinom(x = size, size = size, prob = p)))
cut_point50 <- qgeom(0.50, prob = dbinom(x = size, size = size, prob = p))
cut_point95 <- qgeom(0.95, prob = dbinom(x = size, size = size, prob = p))
ggplot(geom_dist, aes(x = x, y = y)) +
    geom_polygon(data = rbind(data.frame(x = 0, y = 0),
                              geom_dist[geom_dist$x < cut_point95,],
                              data.frame(x = cut_point95, y = 0)),
                 fill = 'grey70') +
    geom_polygon(data = rbind(data.frame(x = 0, y = 0),
                              geom_dist[geom_dist$x < cut_point50,],
                              data.frame(x = cut_point50, y = 0)),
                 fill = 'grey50') +
    geom_path(stat = 'identity', color = 'blue')

p_pass <- dbinom(x = 4:5, size = size, prob = p) |> sum()
1 / p_pass

simulations2 <- integer(1000)
for(i in 1:length(simulations2)) {
    simulations2[i] <- simulate_test(size = size, p = p, stop_score = 0.8)
}
mean(simulations2)

median(simulations2)

qgeom(0.50, prob = p_pass)

qgeom(0.95, prob = p_pass)