T <- 0:4
Y1 <- 20 + 5*T
Y2 <- 10 + 3*T
Y3 <- 40 - 2*T
plot(0, 0, type="n", xlim=c(0,4), ylim=c(0, 50), xlab="T", ylab="Y")
lines(T, Y1, col="red")
lines(T, Y2, col="green")
lines(T, Y3, col="blue")

T <- seq(0, 4, 0.01)
Y1 <- 20 + 5*T - 2*T*T
Y2 <- 10 + 3*T + 1.5*T*T
Y3 <- 40 - 2*T + 0.5*T*T
plot(0, 0, type="n", xlim=c(0,4), ylim=c(0, 50), xlab="T", ylab="Y")
lines(T, Y1, col="red")
lines(T, Y2, col="green")
lines(T, Y3, col="blue")

T <- seq(0, 4, 0.01)
Y1 <- 20 + 5*T - 2*T*T + 0.5*T*T*T
Y2 <- 10 + 0*T + 3*T*T - 0.5*T^3
Y3 <- 30 + 2*T + 3*T*T - 1*T^3
plot(0, 0, type="n", xlim=c(0,4), ylim=c(0, 50), xlab="T", ylab="Y")
lines(T, Y1, col="red")
lines(T, Y2, col="green")
lines(T, Y3, col="blue")

T <- 0:4
T1 <- c(0, 1, 2, 2, 2)
T2 <- c(0, 0, 0, 1, 2)
Y1 <- 20 + 5*T1 - 2*T2
Y2 <- 10 + 3*T1 + 2*T2
Y3 <- 40 - 2*T1 - 2*T2
plot(0, 0, type="n", xlim=c(0,4), ylim=c(0, 50), xlab="T", ylab="Y")
lines(T, Y1, col="red")
lines(T, Y2, col="green")
lines(T, Y3, col="blue")


T <- seq(0, 4, 0.01)
T1 <- T
T1[T1 > 2] <- 2
T2 <- T - T1
T <- 0:4
T1 <- c(0, 1, 2, 2, 2)
T2 <- c(0, 0, 0, 1, 2)
Y1 <- 20 + 5*T1 - 2*T1^2 - 3*T2 -4*T2^2
Y2 <- 10 + 3*T1 + 1.5*T1^2 + 7*T2 - 2*T2^2
Y3 <- 40 - 4*T1 + 0.5*T1^2 + 0.5*T2^2
plot(0, 0, type="n", xlim=c(0,4), ylim=c(0, 50), xlab="T", ylab="Y")
lines(T, Y1, col="red")
lines(T, Y2, col="green")
lines(T, Y3, col="blue")


T <- seq(0, 8, 0.01)
T1 <- T
T1[T1 > 3] <- 3
T2 <- T - T1
T2[T2 > 3] <- 3
T3 <- T - T1 - T2

T <- 0:8
T1 <- c(0, 1, 2, 3, 3, 3, 3, 3, 3)
T2 <- c(0, 0, 0, 0, 1, 2, 3, 3, 3)
T3 <- c(0, 0, 0, 0, 0, 0, 0, 1, 2)
Y1 <- 20 + 5*T1 - 2*T1^2 - 3*T2 + 1.5*T2^2 + 4*T3 + 2*T3^2
Y2 <- 10 + 3*T1 + 1.5*T1^2 + 7*T2 - 2*T2^2 - 3*T3 + 0.5*T3^2
Y3 <- 40 - 4*T1 + 0.5*T1^2 + 0.5*T2^2 - 0.5*T3
plot(0, 0, type="n", xlim=c(0,8), ylim=c(0, 50), xlab="T", ylab="Y")
lines(T, Y1, col="red")
lines(T, Y2, col="green")
lines(T, Y3, col="blue")



dat1 <-read.table("lecture9ex1.csv", sep=",", header=TRUE)
psych::describe(dat1)

dat1$timec <- dat1$time - 1
dat1$timebase <- dat1$timec
dat1[dat1$timebase > 4, "timebase"] <- 4
dat1$timetreat <- dat1$timec - dat1$timebase
dat1[dat1$timetreat > 6, "timetreat"] <- 6
dat1$timefollow <- dat1$timec - dat1$timebase - dat1$timetreat

