library(reshape2)

library(ggplot2)

library(ComplexHeatmap)

library(RColorBrewer)

library(viridis)

library(grid)

library(circlize)

library(ggrepel)

#plot settings

source('~/plot/publication_plot_theme.R')

#### some functions

#order features

get_feature_orders=function(topic_matrix,num_topic=100){

  feature_ids=c()#ordered ids

  for (i in 1:num_topic){

    value=topic_matrix[,i]

    feature_ordered=rownames(topic_matrix)[order(value,decreasing = T)]

    feature_ids=cbind(feature_ids,feature_ordered)

  }

  colnames(feature_ids)=paste0('topic',1:num_topic)

  return(feature_ids)

}

##calculate correlation between ranks

cal_corr=function(version,name_index,id_index,rna_topic,adt_topic){

  #read all gene ids and protein ids

  if (strsplit(version,'/')[[1]][1]=='covid'){

    all_genes=read.csv('./data/covid/rna_name.csv',header=T)

    all_proteins=read.csv('./data/covid/protein_name.csv',header=T)

  } else{#nips

    all_genes=read.csv('./data/rna_name.csv',header=T)

    all_proteins=read.csv('./protein_name_complete.csv',header=T)

  }

  #get corresponding ids based on gene name & protein name

  gene_name_id=c()

  for (i in rownames(rna_topic)){

    id=all_genes[all_genes[,name_index]==i,][,id_index]

    if (length(id)!=0){

      for (d in id){#if more than 1 ids

        gene_name_id=rbind(gene_name_id,c(i,d))

      }    

    }

  }

  colnames(gene_name_id)=c('name','id')

  gene_name_id=as.data.frame(gene_name_id)

  protein_name_id=c()

  for (i in rownames(adt_topic)){

    id=all_proteins[all_proteins[,name_index]==i,][,id_index]

    if (length(id)!=0){

    protein_name_id=rbind(protein_name_id,c(i,id))}

  }

  colnames(protein_name_id)=c('name','id')

  protein_name_id=as.data.frame(protein_name_id)

  ### select common ids 

  common_ids=intersect(gene_name_id$id,protein_name_id$id)

  ## get corresponding names

  common_pro=c()

  for (i in common_ids){

    common_pro=c(common_pro,protein_name_id[protein_name_id$id==i,]$name)

  }

  common_genes=c()

  for (i in common_ids){

    common_genes=c(common_genes,gene_name_id[gene_name_id$id==i,]$name)

  }

  ##make sure the same order

  protein_topic_sub=adt_topic[common_pro,]

  rna_topic_sub=rna_topic[common_genes,]

  ## under each topic, get rank, calculate correlation for each topic

  rank_gene=apply(rna_topic_sub,2,rank)#column

  rank_pro=apply(protein_topic_sub,2,rank)#column

  corr=c()

  p=c()#wilcox paired test

  for (i in 1:num_topic){

    x=rank_gene[,i]

    y=rank_pro[,i]

    c=cor(x,y)

    corr=c(corr,c)

    w=wilcox.test(x, y, paired = TRUE, alternative = "two.sided")

    p=c(p,w$p.value)

  }

  p_adjust=p.adjust(p,'BH')

  plot_corr=data.frame(cor=corr,p_value=p,p_value_adj=p_adjust)

  return(plot_corr)

}

#plot heatmap for top genes and protein

