__author__ = 'Alena Zablotskaya. Laboratory for Cytogenetics and Genome Research (Prof. Dr. Joris R. Vermeesch), KU Leuven, Leuven, Belgium'
#Usage example: python3.4 /path/to/this/script/reads_byLocus.py -l /path/to/general/TSSVlibrary/file/TSSV_library_file_with_tags-5th_field.csv -lb /path/to/bed/file1/positions_of_TandemRepeats.bed -f /path/to/fasta/file/reads.fasta -sb /path/to/bed/file2/reads.bed -d 300

#!!!IMPORTANT!!! If this script is run in several parallel streams (different fasta files, but same folder and same library file), prior to that old '_TSSVlib_foo.csv', '_plusFlanks.bed' files should be removed.

def make_libr(workingDir, name, left_flank, right_flank, variation, unique_tag):
    repeat_library_file_current = workingDir+'/'+name+'_TSSVlib_foo.csv'
    if not os.path.exists(repeat_library_file_current):
        repeat_library_file = open(repeat_library_file_current, 'w')
        repeat_library_file.write(name+"\t"+left_flank+"\t"+right_flank+"\t"+variation+"\t"+unique_tag)
        repeat_library_file.close()
    return repeat_library_file_current

def make_bed(workingDir, sample_name, name, extension_bed, chr, locus_start, locus_end, repeat_name_bed, seq_bed_format):
    repeat_bed_file = workingDir+'/'+sample_name+'_'+name+'.'+extension_bed
    if os.path.exists(repeat_bed_file):
        os.remove(repeat_bed_file)
    if seq_bed_format == 'X' and chr.startswith('chr'):
        chr = chr[3:]       #chrX, chr22 are turned into X, 22 to match the sequence bed file format
    elif seq_bed_format == 'chrX' and not chr.startswith('chr'):
        chr = 'chr' + chr   #X, 22 are turned into chrX, chr22 to match the sequence bed file format
    repeat_bed = open(repeat_bed_file, 'a')
    repeat_bed.write(chr+"\t"+str(locus_start)+"\t"+str(locus_end)+"\t"+repeat_name_bed)
    repeat_bed.close()
    return repeat_bed_file

def make_fasta(repeat_intersect_bed_file, repeat_intersect_fasta_file, all_seq_fasta_cached):
    seq_names_intersect = []
    with open(repeat_intersect_bed_file) as seq_intersect_bed:
        for seq_intersect in seq_intersect_bed:
            a, b, c, seq_name_intersect, d, e = seq_intersect.split("\t")
            if seq_name_intersect not in seq_names_intersect:
                seq_names_intersect.append(seq_name_intersect)

    repeat_intersect_fasta = open(repeat_intersect_fasta_file, 'a')
    TO_SKIP = 0
    TO_PRINT = 1

    for seq_name in seq_names_intersect:
        fasta_name_needed = ">" + seq_name
        for line in all_seq_fasta_cached:
            if line.startswith(">"):
                if line.startswith(fasta_name_needed):
                    status = TO_PRINT
                    repeat_intersect_fasta.write(line)
                else:
                    status = TO_SKIP
                    continue
            elif status == TO_SKIP:
                continue
            elif status == TO_PRINT:
                repeat_intersect_fasta.write(line)
    repeat_intersect_fasta.close()

