library(tidyverse)
library(ggdist)
library(distributional)
library(patchwork) # install.packages("patchwork")

# implementation of https://docs.growthbook.io/statistics/details
calculate_posterior <- function(metric_stats, mu_rel_prior = 0, s2_rel_prior = 0.09) {
    # extract sample sizes:
    n_C <- metric_stats['control', 'sample_size']
    n_T <- metric_stats['treatment', 'sample_size']

    # extract sample means:
    mu_C <- metric_stats['control', 'sample_mean']
    mu_T <- metric_stats['treatment', 'sample_mean']

    # extract sample variances:
    s2_C <- metric_stats['control', 'sample_variance']
    s2_T <- metric_stats['treatment', 'sample_variance']

    # calculate lift (relative effect)
    delta_rel <- (mu_T - mu_C)/mu_C
    s2_delta_rel <- ((s2_C * mu_T**2)/(mu_C**4 * n_C)) + ((s2_T)/(mu_C**2 * n_T))

    # mean and variance for the posterior of relative effect:
    mu_rel_posterior <- (
        # Numerator:
        ((mu_rel_prior/s2_rel_prior) + (delta_rel/s2_delta_rel)) /
            # Denominator:
            ((1/s2_rel_prior) + (1/s2_delta_rel))
    )
    s2_rel_posterior <- 1 / ((1/s2_rel_prior) + (1/s2_delta_rel))

    delta_rel_posterior <- dist_normal(mu_rel_posterior, sqrt(s2_rel_posterior))

    return(list(
        quantities = list(
            sample_sizes = c(control = n_C, treatment = n_T),
            sample_means = c(control = mu_C, treatment = mu_T),
            sample_variances = c(control = s2_C, s2_T),
            lift = c(estimate = delta_rel, variance = s2_delta_rel)
        ),
        distribution = delta_rel_posterior
    ))
}

example_stats <- matrix(
    c(
        # control:
        1e3, # sample size
        0.13, # sample mean
        0.13 * (1 - 0.13), # sample variance
        # treatment:
        1e3, # sample size
        0.145, # sample mean (10% lift from control)
        0.145 * (1 - 0.145) # sample variance
    ),
    nrow = 2, ncol = 3, byrow = TRUE,
    dimnames = list(
        c("control", "treatment"),
        c("sample_size", "sample_mean", "sample_variance")
    )
)

example_posterior <- calculate_posterior(example_stats)

theme_set(theme_ggdist(base_family = "Arial"))

# prior
p1 <- tibble(dist = dist_normal(0, 0.3))