library(mgcv)
library(ggplot2)
library(itsadug)
library(tidyverse)

subcortex <- read.table("subcortex_clusters.txt", sep="\t", header=TRUE)

# As factor
subcortex$SubjID <- as.factor(subcortex$SubjID)

subcortex$FSIQ_high2_F <- as.factor(subcortex$FSIQ_high2)
subcortex$FSIQ_high2_OF <- as.ordered(subcortex$FSIQ_high2)

subcortex_kids <- subcortex[subcortex$Age < 20,]
subcortex_adults <- subcortex[subcortex$Age > 19.9,]

cluster1_model <- cluster_1_kids ~  s(Age) + Sex
fit_cluster1_model <- gamm(cluster1_model, data = subcortex, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster1_model$gam)
plot(fit_cluster1_model$gam)

cluster2_model <- cluster_2_kids ~  s(Age) + Sex
fit_cluster2_model <- gamm(cluster2_model, data = subcortex, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster2_model$gam)
plot(fit_cluster2_model$gam)

cluster3_model <- cluster_3_kids ~  s(Age) + Sex
fit_cluster3_model <- gamm(cluster3_model, data = subcortex, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster3_model$gam)
plot(fit_cluster3_model$gam)

cluster4_model <- cluster_4_kids ~  s(Age) + Sex
fit_cluster4_model <- gamm(cluster4_model, data = subcortex, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster4_model$gam)
plot(fit_cluster4_model$gam)

cluster5_model <- cluster_5_kids ~  s(Age) + Sex
fit_cluster5_model <- gamm(cluster5_model, data = subcortex, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster5_model$gam)
plot(fit_cluster5_model$gam)



cluster1_kidsmodel <- cluster_1_kids ~  s(Age) + Sex
fit_cluster1_kidsmodel <- gamm(cluster1_kidsmodel, data = subcortex_kids, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster1_kidsmodel$gam)
plot(fit_cluster1_kidsmodel$gam)

cluster2_kidsmodel <- cluster_2_kids ~  s(Age) + Sex
fit_cluster2_kidsmodel <- gamm(cluster2_kidsmodel, data = subcortex_kids, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster2_kidsmodel$gam)
plot(fit_cluster2_kidsmodel$gam)

cluster3_kidsmodel <- cluster_3_kids ~  s(Age) + Sex
fit_cluster3_kidsmodel <- gamm(cluster3_kidsmodel, data = subcortex_kids, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster3_kidsmodel$gam)
plot(fit_cluster3_kidsmodel$gam)

cluster4_kidsmodel <- cluster_4_kids ~  s(Age) + Sex
fit_cluster4_kidsmodel <- gamm(cluster4_kidsmodel, data = subcortex_kids, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster4_kidsmodel$gam)
plot(fit_cluster4_kidsmodel$gam)

cluster5_kidsmodel <- cluster_5_kids ~  s(Age) + Sex
fit_cluster5_kidsmodel <- gamm(cluster5_kidsmodel, data = subcortex_kids, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster5_kidsmodel$gam)
plot(fit_cluster5_kidsmodel$gam)


cluster1_adultsmodel <- cluster_1_kids ~  s(Age) + Sex
fit_cluster1_adultsmodel <- gamm(cluster1_adultsmodel, data = subcortex_adults, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster1_adultsmodel$gam)
plot(fit_cluster1_adultsmodel$gam)

cluster2_adultsmodel <- cluster_2_kids ~  s(Age) + Sex
fit_cluster2_adultsmodel <- gamm(cluster2_adultsmodel, data = subcortex_adults, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster2_adultsmodel$gam)
plot(fit_cluster2_adultsmodel$gam)

cluster3_adultsmodel <- cluster_3_kids ~  s(Age) + Sex
fit_cluster3_adultsmodel <- gamm(cluster3_adultsmodel, data = subcortex_adults, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster3_adultsmodel$gam)
plot(fit_cluster3_adultsmodel$gam)

cluster4_adultsmodel <- cluster_4_kids ~  s(Age) + Sex
fit_cluster4_adultsmodel <- gamm(cluster4_adultsmodel, data = subcortex_adults, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster4_adultsmodel$gam)
plot(fit_cluster4_adultsmodel$gam)

cluster5_adultsmodel <- cluster_5_kids ~  s(Age) + Sex
fit_cluster5_adultsmodel <- gamm(cluster5_adultsmodel, data = subcortex_adults, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster5_adultsmodel$gam)
plot(fit_cluster5_adultsmodel$gam)




