##11/04/2018
##@@author: Simon (Chong) Chu, DBMI, Harvard Medical School
##@@contact: chong_chu@hms.harvard.edu

#The original input is generated by: filtering -> gene-annotation -> genotyping --> convert-to-gvcf
#two major classes: gVCF

#To-do-list:1. The structure of the raw output is defined in x_clip_disc_filter.py, need to merge them
#           2. Insertion strand is not added (done 072919)
#           3. Merge different type of repeats to one VCF (in x_vcf_merger.py now 08/30/19)
#           4. Genotype probability

import pysam
from x_TSD import *
from x_rep_type import *
from x_reference import *
from cmd_runner import *



#gVCF: convert xTEA raw output to gVCF
class gVCF():
####
    #these are not used for now:
    # ##INFO=<ID=PROTEIN_CODING__LOF,Number=.,Type=String,Description="Gene(s) on which the SV is predicted to have a loss-of-function effect.">
    # ##INFO=<ID=PROTEIN_CODING__NEAREST_TSS,Number=.,Type=String,Description="Nearest transcription start site to intragenic variants.">
    # ##INFO=<ID=PROTEIN_CODING__PROMOTER,Number=.,Type=String,Description="Genes whose promoter sequence (1 kb) was disrupted by TE insertion.">
    # ##INFO=<ID=PROTEIN_CODING__UTR,Number=.,Type=String,Description="Gene(s) for which the SV is predicted to disrupt a UTR.">
    # ##INFO=<ID=PROTEIN_CODING__INTRONIC,Number=.,Type=String,Description="Gene(s) where the TE insertion was found to lie entirely within an intron.">
    # ##INFO=<ID=INS_DEL,Number=1,Type=String,Description="Insertion with deletion">

    # ##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">
    def gnrt_vcf_head(self, s_sample_id, sf_bam):
        #use the first bam (if there are several bams of the same individual)
        samfile = pysam.AlignmentFile(sf_bam, "r")
        # parse out the header and contig information
        m_header = samfile.header
        shead = "##fileformat=VCFv4.2\n"
        if "SQ" not in m_header:#
            print(("No head information in bam file: {0}".format(sf_bam)))
            shead+="##reference=GCA_000001405.15_GRCh38_no_alt_analysis_set.fna\n"
            shead+=prepare_head_hg38(True)
        else:
            sf_ref=""
            l_chrm=[]
            for rcd in m_header['SQ']:#
                sf_ref="NULL"
                if 'UR' in rcd:
                    sf_ref=rcd['UR']
                i_chrm_len=-1
                if 'LN' in rcd:
                    i_chrm_len=rcd['LN']
                s_chrm="NULL"
                if 'SN' in rcd:
                    s_chrm=rcd['SN']
                l_chrm.append((s_chrm, i_chrm_len))
            fields = sf_ref.split("/")
            shead += ("##reference=" + fields[-1] + "\n")
            for rcd in l_chrm:
                if ("HLA" in rcd[0]) or ("Un" in rcd[0]) or ("_" in rcd[0]):
                    continue
                shead+=("##contig=<ID={0},length={1}>\n".format(rcd[0], rcd[1]))
        s_other_head = self.prepare_head_other()
        shead += s_other_head
        ##CHROM  POS     ID      REF     ALT     QUAL    FILTER  INFO   HG002_NA24385_son_2_GRCh38
        shead+=("#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\t"+s_sample_id+"\n")
        return shead

