#### Modelling pipeline for pial arteries:
## with radius as covariate (for amp and period)
library(dplyr)
library(glmmTMB) #package for fitting mixed models (there are many other packages that can do this)
library(DHARMa) #package for "nice" residual plots
library(emmeans) #package for computing contrasts
library(effects) #package for making "effect"-plots
library(ggplot2)
# Other libraries that need to be installed (but not necessarily loaded): tidyr, readr, ggrepel, plotly, RColorBrewer

setwd("~/Astrocytter/Bloodvessels_sleep/") #Set working directory 

## Read data: Quite time-consuming, so only take the one you need
dataset <- "PialsWT9.csv"
ddr <- read.csv(dataset)


## Filter away noisy measurements
rm(list=setdiff(ls(), c("ddr","dataset")))
source("pipeline_func.R")
dd <- ddr

## Arrange column names and stuff:
dd=rename(dd,mouse=mouse.number)
dd=rename(dd,episode=sequence)
dd$mouse <- factor(dd$mouse)

dd$trialID <- factor(interaction(dd$mouse,dd$trial))

dd$vesselTrial <- factor(interaction(dd$vesselID,dd$trialID,drop=T))
# here I define a new factor which is the combination of each vesselID and trialID. The idea is that
# the vessel numbers which are investigated in different trials are actually *different* vessels 
# (or that they will be sufficiently different to be treated as such)
# This is an important choice!

## What do you want to analyse?
response_values <- c("amp","period","medianepisode")
mode_values <- c("sleep","wake")
linescan_values <- c("lumen")
band_values <- c("VLF", "LF","cardiac","resp")
aggregate.by = mean # mean or median (does not matter if response==meanepisode)