library(lme4)
out1a <- lmer(stress ~ 1 + (1|pid), data=dat1, REML=FALSE)
summary(out1a)
out1b <- lmer(stress ~ 1 + timec + (1|pid), data=dat1, REML=FALSE)
summary(out1b)
out1c <- lmer(stress ~ 1 + timebase + timetreat 
              + timefollow + (1|pid), data=dat1, REML=FALSE)
summary(out1c)
anova(out1b, out1c)

datl <- split(dat1, dat1$pid)
stresspred <- lapply(datl, function(x) { predict(lm(stress ~ timebase + timetreat + timefollow, data=x)) })
dat1$stress2 <- do.call(c, stresspred)
dat1_2 <- dat1[dat1$pid%%20 == 0,]
dat1_2$treat <- factor(dat1_2$treat, labels=c("Control", "Treatment"))
library(ggplot2)
ggplot(data=dat1_2, aes(x=time, y=stress, group=pid, colour=treat)) + geom_line()

ggplot(data=dat1_2, aes(x=time, y=stress2, group=pid, colour=treat)) + geom_line()

psi <- coef(summary(out1c))
dat1g1 <- dat1[dat1$pid == 1,]
yhat <- psi[1,1] + psi[2,1]*dat1g1$timebase + psi[3,1]*dat1g1$timetreat + psi[4,1]*dat1g1$timefollow
plot(dat1g1$time, yhat, type="l", xlab="time", ylab="stress")


out1d <- lmer(stress ~ 1 + timebase + timetreat 
              + timefollow + (1 + timebase|pid), 
              data=dat1, REML=FALSE,
              control = lmerControl(optimizer ="Nelder_Mead"))
anova(out1c, out1d)

out1e <- lmer(stress ~ 1 + timebase + timetreat 
              + timefollow + (1 + timetreat|pid), 
              data=dat1, REML=FALSE,
              control = lmerControl(optimizer ="Nelder_Mead"))
anova(out1c, out1e)

out1f <- lmer(stress ~ 1 + timebase + timetreat 
              + timefollow + (1 + timefollow|pid), 
              data=dat1, REML=FALSE,
              control = lmerControl(optimizer ="Nelder_Mead"))
anova(out1c, out1f)

out1g <- lmer(stress ~ 1 + timetreat + timefollow 
              + (1 + timetreat+ timefollow|pid), 
              data=dat1, REML=FALSE,
              control = lmerControl(optimizer ="Nelder_Mead"))
summary(out1g)

tempsd <- c(0.5799, 1.1468, 1.3455)
tempr <- matrix(c(1, 0.2, -0.67,
           0.2, 1, -0.31,
           -0.67, -0.31, 1), 3, 3)
tempcov <- diag(tempsd) %*% tempr %*% diag(tempsd)
temp <- MASS::mvrnorm(1000, c(52.41, -0.88, 0.587), tempcov)

plot(temp[,1], temp[,2], xlab="Baseline Stress", ylab="Stress Change in Treatment Period")
plot(temp[,1], temp[,3], xlab="Baseline Stress", ylab="Stress Change in Follow-up Period")
plot(temp[,2], temp[,3], xlab="Stress Change in Treatment Period", ylab="Stress Change in Follow-up Period")


dat1$sleephrs <- dat1$sleep / 60
dat1$avesleep <- ave(dat1$sleephrs, dat1$pid)
dat1$diffsleep <- dat1$sleephrs - dat1$avesleep

dat1a <- dat1[!duplicated(dat1$pid),]
round(apply(dat1a, 2, mean), 3)

out1h <- lmer(stress ~ 1 + timetreat + timefollow + diffsleep 
              + I(avesleep - 7.075) + I(treat - 0.5) + I(age - 36.638)
              + (1 + timetreat+ timefollow|pid), data=dat1, REML=FALSE,
              control = lmerControl(optimizer ="Nelder_Mead"))
summary(out1h)

