#setwd("D:/scripts/IDTWES_CNV/plot_20201120/")
rm(list=ls())
library(RColorBrewer)
library(ggplot2)
library(plyr)
library(gridExtra)

args<- commandArgs(TRUE)
sample<-args[1]
#sample<-"CG316V3459-Lib7839"
max_ratio<-3.4
chrOrder<-c(1:22,'X','Y')
col= brewer.pal(9,"Greys")[c(7,4)]

samdata<-read.table(file=paste(sample,"_plot.tsv",sep=''),sep="\t",header=TRUE,as.is=TRUE)
names(samdata)<-c('chrom','start','end','Name','PTN','TN','Seg')
samdata$nPTN<-2**samdata$PTN
samdata$nPTN[samdata$nPTN> max_ratio]<-max_ratio
samdata$nTN<-2**samdata$TN
samdata$nSeg<-samdata$Seg
samdata$nSeg[samdata$nSeg> max_ratio]<-max_ratio
samdata$chrom<-sub('chr','',samdata$chrom,ignore.case=TRUE)
samdata$chrom<-factor(samdata$chrom,levels=chrOrder)
data12<-arrange(samdata,chrom,start)
chr_probe<-data.frame(chr=names(table(samdata$chrom)),
                      counts=as.vector(table(samdata$chrom)) )
chr_probe<-chr_probe[chr_probe$counts>0,]
count<-chr_probe$counts
num<-dim(data12)[1]

snpdata<-read.table(file=paste(sample,".ps.tab.xls",sep=''),sep="\t",header=TRUE,as.is=TRUE)
names(snpdata)<-c('chrom','Pos','Ref','Alt','Other','AF','VAFadj','ROH')
snpdata<-snpdata[snpdata$Ref+snpdata$Alt>40,]
snpdata<-snpdata[snpdata$AF>0.3 & snpdata$AF<0.7,]
#snpdata$VAF=snpdata$Alt/(snpdata$Alt+snpdata$Ref+snpdata$Other)
snpdata$VAF=round(snpdata$Alt/(snpdata$Alt+snpdata$Ref), 3)
snpdata$chrom<-sub('chr','',snpdata$chrom,ignore.case=TRUE)
snpdata$chrom<-factor(snpdata$chrom,levels=chrOrder)
snpdata<-arrange(snpdata,chrom,Pos)
chr_snp<-data.frame(chr=names(table(snpdata$chrom)),
                    counts=as.vector(table(snpdata$chrom)) )
chr_snp<-chr_snp[chr_snp$counts>0,]
snp_cnt<-chr_snp$counts
snp_num<-dim(snpdata)[1]

roh<-factor(snpdata$ROH,levels=unique(snpdata$ROH))
roh_snp<-data.frame(ROH=unique(roh),
                    counts=as.vector(table(roh)) )
roh_cnt<-roh_snp$counts

cum<-function(x){
  if(length(x)==1){
    return(x[1])
  }
  a<-as.vector(x[1])
  for (i in 2:length(x)){
    a[i]<-a[i-1]+x[i]
  }
  return(a)
}
cum0<-function(x){
  if(length(x)==1){
    return(1)
  }
  a<-(1)
  for (i in 2:length(x)){
    a[i]<-a[i-1]+x[i-1]
  }
  return(a)
}
cum_mid<-function(x){
  if(length(x)==1){
    return(round(x[1]/2))
  }
  a<-as.vector(round(x[1]/2))
  for ( i in 2:length(x)){
    a[i]<-cum(x)[i-1] + round(x[i]/2)
  }
  return(a)
}
gene_xys<-function(x){
  x_genes<-c()
  y_genes<-c()
  for (i in 1:length(x)){
    name<-x[i]
    x_genes[i]<-max( as.numeric(row.names(gene[gene$Name==name,])) )
    ymin<-min(gene[gene$Name==name,'Seg'])
    if(ymin < 0.6){
      y_genes[i]<-max(ymin-0.2,0)
    }else{
      y_genes[i]<-min(ymin+0.2,1.9)
    }
  }
  return(data.frame(xpos=x_genes,ypos=y_genes))
}
reg.data<-data.frame(
  x=cum0(count),
  xend=cum(count),
  y=rep(-0.05,length(count)),
  yend=rep(-0.05,length(count)),
  alpha=rep(c(0.5,0),length(count))[1:length(count)],
  stringsAsFactors = FALSE
)
cols<-rep(col,length(count))[1:length(count)]

