source('https://myweb.uiowa.edu/pbreheny/7110/f21/notes/fun.R')
library(magrittr)

# Neyman-Scott
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))
}
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
}

# 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)
analyze(x, y, s, 'naive')

# Analysis
Tab <- rbind(analyze(x, y, s, 'mle'),
             analyze(x, y, s, 'naive'),
             analyze(x, y, s, 'marginal'),
             analyze(x, y, s, 'diff'),
             analyze(x, y, s, 'oracle', u))
knitr::kable(Tab, 'latex', booktabs=TRUE, digits=2)

# Illustration of marginal idea
X <- cbind(1, x)
V <- Matrix::bandSparse(n*m, k=0:1, diagonals=list(rep(2, n*m), c(rep(1:0, n), 1)), symmetric=TRUE)
SVD <- svd(V)
A <- SVD$u %*% diag(SVD$d^(-1/2)) %*% t(SVD$v)
yy <- A %*% y
XX <- A %*% X
lm(yy ~ XX) %>% summary()

# 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)
Tab <- rbind(analyze(x, y, s, 'mle'),
             analyze(x, y, s, 'naive'),
             analyze(x, y, s, 'marginal'),
             analyze(x, y, s, 'diff'),
             analyze(x, y, s, 'oracle', u))
knitr::kable(Tab, 'latex', booktabs=TRUE, digits=2)

# Nonlinear ---------------------------------------------------------------

# 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)))
var(m)
pi/(2*n)
f <- function(x) {
  n * choose(n-1, 5) * x^2 * pnorm(x)^5 * (1-pnorm(x))^5
}
GH <- lme4::GHrule(20, FALSE)
sum(GH$w * f(GH$z))
GH <- lme4::GHrule(100, FALSE)
sum(GH$w * f(GH$z))

# 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',]

# Table
Tab <- rbind(b1[1:2], b2[1:2], b3[c(1, 3)], b4[1:2])
rownames(Tab) <- c('Naive', 'Profile', 'Conditional', 'Marginal')
knitr::kable(Tab, 'latex', booktabs=TRUE, digits=2)