out1i <- lmer(stress ~ 1 + timetreat + timefollow + diffsleep 
              + timetreat*I(avesleep - 7.075) + timetreat*I(treat - 0.5) 
              + timetreat*I(age - 36.638) + timefollow*I(avesleep - 7.075) 
              + timefollow*I(treat - 0.5) + timefollow*I(age - 36.638)
              + (1 + timetreat+ timefollow|pid), data=dat1, REML=FALSE,
              control = lmerControl(optimizer ="Nelder_Mead", calc.derivs = FALSE))
summary(out1i)


psi <- coef(summary(out1i))
ytreat <- 52.499 - 1.812*dat1g1$timetreat + 1.185*dat1g1$timefollow
yctrl <- 52.389 + 0.052*dat1g1$timetreat - 0.005*dat1g1$timefollow

plot(dat1g1$time, ytreat, type="n", xlab="time", ylab="stress", ylim=c(40,55))
lines(dat1g1$time, ytreat, col="red")
lines(dat1g1$time, yctrl, col="green")
legend(1, 43, c("Treatment", "Control"), lty=1, col=c("red", "green"))

sumout1i <- summary(out1i)
coef(sumout1i)

psi <- coef(sumout1i)[,1]
vpsi <- vcov(sumout1i)
cvec1 <- c(1, 0, 0, 0, 0, 0.5, 0, 0, 0, 0, 0, 0, 0) 
cvec2 <- c(1, 0, 0, 0, 0, -0.5, 0, 0, 0, 0, 0, 0, 0)   
cvec3 <- c(0, 1, 0, 0, 0, 0, 0, 0, 0.5, 0, 0, 0, 0) 
cvec4 <- c(0, 1, 0, 0, 0, 0, 0, 0, -0.5, 0, 0, 0, 0) 
cvec5 <- c(0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0.5, 0) 
cvec6 <- c(0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, -0.5, 0)   
cmat <- cbind(cvec1, cvec2, cvec3, cvec4, cvec5, cvec6)
est <- t(cmat) %*% psi
se <- sqrt(diag(t(cmat) %*% vpsi %*% cmat))
zval <- est/se
pval <- round(pnorm(-abs(zval)) * 2, 6)
cbind(est, se, zval, pval)

cvec1 <- c(0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0) 
cvec2 <- c(0, 0, 0, 0, 0, 1, 0, 0, 1, 0, 0, 0, 0)   
cvec3 <- c(0, 0, 0, 0, 0, 1, 0, 0, 2, 0, 0, 0, 0)  
cvec4 <- c(0, 0, 0, 0, 0, 1, 0, 0, 3, 0, 0, 0, 0)  
cvec5 <- c(0, 0, 0, 0, 0, 1, 0, 0, 4, 0, 0, 0, 0)  
cvec6 <- c(0, 0, 0, 0, 0, 1, 0, 0, 5, 0, 0, 0, 0) 
cvec7 <- c(0, 0, 0, 0, 0, 1, 0, 0, 6, 0, 0, 0, 0) 
cvec8 <- c(0, 0, 0, 0, 0, 1, 0, 0, 6, 0, 0, 1, 0) 
cvec9 <- c(0, 0, 0, 0, 0, 1, 0, 0, 6, 0, 0, 2, 0) 
cvec10 <- c(0, 0, 0, 0, 0, 1, 0, 0, 6, 0, 0, 3, 0) 
cmat <- cbind(cvec1, cvec2, cvec3, cvec4, cvec5, cvec6, cvec7, cvec8, cvec9, cvec10)
est <- t(cmat) %*% psi
se <- sqrt(diag(t(cmat) %*% vpsi %*% cmat))
zval <- est/se
pval <- round(pnorm(-abs(zval)) * 2, 6)
cbind(est, se, zval, pval)

# Probing interaction between time and tenure0

plot(dat1$time, dat1$jsat, type="n", xlab="Year of Measurement", ylab="Job Satisfaction")
abline(a = 52.06, b = -2.89, col="red")
abline(a = 49.34, b = -0.81, col="green")
abline(a = 46.62, b = 1.27, col="blue")
legend(0, 90, c(0, 2, 4), col=c("red", "green", "blue"), 
       lty=1, title="Tenure at Year 0")