pdf("cluster_graphs_lifespan.pdf")
par(mfrow=c(3,2),
    oma = c(2,2,0,0),              # two rows of text at the outer left and bottom margin
    mar = c(1,1,3,3) ,             # spave for one row of text at ticks and to separate plots
    mgp = c(2, 0.5, 0),              # axis label at 2 rows distance, tick labels at 0 row
    xpd = NA)  
plot(fit_cluster1_model$gam, main = "", xlab = "", ylab = "", col = "royalblue4" , shade = TRUE, rug = FALSE, bty="n", shade.col = "cornflowerblue", cex.lab=1.5, cex.axis=1.5)
plot(fit_cluster2_model$gam, main = "", xlab = "", ylab = "", col = "lightseagreen", shade = TRUE, rug = FALSE, bty="n", shade.col = "aquamarine", cex.lab=1.5, cex.axis=1.5)
plot(fit_cluster3_model$gam, main = "", xlab = "", ylab = "", col = "mediumorchid4", shade = TRUE, rug = FALSE, bty="n", shade.col = "darkorchid1", cex.lab=1.5, cex.axis=1.5)
plot(fit_cluster4_model$gam, main = "", xlab = "", ylab = "", col = "steelblue4", shade = TRUE, rug = FALSE,bty="n", shade.col = "lightskyblue4", cex.lab=1.5, cex.axis=1.5)
plot(fit_cluster5_model$gam, main = "", xlab = "", ylab = "", col = "lightsalmon3", shade = TRUE, rug = FALSE,bty="n", shade.col = "sandybrown", cex.lab=1.5, cex.axis=1.5)
dev.off()

pdf("cluster_graphs_kids2.pdf")
par(mfrow=c(3,2),
    oma = c(2,2,0,0),              # two rows of text at the outer left and bottom margin
    mar = c(1,1,3,3) ,             # spave for one row of text at ticks and to separate plots
    mgp = c(2, 0.5, 0),              # axis label at 2 rows distance, tick labels at 0 row
    xpd = NA)  
plot(fit_cluster1_kidsmodel$gam, main = "", xlab = "", ylab = "", col = "royalblue4" , shade = TRUE, rug = FALSE, bty="n", shade.col = "cornflowerblue", cex.lab=1.5, cex.axis=1.5)
plot(fit_cluster2_kidsmodel$gam, main = "", xlab = "", ylab = "", col = "lightseagreen", shade = TRUE, rug = FALSE, bty="n", shade.col = "aquamarine", cex.lab=1.5, cex.axis=1.5)
plot(fit_cluster3_kidsmodel$gam, main = "", xlab = "", ylab = "", col = "mediumorchid4", shade = TRUE, rug = FALSE, bty="n", shade.col = "darkorchid1", cex.lab=1.5, cex.axis=1.5)
plot(fit_cluster4_kidsmodel$gam, main = "", xlab = "", ylab = "", col = "steelblue4", shade = TRUE, rug = FALSE,bty="n", shade.col = "lightskyblue4", cex.lab=1.5, cex.axis=1.5)
plot(fit_cluster5_kidsmodel$gam, main = "", xlab = "", ylab = "", col = "lightsalmon3", shade = TRUE, rug = FALSE,bty="n", shade.col = "sandybrown", cex.lab=1.5, cex.axis=1.5)
dev.off()


pdf("cluster_graphs_adults.pdf")
par(mfrow=c(3,2),
    oma = c(2,2,0,0),              # two rows of text at the outer left and bottom margin
    mar = c(1,1,3,3) ,             # spave for one row of text at ticks and to separate plots
    mgp = c(2, 0.5, 0),              # axis label at 2 rows distance, tick labels at 0 row
    xpd = NA)  
plot(fit_cluster1_adultsmodel$gam, main = "", xlab = "", ylab = "", col = "royalblue4" , shade = TRUE, rug = FALSE, bty="n", shade.col = "cornflowerblue", cex.lab=1.5, cex.axis=1.5)
plot(fit_cluster2_adultsmodel$gam, main = "", xlab = "", ylab = "", col = "lightseagreen", shade = TRUE, rug = FALSE, bty="n", shade.col = "aquamarine", cex.lab=1.5, cex.axis=1.5)
plot(fit_cluster3_adultsmodel$gam, main = "", xlab = "", ylab = "", col = "mediumorchid4", shade = TRUE, rug = FALSE, bty="n", shade.col = "darkorchid1", cex.lab=1.5, cex.axis=1.5)
plot(fit_cluster4_adultsmodel$gam, main = "", xlab = "", ylab = "", col = "steelblue4", shade = TRUE, rug = FALSE,bty="n", shade.col = "lightskyblue4", cex.lab=1.5, cex.axis=1.5)
plot(fit_cluster5_adultsmodel$gam, main = "", xlab = "", ylab = "", col = "lightsalmon3", shade = TRUE, rug = FALSE,bty="n", shade.col = "sandybrown", cex.lab=1.5, cex.axis=1.5)
dev.off()





