#!/usr/bin/env python
##coding=utf-8

import os
import argparse
import glob

parser = argparse.ArgumentParser( description="")
parser.add_argument('-s', '--sample', help = 'the sample information file:sample\tdata_pathway', required = True)
parser.add_argument('-o', '--outputway', help = 'the release pathway', required = True)
parser.add_argument('--str',help="str analysis program ",choices=['y','n'],default='y')
parser.add_argument('--hla',help="str analysis program ",choices=['y','n'],default='y')
parser.add_argument('--sv',help="sv variant analysis program ",choices=['y','n'],default='y')
argv = vars(parser.parse_args())
outpath = os.path.abspath(argv['outputway'].strip())
samc= os.path.abspath(argv['sample'].strip())

def checkDir(Dir):
    if not os.path.exists(Dir):
        os.makedirs(Dir)

def delln(datalist,dir):
    for datai in datalist:
        if not os.path.exists(datai):
            print("ERROR : %s not exit"%(datai))
            break
        else:
            os.system("ln -s %s %s "%(datai,dir))

inputform=""
def trimsample(paw,samplef):
    slist={}
    sc=open(samplef,'r')
    for sci in sc:
        scl=sci.strip().split('\t')
        samplename=str(scl[0])
        if samplename not in slist.keys():
            slist[samplename]=""    
        samplepath=scl[1]
        if samplepath.endswith('fq.gz') or samplepath.endswith('fastq.gz'):
            inputform="fastq"
            slist[samplename]="%s/01.fastq/%s/%s.fastq.gz"%(paw,samplename,samplename)
        elif samplepath.endswith('bam'):
            inputform="bam"
            slist[samplename]=samplepath
    return slist

class release_hifi:
    def __init__(self,minipaw,sdict):
        self.minipaw=minipaw
        self.dict=sdict

    def rawdata(self):
        checkDir("%s/release/00.rawdata"%(self.minipaw))
        for s in self.dict.keys():
            dataline=self.dict[s]
            datadir="./00.rawdata/"+s+"/"
            checkDir(datadir)
            delln(dataline,datadir)

    def fasta(self):
        if inputform=="bam":
            checkDir("%s/release/01.fasta"%(self.minipaw))
            for s in self.dict.keys():
                datadir="%s/release/01.fasta/%s/"%(self.minipaw,s)
                checkDir(datadir)
                os.path.join(self.minipaw,"01.BWA",'%s.fastq.gz')%(s)
                dataline="%s/01.fastq/%s/%s.fastq.gz"%(self.minipaw,s,s)
                datalist=[dataline]
                delln(dataline,datadir)

    def mapping(self):
        checkDir("%s/release/02.mapping"%(self.minipaw))
        for s in self.dict.keys():
            datadir="%s/release/02.mapping/%s/"%(self.minipaw,s)
            checkDir(datadir)
            dataline="%s/02.mapping/%s/%s.sort.bam"%(self.minipaw,s,s)
            dataline1="%s/02.mapping/%s/%s.sort.bam.bai"%(self.minipaw,s,s)
            datalist=[dataline,dataline1]
            delln(datalist,datadir)

    
    def mutation(self):
        checkDir("%s/release/03.mutation"%(self.minipaw))
        for s in self.dict.keys():
            datadir="%s/release/03.mutation/%s/"%(self.minipaw,s)
            checkDir(datadir)  
            dataline="%s/03.mutation/%s/%s.vcf.gz"%(self.minipaw,s,s)
            dataline1="%s/03.mutation/%s/%s.vcf.gz.tbi"%(self.minipaw,s,s)
            dataline2="%s/03.mutation/%s/%s.filter.vcf.gz"%(self.minipaw,s,s)
            dataline3="%s/03.mutation/%s/%s.filter.vcf.gz.tbi"%(self.minipaw,s,s)
            dataline4="%s/03.mutation/%s/%s.visual_report.html"%(self.minipaw,s,s)
            datalist=[dataline,dataline1,dataline2,dataline3,dataline4]
            delln(datalist,datadir)

    def haplo(self):
        checkDir("%s/release/04.haplo"%(self.minipaw))
        for s in self.dict.keys():
            datadir="%s/release/04.haplo/%s/"%(self.minipaw,s)
            checkDir(datadir)
            dataline="%s/04.haplo/%s/%s.haplot.vcf.gz"%(self.minipaw,s,s)
            dataline1="%s/04.haplo/%s/%s.haplot.vcf.gz.csi"%(self.minipaw,s,s)
            dataline2="%s/04.haplo/%s/%s.haplot.bam"%(self.minipaw,s,s)
            dataline3="%s/04.haplo/%s/%s.haplot.bam.bai"%(self.minipaw,s,s)
            dataline4="%s/04.haplo/%s/%s.haplot.tsv"%(self.minipaw,s,s)
            datalist=[dataline,dataline1,dataline2,dataline3,dataline4]
            delln(datalist,datadir)
    
    def SV(self):
        checkDir("%s/release/05.SV"%(self.minipaw))
        for s in self.dict.keys():
            datadir="%s/release/05.SV/%s/"%(self.minipaw,s)
            checkDir(datadir)
            dataline="%s/05.SV/%s/%s.sv.vcf"%(self.minipaw,s,s)
            datalist=[dataline]
            delln(datalist,datadir)


    def STR(self):
        checkDir("%s/release/06.STR"%(self.minipaw))
        for s in self.dict.keys():
            datadir="%s/release/06.STR/%s/"%(self.minipaw,s)
            checkDir(datadir)
            datalist=glob.glob(os.path.join(self.minipaw,'06.STR',s,'logbam','RepBAM_all*.log'))
            delln(datalist,datadir)

    def HLA(self):
        checkDir("%s/release/07.HLA"%(self.minipaw))
        for s in self.dict.keys():
            datadir="%s/release/07.HLA/%s/"%(self.minipaw,s)
            checkDir(datadir)
            dataline="%s/07.HLA/%s/hla/R1_bestguess_G.txt"%(self.minipaw,s)
            datalist=[dataline]
            delln(datalist,datadir)

checkDir("%s/release"%(outpath))
os.system('cd %s/release'%(outpath))
sampledict=trimsample(outpath,samc)
rehifi=release_hifi(outpath,sampledict)
rehifi.fasta()
rehifi.mapping()
rehifi.mutation()
rehifi.haplo()
if argv['sv'].strip()=='y':
    rehifi.SV()
if argv['hla'].strip()=='y':
    rehifi.HLA()
if argv['str'].strip()=='y':
    rehifi.STR()