def process_input(library_file, lib_bed_file, fasta_file, seq_bed_file, distance):

    #==================DETERMINE_PATHS_TO_PROGRAM_FILES==================
    bedtools_path = '/cm/shared/apps/bedtools/current/bin/'				#bamToBed, intersectBed
    TSSV_path = '/home/azablo0/TSSV/tssv-0.2.0/tssv/'
    TSSVintoGenotype_path = '/home/azablo0/myScripts/'
    #====================================================================

    workingDir = os.path.dirname(seq_bed_file)
    os.chdir(workingDir)
    sample_name, extension_bed = os.path.basename(seq_bed_file).split('.')
    lib_bed_file_name, extension_bed = os.path.basename(lib_bed_file).split('.')

    #Checking format of the sequence bed file ('chrX' or 'X')
    seq_bed = open(seq_bed_file, 'r')
    first_line = seq_bed.readline()
    if first_line.startswith('chr'):
        seq_bed_format = 'chrX'
    else:
        seq_bed_format = 'X'
    seq_bed.close()

    #Cleaning '*_CCSs.fasta' files
    for CCSs_fasta_file in glob.glob(sample_name+'*_CCSs.fasta'):
        os.remove(CCSs_fasta_file)

    bed_withFlanks_file = workingDir+'/'+sample_name+'_'+lib_bed_file_name+'_plusFlanks.'+extension_bed
    if os.path.exists(bed_withFlanks_file):
        os.remove(bed_withFlanks_file)
    bed_withFlanks = open(bed_withFlanks_file, 'w')			#throws an error

    TSSV_jobs_file = workingDir+'/'+sample_name+'_TSSVjobs.txt'
    if os.path.exists(TSSV_jobs_file):
        os.remove(TSSV_jobs_file)
    TSSV_jobs = open(TSSV_jobs_file, 'a')

    all_seq_fasta_cached = []
    with open(fasta_file) as all_seq_fasta:
        for line in all_seq_fasta:
            all_seq_fasta_cached.append(line)

    unique_repeats = []
    with open(library_file) as library:
        for repeat_library in library:
            name, left_flank, right_flank, variation, unique_tag = repeat_library.split("\t")
            left_flank_length = len(left_flank)
            right_flank_length = len(right_flank)

            with open(lib_bed_file) as bed:
                for bed_entry in bed:
                    chr, repeat_start, repeat_end, repeat_name_bed = bed_entry.split("\t")
                    if repeat_name_bed.startswith(name):
                        left_flank_start = int(repeat_start) - left_flank_length
                        right_flank_end = int(repeat_end) + right_flank_length
                        if unique_tag not in unique_repeats:
                            bed_withFlanks.write(chr+"\t"+str(left_flank_start)+"\t"+str(right_flank_end)+"\t"+repeat_name_bed)
                            TSSV_report_file = workingDir+'/perLocus/'+name+'_report.txt'
                            TSSV_report_dir = workingDir+'/perLocus/'+name
                            TSSV_log_file = workingDir+'/perLocus/'+name+'_report.log'

                            if unique_tag.startswith('cluster'):
                                #First member of a cluster -> to be processed
                                repeat_intersect_fasta_file = workingDir+'/'+sample_name+'_'+chr+'_CCSs.fasta'

                                #Making a separate library file for the current repeat
                                repeat_library_file_current = make_libr(workingDir, name, left_flank, right_flank, variation, unique_tag)

                                if not os.path.exists(repeat_intersect_fasta_file):
                                    #Making a temporary bed-file for the current chromosome
                                    chr_start = 0
                                    chr_end = 999999999
                                    repeat_bed_file = make_bed(workingDir, sample_name, name, extension_bed, chr, chr_start, chr_end, repeat_name_bed, seq_bed_format)

                                    #Obtaining bed-file with reads intersecting the current locus
                                    repeat_intersect_bed_file = workingDir+'/'+sample_name+'_'+name+'_intersect.'+extension_bed
                                    command = bedtools_path+"intersectBed -a "+seq_bed_file+" -b "+repeat_bed_file+" -wa > "+repeat_intersect_bed_file
                                    subprocess.call(command, shell=True)
                                    os.chmod(repeat_intersect_bed_file, 0o755)
                                    os.remove(repeat_bed_file)

                                    #Making a separate fasta file with the reads intersecting the current locus
                                    make_fasta(repeat_intersect_bed_file, repeat_intersect_fasta_file, all_seq_fasta_cached)

                                #Appending a corresponding TSSV job to the file with all jobs
                                TSSV_jobs.write(TSSV_path+'tssv.py '+repeat_intersect_fasta_file+' '+repeat_library_file_current+' -r '+TSSV_report_file+' -d '+TSSV_report_dir+'/ 2>&1 | tee '+TSSV_log_file+"\n")
                                unique_repeats.append(unique_tag)
                                break
                            else:
                                #Non-clustered repeat -> to be processed

                                #Making a separate library file for the current repeat
                                repeat_library_file_current = make_libr(workingDir, name, left_flank, right_flank, variation, unique_tag)

                                #Making a temporary bed-file for the current locus (with N bp frame around the repeat)
                                left_distance = left_flank_start - int(distance)
                                right_distance = right_flank_end + int(distance)
                                repeat_bed_file = make_bed(workingDir, sample_name, name, extension_bed, chr, left_distance, right_distance, repeat_name_bed, seq_bed_format)

                                #Obtaining bed-file with reads intersecting the current locus
                                repeat_intersect_bed_file = workingDir+'/'+sample_name+'_'+name+'_intersect.'+extension_bed
                                command = bedtools_path+"intersectBed -a "+seq_bed_file+" -b "+repeat_bed_file+" -wa > "+repeat_intersect_bed_file
                                subprocess.call(command, shell=True)
                                os.chmod(repeat_intersect_bed_file, 0o755)
                                os.remove(repeat_bed_file)

                                #Making a separate fasta file with the reads intersecting the current locus
                                repeat_intersect_fasta_file = workingDir+'/'+sample_name+'_'+name+'_CCSs.fasta'
                                make_fasta(repeat_intersect_bed_file, repeat_intersect_fasta_file, all_seq_fasta_cached)

                                #Appending a corresponding TSSV job to the file with all jobs
                                TSSV_jobs.write(TSSV_path+'tssv.py '+repeat_intersect_fasta_file+' '+repeat_library_file_current+' -r '+TSSV_report_file+' -d '+TSSV_report_dir+'/ 2>&1 | tee '+TSSV_log_file+"\n")
                                unique_repeats.append(unique_tag)
                                break
                        else:
                            #One of the repeats from a cluster whose member has been already processed -> to skip
                            if unique_tag.startswith('cluster'):
                                bed_withFlanks.write(chr+"\t"+str(left_flank_start)+"\t"+str(right_flank_end)+"\t"+repeat_name_bed)
                            break
            if os.path.exists(repeat_intersect_bed_file):
                os.remove(repeat_intersect_bed_file)

    bed_withFlanks.close()
    TSSV_jobs.close()

import sys
import argparse
import os
import subprocess
import glob

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("-l", "--library", help="Path to input TSSV library file.", default="")
    parser.add_argument("-lb", "--libbed", help="Path to input BED file with tandem repeat locations.", default="")
    parser.add_argument("-f", "--fa", help="Path to input FASTA file with all sequencing reads.", default="")
    parser.add_argument("-sb", "--seqbed", help="Path to input BED file with mapping locations for sequencing reads.", default="")
    parser.add_argument("-d", "--distance", help="Mapping distance from a tandem repeat to select seq reads for further analysis by TSSV.", default="0")
    args = parser.parse_args()
    print(args)
    process_input(args.library, args.libbed, args.fa, args.seqbed, args.distance)

if __name__ == "__main__":
    sys.exit(main())