####
    def prepare_head_hg19(self, b_chr):
        s_chr=""
        if b_chr==True:
            s_chr="chr"
        sinfo = "##contig=<ID={0}1,length=249250621>\n".format(s_chr)
        sinfo += "##contig=<ID={0}2,length=243199373>\n".format(s_chr)
        sinfo += "##contig=<ID={0}3,length=198022430>\n".format(s_chr)
        sinfo += "##contig=<ID={0}4,length=191154276>\n".format(s_chr)
        sinfo += "##contig=<ID={0}5,length=180915260>\n".format(s_chr)
        sinfo += "##contig=<ID={0}6,length=171115067>\n".format(s_chr)
        sinfo += "##contig=<ID={0}7,length=159138663>\n".format(s_chr)
        sinfo += "##contig=<ID={0}8,length=146364022>\n".format(s_chr)
        sinfo += "##contig=<ID={0}9,length=141213431>\n".format(s_chr)
        sinfo += "##contig=<ID={0}10,length=135534747>\n".format(s_chr)
        sinfo += "##contig=<ID={0}11,length=135006516>\n".format(s_chr)
        sinfo += "##contig=<ID={0}12,length=133851895>\n".format(s_chr)
        sinfo += "##contig=<ID={0}13,length=115169878>\n".format(s_chr)
        sinfo += "##contig=<ID={0}14,length=107349540>\n".format(s_chr)
        sinfo += "##contig=<ID={0}15,length=102531392>\n".format(s_chr)
        sinfo += "##contig=<ID={0}16,length=90354753>\n".format(s_chr)
        sinfo += "##contig=<ID={0}17,length=81195210>\n".format(s_chr)
        sinfo += "##contig=<ID={0}18,length=78077248>\n".format(s_chr)
        sinfo += "##contig=<ID={0}19,length=59128983>\n".format(s_chr)
        sinfo += "##contig=<ID={0}20,length=63025520>\n".format(s_chr)
        sinfo += "##contig=<ID={0}21,length=48129895>\n".format(s_chr)
        sinfo += "##contig=<ID={0}22,length=51304566>\n".format(s_chr)
        sinfo += "##contig=<ID={0}X,length=155270560>\n".format(s_chr)
        sinfo += "##contig=<ID={0}Y,length=59373566>\n".format(s_chr)
        return sinfo

