# 1. Load functions and Libraries ----------------------------------------------
source("./R/graphon_distribution.R")

# 2. Set constants -------------------------------------------------------------
withr::with_seed(42) # set seed in a local environment

a0 <- c(0, 0)
a1 <- c(1.0)
a2 <- c(0.5, -0.3)

N <- 400 # number of samples for X

X1 = matrix(rnorm(N))
X2 = matrix(rnorm(2 * N), ncol = 2)


# 3. Normally distributed v's --------------------------------------------------
## 3.1 Plot a = 0 ==============================================================
# In this section we plot the conditional density p(u | x) for v ~ N(0,1) and
# a = (0, 0)
p1 <- create_cond_density(a0, dnorm, pnorm, X2)
givenX <- c(0, 0)
p1_plot <- \(x) p1(x, givenX)
plot.function(p1_plot, xlab="u", ylab="p(u |x) ")
title(main=expression("Conditional density for" ~ v %~% N(0, 1)), line=2)
title(
  main = bquote(
    a == (.(paste(a0, collapse = ", "))) ~ "," ~
      N == .(N) ~ "," ~
      x == (.(paste(givenX, collapse = ", ")))
  ),
  line = 1,
  cex = 0.65
)

## 3.2 Plot a != 0, X != 0 =====================================================
p2 <- create_cond_density(a2, dnorm, pnorm, X2)
givenX <- c(-0.5, 0.5)
p2_plot <- \(x) p2(x, givenX)
plot.function(p2_plot, xlab="u", ylab="p(u |x) ")
title(main=expression("Conditional density for" ~ v %~% N(0, 1)), line=2)
title(
  main = bquote(
    a == (.(paste(a2, collapse = ", "))) ~ "," ~
      N == .(N) ~ "," ~
      x == (.(paste(givenX, collapse = ", ")))
  ),
  line = 1,
  cex = 0.65
)

## 3.3 Plot one dimensional X_i ================================================
p3 <- create_cond_density(a1, dnorm, pnorm, X1)
givenX <- c(-0.5)
p3_plot <- \(x) p3(x, givenX)
plot.function(p3_plot, xlab="u", ylab="p(u |x) ")
title(main=expression("Conditional density for" ~ v %~% N(0, 1)), line=2)
title(
  main = bquote(
    a == (.(paste(a1, collapse = ", "))) ~ "," ~
      N == .(N) ~ "," ~
      x == (.(paste(givenX, collapse = ", ")))
  ),
  line = 1,
  cex = 0.65
)

## 3.4 Plot v ~ N(0, 2) ========================================================
dens_func <-\(x) dnorm(x, sd=sqrt(2))
cdf <- \(x) pnorm(x, sd=sqrt(2))
p4 <- create_cond_density(a1, dens_func, cdf, X1)
givenX <- c(-0.5)
p4_plot <- \(x) p4(x, givenX)
plot.function(p4_plot, xlab="u", ylab="p(u |x) ")
title(main=expression("Conditional density for" ~ v %~% N(0, 2)), line=2)
title(
  main = bquote(
    a == (.(paste(a1, collapse = ", "))) ~ "," ~
      N == .(N) ~ "," ~
      x == (.(paste(givenX, collapse = ", ")))
  ),
  line = 1,
  cex = 0.65
)

# 4. Expontial distribution for v's --------------------------------------------
## 4.1 One dimensional X_i =====================================================
# This example has a jump in [0.14, 0.15].
# The zero can not be explained and it's not a probability as:
# integrate(p5_plot, 0.15, 1, subdivisions = 500)
# => 0.952492 with absolute error < 0.00011
lambda <- 1.0
p5 <- create_cond_density(a1, dexp, pexp, X1)
givenX <- c(-0.5)
p5_plot <- \(x) p5(x, givenX)
plot.function(p5_plot, xlab="u", ylab="p(u |x) ")
title(main=expression("Conditional density for" ~ v %~% Exp(1)), line=2)
title(
  main = bquote(
    a == (.(paste(a1, collapse = ", "))) ~ "," ~
      N == .(N) ~ "," ~
      x == (.(paste(givenX, collapse = ", ")))
  ),
  line = 1,
  cex = 0.65
)

## 4.2 Two dimensional X_i =====================================================
lambda <- 1.0
p6 <- create_cond_density(a2, dexp, pexp, X2)
givenX <- c(-0.5)
p6_plot <- \(x) p6(x, givenX)
plot.function(p6_plot, xlab="u", ylab="p(u |x) ")
title(main=expression("Conditional density for" ~ v %~% Exp(1)), line=2)
title(
  main = bquote(
    a == (.(paste(a2, collapse = ", "))) ~ "," ~
      N == .(N) ~ "," ~
      x == (.(paste(givenX, collapse = ", ")))
  ),
  line = 1,
  cex = 0.65
)

## 4.3 Higher Rate of exponential distribtion ==================================
lambda <- 2.0
dens_func <- \(x) dexp(x, rate=lambda)
cdf <- \(x) pexp(x, rate=lambda)
p7 <- create_cond_density(a2, dens_func, cdf, X2)
givenX <- c(-0.5, 0.5)
p7_plot <- \(x) p7(x, givenX)
plot.function(p7_plot, xlab="u", ylab="p(u |x) ")
title(main=expression("Conditional density for" ~ v %~% Exp(2.0)), line=2)
title(
  main = bquote(
    a == (.(paste(a2, collapse = ", "))) ~ "," ~
      N == .(N) ~ "," ~
      x == (.(paste(givenX, collapse = ", ")))
  ),
  line = 1,
  cex = 0.65
)

# 5. Uniform Distribution ------------------------------------------------------

## Two dimensional X_i =========================================================
p8 <- create_cond_density(a2, dunif, punif, X2)
givenX <- c(-0.5, 0.5)
p8_plot <- \(x) p8(x, givenX)
plot.function(p8_plot, xlab="u", ylab="p(u |x) ")
title(main=expression("Conditional density for" ~ v %~% U(0,1)), line=2)
title(
  main = bquote(
    a == (.(paste(a2, collapse = ", "))) ~ "," ~
      N == .(N) ~ "," ~
      x == (.(paste(givenX, collapse = ", ")))
  ),
  line = 1,
  cex = 0.65
)

## 5.2 One dimensional X_i =====================================================
p9 <- create_cond_density(a1, dunif, punif, X1)
givenX <- c(-0.5)
p9_plot <- \(x) p9(x, givenX)
plot.function(p9_plot, xlab="u", ylab="p(u |x) ")
title(main=expression("Conditional density for" ~ v %~% U(0,1)), line=2)
title(
  main = bquote(
    a == (.(paste(a1, collapse = ", "))) ~ "," ~
      N == .(N) ~ "," ~
      x == (.(paste(givenX, collapse = ", ")))
  ),
  line = 1,
  cex = 0.65
)