##### Interactions with IQ

	

cluster1_IQmodel0 <- cluster_1_kids ~  s(Age) + FSIQ_high2_F + Sex
fit_cluster1_IQmodel0 <- gamm(cluster1_IQmodel0, data = subcortex, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster1_IQmodel0$gam)

cluster1_IQmodel0x <- cluster_1_kids ~  s(Age) + FSIQ_high2_F + Sex + EstimatedTotalIntraCranialVol
fit_cluster1_IQmodel0x <- gamm(cluster1_IQmodel0x, data = subcortex, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster1_IQmodel0x$gam)

cluster1_IQmodel <- cluster_1_kids ~  ti(Age) + FSIQ_high2_OF + ti(Age, by=FSIQ_high2_OF) + Sex
fit_cluster1_IQmodel <- gamm(cluster1_IQmodel, data = subcortex, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster1_IQmodel$gam)

cluster1_IQmodel <- cluster_1_kids ~  ti(Age) + FSIQ_high2_OF + ti(Age, by=FSIQ_high2_OF) + Sex
fit_cluster1_IQmodel <- gamm(cluster1_IQmodel, data = subcortex, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster1_IQmodel$gam)


cluster2_IQmodel0 <- cluster_2_kids ~  s(Age) + FSIQ_high2_F + Sex
fit_cluster2_IQmodel0 <- gamm(cluster2_IQmodel0, data = subcortex, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster2_IQmodel0$gam)

cluster2_IQmodel <- cluster_2_kids ~  ti(Age) + FSIQ_high2_OF + ti(Age, by=FSIQ_high2_OF) + Sex
fit_cluster2_IQmodel <- gamm(cluster2_IQmodel, data = subcortex, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster2_IQmodel$gam)


cluster3_IQmodel0 <- cluster_3_kids ~  s(Age) + FSIQ_high2_OF + Sex
fit_cluster3_IQmodel0 <- gamm(cluster3_IQmodel0, data = subcortex, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster3_IQmodel0$gam)

cluster3_IQmodel0x <- cluster_3_kids ~  s(Age) + FSIQ_high2_F + Sex + EstimatedTotalIntraCranialVol
fit_cluster3_IQmodel0x <- gamm(cluster3_IQmodel0x, data = subcortex, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster3_IQmodel0x$gam)

cluster3_IQmodel <- cluster_3_kids ~  ti(Age) + FSIQ_high2_OF + ti(Age, by=FSIQ_high2_OF) + Sex
fit_cluster3_IQmodel <- gamm(cluster3_IQmodel, data = subcortex, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster3_IQmodel$gam)



cluster4_IQmodel0 <- cluster_4_kids ~  s(Age) + FSIQ_high2_F + Sex
fit_cluster4_IQmodel0 <- gamm(cluster4_IQmodel0, data = subcortex, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster4_IQmodel0$gam)

cluster4_IQmodel0x <- cluster_4_kids ~  s(Age) + FSIQ_high2_F + Sex + EstimatedTotalIntraCranialVol
fit_cluster4_IQmodel0x <- gamm(cluster4_IQmodel0x, data = subcortex, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster4_IQmodel0x$gam)

cluster4_IQmodel <- cluster_4_kids ~  ti(Age) + FSIQ_high2_OF + ti(Age, by=FSIQ_high2_OF) + Sex
fit_cluster4_IQmodel <- gamm(cluster4_IQmodel, data = subcortex, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster4_IQmodel$gam)


cluster5_IQmodel0 <- cluster_5_kids ~  s(Age) + FSIQ_high2_F + Sex
fit_cluster5_IQmodel0 <- gamm(cluster5_IQmodel0, data = subcortex, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster5_IQmodel0$gam)

cluster5_IQmodel <- cluster_5_kids ~  ti(Age) + FSIQ_high2_OF + ti(Age, by=FSIQ_high2_OF) + Sex
fit_cluster5_IQmodel <- gamm(cluster5_IQmodel, data = subcortex, family = "gaussian", discrete = TRUE, random=list(SubjID=~1))
summary(fit_cluster5_IQmodel$gam)


#### Produce plots