plot_top_feature_in_selected_topic=function(type,top_feature_num,

                                            all_ids,topic_matrix,

                                            selected_topic,

                                            save_path,name='',

                                            width=5,height=10){

    top=all_ids[1:top_feature_num,selected_topic]

    all_names=melt(top)[,3]#all top gene/protein names

    m=c()#values

    for ( i in selected_topic){

      m=cbind(m,topic_matrix[all_names,i])

    }

    rownames(m)=all_names

    colnames(m)=as.character(selected_topic)

    m_scale=apply(m,2,function(x) x/max(abs(x)))#scale by column(within each topic)

    m_plot=melt(m_scale)

    m_plot$Var2=factor(m_plot$Var2,levels=unique(m_plot$Var2))

    plot=ggplot(m_plot,aes(x=Var2,y=Var1,fill=value))+

         geom_tile() +

         scale_fill_gradient2(low = "blue", mid = "white", high = "red",midpoint = 0)+

         labs(title = "",x='topic',y=type) +

         theme_bw() +

         theme(axis.text.x = element_text(angle = -90, hjust = 0,size=10),

               axis.text.y = element_text(size=10))

    save_name=paste0(save_path,'top',top_feature_num,type,name,'.png')

    ggsave(filename=save_name,plot=plot,dpi=320,width=width,height=height)

}

### plot sample by topic

plot_sample_by_topic=function(plot_cell_topic,version,coll,num_select_topic,save_name,is_cluster=TRUE){

# Create the heatmap annotation

  if (version=='covid'){

    ha <- HeatmapAnnotation(

      cellType=plot_cell_topic$initial_clustering,

      severity=plot_cell_topic$Status_on_day_collection_summary,

      status=plot_cell_topic$Status,

      col = coll)

  } else if(version=='rna_atac'){#atac

    ha <- HeatmapAnnotation(

      cellType=plot_cell_topic$cellType,

      col = coll,

      annotation_name_gp= gpar(fontsize = 20))

  } else{

    ha <- HeatmapAnnotation(

      cellType1=plot_cell_topic$cellType1,

      cellType2=plot_cell_topic$cellType2,

      col = coll,

      annotation_name_gp= gpar(fontsize = 20))

  }

  #prepare plot data

  plot_matrix=as.matrix(t(plot_cell_topic[,1:num_select_topic]))

  col_fun = colorRamp2(c(min(plot_matrix), 0, max(plot_matrix)), c("blue", "white", "red"))

  # Combine the heatmap and the annotation

  # !!!!! check dimension of cell_topic_select, make sure all are numeric !!!!

  png(file=save_name,

      width = 1500, height = 1000,units='px',bg = "transparent",res=100)

  h=Heatmap(plot_matrix, col=col_fun,

            show_column_names = FALSE,

            cluster_columns = is_cluster,

            cluster_rows = TRUE,

            top_annotation = ha,

            row_names_gp = grid::gpar(fontsize = 20)

            )

  draw(h)

  dev.off()

}

## plot correlation

plot_correlation=function(corr_plot,save_name){

  plot_corr=ggplot(corr_plot,aes(x=x,y=y,color=color))+

    geom_point(size=2)+

    theme_Publication(base_family='Arial')+

    theme(panel.grid = element_blank(),

          panel.grid.major=element_line(colour=NA),

          legend.position='none',

          axis.title.x=element_text(size=20),

          axis.title.y=element_text(size=20),

          axis.text.y = element_text(size = 20),

          axis.text.x = element_text(size = 20))+

    labs(x='',y='correlation',title='Correlation among genes and proteins')+

    geom_hline(yintercept = 0,linetype = 'dashed')

  ggsave(filename=save_name,plot=plot_corr,dpi=320,width=10,height=3)

}

## plot q values

plot_q_value=function(plot_data,save_name){

  plot_q=ggplot(plot_data,aes(x=x,y=y,color=color,label=label))+

    geom_point(alpha=0.3,size=1.5)+

    #geom_text_repel()+

    theme_Publication(base_family='Arial')+

    theme(panel.grid=element_blank(),

          panel.grid.major=element_line(colour=NA),

          legend.position='none',

          axis.title.x=element_text(size=20),

          axis.title.y=element_text(size=20),

          axis.text.y = element_text(size = 20),

          axis.text.x = element_text(size = 20))+

    labs(x='',y='-ln(q value)',title='')+

    geom_hline(yintercept = 3,linetype = 'dashed',col='red')

  ggsave(filename=save_name,plot=plot_q,dpi=320,width=10,height=3)

}