for (i in 1:length(response_values)){
  for (j in 1:length(mode_values)){
    for (k in 1:length(linescan_values)){
      
      # Current analysis
      response = response_values[i]
      mode = mode_values[j]
      linescan=linescan_values[k]
      
      if (response=="medianepisode"){
        band_values = "no"
      }else{
        band_values <- c("VLF", "LF", "cardiac", "resp")
      }
      
      for (l in 1:length(band_values)){
        if (exists("mod")) rm(mod)
        if (exists("mod2")) rm(mod2)
        # Current band:
        band = band_values[l]
        
        # Filter by correlation
        if (linescan=="PVS"){
          if (band=="VLF" | band=="LF"| band=="no"){ #"no"=medianepisode
            ddf <- filter(dd,corrlumen>-1 & correndfoot>-1)
          }else{ #resp, cardiac
            ddf <- filter(dd,corrlumen>0.8 & correndfoot>0.7) #filter(dd,corrlumen>0.8 & correndfoot>0.7) filter(dd,corrlumen>0.9 & correndfoot>0.8)
          }
        }else if(linescan=="lumen"){
          if (band=="VLF" | band=="LF"| band=="no"){ #"no"=medianepisode
            ddf <- filter(dd,corrlumen>-1)
          }else{ #resp, cardiac 
            ddf <- filter(dd,corrlumen>0.8) #filter(dd,corrlumen>0.9)
          }
        }else if(linescan=="endfoot"){
          if (band=="VLF" | band=="LF"| band=="no"){ #"no"=medianepisode
            ddf <- filter(dd,correndfoot>-1)
          }else{ #resp, cardiac 
            ddf <- filter(dd,correndfoot>0.7) #filter(dd,correndfoot>0.7)
          }
        }
        
        ## Subset dataset according to the analysis:
        if (mode=="sleep"){
          ddi = ddf[which(ddf$stage %in% c("Baseline","Awakening","NREM","IS","REM")),]
          ddi$stage = factor(ddi$stage,levels=c("Baseline","NREM","IS","REM","Awakening")) # changes ordering of levels (mostly for plotting)
          levels(ddi$stage)[levels(ddi$stage)=="Baseline"] <- "Quiet" #change name of "Baseline" stage
        }else if (mode=="wake"){
          ddi <- dd[which(dd$stage %in% c("Quiet","Whisking","Locomotion")),]
          ddi$stage <- factor(ddi$stage,levels=c("Quiet","Whisking","Locomotion")) # changes ordering of levels (mostly for plotting)
        }
        
        if (response=="medianepisode"){
          dds = ddi[which(ddi$linescan==linescan),]
        }else{
          dds = ddi[which(ddi$linescan==linescan & ddi$bandname==band),]
        }
        
        ## Aggregate data per episode
        dda <- dds %>% group_by(mouse,vesselTrial,stage,episode) %>% 
          summarise(yy=aggregate.by(!!sym(response),na.rm=T),n=n(),radius=mean(medianlumen))
  
        
        ## Filter away vessels without enough data:
        ddv <- table(dda$vesselTrial,dda$stage) 
        idv = rownames(ddv)[which(rowSums(ddv)< 2 | rowSums(ddv>0)<2)] #idv = rownames(ddv)[which(rowSums(ddv)< length(unique(dda$stage)) | rowSums(ddv>0)<2)]
        if (length(which(dda$vesselTrial%in%idv))>0) dda <- dda[-which(dda$vesselTrial%in%idv),]
        
        ## Aggregate data per vesselTrial and stage (for source files)
        mvessel <- aggregate(dda$yy,by=list(dda$mouse,dda$vesselTrial,dda$stage),FUN=median,na.rm=T)
        vesselorder <- order(mvessel$Group.1,mvessel$Group.2,mvessel$Group.3)
        mvessel <- mvessel[vesselorder,]
        colnames(mvessel)[1:4] <- c("mouse","vessel","stage","median")
        mvessel_wide = mvessel %>% tidyr::spread(stage, median)
        
        if (mode=="sleep"){
          dda[,"state2"] <- dda[,"stage"]
          levels(dda$state2) <- c("awake","sleep","sleep","sleep","awake")
        }else if (mode=="wake"){
          dda[,"state2"] <- dda[,"stage"]
          levels(dda$state2) <- c("quiet","active","active")
        }
        
        
        ### Testing simpler and simpler model until one is able to fit one without error/warning:
        if (response=="medianepisode"){
          mod <- tryCatch({mod2 <- glmmTMB(log(yy) ~ stage + (1+state2|vesselTrial),
                                           data=dda,dispformula = ~stage+mouse)
          if (sum(is.na(summary(mod2)$vcov$cond))>0){
            warning("bad convergence")
          }},warning=function(cond){
            tryCatch({mod <- glmmTMB(log(yy) ~ stage + (1+state2|vesselTrial) ,data=dda,
                                     dispformula = ~stage)
            if (sum(is.na(summary(mod)$vcov$cond))>0){
              warning("bad convergence")
            }
            return(mod)},warning=function(cond){
              tryCatch({mod <- glmmTMB(log(yy) ~ stage + (1+state2|vesselTrial) ,data=dda,
                                       dispformula = ~mouse)
              if (sum(is.na(summary(mod)$vcov$cond))>0){
                warning("bad convergence")
              }
              return(mod)},warning=function(cond){
                tryCatch({mod <- glmmTMB(log(yy) ~ stage + (1+state2|vesselTrial) ,data=dda)
                if (sum(is.na(summary(mod)$vcov$cond))>0){
                  warning("bad convergence")
                }
                return(mod)},warning=function(cond){
                  tryCatch({mod <- glmmTMB(log(yy) ~ stage + (1|vesselTrial) ,data=dda,
                                           dispformula = ~stage+mouse)
                  if (sum(is.na(summary(mod)$vcov$cond))>0){
                    warning("bad convergence")
                  }
                  return(mod)},
                  warning=function(cond){
                    mod <- glmmTMB(log(yy) ~ stage + (1|vesselTrial) ,data=dda)
                    return(mod)})
                })
              })
            })}
          )
        }else{
          dda$radius = dda$radius-mean(dda$radius,na.rm=T)

          mod <- tryCatch({mod2 <- glmmTMB(log(yy) ~ stage + radius + I(radius^2)  + (1+state2|vesselTrial),
                                           data=dda,dispformula = ~stage+mouse)
          if (sum(is.na(summary(mod2)$vcov$cond))>0){
            warning("bad convergence")
          }},
          warning=function(cond){
            tryCatch({mod <- glmmTMB(log(yy) ~ stage + radius + I(radius^2)  + (1+state2|vesselTrial) ,data=dda,
                                     dispformula = ~stage)
            if (sum(is.na(summary(mod)$vcov$cond))>0){
              warning("bad convergence")
            }
            return(mod)},
            warning=function(cond){
              tryCatch({mod <- glmmTMB(log(yy) ~ stage + radius + I(radius^2)  + (1+state2|vesselTrial) ,data=dda,
                                       dispformula = ~mouse)
              if (sum(is.na(summary(mod)$vcov$cond))>0){
                warning("bad convergence")
              }
              return(mod)},
              warning=function(cond){
                tryCatch({mod <- glmmTMB(log(yy) ~ stage + radius + I(radius^2)  + (1+state2|vesselTrial),
                                         data=dda)
                if (sum(is.na(summary(mod)$vcov$cond))>0){
                  warning("bad convergence")
                }
                return(mod)},
                warning=function(cond){
                  tryCatch({mod <- glmmTMB(log(yy) ~ stage + radius + I(radius^2)  + (1|vesselTrial) ,data=dda,
                                           dispformula = ~stage+mouse)
                  if (sum(is.na(summary(mod)$vcov$cond))>0){
                    warning("bad convergence")
                  }
                  return(mod)},
                  warning=function(cond){
                    mod <- glmmTMB(log(yy) ~ stage + radius + I(radius^2)  + (1|vesselTrial) ,
                                   data=dda)
                    return(mod)})
                })
              })
            })}
          )
        }
        if (exists("mod2")){
          if (sum(is.na(summary(mod2)$vcov$cond))==0){
            mod <- mod2
          }
        }

        simres <- simulateResiduals(mod)
        
        ## Writing result files:
        if (response=="medianepisode"){
          file_intr <- paste(strsplit(dataset,".",fixed=T)[[1]][1],mode,linescan,response,sep="_")
        }else{
          file_intr <- paste(strsplit(dataset,".",fixed=T)[[1]][1],mode,linescan,band,response,sep="_")
        }
        print(file_intr)
        file_RE <- paste(file_intr,"randomEffects.txt",sep="_")

        if (mode=="sleep"){
          res <- REests(mod,table(dda$mouse,dda$vesselTrial))
        }else if (mode=="wake"){
          res <- REests_wake(mod,table(dda$mouse,dda$vesselTrial))
        }
        write.table(res,file=file_RE)

        ss <- summary(mod)
        file_FE <- paste(file_intr,"fixedEffects.txt",sep="_")

        write.table(ss$coefficients$cond[,c(1,2,4)],file=file_FE)

        emv <- emmeans(mod,~ stage,type="response")
        pp <- summary(pairs(emv,adjust="tukey",reverse=T))
        pp <- cbind(gsub(" ", "", pp[,"contrast"], fixed = TRUE),pp[,c("ratio","SE","p.value")]) #c("estimate","SE","p.value")
        colnames(pp)[1] <- "contrast"
        file_contr <-  paste(file_intr,"contrasts.txt",sep="_")

        write.table(pp,file=file_contr)
        
        source_file <- paste(file_intr,"source_file.csv",sep="_")
        readr::write_excel_csv(mvessel_wide,file=source_file)
        
        ## Saving plots:
        if (response=="medianepisode"){
          labelr <- paste(linescan,response, sep=", ")
        }else{
          labelr <- paste(linescan, band, response, sep=", ")
        }
        lims <- c(min(dda$yy)-0.1*mean(dda$yy),
                  max(dda$yy)+0.1*mean(dda$yy))

        vv1 <- vplot_mve_modfit(response,dda,ylab=labelr,ylims=lims,modfit=mod)
        vv2 <- vplot_mve_modfit(response,dda,ylab=labelr,ylims=lims,modfit=mod,transp=1)

        setEPS()
        postscript(paste(file_intr,".eps",sep=""),width=11,height=7)
        print(vv2)
        dev.off()

        pdf(paste(file_intr,".pdf",sep=""),width=11,height=7)
        print(vv1)
        dev.off()

        if (response!="medianepisode"){
          pdf(paste(file_intr,"_radiusEffect.pdf",sep=""),width=11,height=7)
          plot(predictorEffects(mod,~radius),axes=list(y=list(transform=exp, lab=response)))
          dev.off()
        }

        pdf(paste(file_intr,"_res_plot.pdf",sep=""),width=11,height=7)
        plot(simres)
        dev.off()
        
      }
    }
  }
}