#!/usr/bin/env python3
import singlecellmultiomics.utils
import pysam
import collections
import numpy as np
import re
import more_itertools
import gzip
import argparse
import uuid
import os
from singlecellmultiomics.utils import BlockZip


if __name__ == '__main__':
    argparser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="""In silico digest genome""")
    argparser.add_argument(
        'fasta',
        metavar='fastafile',
        type=str,
        help='Fasta file to digest')
    argparser.add_argument(
        '-minlen',
        type=int,
        default=20,
        help='minimum read length')
    argparser.add_argument(
        '-maxlen',
        type=int,
        default=69,
        help='maximum read length')
    argparser.add_argument('-digest_sequence', type=str, default='CATG')
    argparser.add_argument('-head', type=int)
    argparser.add_argument(
        '-contigs',
        type=str,
        default=None,
        help='Comma separated contigs to run analysis for, when None specified all contigs are used')
    args = argparser.parse_args()

    r1_read_length = args.maxlen
    minlen = args.minlen

    ref = pysam.FastaFile(args.fasta)

    def seq_to_fastq(seq, header):
        return f'@{header}\n{seq}\n+\n{"E"*len(seq)}\n'

    selected_contigs = None if args.contigs is None else set(
        args.contigs.split(','))

    fastq_path = f'simulated_{args.digest_sequence}_single_{r1_read_length}_{uuid.uuid4()}.fastq.gz'

    outbam = f'simulated_{args.digest_sequence}_single_{r1_read_length}.bam'

    with gzip.open(fastq_path, 'wt') as out:
        processed = 0

        for contig, contig_len in zip(ref.references, ref.lengths):
            if selected_contigs is not None and contig not in selected_contigs:
                continue
            print(f'{contig}\t{contig_len}')
            seq = ref.fetch(contig).upper()
            frag_locations = np.diff(
                [m.start() for m in re.finditer(args.digest_sequence, seq)])
            # Generate fragments:
            for i, (start, end_excluding_catg) in enumerate(more_itertools.windowed(
                    (m.start() for m in re.finditer(args.digest_sequence, seq)), 2)):
                if end_excluding_catg is None or args.head is not None and i >= (
                        args.head - 1):
                    continue
                end = end_excluding_catg + len(args.digest_sequence)

                forward_read = seq[start:end][:r1_read_length]
                if len(forward_read) >= minlen:
                    out.write(
                        seq_to_fastq(
                            forward_read,
                            f'DR:fwd;ct:{contig};ds:{start};qn:{len(forward_read)}'))
                reverse_read = singlecellmultiomics.utils.reverse_complement(seq[start:end])[
                    :r1_read_length]
                if len(reverse_read) >= minlen:
                    out.write(
                        seq_to_fastq(
                            reverse_read,
                            f'DR:rev;ct:{contig};ds:{end_excluding_catg};qn:{len(reverse_read)}'))

            processed += contig_len

    # Map the fastq file:
    print("Now mapping ...")
    os.system(f"bwa mem -t 4 {args.fasta} {fastq_path} | samtools view -b - > ./{outbam}.unsorted.bam; samtools sort -T ./temp_sort -@ 4 ./{outbam}.unsorted.bam > ./{outbam}.unfinished.bam ; mv ./{outbam}.unfinished.bam ./{outbam} ; samtools index ./{outbam} & rm {fastq_path}")

    print("Creating site database ...")
    # Create site dictionary:
    sites = {}  # site-> wrongly_assinged_to, correctly_assigned, lost

    ref_reads = pysam.AlignmentFile(outbam)
    # Perform site detection:
    for read in ref_reads:

        if read.is_supplementary:
            continue
        # Check if there is reads mapping to this site

        kv = {kv.split(':')[0]: kv.split(':')[1]
              for kv in read.query_name.split(';')}

        site_pos = int(kv['ds'])
        site_chrom = kv['ct']
        strand = (kv['DR'] == 'rev')

        read_site = read.reference_end - 4 if read.is_reverse else read.reference_start
        read_contig = read.reference_name

        key = (site_chrom, site_pos, strand)
        key_mapped = (read_contig, read_site, read.is_reverse)

        if key not in sites:
            sites[key] = {'correct': 0, 'wrong_gain': 0, 'lost': 0}

        if read_site == site_pos and site_chrom == read_contig and read.is_reverse == strand:  # Correct assignment
            sites[key]['correct'] += 1
        else:
            # Assign a missing count to the origin site:
            sites[key]['lost'] += 1

            # Assing a wrong annotation to the target site:
            if key_mapped not in sites:
                sites[key_mapped] = {'correct': 0, 'wrong_gain': 0, 'lost': 0}
            sites[key_mapped]['wrong_gain'] += 1

    print("Writing site statistics ...")
    outtab = f'simulated_{args.digest_sequence}_single_{r1_read_length}.mappability.stats.bgzf'
    outtabsafe = f'simulated_{args.digest_sequence}_single_{r1_read_length}.mappability.safe.bgzf'

    keys = sorted(
    [(contig, pos, strand)
     for  (contig, pos, strand) in sites.keys()
     if contig is not None])

    with BlockZip(outtab, 'w') as stats, BlockZip(outtabsafe, 'w') as safe:
        for (contig, pos, strand) in keys:
            measured= sites[(contig, pos, strand)]
            stats.write(
                contig,
                pos,
                strand,
                f'{measured["correct"]}\t{measured["lost"]}\t{measured["wrong_gain"]}')

            if measured['wrong_gain'] == 0 and measured['lost'] == 0 and measured['correct'] == 1:
                safe.write(contig, pos, strand, 'ok')