# load data
args = commandArgs(trailingOnly=TRUE)

mapfile = args[1]
divfiledir = args[2]
tag = args[3]

if(T){
  mapda<-read.csv(file=mapfile, header=F,sep=",")
  rownames(mapda) <- mapda$V1

  hlagenes <- c("HLA-F","HLA-V","HLA-P","HLA-G","HLA-H",
                "HLA-T","HLA-K","HLA-U","HLA-A","HLA-W","HLA-Y",
                "HLA-J","HLA-L","HLA-N","HLA-E","HLA-C","HLA-B",
                "HLA-S","MICA","MICB", "C4A", "C4B", "HLA-DRA","HLA-DRB3",
                "HLA-DRB4","HLA-DRB5","HLA-DRB1","HLA-DQA1",
                "HLA-DQB1","HLA-DQB2", "HLA-DQA2","HLA-DOB","TAP2","TAP1",
                "HLA-DMB","HLA-DMA","HLA-DOA", "HLA-DPA1",
                "HLA-DPB1","HLA-DPA2","HLA-DPB2")

  kirgenes <- c("KIR3DL3", "KIR2DS2", "KIR2DL3","KIR2DL2",
                "KIR2DL5B", "KIR2DP1","KIR2DL1","KIR3DP1",
                "KIR2DL4","KIR3DS1","KIR2DL5A","KIR2DS3",
                "KIR2DS5","KIR3DL1", "KIR2DS1", "KIR2DS4","KIR3DL2")


  hladiv <- matrix(NA, nrow=length(hlagenes), ncol=3)
  kirdiv <- matrix(NA, nrow=length(kirgenes), ncol=3)
  rownames(hladiv) <- hlagenes
  rownames(kirdiv) <- kirgenes
}

# add diversity score
if(T){

  for(gene in hlagenes){
    tid = mapda[gene,4]
    divfile=paste0(divfiledir, '/', tid, ".cds.div.txt.gz")
    if(file.info(divfile)$size < 40) next
    divval <- read.csv(divfile, header=F, sep=" ")

    hladiv[gene, 1] = mean(divval$V4[divval$V1 == 'hprc-afr'])
    hladiv[gene, 2] = mean(divval$V4[divval$V1 == 'cpc'])
    hladiv[gene, 3] = mean(divval$V4[divval$V1 == 'hprc-afr_cpc'])
  }

  colnames(hladiv) <- c("hprcafr", "cpc", "hprc_cpc")

  for(gene in kirgenes){
    tid = mapda[gene,4]
    divfile=paste0(divfiledir, '/', tid, ".cds.div.txt.gz")
    if(file.info(divfile)$size < 40) next
    divval <- read.csv(divfile, header=F, sep=" ")

    kirdiv[gene, 1] = mean(divval$V4[divval$V1 == 'hprc-afr'])
    kirdiv[gene, 2] = mean(divval$V4[divval$V1 == 'cpc'])
    kirdiv[gene, 3] = mean(divval$V4[divval$V1 == 'hprc-afr_cpc'])
  }

  colnames(kirdiv) <- c("hprcafr", "cpc", "hprc_cpc")

  outfile=paste0(tag, ".hla.cds.div.txt")
  write.table(hladiv, file=outfile, col.names=T, quote=F)

  outfile=paste0(tag, ".kir.cds.div.txt")
  write.table(kirdiv, file=outfile, col.names=T, quote=F)

}