gene<-data12[data12$Seg>1.4 |data12$Seg< 0.6,]

if(dim(gene)[1]==0){
  gene_label<-c('')
  gene_pos<-data.frame(xpos=c(0),ypos=c(0))
}else{
  gene_label<-names(table(gene$Name))
  gene_pos<-gene_xys(gene_label)
}

reg.snp<-data.frame(
  x=cum0(snp_cnt),
  xend=cum(snp_cnt),
  y=rep(-0.03,length(snp_cnt)),
  yend=rep(-0.03,length(snp_cnt)),
  alpha=rep(c(0.5,0),length(snp_cnt))[1:length(snp_cnt)],
  stringsAsFactors = FALSE
)
reg.roh<-data.frame(
  x=cum0(roh_cnt),
  xend=cum(roh_cnt),
  y=rep(-0.03,length(roh_cnt)),
  yend=rep(0,length(roh_cnt)),
  alpha=ifelse(substr(roh_snp$ROH,1,3)=='ROH',1,0),
  col_roh<-brewer.pal(9,"RdBu")[ifelse(substr(roh_snp$ROH,1,3)=='ROH',2,6)],
  stringsAsFactors = FALSE
)
snp_cols<-rep(col,length(snp_cnt))[1:length(snp_cnt)]

pdf(file=paste(sample,'.CR_VAF.pdf',sep=''),width = 8,height = 4.5)
p1<-ggplot(data=data12 ) +
  geom_rect(data = reg.data,aes(xmin=x, xmax=xend, alpha=alpha), ymin= -Inf,
            ymax= Inf, fill= brewer.pal(9,"Greys")[2],color= "NA",size= 2) +
  geom_point(stat="identity",size=0.05,aes(x=1:num,y=nPTN)) +
  geom_hline(aes(yintercept=1),color='steelblue2',size=1.5) +
  ylim(-0.1,max(data12$nPTN,1.8)+0.1) +xlab("Probe")+ylab("Copy Ratio")+
  scale_x_continuous(breaks=cum_mid(count),labels=chr_probe$chr,
                    expand = c(0.002,0.002) ) +
  geom_point(aes(x=1:num,y=nSeg),size=0.1,colour="red") +
  geom_segment(aes(x=x,xend=xend,y=y,yend=yend),data=reg.data,color=cols,size=1) +
 # annotate("text",x=gene_pos$xpos,y=gene_pos$ypos,label=gene_label,size=2,color='steelblue1')+
  theme(axis.text.x=element_text(angle=330,size=8),
        axis.text.y=element_text(size=8),
        axis.ticks = element_line(size = 0.2),
        axis.ticks.length=unit(.05, "cm"),
        panel.background = element_rect(fill = "transparent",colour = NA),
        legend.position = "none")

p2<-ggplot(data=snpdata ) +
  geom_rect(data = reg.snp,aes(xmin=x, xmax=xend, alpha=alpha), ymin= -Inf,
            ymax= Inf, fill= brewer.pal(9,"Greys")[3],color= "NA",size= 2) +
  geom_rect(data = reg.roh,aes(xmin=x, xmax=xend, alpha=alpha,color=col_roh), ymin= -Inf,
            ymax= Inf, fill= col_roh,color= "NA",size= 2) +
  geom_point(stat="identity",size=0.05,aes(x=1:snp_num,y=VAF)) +
  ylim(-0.03,max(snpdata$VAF,1)+0.02) +xlab("SNP")+ylab("VAF")+
  scale_x_continuous(breaks=cum_mid(snp_cnt),labels=chr_snp$chr,
                     expand = c(0.002,0.002) ) +
  theme(axis.text.x=element_text(angle=330,size=8),
        axis.text.y=element_text(size=8),
        axis.ticks = element_line(size = 0.2),
        axis.ticks.length=unit(.05, "cm"),
        panel.background = element_rect(fill = "transparent",colour = NA),
        legend.position = "none")
#geom_segment(aes(x=x,xend=xend,y=y,yend=yend),data=reg.snp,color=snp_cols,size=1) +
grid.arrange(p2, p1)
dev.off()