library(tidyverse)
library(samplr)
library(patchwork)

theme_set(theme_minimal(20))
red <- "#8C1A24"
blue <- "#375C8C"


# Exercise 1: generate samples from a Normal distribution -----------

sampler_mh(
  start = , # fill this field
  distr_name = , # fill this field
  distr_params = # fill this field
)


# Hint: 
# run vignette(topic="supported-distributions", package="samplr")
# to see which distributions are supported and which parameters they accept



# Exercise 2: How does the acceptance rate change? ------------------------
small_proposal_width <- sampler_mh(
  start = , # fill this field
  distr_name = , # fill this field
  distr_params = , # fill this field
  sigma_prop = # fill this field
)$`Acceptance Ratio`


large_proposal_width <- sampler_mh(
  start = , # fill this field
  distr_name = , # fill this field
  distr_params = , # fill this field
  sigma_prop = # fill this field
)$`Acceptance Ratio`

print(small_proposal_width, large_proposal_width)

# Exercise 3: Sample from a Beta(4,4) ------------------------------------------

sampler_mh(
  start = , # fill this field
  distr_name = , # fill this field
  distr_params = # fill this field
)


# Hint: 
# run vignette(topic="supported-distributions", package="samplr")
# to see which distributions are supported and which parameters they accept


# Extra: some plots from slides -------------------------------------------

# In the slides we print the proposal width / acceptance relationship:
a <- c()

s_props <- seq(.01, 100, .5)
# iterate over several proposal widths...
for (s_prop in s_props){
  set.seed(1)
  # ... get acceptance ratio for each width
  a <- c(a, sampler_mh(
    start = 0, 
    distr_name = "norm", 
    distr_params = c(3, sqrt(5)), 
    sigma_prop = s_prop
  )$`Acceptance Ratio`)
}

tibble(s = s_props, a) %>% 
  ggplot(aes(s,a)) + 
  geom_path() +
  xlab("Proposal Width") + 
  ylab("Acceptance Ratio") + 
  theme(axis.title = element_text(size = 20))

# We also show how high acceptance might mean low coverage
set.seed(1)
tibble(x = sampler_mh(
  start = 0, 
  distr_name = "norm", 
  distr_params = c(3, sqrt(5)), 
  sigma_prop = s_props[1]
)$Samples[,1]) %>% 
  ggplot(aes(x)) + 
  geom_histogram(aes(y=after_stat(density))) + 
  geom_function(fun=\(x){dnorm(x, 3, sqrt(5))}, linewidth=1.5) + 
  xlim(-4, 10)

set.seed(1)
tibble(x = sampler_mh(
  start = 0, 
  distr_name = "norm", 
  distr_params = c(3, sqrt(5)), 
  sigma_prop = s_props[100]
)$Samples[,1]) %>% 
  ggplot(aes(x)) + 
  geom_histogram(aes(y=after_stat(density))) + 
  geom_function(fun=\(x){dnorm(x, 3, sqrt(5))}, linewidth=1.5) + 
  xlim(-4, 10)


# Exercise 4: Anchoring ---------------------------------------------------
set.seed(1)
starts <- c(60, 320)

# Feel free to change these values! ⬇️
target_sd <- 30
proposal_variance <- 20
n.max <- 200 # maximum number of samples we're going to test
# Feel free to change these values! ⬆️


# Initialise vector
responses <- c()

# Loop over starting points
for (start in starts){
  # loop over possible n values
  for (n in 1:n.max){
    samples <- sampler_mh(
      start = start, 
      distr_name = "norm", 
      distr_params = c(157, target_sd), 
      sigma_prop = proposal_variance, 
      iterations = n
    )$Samples[n] # choose the nth sample
    
    # add samples to vector
    responses <- c(responses, samples)
  }
}

# Dataframe of responses
df <- tibble(
  r=responses, s = rep(starts, each=n.max), n=rep(1:n.max, 2)
) 