plot(dat1$tenure0, dat1$jsat, type="n", xlab="Tenure at Year 0", ylab="Job Satisfaction")
abline(a = 49.27 - (2.054*-1.36), b = -1.36, col="red")
abline(a = 48.52 - (2.054*-0.32), b = -0.32, col="green")
abline(a = 47.77 - (2.054*0.72), b = 0.72, col="blue")
abline(a = 47.02 - (2.054*1.76), b = 1.76, col="black")
legend(0, 90, c("0", "1", "2", "3"), col=c("red", "green", "blue", "black"), 
       lty=1, title="Year of Measurement")

cvec1 <- c(1, 0, 0, 0, -2.054, 0, 0, 0, 0, 0, 0, 0) 
cvec2 <- c(1, 0, 0, 0, -0.054, 0, 0, 0, 0, 0, 0, 0)  
cvec3 <- c(1, 0, 0, 0, 1.946, 0, 0, 0, 0, 0, 0, 0) 
cvec4 <- c(0, 1, 0, 0, 0, 0, 0, 0, 0, -2.054, 0, 0) 
cvec5 <- c(0, 1, 0, 0, 0, 0, 0, 0, 0, -0.054, 0, 0) 
cvec6 <- c(0, 1, 0, 0, 0, 0, 0, 0, 0, 1.946, 0, 0)  
cmat <- cbind(cvec1, cvec2, cvec3, cvec4, cvec5, cvec6)
est <- t(cmat) %*% psi
se <- sqrt(diag(t(cmat) %*% vpsi %*% cmat))
zval <- est/se
pval <- round(pnorm(-abs(zval)) * 2, 6)
cbind(est, se, zval, pval)

cvec1 <- c(1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0) 
cvec2 <- c(1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0)  
cvec3 <- c(1, 2, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0) 
cvec4 <- c(1, 3, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0) 
cvec5 <- c(0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0) 
cvec6 <- c(0, 0, 0, 0, 1, 0, 0, 0, 0, 1, 0, 0)  
cvec7 <- c(0, 0, 0, 0, 1, 0, 0, 0, 0, 2, 0, 0)  
cvec8 <- c(0, 0, 0, 0, 1, 0, 0, 0, 0, 3, 0, 0)  
cmat <- cbind(cvec1, cvec2, cvec3, cvec4, cvec5, cvec6, cvec7, cvec8)
est <- t(cmat) %*% psi
se <- sqrt(diag(t(cmat) %*% vpsi %*% cmat))
zval <- est/se
pval <- round(pnorm(-abs(zval)) * 2, 6)
cbind(est, se, zval, pval)


###################################

set.seed(123321)
k <- 10
n <- 6
ar1 <- 0.9

test1 <- NULL
test2 <- NULL
for(i in 1:k) {
        S <- matrix(ar1, n, n)
        for(a in 1:n) {
                for(b in 1:n) {
                     S[a, b] <- ar1^abs(a - b)   
                }
        }
        temp1 <- rnorm(n, 0, 1)
        temp1a <- cbind(i, 1:n, temp1)
        test1 <- rbind(test1, temp1a)
        temp2 <- MASS::mvrnorm(1, rep(0, n), S)
        temp2a <- cbind(i, 1:n, temp2)
        test2 <- rbind(test2, temp2a)
}
colnames(test1) <- c("id", "t", "y")
colnames(test2) <- c("id", "t", "y")
test1 <- data.frame(test1)
test2 <- data.frame(test2)

library(ggplot2)
ggplot(data=test1, aes(x=t, y=y, group=id)) + geom_line()
ggplot(data=test2, aes(x=t, y=y, group=id)) + geom_line()

library(nlme)
outtest1 <- lme(y ~ 1 + t, random = ~1|id, data=test1, method="ML", na.action=na.omit)
outtest1a <- lme(y ~ 1 + t, random = ~1|id, data=test1, method="ML", na.action=na.omit, correlation=corARMA(p = 1))
anova(outtest1, outtest1a)

