# Neyman Scott
library(breheny)
library(magrittr)
ll <- function(v, Y, m) {
  n <- length(m)
  -n*log(v) - 1/(2*v)*sum((Y[,1] -m)^2 + (Y[,2] -m)^2)
}
llv <- function(s2, Y) {
  n <- nrow(Y)
  v <- apply(Y, 1, diff)/sqrt(2)
  -(n/2)*log(s2) - (1/(2*s2))*sum(v^2)
}
lik_plots <- function(n, seed=3) {
  set.seed(seed)
  m <- rnorm(n)
  Y <- matrix(rnorm(2*n, rep(m, 2)), n, 2)
  vv <- seq(0.01, 2, len=299)
  oll <- ll(vv, Y, m)
  pll <- ll(vv, Y, apply(Y, 1, mean))
  mll <- llv(vv, Y)
  plotL(vv, exp(pll-max(pll)), col=pal(3)[1], xlab=expression(sigma^2), ylab=expression("L"*(sigma^2)))
  plotL(vv, exp(mll-max(mll)), col=pal(3)[2], add=TRUE)
  plotL(vv, exp(oll-max(oll)), col=pal(3)[3], add=TRUE)
  mtext(paste0('n = ', n))
  rightlegend(legend=c('Profile', 'Marginal', 'Oracle'), lwd=3, col=pal(3))
}
par(mar=c(4.5, 4.5, 2, 6))
lik_plots(20)

# Linear
set.seed(NULL)
analyze <- function(x, y, s, method, u) {
  if (method == 'mle') {
    fit <- lm(y ~ x + as.factor(s))
    v <- sum(fit$residuals^2)/(n*m)
    b <- coef(summary(fit))['x',]
  } else if (method == 'marginal') {
    fit <- lme4::lmer(y ~ x + (1|s))
    v <- summary(fit)$sigma^2
    b <- coef(summary(fit))['x',]
  } else if (method == 'oracle') {
    yy <- y-u
    fit <- lm(yy ~ x)
    v <- sum(fit$residuals^2)/(n*m)
    b <- coef(summary(fit))['x',]
  } else if (method == 'naive') {
    fit <- lm(y ~ x)
    v <- summary(fit)$sigma^2
    b <- coef(summary(fit))['x',]
  } else if (method == 'diff') {
    n <- length(y)
    yy <- y[2*(1:n)] - y[2*(1:n)-1]
    xx <- x[2*(1:n)] - x[2*(1:n)-1]
    fit <- lm(yy ~ 0 + xx)
    v <- sum(fit$residuals^2)/(n)
    b <- coef(summary(fit))['xx',]
  }
  out <- data.frame(Estimate=b['Estimate'], SE=b['Std. Error'], Variance=v)
  rownames(out) <- method
  out
}
set_rownames <- function(df, rownames) {
  rownames(df) <- rownames
  df
}

# Generate data
set.seed(4)
n <- 100
m <- 2
x <- runif(n*m)
u <- rep(rnorm(n), each=m)
y <- rnorm(m*n, x+u)
s <- rep(1:n, each=m)

# Analysis
rbind(analyze(x, y, s, 'oracle', u),
      analyze(x, y, s, 'marginal'),
      analyze(x, y, s, 'diff'),
      analyze(x, y, s, 'mle'),
      analyze(x, y, s, 'naive')) |>
  set_rownames(c('Oracle', 'Marginal', 'Differencing', 'Profile', 'Naive')) |>
  knitr::kable('latex', booktabs=TRUE, digits=2) |>
  kableExtra::kable_styling(position = "center")

# What if u is correlated with x
set.seed(4)
n <- 1000
m <- 2
u <- rep(rnorm(n), each=m)
x <- rnorm(m*n, u)
y <- rnorm(m*n, x+u)
s <- rep(1:n, each=m)
rbind(analyze(x, y, s, 'oracle', u),
      analyze(x, y, s, 'marginal'),
      analyze(x, y, s, 'diff'),
      analyze(x, y, s, 'mle'),
      analyze(x, y, s, 'naive')) |>
  set_rownames(c('Oracle', 'Marginal', 'Differencing', 'Profile', 'Naive')) |>
  knitr::kable('latex', booktabs=TRUE, digits=2) |>
  kableExtra::kable_styling(position = "center")

# Quadrature: Wacky example
z <- rnorm(1000000)
mean(sqrt(abs(z)+abs(z^3)))
GH <- lme4::GHrule(20, FALSE)
sum(GH$w * sqrt(abs(GH$z)+abs(GH$z^3)))
GH <- lme4::GHrule(100, FALSE)
sum(GH$w * sqrt(abs(GH$z)+abs(GH$z^3)))

# Quadrature: Variance of median
n <- 11
m <- replicate(100000, median(rnorm(n)))
f <- function(x) {
n * choose(n-1, 5) * x^2 * pnorm(x)^5 * (1-pnorm(x))^5
}
GH20 <- lme4::GHrule(20, FALSE)
GH100 <- lme4::GHrule(100, FALSE)
data.frame(Variance=c(
var(m),
pi / (2*n),
sum(GH20$w * f(GH20$z)),
sum(GH100$w * f(GH100$z))),
row.names=c(
  'Monte Carlo ($N=100,000$)',
  'Asymptotic',
  'Gauss-Hermite ($K=20$)',
  'Gauss-Hermite ($K=100$)')) |>
knitr::kable('latex', booktabs=TRUE, digits=4, escape=FALSE) |>
kableExtra::kable_styling(position = "center")

# Logistic
set.seed(3)
n <- 1000
m <- 2
x <- rnorm(n*m)
u <- rep(rnorm(n), m)
f <- function(i, n, x, u) {
  ind <- n*(i-1) + 1:n
  rbinom(n, 1, binomial()$linkinv(x[ind]+u[ind]))
}
y <- sapply(1:m, f, n=n, x=x, u=u) %>% as.integer
s <- rep(1:n, m)

# Naive
fit <- glm(y ~ x, family='binomial')
b1 <- coef(summary(fit))['x',]

# MLE/GLM
fit <- glm(y ~ x + as.factor(s), family='binomial')
b2 <- coef(summary(fit))['x',]

# Conditional
library(survival)
fit <- clogit(y ~ x + strata(s), data.frame(y=y, x=x, s=s))
b3 <- coef(summary(fit))

# Marginal
library(lme4)
fit <- glmer(y ~ x + (1|s), family='binomial')
b4 <- coef(summary(fit))['x',]

# Display
rbind(b1[1:2], b2[1:2], b3[c(1, 3)], b4[1:2]) |>
  set_rownames(c('Naive', 'Profile', 'Conditional', 'Marginal')) |>
  knitr::kable('latex', booktabs=TRUE, digits=2) |>
  kableExtra::kable_styling(position = "center")