## Plotting functions for app
pall <- c("M10"=2,"M11"=3,"M12"=4,"M21"=5,"M22"=6,
          "M13"=8,"M14"=7,"M16"=6,"M17"=4,"M19"=3,
          "M24"=2)
library(tidyverse)
library(ggplot2)
library(ggrepel)
library(plotly)
library(RColorBrewer)

# Plot with genotype, mouse and trial level
hhplot_gmt <- function(yy,df,xlab,fun=median,type){
  # yy: string with response name
  # df: a matrix with a column with the response and with columns "mouse_geno",
  # "mouse_name_short", "ts_day", "ts_name", "ts_stage_z_um","quality_category","max_speed_pre" (for puff data)
  idy <- which(colnames(df)==yy)
  colnames(df)[idy] <- "yy"
  
  # Remove low quality trials
  df <- filter(df,quality_category=="OK")
  idna <- which(is.na(df$yy))
  if (length(idna)>0) df <- df[-idna,]
  
  if(type=="Startle"){
    # Remove traces where the mouse runs before puff
    df <- filter(df,max_speed_pre<=10)
  }
  
  mgeno <- aggregate(df$yy,by=list(df$mouse_geno),FUN=fun,na.rm=T)
  colnames(mgeno)[1] <- c("geno")
  
  mmouse <- aggregate(df$yy,by=list(df$mouse_geno,df$mouse_name_short),FUN=fun,na.rm=T)
  morder <- order(mmouse$Group.1)
  mmouse <- mmouse[morder,]
  colnames(mmouse)[1:2] <- c("geno","mouse")
  
  l1 <- 3
  l2 <- 2
  l3 <- 1
  
  nm <- dim(mmouse)[1]
  
  mouselines <- matrix(NA,nm,2)
  mouselines[,1] <- rep(mgeno$x,times=rle(mmouse$geno)[1]$lengths)
  mouselines[,2] <- mmouse$x
  
  mouselines <- as.data.frame(mouselines)
  colnames(mouselines) <- c("x","xend")
  mouselines$y <- rep(l1,nm)
  mouselines$yend <- rep(l2,nm)
  mouselines$geno <- mmouse$geno
  mouselines$mouse <- mmouse$mouse
  
  mtrial <- aggregate(cbind(df$yy,df$ts_stage_z_um),by=list(df$mouse_geno,df$mouse_name_short,df$ts_name),FUN=fun,na.rm=T)
  mtrial <- mtrial[order(mtrial$Group.1,mtrial$Group.2),]
  colnames(mtrial)[1:5] <- c("geno","mouse","trial","x","depth")
  mtrial$depth <- round(scale(mtrial$depth),2)
  
  nt <- dim(mtrial)[1]
  
  tlines <- matrix(NA,nt,2)
  tlines[,1] <- rep(mmouse$x,times=rle(mtrial$mouse)[1]$lengths)
  tlines[,2] <- mtrial$x
  tlines <- as.data.frame(tlines)
  colnames(tlines) <- c("x","xend")
  tlines$y <- rep(l2,nt)
  tlines$yend <- rep(l3,nt)
  tlines$geno <- mtrial$geno
  tlines$mouse <- mtrial$mouse
  
  hhplot2 <- ggplot(mgeno,aes(x,c(l1,l1),colour=geno)) + theme(legend.position="none") +
    geom_point(size=3.5) + scale_color_brewer(palette="Set1") + 
    scale_y_continuous(breaks=c(l1,l2,l3),labels=c("genotype", "mouse","trial")) + 
    #scale_x_continuous(limits=c(0,1)) +
    geom_point(aes(x,rep(l2,nm),text=mouse),mmouse,shape=21,stroke=1.5,alpha=1,size=2) +
    geom_segment(data=mouselines,aes(x=x,y=y,xend=xend,yend=yend)) +
    geom_point(aes(x,rep(l3,nt),text=mouse),mtrial,shape=4,alpha=1)   +
    geom_segment(data=tlines,aes(x=x,y=y,xend=xend,yend=yend),size=0.1) +
    labs(x=xlab,y=" ") +
    annotate("text",label=paste0("n = ", dim(df)[1],paste(" ",type,", ",sep=""),
                                length(unique(df$mouse_name_short[df$mouse_geno=="wt/wt"])), " WT, ",
                                length(unique(df$mouse_name_short[df$mouse_geno=="tg/wt"])), " TG."),
             x=max(df$yy,na.rm=T)-0.2,xend=max(df$yy,na.rm=T),y=l1+0.1)
    #+ xlim(xlims[1], xlims[2])
  #hh <- ggplotly(hhplot2,tooltip=c("text")) 
  return(hhplot2)
}

# x-y plot
xyplot_pupil <- function(yy,df,xlab,ylab,type){
  # yy: string with response name
  # df: a matrix with a column with the response and with columns "pup_diff", "mouse_geno",
  # "mouse_name_short", "ts_day", "ts_name", "ts_stage_z_um","quality_category","max_speed_pre" (for puff data)
  idy <- which(colnames(df)==yy)
  colnames(df)[idy] <- "yy"
  
  # Remove low quality trials
  df <- filter(df,quality_category=="OK")
  idna <- which(is.na(df$yy))
  if (length(idna)>0) df <- df[-idna,]
  
  if(type=="Startle"){
    # Remove traces where the mouse runs before puff
    df <- filter(df,max_speed_pre<=10)
  }
  
  xlims <- c(min(df$pup_diff),max(df$pup_diff))
  
  gg <- ggplot(df,aes(x=pup_diff,y=yy)) + geom_point(aes(colour=mouse_name_short,text=mouse_name_short),size=2) + 
    theme_bw() +
    geom_smooth(aes(group=1),method='lm', formula= y~x,se=F,colour="black",size=1.5) + facet_wrap(~mouse_geno,ncol=4) + 
    geom_smooth(aes(colour=mouse_name_short),method='lm', formula= y~x,se=F,linetype=2,size=0.5) +
    theme(legend.position="none")   + xlim(xlims[1],xlims[2]) +
    labs(x=xlab,y=ylab) +  
    scale_color_manual(name="Mouse",values=pall) +
    annotate("text",label=paste0("n = ", dim(df)[1],paste(" ",type,", ",sep=""),
                                 length(unique(df$mouse_name_short[df$mouse_geno=="wt/wt"])), " WT, ",
                                 length(unique(df$mouse_name_short[df$mouse_geno=="tg/wt"])), " TG."),
             x=max(df$yy,na.rm=T)-0.2,xend=max(df$yy,na.rm=T),y=-0.1)
  #gg2 <- ggplotly(gg,tooltip=c("text"))
  return(gg)
}