#!/usr/bin/env Rscript
library(rtracklayer) 
library(BSgenome) 
library(BiocGenerics) 
library(Biostrings) 
library(GenomeInfoDb)
library(IRanges)
library(S4Vectors)

# 自定义函数，接受VCF文件路径作为参数
analyze_vcf <- function(vcf_file, child, father, mother, output_file1, output_file2) {
  # 读取VCF文件的所有行
  vcf <- read.table(vcf_file, header=T, stringsAsFactors = F, comment.char = "#")

  # 使用get()函数来动态获取列
  child_data <- as.data.frame(t(as.data.frame(sapply(vcf[[child]], strsplit, split="|"))))
  father_data <- as.data.frame(t(as.data.frame(sapply(vcf[[father]], strsplit, split="|"))))
  mother_data <- as.data.frame(t(as.data.frame(sapply(vcf[[mother]], strsplit, split="|"))))
  # 提取父母的基因型
  child <- as.data.frame(t(as.data.frame(sapply(vcf[[child]], strsplit, split="|"))))
  father <- as.data.frame(t(as.data.frame(sapply(vcf[[father]], strsplit, split="|"))))
  mother <- as.data.frame(t(as.data.frame(sapply(vcf[[mother]], strsplit, split="|"))))
  
  # 根据参考和替代等位基因设置基因型
  child[,1] <- ifelse(child[,1]=="0", vcf$REF, vcf$ALT)
  child[,3] <- ifelse(child[,3]=="0", vcf$REF, vcf$ALT)
  mother[,1] <- ifelse(mother[,1]=="0", vcf$REF, vcf$ALT)
  mother[,3] <- ifelse(mother[,3]=="0", vcf$REF, vcf$ALT)
  father[,1] <- ifelse(father[,1]=="0", vcf$REF, vcf$ALT)
  father[,3] <- ifelse(father[,3]=="0", vcf$REF, vcf$ALT)
  
  # 创建单倍型字符串
  fa.hap1 <- paste0(father[,1], collapse = "")
  fa.hap2 <- paste0(father[,3], collapse = "")
  mo.hap1 <- paste0(mother[,1], collapse = "")
  mo.hap2 <- paste0(mother[,3], collapse = "")
  ch.hap1 <- paste0(child[,1], collapse = "")
  ch.hap2 <- paste0(child[,3], collapse = "")
  
  # 计算身份一致性
  child.seq <- c(ch.hap1 = ch.hap1, ch.hap2 = ch.hap2)
  parent.seq <- c(fa.hap1 = fa.hap1, fa.hap2 = fa.hap2, mo.hap1 = mo.hap1, mo.hap2 = mo.hap2)
  
  numIdentiySet1 <- sapply(parent.seq, function(seq) {
    sum(as.numeric(strsplit(ch.hap1, "")[[1]] == strsplit(as.character(seq), "")[[1]])) / nchar(ch.hap1)
  })
  
  numIdentiySet2 <- sapply(parent.seq, function(seq) {
    sum(as.numeric(strsplit(ch.hap2, "")[[1]] == strsplit(as.character(seq), "")[[1]])) / nchar(ch.hap2)
  })
  
  # 找出与孩子单倍型最匹配的父母单倍型，并记录差异
  for(i in 1:4) {
    if (which.max(numIdentiySet1) == i) {
      difference1 <- as.numeric(strsplit(ch.hap1, "")[[1]] != strsplit(as.character(parent.seq[i]), "")[[1]])
      res1 <- data.frame(chr = vcf$CHROM, pos = vcf$POS, diffORNot = difference1)
      write.table(res1, file = output_file1, quote = FALSE, row.names = FALSE, sep = '\t')
    }
  }
  
  for(i in 1:4) {
    if (which.max(numIdentiySet2) == i) {
      difference2 <- as.numeric(strsplit(ch.hap2, "")[[1]] != strsplit(as.character(parent.seq[i]), "")[[1]])
      res2 <- data.frame(chr = vcf$CHROM, pos = vcf$POS, diffORNot = difference2)
      write.table(res2, file = output_file2, quote = FALSE, row.names = FALSE, sep = '\t')
    }
  }
}

# 获取命令行参数
args <- commandArgs(trailingOnly = TRUE)

# 检查参数数量
if (length(args) < 6) {
  stop("Usage: Rscript alignment_alleleDiff.R <vcf_file> <output_file1> <output_file2>", call. = FALSE)
}

# 调用函数
analyze_vcf(args[1], args[2], args[3], args[4], args[5], args[6])