####
    def prepare_head_hg38(self, b_chr):
        s_chr = ""
        if b_chr == True:
            s_chr = "chr"
        sinfo = "##contig=<ID={0}1,length=248956422>\n".format(s_chr)
        sinfo += "##contig=<ID={0}2,length=242193529>\n".format(s_chr)
        sinfo += "##contig=<ID={0}3,length=198295559>\n".format(s_chr)
        sinfo += "##contig=<ID={0}4,length=190214555>\n".format(s_chr)
        sinfo += "##contig=<ID={0}5,length=181538259>\n".format(s_chr)
        sinfo += "##contig=<ID={0}6,length=170805979>\n".format(s_chr)
        sinfo += "##contig=<ID={0}7,length=159345973>\n".format(s_chr)
        sinfo += "##contig=<ID={0}8,length=145138636>\n".format(s_chr)
        sinfo += "##contig=<ID={0}9,length=138394717>\n".format(s_chr)
        sinfo += "##contig=<ID={0}10,length=133797422>\n".format(s_chr)
        sinfo += "##contig=<ID={0}11,length=135086622>\n".format(s_chr)
        sinfo += "##contig=<ID={0}12,length=133275309>\n".format(s_chr)
        sinfo += "##contig=<ID={0}13,length=114364328>\n".format(s_chr)
        sinfo += "##contig=<ID={0}14,length=107043718>\n".format(s_chr)
        sinfo += "##contig=<ID={0}15,length=101991189>\n".format(s_chr)
        sinfo += "##contig=<ID={0}16,length=90338345>\n".format(s_chr)
        sinfo += "##contig=<ID={0}17,length=83257441>\n".format(s_chr)
        sinfo += "##contig=<ID={0}18,length=80373285>\n".format(s_chr)
        sinfo += "##contig=<ID={0}19,length=58617616>\n".format(s_chr)
        sinfo += "##contig=<ID={0}20,length=64444167>\n".format(s_chr)
        sinfo += "##contig=<ID={0}21,length=46709983>\n".format(s_chr)
        sinfo += "##contig=<ID={0}22,length=50818468>\n".format(s_chr)
        sinfo += "##contig=<ID={0}X,length=156040895>\n".format(s_chr)
        sinfo += "##contig=<ID={0}Y,length=57227415>\n".format(s_chr)
        return sinfo

    def prepare_head_other(self):#
        sinfo=""
        sinfo += "##ALT=<ID=INS:ME:SVA,Description=\"Insertion of SVA element\">\n"
        sinfo += "##ALT=<ID=INS:ME:LINE1,Description=\"Insertion of LINE1 element\">\n"
        sinfo += "##ALT=<ID=INS:ME:ALU,Description=\"Insertion of ALU element\">\n"
        sinfo += "##ALT=<ID=INS:ME:HERV-K,Description=\"Insertion of HERV-K element\">\n"
        sinfo += "##ALT=<ID=INS:MT,Description=\"Mitochondrial insertion\">\n"
        sinfo += "##ALT=<ID=INS:PSDGN,Description=\"Pseudogene insertion\">\n"
        sinfo += "##FILTER=<ID=PASS,Description=\"All filters passed\">\n"
        sinfo += "##FILTER=<ID=LowConfident,Description=\"All filters passed\">\n"
        sinfo += "##INFO=<ID=SVTYPE,Number=1,Type=String,Description=\"Type of structural variant\">\n"
        sinfo += "##INFO=<ID=SVLEN,Number=1,Type=Integer,Description=\"Insertion length\">\n"
        sinfo += "##INFO=<ID=END,Number=1,Type=Integer,Description=\"End coordinate of this variant\">\n"
        sinfo += "##INFO=<ID=TSD,Number=1,Type=String,Description=\"Precise Target Site Duplication for bases, if unknown, value will be NULL\">\n"
        sinfo += "##INFO=<ID=TSDLEN,Number=1,Type=Integer,Description=\"TSD (dup/del) length, if unknown, then -1\">\n"
        sinfo += "##INFO=<ID=SUBTYPE,Number=1,Type=String,Description=\"Subtype based on TSD/PolyA/side-info\">\n"
        sinfo += "##INFO=<ID=TD_SRC,Number=1,Type=String,Description=\"Transduction source\">\n"
        sinfo += "##INFO=<ID=STRAND,Number=1,Type=String,Description=\"Insertion orientation (+/- for sense/antisense)\">\n"
        sinfo += "##INFO=<ID=AF,Number=A,Type=Float,Description=\"Estimated allele frequency in the range (0,1)\">\n"
        sinfo += "##INFO=<ID=LCLIP,Number=1,Type=Integer,Description=\"Number of left clipped reads (aligned to copies)\">\n"
        sinfo += "##INFO=<ID=RCLIP,Number=1,Type=Integer,Description=\"Number of right clipped reads (aligned to copies)\">\n"
        sinfo += "##INFO=<ID=LDISC,Number=1,Type=Integer,Description=\"Number of left discordant reads (aligned to copies)\">\n"
        sinfo += "##INFO=<ID=RDISC,Number=1,Type=Integer,Description=\"Number of right discordant reads (aligned to copies)\">\n"
        sinfo += "##INFO=<ID=LPOLYA,Number=1,Type=Integer,Description=\"Number of left polyA reads\">\n"
        sinfo += "##INFO=<ID=RPOLYA,Number=1,Type=Integer,Description=\"Number of right polyA reads\">\n"
        sinfo += "##INFO=<ID=LRAWCLIP,Number=1,Type=Integer,Description=\"Number of left raw clip reads\">\n"
        sinfo += "##INFO=<ID=RRAWCLIP,Number=1,Type=Integer,Description=\"Number of right raw clip reads\">\n"
        sinfo += "##INFO=<ID=AF_CLIP,Number=1,Type=Integer,Description=\"Number of effective clip reads (for AF)\">\n"
        sinfo += "##INFO=<ID=AF_FMAP,Number=1,Type=Integer,Description=\"Number of effective fully mapped reads (for AF)\">\n"
        sinfo += "##INFO=<ID=AF_DISC,Number=1,Type=Integer,Description=\"Number of effective discordant pairs (for AF)\">\n"
        sinfo += "##INFO=<ID=AF_CONCORDNT,Number=1,Type=Integer,Description=\"Number of effective concordant pairs (for AF)\">\n"
        sinfo += "##INFO=<ID=LDRC,Number=1,Type=Integer,Description=\"Number of left discordant reads are reverse-complementary\">\n"
        sinfo += "##INFO=<ID=LDNRC,Number=1,Type=Integer,Description=\"Number of left discordant reads are not reverse-complementary\">\n"
        sinfo += "##INFO=<ID=RDRC,Number=1,Type=Integer,Description=\"Number of right discordant reads are reverse-complementary\">\n"
        sinfo += "##INFO=<ID=RDNRC,Number=1,Type=Integer,Description=\"Number of right discordant reads are not reverse-complementary\">\n"
        sinfo += "##INFO=<ID=LCOV,Number=1,Type=Float,Description=\"Left focal coverage\">\n"
        sinfo += "##INFO=<ID=RCOV,Number=1,Type=Float,Description=\"Right focal coverage\">\n"
        sinfo += "##INFO=<ID=LD_AKR_RC,Number=1,Type=Integer,Description=\"Number of left discordant anchor reads are reverse-complementary\">\n"
        sinfo += "##INFO=<ID=LD_AKR_NRC,Number=1,Type=Integer,Description=\"Number of left discordant anchor reads are not reverse-complementary\">\n"
        sinfo += "##INFO=<ID=RD_AKR_RC,Number=1,Type=Integer,Description=\"Number of right discordant anchor reads are reverse-complementary\">\n"
        sinfo += "##INFO=<ID=RD_AKR_NRC,Number=1,Type=Integer,Description=\"Number of right discordant anchor reads are not reverse-complementary\">\n"
        sinfo += "##INFO=<ID=LC_CLUSTER,Number=1,Type=String,Description=\"Left clip formed cluster in consensus\">\n"
        sinfo += "##INFO=<ID=RC_CLUSTER,Number=1,Type=String,Description=\"Right clip formed cluster in consensus\">\n"
        sinfo += "##INFO=<ID=LD_CLUSTER,Number=1,Type=String,Description=\"Left discordant formed cluster in consensus\">\n"
        sinfo += "##INFO=<ID=RD_CLUSTER,Number=1,Type=String,Description=\"Right discordant formed cluster in consensus\">\n"
        sinfo += "##INFO=<ID=NINDEL,Number=1,Type=Integer,Description=\"Number of large indel (or with several mismatches) reads\">\n"
        sinfo += "##INFO=<ID=CLIP_LEN,Number=1,Type=String,Description=\"Clip reads length\">\n"
        sinfo += "##INFO=<ID=INS_INV,Number=1,Type=String,Description=\"5 prime inversion\">\n"
        sinfo += "##INFO=<ID=REF_REP,Number=1,Type=String,Description=\"Fall in reference repeat copy or not\">\n"
        sinfo += "##INFO=<ID=GENE_INFO,Number=1,Type=String,Description=\"Fall in gene region\">\n"
        sinfo += "##FORMAT=<ID=GT,Number=1,Type=String,Description=\"Genotype\">\n"
        return sinfo