## plot p values for covid topics

plot_topic_p_value=function(covid_p_plot,save_name){

  plot=ggplot(covid_p_plot,aes(x=topic,y=-log(p++exp(-20)),color=color,label=label))+

    geom_point(size=2)+

    geom_text_repel()+

    theme_Publication(base_family='Arial')+

    theme(panel.grid = element_blank(),

          panel.grid.major=element_line(colour=NA),

          legend.position='none',

          axis.title.x=element_text(size=20),

          axis.title.y=element_text(size=20),

          axis.text.y = element_text(size = 20),

          axis.text.x = element_text(size = 20))+

    labs(x='',y='-log(p)',title='')#+

    #geom_hline(yintercept = 0.05,linetype = 'dashed')

  ggsave(filename=save_name,plot=plot,dpi=320,width=10,height=3)

}

## Differential analysis of topic expression

diff_topic=function(label,cell_info,cell_topic,num_topic=100,alter='greater'){

  label_index=which(names(cell_info)==label)

  label_type=unique(cell_info[,label_index])

  topic_label_p=c()# topic x full_cell_type  matrix storing p values

  topic_label_mean=c() # topic x full_cell_type  matrix storing mean difference 

  for (i in label_type){

    group1=cell_topic[cell_info[,label_index]==i,]#topic values with label

    group2=cell_topic[cell_info[,label_index]!=i,]#topic values without label

    topic_p=c()

    topic_mean=c()

    for (j in 1:num_topic){

      topic_p=c(topic_p,t.test(group1[,j],group2[,j],alternative=alter)$p.value)#upregulated

      topic_mean=c(topic_mean,mean(group1[,j])-mean(group2[,j]))

    }

    #topic_p_adj=p.adjust(topic_p)

    topic_label_p=cbind(topic_label_p,topic_p)

    topic_label_mean=cbind(topic_label_mean,topic_mean)

  }

  #

  topic_label_adj=p.adjust(topic_label_p)

  topic_label_p=matrix(topic_label_adj,nrow=num_topic)

  topic_label_p=as.data.frame(topic_label_p)

  colnames(topic_label_p)=label_type

  #

  topic_label_mean=as.data.frame(topic_label_mean)

  colnames(topic_label_mean)=label_type

  return(list('topic_label_p'=topic_label_p,'topic_label_mean'=topic_label_mean))

}

## plot differntially expressed topics 

plot_diff_topic=function(topic_label_list,save_name,width=1000,topic_select=c(1:100)){

  topic_label_mean=as.matrix(topic_label_list$topic_label_mean)

  rownames(topic_label_mean)=1:100

  topic_label_p=as.matrix(topic_label_list$topic_label_p)

  topic_label_mean_select=topic_label_mean[topic_select,]

  topic_label_p_select=topic_label_p[topic_select,]

  col_fun = colorRamp2(c(min(topic_label_mean_select), 0, max(topic_label_mean_select)), c("blue", "white", "red"))

  #col_fun = colorRamp2(c(-0.6124488, 0, 0.7623149), c("blue", "white", "red"))#confounder same range

  png(file=save_name,width = width, height = 1000,units='px',bg = "transparent",res=120)

  h=Heatmap(topic_label_mean_select, col=col_fun,cluster_rows = F,cluster_columns = T,

            row_names_gp = gpar(fontsize = 20),column_names_gp = gpar(fontsize = 20),

          cell_fun = function(j, i, x, y, w, h, fill) {

            if(topic_label_p_select[i, j] < 0.001) {

              gb = textGrob("*")

              gb_w = convertWidth(grobWidth(gb), "mm")

              gb_h = convertHeight(grobHeight(gb), "mm")

              grid.text("*", x, y - gb_h*0.5 + gb_w*0.4,gp = gpar(fontsize = 20))

            } 

          }

          )

  draw(h)

  dev.off()

}