# ctrl <- lmeControl(maxIter=1000, msMaxIter = 1000)

outtest2 <- lme(y ~ 1 + t, random = ~1|id, data=test2, method="ML", na.action=na.omit)
outtest2a <- lme(y ~ 1 + t, random = ~1|id, data=test2, method="ML", na.action=na.omit, correlation=corARMA(p = 1))
anova(outtest2, outtest2a)

###########################

dat3 <- read.table("lecture9ex2.csv", sep=",", header=TRUE)
out3a <- lmer(test ~ 1 + I(year-2017) + (1|timeid), data=dat3, REML=FALSE)
summary(out3a)

dat4 <- read.table("lecture9ex3.csv", sep=",", header=TRUE)
out4a <- lmer(test ~ 1 + I(year-2017) + (1|timeid), data=dat4, REML=FALSE)
out4b <- lmer(test ~ 1 + I(year-2017) + (1|timeid) + (1|schoolid), data=dat4, REML=FALSE)
anova(out4a, out4b)
summary(out4b)

out4c <- lmer(test ~ 1 + I(year-2017) + (1|timeid) + (1 + I(year-2017)|schoolid), 
              data=dat4, REML=FALSE, control = lmerControl(optimizer ="Nelder_Mead"))
anova(out4b, out4c)
summary(out4c)

tempsd <- c(4.818, 2.646)
tempr <- matrix(c(1, 0.31,
                  0.31, 1), 2, 2)
tempcov <- diag(tempsd) %*% tempr %*% diag(tempsd)
temp <- MASS::mvrnorm(1000, c(47.79, 2.2), tempcov)

plot(temp[,1], temp[,2], xlab="School Averaged Test Score in 2017", ylab="School Change in Test Score per Year")

out4d <- lmer(test ~ 1 + I(year-2017)*public + (1|timeid) + (1 + I(year-2017)|schoolid), 
              data=dat4, REML=FALSE, control = lmerControl(optimizer ="Nelder_Mead"))
summary(out4d)


sumout4d <- summary(out4d)
coef(sumout4d)

psi <- coef(sumout4d)[,1]
vpsi <- vcov(sumout4d)
cvec1 <- c(1, 0, 1, 0) 
cvec2 <- c(1, 0, 0, 0)   
cvec3 <- c(0, 1, 0, 1) 
cvec4 <- c(0, 1, 0, 0) 
cmat <- cbind(cvec1, cvec2, cvec3, cvec4)
est <- t(cmat) %*% psi
se <- sqrt(diag(t(cmat) %*% vpsi %*% cmat))
zval <- est/se
pval <- round(pnorm(-abs(zval)) * 2, 6)
cbind(est, se, zval, pval)


psi <- coef(sumout4d)[,1]
vpsi <- vcov(sumout4d)
cvec1 <- c(1, 0, 0, 0) 
cvec2 <- c(1, 1, 0, 0)   
cvec3 <- c(1, 2, 0, 0) 
cvec4 <- c(0, 0, 1, 0) 
cvec5 <- c(0, 0, 1, 1) 
cvec6 <- c(0, 0, 1, 2) 
cmat <- cbind(cvec1, cvec2, cvec3, cvec4, cvec5, cvec6)
est <- t(cmat) %*% psi
se <- sqrt(diag(t(cmat) %*% vpsi %*% cmat))
zval <- est/se
pval <- round(pnorm(-abs(zval)) * 2, 6)
cbind(est, se, zval, pval)


##########################



set.seed(123321)
n <- 10
time <- 6
start <- rnorm(n, 0, 1)
rho0 <- matrix(NA, n, time)
rho3 <- matrix(NA, n, time)
rho5 <- matrix(NA, n, time)
rho7 <- matrix(NA, n, time)
rho9 <- matrix(NA, n, time)
rho0[,1] <- start
rho3[,1] <- start
rho5[,1] <- start
rho7[,1] <- start
rho9[,1] <- start