####

    def sort_raw_rslt(self, sf_raw_rslt, sf_sorted_rslt):
        cmd="sort -k1,1V -k2,2n -o {0} {1}".format(sf_sorted_rslt, sf_raw_rslt)
        cmd_runner = CMD_RUNNER()
        cmd_runner.run_cmd_small_output(cmd)

    ####convert raw results to gVCF records
    def cvt_raw_rslt_to_gvcf(self, s_sample_id, sf_bam, sf_raw_rslt, i_rep_type, sf_ref, sf_vcf):
        #first sort the raw rslt
        sf_raw_rslt_sorted=sf_raw_rslt+".sorted"
        self.sort_raw_rslt(sf_raw_rslt, sf_raw_rslt_sorted)

        l_rslts=[]
        l_sites = []
        l_pos=[]#this is the exact position
        with open(sf_raw_rslt_sorted) as fin_rslt:
            for line in fin_rslt:
                fields=line.split()
                l_rslts.append(fields)
                l_sites.append((fields[0], int(fields[2]), int(fields[3])))
                lclip_pos = int(fields[2])
                rclip_pos = int(fields[3])
                i_begin = lclip_pos
                i_end = rclip_pos
                if lclip_pos > rclip_pos:
                    i_begin = rclip_pos
                    i_end = lclip_pos
                if i_begin <= 0:
                    i_begin = i_end
                l_pos.append((fields[0], i_begin, i_begin+1))