# Plot
df %>% 
  ggplot(aes(n, r)) + 
  geom_line(aes(group=s, color=factor(s))) + 
  geom_hline(yintercept = 273, color=red) +
  geom_hline(yintercept = 89, color=blue) +
  scale_color_manual(values = c(blue, red)) + 
  theme(axis.text = element_text(size = 20), axis.title = element_text(size=30)) +
  geom_hline(yintercept = 157, linetype = "dashed", color="black", linewidth=1.1) +
  geom_point(x=0, y=320, color=red, size=5) +
  geom_point(x=0, y=60, color=blue, size=5) + 
  xlab("N Samples") + ylab("Response") + 
  theme(legend.position = "none") + 
  annotate(geom="text", x=150, y=280, label="Mean High Anchor", color=red) + 
  annotate(geom="text", x=150, y=75, label="Mean Low Anchor", color=blue) 



# Exercise 5 Sampling from other algorithms --------------------------------------------------------------
## Setup - do not edit..
# Setup: names, parameters, weights
names <- rep("norm", 2)
parameters <- list(c(-2, .25), c(2, .25))
weights <- rep(.5, 2)

# Setup: some convenience functions/variables for plotting purposes
pdf <- function(x){
  .5 * dnorm(x, -2, .25) + .5 * dnorm(x, 2, .25)
}

  # We make a side plot with what the density should look like
side.plot <- tibble(x = seq(-3, 3, .01)) %>% 
  mutate(y = pdf(x)) %>% 
  ggplot() + 
  geom_line(aes(x, y)) + 
  coord_flip() + 
  xlim(-3, 3) + 
  theme_void()

hist_side_plot <- function(df){
  side.plot + 
    geom_histogram(data = df, mapping=aes(x=y, y=after_stat(density)))
}

main_plot <- function(df){
  df %>% 
    ggplot(aes(x, y)) + 
    geom_point() + 
    geom_line() + 
    ylim(-3, 3) + 
    xlab("Iteration") + ylab("Sample")
}

full_plot <- function(df){
  main_plot(df) + hist_side_plot(df) + plot_layout(design = "aaaab")
}


## Sampling (edit here ⬇️): 


# hint: don't know what a parameter does? Get the documentation by prepending a function with ?
# e.g. ?sampler_hmc

# Metropolis Hastings
set.seed(1)
y <- sampler_mh(
  start = -2, 
  distr_name = names, 
  distr_params = parameters, 
  weights = weights,
  sigma_prop = .5
)$Samples[,1]

df <- tibble(y) %>% 
  mutate(x = 1:n()) %>% 
  mutate(d = pdf(y))

full_plot(df)

# MC3
set.seed(1)
y <- sampler_mc3(
  start = -2, 
  distr_name = names, 
  distr_params = parameters, 
  weights = weights,
  sigma_prop = , # fill this field
  nChains = # fill this field
)$Samples[,,1] # notice we take 1st slice (= cold chain)

df <- tibble(y) %>% 
  mutate(x = 1:n()) %>% 
  mutate(d = pdf(y))

full_plot(df)

# HMC
y <- sampler_hmc(
  start = -2, 
  distr_name = names, 
  distr_params = parameters, 
  weights = weights,
  epsilon = ,  # fill this field
  L =  # fill this field
)$Samples[,1]

df <- tibble(y) %>% 
  mutate(x = 1:n()) %>% 
  mutate(d = pdf(y))

full_plot(df)

# REC
set.seed(1)
y <- sampler_rec(
  start = -2, 
  distr_name = names, 
  distr_params = parameters, 
  weights = weights,
  epsilon = ,  # fill this field
  L = , 
  alpha =  # fill this field
)$Samples[,1]

df <- tibble(y) %>% 
  mutate(x = 1:n()) %>% 
  mutate(d = pdf(y))

full_plot(df)

# MCREC
set.seed(1)
y <- sampler_mcrec(
  start = -2, 
  distr_name = names, 
  distr_params = parameters, 
  weights = weights,
  nChains = , # fill this field
  epsilon = ,  # fill this field
  L = , # fill this field
  alpha =  # fill this field
)$Samples[,,1] # notice we take 1st slice (= cold chain)

df <- tibble(y) %>% 
  mutate(x = 1:n()) %>% 
  mutate(d = pdf(y))

full_plot(df)