# Cluster 1 
grid <- crossing( 
  Age = seq(from = 4, to = 90, by = 1), 
  FSIQ_high2_OF = ordered(c("High", "Low")), 
  Sex = "F"
) 

fit <- predict(fit_cluster1_IQmodel$gam, newdata = grid, se.fit = TRUE) 

grid <- grid %>% 
  mutate(fit2 = fit$fit, se = fit$se.fit) %>%
  mutate(
    ymin = fit2 + qnorm(.025) * se,
    ymax = fit2 + qnorm(.975) * se
)

cluster1_plot <- ggplot(grid, aes(x = Age, y = fit2, group = FSIQ_high2_OF, color = FSIQ_high2_OF), size=1.5) + 
  geom_line(size=1.25) + labs (x = "Age (years)", y = "Volume (mm3)") + theme(axis.title.x = element_text(size = 12),
                                                                                 axis.title.y = element_text(size = 12),
                                                                                 axis.line = element_line(size=1.25),
                                                                                 axis.text.x = element_text(size = 12, color="black"),
                                                                                 axis.text.y = element_blank(),
                                                                              plot.title = element_text(hjust = 0.5, size=14),
                                                                              legend.position = "none") + ggtitle("Cluster 1") +
  geom_line(aes(y = ymin), linetype = "dashed") + 
  geom_line(aes(y = ymax), linetype = "dashed") +
  theme(panel.grid.major = element_blank(), panel.grid.minor = element_blank(), panel.background = element_blank(), axis.line = element_line(colour = "black"))


# Cluster 3

grid <- crossing( 
  Age = seq(from = 4, to = 90, by = 1), 
  FSIQ_high2_OF = ordered(c("High", "Low")), 
  Sex = "F"
) 

fit <- predict(fit_cluster3_IQmodel$gam, newdata = grid, se.fit = TRUE) 

grid <- grid %>% 
  mutate(fit2 = fit$fit, se = fit$se.fit) %>%
  mutate(
    ymin = fit2 + qnorm(.025) * se,
    ymax = fit2 + qnorm(.975) * se
  )


cluster3_plot <- ggplot(grid, aes(x = Age, y = fit2, group = FSIQ_high2_OF, color = FSIQ_high2_OF), size=1.5) + 
  geom_line(size=1.25) + labs (x = "Age (years)") + theme(axis.title.x = element_text(size = 12),
                                                                                 axis.title.y = element_blank(),
                                                                                 axis.line = element_line(size=1.25),
                                                                                 axis.text.x = element_text(size = 12, color="black"),
                                                                                 axis.text.y = element_blank(),
                                                          plot.title = element_text(hjust = 0.5, size=14),
                                                          legend.position = "none") + ggtitle("Cluster 3") +
  geom_line(aes(y = ymin), linetype = "dashed") + 
  geom_line(aes(y = ymax), linetype = "dashed") +
  theme(panel.grid.major = element_blank(), panel.grid.minor = element_blank(), panel.background = element_blank(), axis.line = element_line(colour = "black"))

# Cluster 4

grid <- crossing( 
  Age = seq(from = 4, to = 90, by = 1), 
  FSIQ_high2_OF = ordered(c("High", "Low")), 
  Sex = "F"
) 

fit <- predict(fit_cluster4_IQmodel$gam, newdata = grid, se.fit = TRUE) 

grid <- grid %>% 
  mutate(fit2 = fit$fit, se = fit$se.fit) %>%
  mutate(
    ymin = fit2 + qnorm(.025) * se,
    ymax = fit2 + qnorm(.975) * se
  )

cluster4_plot <- ggplot(grid, aes(x = Age, y = fit2, group = FSIQ_high2_OF, color = FSIQ_high2_OF), size=1.5) + 
  geom_line(size=1.25) + labs (x = "Age (years)") + theme(axis.title.x = element_text(size = 12),
                                                                                 axis.title.y = element_blank(),
                                                                                 axis.line = element_line(size=1.25),
                                                                                 axis.text.x = element_text(size = 12, color="black"),
                                                                                 axis.text.y = element_blank(),
                                                          plot.title = element_text(hjust = 0.5, size=14),
                                                          legend.position = "none") + ggtitle("Cluster 4") +
  geom_line(aes(y = ymin), linetype = "dashed") + 
  geom_line(aes(y = ymax), linetype = "dashed") +
  theme(panel.grid.major = element_blank(), panel.grid.minor = element_blank(), panel.background = element_blank(), axis.line = element_line(colour = "black"))

multiplot(cluster1_plot,cluster4_plot,cluster3_plot, cols=3, nrow = 2)