for(i in 1:(time-1)) {
        rho0[,i+1] <- 0*rho0[,i] + rnorm(n, 0, sqrt(1 - 0))
        rho3[,i+1] <- 0.3*rho3[,i] + rnorm(n, 0, sqrt(1 - 0.3^2))
        rho5[,i+1] <- 0.5*rho5[,i] + rnorm(n, 0, sqrt(1 - 0.5^2))
        rho7[,i+1] <- 0.7*rho7[,i] + rnorm(n, 0, sqrt(1 - 0.7^2))
        rho9[,i+1] <- 0.9*rho9[,i] + rnorm(n, 0, sqrt(1 - 0.9^2))
}

rho0 <- data.frame(id=rep(1:n, each=time), t=rep(1:time, n), y=as.vector(t(rho0)))
rho3 <- data.frame(id=rep(1:n, each=time), t=rep(1:time, n), y=as.vector(t(rho3)))
rho5 <- data.frame(id=rep(1:n, each=time), t=rep(1:time, n), y=as.vector(t(rho5)))
rho7 <- data.frame(id=rep(1:n, each=time), t=rep(1:time, n), y=as.vector(t(rho7)))
rho9 <- data.frame(id=rep(1:n, each=time), t=rep(1:time, n), y=as.vector(t(rho9)))


library(ggplot2)
ggplot(data=rho0, aes(x=t, y=y, group=id)) + geom_line()
ggplot(data=rho3, aes(x=t, y=y, group=id)) + geom_line()
ggplot(data=rho5, aes(x=t, y=y, group=id)) + geom_line()
ggplot(data=rho7, aes(x=t, y=y, group=id)) + geom_line()
ggplot(data=rho9, aes(x=t, y=y, group=id)) + geom_line()

set.seed(123321)
x <- rep(NA, 100)
x[1] <- 1
x[2] <- -0.5
time <- 100
for(i in 3:time) {
        x[i] <- 0.4*x[i-1] - 0.1*x[i-2] + rnorm(1, 0, sqrt(0.5))
}
plot(x, type="l", xlab="Time", ylab="e")


dat1 <-read.table("lecture9ex1.csv", sep=",", header=TRUE)
dat1$timec <- dat1$time - 1
dat1$timebase <- dat1$timec
dat1[dat1$timebase > 4, "timebase"] <- 4
dat1$timetreat <- dat1$timec - dat1$timebase
dat1[dat1$timetreat > 6, "timetreat"] <- 6
dat1$timefollow <- dat1$timec - dat1$timebase - dat1$timetreat

library(lme4)
out1a <- lmer(stress ~ 1 + timetreat + timefollow 
               + (1|pid), 
               data=dat1, REML=FALSE,
               control = lmerControl(optimizer ="Nelder_Mead"))
logLik(out1a)
out1a2 <- lmer(stress ~ 1 + timetreat + timefollow 
              + (1 + timetreat+ timefollow|pid), 
              data=dat1, REML=FALSE,
              control = lmerControl(optimizer ="Nelder_Mead"))
logLik(out1a2)

library(nlme)
out1n <- lme(stress ~ 1 + timetreat + timefollow, 
             random = ~1|pid, data=dat1, method="ML", 
             na.action=na.omit)
logLik(out1n)
out1n2 <- lme(stress ~ 1 + timetreat + timefollow, 
              random = ~1 + timetreat + timefollow|pid, 
              data=dat1, method="ML", na.action=na.omit)
logLik(out1n2)


out1n1 <- lme(stress ~ 1 + timetreat + timefollow, 
              random = ~1|pid, data=dat1, method="ML", 
              na.action=na.omit, 
              correlation=corARMA(form = ~ 1 + timec | pid, p = 1))
summary(out1n1)
anova(out1n, out1n1)

out1n3 <- lme(stress ~ 1 + timetreat + timefollow, 
              random = ~1 + timetreat + timefollow|pid, 
              data=dat1, method="ML", na.action=na.omit, 
              correlation=corARMA(form = ~ 1 + timec | pid, p = 1))
summary(out1n3)
anova(out1n2, out1n3)