####
        #get the repeat type
        rep_type=RepType()
        s_rep_type=rep_type.get_rep_vcf_alt_type(i_rep_type)
        #get the TSD sequences
        xTSD=XTSD()
        l_sTSD=xTSD.get_TSD_seqs(sf_ref, l_sites)
        #get the reference site sequences
        #l_ref_seqs=xTSD.get_TSD_seqs(sf_ref, l_pos)
        xref=XReference()
        l_ref_seqs=xref.get_ref_seqs_of_sites(sf_ref, l_pos)

        with open(sf_vcf, "w") as fout_vcf:
            s_head = self.gnrt_vcf_head(s_sample_id, sf_bam)
            fout_vcf.write(s_head)
            #for rcd in l_rslts:
            if len(l_rslts) != len(l_sites) or len(l_sites) != len(l_ref_seqs) or len(l_rslts) != len(l_ref_seqs):
                print("[Error in x_vcf.py]: Arrzy size is different!\n")
            for (l_fields, s_tsd, s_site_seq) in zip(l_rslts, l_sTSD, l_ref_seqs):
                sid="."
                s_rcd=self._cvt_to_vcf_rcd(l_fields, s_rep_type, s_tsd, sid, s_site_seq)
                fout_vcf.write(s_rcd)
####
    # convert to vcf records
    def _cvt_to_vcf_rcd(self, l_fields, s_type, s_tsd, s_id, s_ref_seq):
        s_rcd=""
        ins_chrm=l_fields[0]
        #ins_pos=int(l_fields[1])
        lclip_pos=int(l_fields[2])
        rclip_pos=int(l_fields[3])
        i_tsd_len=int(l_fields[4])
        n_lclip_cns=int(l_fields[5])
        n_rclip_cns=int(l_fields[6])
        n_ldisc_cns=int(l_fields[7])
        n_rdisc_cns=int(l_fields[8])
        n_lpolyA=int(l_fields[9])
        n_rpolyA=int(l_fields[10])
        f_lcov=float(l_fields[11])
        f_rcov=float(l_fields[12])
        slc_cluster = l_fields[19]
        src_cluster = l_fields[20]
        sld_cluster = l_fields[21]
        srd_cluster = l_fields[22]
        s_transduction=l_fields[23]
        n_lpe_same=int(l_fields[27])#PE of same oritentation(anchor oritentaion and mate on cns oritentation)
        n_lpe_diff=int(l_fields[28])
        n_rpe_same=int(l_fields[29])
        n_rpe_diff=int(l_fields[30])
        s_inv=l_fields[31]
        s_sub_type=l_fields[32]

        n_AF_clip=int(l_fields[35])
        n_AF_full=int(l_fields[36])
        f_AF=0
        if (n_AF_clip+n_AF_clip)>0:
            f_AF=float(n_AF_clip)/float(n_AF_clip+n_AF_full)
        n_raw_lclip=int(l_fields[37])
        n_raw_rclip=int(l_fields[38])
        n_AF_disc=int(l_fields[39])
        n_AF_concod=int(l_fields[40])
        n_indel=int(l_fields[41])
        s_clip_lens=l_fields[42]

        n_lanchor_rc=int(l_fields[43])
        n_lanchor_nrc = int(l_fields[44])
        n_ranchor_rc = int(l_fields[45])
        n_ranchor_nrc = int(l_fields[46])

        i_ins_len=int(l_fields[47])
        s_fall_in_ref_copy=l_fields[48]
        s_gene_info="."
        if len(l_fields)>49:
            s_gene_info=l_fields[49]
        s_gntp="./."
        if len(l_fields)>50:
            s_gntp=l_fields[50]#genotype

        s_strand="."
        if n_lpolyA>0 and n_rpolyA==0:#sense, left clipped reads should be polyA
            s_strand="+"
        elif n_lpolyA==0 and n_rpolyA>0:#antisense, right clipped reads should be polyT
            s_strand="-"

        ##CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO
        s_rcd+=ins_chrm+"\t"#chrom
        i_begin=lclip_pos
        i_end=rclip_pos
        if lclip_pos>rclip_pos:
            i_begin=rclip_pos
            i_end=lclip_pos
        if i_begin<=0:
            i_begin=i_end
        s_rcd+=(str(i_begin)+"\t")#pos
        s_rcd+=(s_id+"\t")#ID
        s_rcd+=(s_ref_seq+"\t")#REF
        s_rcd+=("<"+s_type+">\t")#ALT
        s_rcd+=(".\t")#QUAL
        s_rcd+=("PASS\t")#FILTER

        #INFO part
        sinfo=""
        sinfo += ("SVTYPE="+s_type+";")
        sinfo += ("SVLEN="+str(abs(i_ins_len))+";")
        sinfo += ("END="+str(i_end)+";")
        sinfo += ("TSD="+s_tsd+";")
        sinfo += ("TSDLEN="+str(i_tsd_len)+";")
        sinfo += ("SUBTYPE="+s_sub_type+";")
        sinfo += ("TD_SRC="+s_transduction+";")
        sinfo += ("STRAND="+s_strand+";") #########################################################
        sinfo += ("AF="+str(f_AF)+";")#this is from clip/(clip+full-map), may not that accuracy!!!#
        sinfo += ("LCLIP="+str(n_lclip_cns)+";")
        sinfo += ("RCLIP="+str(n_rclip_cns)+";")
        sinfo += ("LDISC="+str(n_ldisc_cns)+";")
        sinfo += ("RDISC="+str(n_rdisc_cns)+";")
        sinfo += ("LPOLYA="+str(n_lpolyA)+";")
        sinfo += ("RPOLYA=" +str(n_rpolyA)+";")
        sinfo += ("LRAWCLIP=" + str(n_raw_lclip) +";")
        sinfo += ("RRAWCLIP=" + str(n_raw_rclip) +";")
        sinfo += ("AF_CLIP=" + str(n_AF_clip) +";")
        sinfo += ("AF_FMAP=" + str(n_AF_full) +";")
        sinfo += ("AF_DISC=" + str(n_AF_disc) +";")
        sinfo += ("AF_CONCORDNT=" + str(n_AF_concod) +";")
        sinfo += ("LDRC=" + str(n_lpe_same) +";")
        sinfo += ("LDNRC=" + str(n_lpe_diff) +";")
        sinfo += ("RDRC=" + str(n_rpe_same) +";")
        sinfo += ("RDNRC=" + str(n_rpe_diff) +";")
        sinfo += ("LCOV=" + str(f_lcov) +";")
        sinfo += ("RCOV=" + str(f_rcov) +";")
        sinfo += ("LD_AKR_RC=" + str(n_lanchor_rc) + ";")
        sinfo += ("LD_AKR_NRC=" + str(n_lanchor_nrc) + ";")
        sinfo += ("RD_AKR_RC=" + str(n_ranchor_rc) + ";")
        sinfo += ("RD_AKR_NRC=" + str(n_ranchor_nrc) + ";")
        sinfo += ("LC_CLUSTER=" + slc_cluster +";")
        sinfo += ("RC_CLUSTER=" + src_cluster +";")
        sinfo += ("LD_CLUSTER=" + sld_cluster +";")
        sinfo += ("RD_CLUSTER=" + srd_cluster +";")
        sinfo += ("NINDEL=" + str(n_indel)+";")
        sinfo += ("CLIP_LEN=" + s_clip_lens +";")
        sinfo += ("INS_INV=" + s_inv +";")
        sinfo += ("REF_REP=" + s_fall_in_ref_copy + ";")
        sinfo += ("GENE_INFO=" + s_gene_info+"\t")
        #genotype part
        sinfo += ("GT\t"+s_gntp+"\n")
        s_rcd+=sinfo
        return s_rcd

####