require(reshape);
require(ggplot2)

plotQualityHeatmap=function(assignment, cells, minlim=5, maxlim=150) {
  if(sum(assignment$QUAL>150)>0) {
    assignment[assignment$QUAL>150,]$QUAL = 150
  }
  qual=melt(assignment[assignment$Barcode %in% cells,][,c(1,3)],id.vars="Barcode")
  qual$value=as.numeric(qual$value)
  qual=qual[order(qual$value,decreasing=T),]
  qual_reorder=qual
  qual_reorder$Barcode = c(1:nrow(qual_reorder))
  p = ggplot(qual_reorder, aes(x = Barcode,y=variable)) +
    geom_raster(aes(fill = value), interpolate=TRUE) +
    #scale_fill_gradient2(low="grey26", mid="grey85", high="red", 
    #                     midpoint=10, limits=c(minlim, maxlim)) +
    scale_fill_gradientn(colours = c("#CEEBFA","#EEDCC5","#F4A473","#E67A56","#CD473E","#AE1D3E"), limits=c(minlim, maxlim))+
    theme_classic()
  print(p)
}
convertCellVariantMatrixSortCellByQual = function(data,cells,variant_range=NA,assignment){
  qual=assignment[assignment$Barcode %in% cells,][,c(1,3)]
  order = order(qual$QUAL,decreasing=T)
  
  if (!is.na(variant_range)) {
    matrix = data[variant_range,][,colnames(data) %in% cells]
    matrix = cbind(data[variant_range,]$variant,matrix[,order])
  } else {
    matrix = data[,colnames(data) %in% cells]
    matrix = cbind(data$variant,matrix[,order])
  }
  #print(head(matrix[,c(1:5)]))
  colnames(matrix) = c("variant", 1:(ncol(matrix)-1))
  matrix$variant = c(1:nrow(matrix))
  #print(head(matrix[,c(1:5)]))
  df_heatmap=melt(matrix,id.vars="variant")
  #print(head(df_heatmap))
  df_heatmap$variant = as.numeric(df_heatmap$variant)
  df_heatmap$variable = as.numeric(df_heatmap$variable)
  df_heatmap$value=as.numeric(df_heatmap$value)
  #print(head(df_heatmap))
  return(df_heatmap)
}
convertCellVariantMatrix = function(data,cells,variant_range=NA){
  if (!is.na(variant_range)) {
    matrix = data[variant_range,][,colnames(data) %in% cells]
    matrix = cbind(data[variant_range,]$variant,matrix)
  } else {
    matrix = data[,colnames(data) %in% cells]
    matrix = cbind(data$variant,matrix)
  }
  #print(head(matrix[,c(1:5)]))
  colnames(matrix) = c("variant", 1:(ncol(matrix)-1))
  matrix$variant = c(1:nrow(matrix))
  #print(head(matrix[,c(1:5)]))
  df_heatmap=melt(matrix,id.vars="variant")
  #print(head(df_heatmap))
  df_heatmap$variant = as.numeric(df_heatmap$variant)
  df_heatmap$variable = as.numeric(df_heatmap$variable)
  df_heatmap$value=as.numeric(df_heatmap$value)
  #print(head(df_heatmap))
  return(df_heatmap)
}

sample="atac"
data=read.table(paste("../",sample,".scaf.tsv",sep=""),sep="\t",header=T)
assignment=read.table(paste("../expected_results/",sample,".assign_results.tsv",sep=""),header=T)
rownames(assignment) = assignment$Barcode
assignment$Barcode = gsub("-",".",assignment$Barcode)
rownames(assignment) = assignment$Barcode
assignment$QUAL[assignment$ASIG=="UNASSIGN"] = NA
assignment$QUAL=as.numeric(as.character(assignment$QUAL))
assignment$QUAL[is.na(assignment$QUAL)] = 0

if (sum(assignment$QUAL == Inf)>0){
  assignment[assignment$QUAL == Inf,]$QUAL = max(assignment[assignment$QUAL!=Inf,]$QUAL)
}
max(assignment$QUAL)
head(assignment)

cells = rownames(assignment[assignment$ASIG == "tumor",])
df1 = convertCellVariantMatrixSortCellByQual(data,cells,assignment=assignment)
pdf(paste(sample,"_assigntotumor_cell_qual_sort.pdf",sep=""),width=4,height=2)
plotQualityHeatmap(assignment, cells)
dev.off()

cells = rownames(assignment[assignment$ASIG == "normal",])
df2 = convertCellVariantMatrix(data,cells)
df2$variable = df2$variable + max(df1$variable)
pdf(paste(sample,"_assigntonormal_cell_qual_sort.pdf",sep=""),width=4,height=2)
plotQualityHeatmap(assignment, cells)
dev.off()


cells = rownames(assignment[assignment$ASIG == "UNASSIGN",])
df3 = convertCellVariantMatrix(data,cells)
df3$variable = df3$variable + max(df2$variable)


df_combined_1 = rbind(df1,df2,df3)
b = df_combined_1[df_combined_1$value==0,]
c = df_combined_1[df_combined_1$value>0,]
pdf(paste(sample,"_genotype.pdf",sep=""),width=8,height=4)
plot(b$variable, b$variant,pch=16, cex=0.5,col="#00FF0050",ylim=c(1,max(df_combined_1$variant)),xlim=c(1,max(df_combined_1$variable)),
     xlab="Cell index",ylab="Variant index", main="Genotype")
points(c$variable, c$variant,pch=16, cex=0.5,col="#FF000050",ylim=c(1,max(df_combined_1$variant)),xlim=c(1,max(df_combined_1$variable)))
abline(v=max(df1$variable) + 0.5)
abline(v=max(df2$variable) + 0.5)
dev.off()
