import logging

# Setup logging
logging.basicConfig(filename=snakemake.output.removed_transcripts,
                    level=logging.DEBUG,
                    format='%(asctime)s %(levelname)s %(name)s %(message)s')
logger = logging.getLogger(__name__)


def format_oxford(abundance):
    """
    Takes oxford abundance file and produces dictionary with counts per
    transcript.

    :param abundance: File containing transcript counts
    :return: Dictionary: key = transcript id, value = counts
    """
    ox_dict = {}
    with open(abundance, 'r') as file:
        # Skip header
        next(file)
        for line in file:
            line = line.strip("\n").split(',')
            transcript_id = line[0]
            # Can contain an empty string instead of 0
            counts = sum([int(count) for count in line[1:] if count != ''])
            ox_dict[transcript_id] = counts
    return ox_dict


def format_talon(abundance):
    """
    Takes talon abundance file and produces dictionary with counts per
    transcript.

    :param abundance: File containing transcript counts
    :return: Dictionary: key = transcript id, value = counts
    """
    talon_dict = {}
    with open(abundance, 'r') as file:
        # Skip header
        next(file)
        for line in file:
            line = line.strip('\n').split('\t')
            transcript_id = line[3]
            # Always has a round number
            counts = sum(map(int, line[11:]))
            talon_dict[transcript_id] = counts
    return talon_dict


def format_flair(abundance):
    """
    Takes flair abundance file and produces dictionary with counts per
    transcript.

    :param abundance: File containing transcript counts
    :return: Dictionary: key = transcript id, value = counts
    """
    flair_dict = {}
    with open(abundance, 'r') as file:
        # Skip header
        next(file)
        for line in file:
            line = line.strip('\n').split('\t')
            transcript_id = line[0]
            # Numbers are always rounded but saves them as x.0
            counts = sum(map(float, line[1:]))
            flair_dict[transcript_id] = counts
    return flair_dict


def combine(gtf_file, counts, output, missing_output):
    """
    Loops through the GTF file and checks the counts dictionary
    to add counts to the line and then writes it to a new file.

    This function is used per GTF file.

    Flair combines the transcript id with the gene id for known
    transcripts using '_' as a delimiter. This can cause problems
    when the names in the annotation file contain '_'.

    Flair can have transcripts present in the GTF file that are missing
    from the abundance file. This is caused by minimap2 mapping
    ambiguity. If there are any missing transcripts these will be
    written to the missing transcript log and extracted to a separate
    GTF.

    Stringtie adds transcripts with 0 count to the GTF file in
    annotation-guided mode. These 0 count transcripts get removed and
    extracted to a separate GTF file. The IDs of the removed transcripts
    are written to the log files.

    :param missing_output: GTF file with the removed transcripts
    :param gtf_file: The GTF file produced by a pipeline.
    :param counts: Dictionary: key = transcript id, value = counts
    :param output: Output directory and file name
    :return: Writes a file to the output param
    """
    # Read the GTF file
    with open(gtf_file) as GTF, open(output, 'w') as outfile, open(missing_output, 'a') as removed_out:
        origin_file = gtf_file.split('/')[-1]
        write = True
        for line in GTF:
            # Check if line is a transcript
            if line.split()[2] == 'transcript':
                try:
                    # Check for transcript id in counts
                    transcript_id = line.split()[11][1:-2]
                    count = counts[transcript_id]
                    # if count is 0 transcript is removed
                    if count == 0:
                        logger.warning(f'ORIGIN: {origin_file} 0 count: {transcript_id}')
                        removed_out.write(f'{line[:-1]}\n')
                        write = False
                        continue
                except KeyError:
                    # Transcript is missing from the abundance file
                    try:
                        # This is for the flair '_' delimiter issue
                        gene_id = line.split()[9][1:-2]
                        new_transcript_id = f'{transcript_id}_{gene_id}'
                        count = counts[new_transcript_id]
                    except KeyError:
                        # transcript is actually missing
                        # Add missing transcript id to missing transcripts.log
                        logger.error(f'ORIGIN: {origin_file} FIRST TRY: {transcript_id} SECOND TRY: {new_transcript_id}')
                        removed_out.write(f'{line[:-1]}\n')
                        write = False
                        continue

                # transcript is not missing from the abundance and has at least 1 count.
                # Add transcript count to end of line and write to file
                # counts are not in TPM, hijacking the column for use with GffCompare
                write = True
                outfile.write(f'{line[:-1]} TPM "{str(count)}";\n')
            else:
                # Line is not a transcript
                if write:
                    # Line is not an exon from a removed transcript
                    outfile.write(line)
                else:
                    # Line is an exon from a removed transcript
                    removed_out.write(line)


def main():
    """
    Takes abundance file and passes it to a format function. The
    resulting dictionary is passed to combine. Combine produces a new
    GTF that has counts for transcripts.
    :return: -
    """
    flair_abundance = snakemake.input.flair_count
    flair_gtf = snakemake.input.flair_transcripts
    flair_dict = format_flair(flair_abundance)
    combine(flair_gtf, flair_dict, snakemake.output.flair_combined, snakemake.output.flair_removed)
    flair_dict.clear()

    ox_abundance = snakemake.input.oxford_count
    ox_gtf = snakemake.input.oxford_transcript
    ox_dict = format_oxford(ox_abundance)
    combine(ox_gtf, ox_dict, snakemake.output.oxford_combined, snakemake.output.oxford_removed)
    ox_dict.clear()

    talon_abundance = snakemake.input.talon_count
    talon_gtf = snakemake.input.talon_transcripts
    talon_dict = format_talon(talon_abundance)
    combine(talon_gtf, talon_dict, snakemake.output.talon_combined, snakemake.output.talon_removed)


main()

from collections import defaultdict
import numpy as np
import matplotlib.pyplot as plt


def tracking(gfftracking):
    """
    Read the data from the GffCompare gffcmp.tracking file and
    format as a dictionary.

    Only takes three-way matches from the file.

    :return: tcons_XXXX (key) : [[transcript_id 1],
                                [transcript_id 2],
                                [transcript_id 3]] (value)
    """
    tcons = {}
    with open(gfftracking) as file:
        for line in file:
            line = line.split()
            transcripts = line[4::]
            temp_list = []
            # '-' means that at least one pipeline did not have a matching transcript
            # only take lines with three transcripts
            if '-' not in transcripts:
                for transcript in transcripts:
                    temp_list.append(transcript.split('|')[1])
                tcons[line[0]] = temp_list
    return tcons


def gtf(file):
    """
    Read transcript data from a GTF file and format as dictionary.
    :param file: GTF file
    :return: transcript_id (key): [(exon_start, exon_end),
                                   (exon_start, exon_end),
                                   (exon_start, exon_end) etc...](value)
    """
    gtf_dict = defaultdict(list)
    with open(file) as GTF:
        for line in GTF:
            # skip comments
            if line.startswith("#"):
                continue
            line = line.split()
            if line[2] == 'exon':
                transcript_id = line[11][1:-2]
                start = int(line[3])
                end = int(line[4])
                exon = (start, end)
                gtf_dict[transcript_id].append(exon)
    return gtf_dict


def sort_gtf(gtf_dict):
    """
    The exons for each transcript in the TALON gtf file are not sorted
    correctly. (issue with talon)
    :param gtf_dict: transcript_id (key): [(exon_start, exon_end),
                                   (exon_start, exon_end),
                                   (exon_start, exon_end) etc...](value)
    :return: gtf_dict but with sorted exons
    """
    for key, value in gtf_dict.items():
        gtf_dict[key] = sorted(value)
    return gtf_dict


def retrieve_terminal_exon(tcons, oxford, flair, talon):
    """
    Places terminal exon positions of transcripts into dictionary by
    matching transcript id.

    :param tcons: tcons_XXXX (key) : [[transcript_id 1 ],
                                      [transcript_id 2 ],
                                      [transcript_id 3]] (value)

    :param oxford: key = transcript_id, value = [(exon_start, exon_end)]
    :param flair: key = transcript_id, value = [(exon_start, exon_end)]
    :param talon: key = transcript_id, value = [(exon_start, exon_end)]

    :return: tcons_XXXX(key):[[(exon_start, exon_end), (exon_start, exon_end)],
                            [(exon_start, exon_end), (exon_start, exon_end)],
                            [(exon_start, exon_end), (exon_start, exon_end)]] (value)
    """
    tcons_exon = {}
    for tcon_id, transcript_ids in tcons.items():
        # no need to select terminal exons for single-exon transcripts
        # data only has three-matches and all transcripts have the same number of exons
        if len(oxford[transcript_ids[0]]) == 1:
            # Get exons that belong to transcript_id
            oxford_exons = oxford[transcript_ids[0]]
            flair_exons = flair[transcript_ids[1]]
            talon_exons = talon[transcript_ids[2]]
            # add exons to tcon_exon dictionary
            tcons_exon[tcon_id] = [oxford_exons, flair_exons, talon_exons]
        # Slice terminal exons for multi-exon transcripts
        else:
            oxford_terminal_exons = [oxford[transcript_ids[0]][0],
                                     oxford[transcript_ids[0]][-1]]
            flair_terminal_exons = [flair[transcript_ids[1]][0],
                                    flair[transcript_ids[1]][-1]]
            talon_terminal_exons = [talon[transcript_ids[2]][0],
                                    talon[transcript_ids[2]][-1]]
            tcons_exon[tcon_id] = [oxford_terminal_exons, flair_terminal_exons, talon_terminal_exons, ]
    return tcons_exon


def calculate_difference(transcript_exons):
    """
    Calculate the differences between exon positions
    :param transcript_exons: tcons_XXXX (key): [[(exon_start, exon_end)],
                                                [(exon_start, exon_end)],
                                                [(exon_start, exon_end)]]
    :return: Differences between terminal exon positions
    """
    oxford_flair = []
    oxford_talon = []
    flair_talon = []
    for key, value in transcript_exons.items():
        value = np.array(value, dtype=object)
        oxford = value[0]
        flair = value[1]
        talon = value[2]
        a = np.subtract(oxford, flair)
        b = np.subtract(oxford, talon)
        c = np.subtract(flair, talon)
        oxford_flair.append(a)
        oxford_talon.append(b)
        flair_talon.append(c)
    return oxford_flair, oxford_talon, flair_talon


def start_end(compare):
    """
    Sorts the calculated differences into 4 dictionaries.
    Start-exon - start position and end position
    End-exon - start position and end position
    single-exon - start position and end position
    :param compare: Dictionary with calculated differences.
    :return: Dictionaries with start and end positions.
    """
    start_exon_start = defaultdict(list)
    start_exon_end = defaultdict(list)
    end_exon_start = defaultdict(list)
    end_exon_end = defaultdict(list)
    single_exon_start = defaultdict(list)
    single_exon_end = defaultdict(list)
    for positions in compare:
        # single exon transcripts
        if len(positions) == 1:
            # start position
            single_exon_start['0'].append(abs(positions[0][0]))
            # end position
            single_exon_end['0'].append(abs(positions[0][1]))
        else:
            # multi exon transcripts
            for number, position in enumerate(positions):
                # start exon
                if number == 0:
                    # start exon - start
                    start_exon_start[number].append(abs(position[0]))
                    # start exon - end
                    start_exon_end[number].append(abs(position[1]))
                    # end exon
                else:
                    # end exon start
                    end_exon_start[number].append(abs(position[0]))
                    # end exon end
                    end_exon_end[number].append(abs(position[1]))
    return start_exon_start, start_exon_end, end_exon_start, end_exon_end, single_exon_start, single_exon_end


def plot(start_exon_start, start_exon_end, end_exon_start, end_exon_end, single_exon_start, single_exon_end, output):
    """
    Plots the exon differences into 3 figures with 6 subplots each.
    """
    plt.rcParams["figure.figsize"] = [15, 15]
    plt.rcParams["figure.autolayout"] = True
    plt.rcParams.update({'font.size': 15})

    main_figure, axs = plt.subplots(nrows=3, ncols=2)
    title = output.split('/')[-1].split('.')[0]
    main_figure.suptitle(f'Comparison of {title} start and end exons from three-way matched transcripts')

    labels_start = single_exon_start.keys()
    data_start =  single_exon_start.values()
    axs[0, 0].boxplot(data_start, showfliers=True)
    axs[0, 0].set(xlabel="Exon number", ylabel="Number of bases")
    axs[0, 0].set_title('Single-exon, start position differences')
    axs[0, 0].set_xticklabels(labels_start, rotation=45)

    labels_start = single_exon_end.keys()
    data_start =  single_exon_end.values()
    axs[0, 1].boxplot(data_start, showfliers=True)
    axs[0, 1].set(xlabel="Exon number", ylabel="Number of bases")
    axs[0, 1].set_title('Single-exon, end position differences')
    axs[0, 1].set_xticklabels(labels_start, rotation=45)

    labels_start = start_exon_start.keys()
    data_start =  start_exon_start.values()
    axs[1, 0].boxplot(data_start, showfliers=True)
    axs[1, 0].set(xlabel="Exon number", ylabel="Number of bases")
    axs[1, 0].set_title('Start exon, start position differences')
    axs[1, 0].set_xticklabels(labels_start, rotation=45)

    labels_start = start_exon_end.keys()
    data_start =  start_exon_end.values()
    axs[1, 1].boxplot(data_start, showfliers=True)
    axs[1, 1].set(xlabel="Exon number", ylabel="Number of bases")
    axs[1, 1].set_title('Start exon, end position differences')
    axs[1, 1].set_xticklabels(labels_start, rotation=45)

    labels_start = end_exon_start.keys()
    data_start =  end_exon_start.values()
    axs[2, 0].boxplot(data_start, showfliers=True)
    axs[2, 0].set(xlabel="Exon number", ylabel="Number of bases")
    axs[2, 0].set_title('End exon, start position differences')
    axs[2, 0].set_xticklabels(labels_start, rotation=45)

    labels_start = end_exon_end.keys()
    data_start =  end_exon_end.values()
    axs[2, 1].boxplot(data_start, showfliers=True)
    axs[2, 1].set(xlabel="Exon number", ylabel="Number of bases")
    axs[2, 1].set_title('End exon, end position differences')
    axs[2, 1].set_xticklabels(labels_start, rotation=45)

    plt.savefig(output, dpi=200)


def main():
    tcons = tracking(snakemake.input.tracking)
    oxford = gtf(snakemake.input.oxford)
    flair = gtf(snakemake.input.flair)
    talon = gtf(snakemake.input.talon)
    talon_sorted = sort_gtf(talon)

    transcript_exons = retrieve_terminal_exon(tcons, oxford, flair, talon_sorted)

    oxford_flair, oxford_talon, flair_talon = calculate_difference(transcript_exons)

    start_exon_start, start_exon_end, end_exon_start, end_exon_end, single_exon_start, single_exon_end = start_end(oxford_flair)
    plot(start_exon_start, start_exon_end, end_exon_start, end_exon_end, single_exon_start, single_exon_end, snakemake.output.oxford_flair)

    start_exon_start, start_exon_end, end_exon_start, end_exon_end, single_exon_start, single_exon_end = start_end(oxford_talon)
    plot(start_exon_start, start_exon_end, end_exon_start, end_exon_end, single_exon_start, single_exon_end, snakemake.output.oxford_talon)

    start_exon_start, start_exon_end, end_exon_start, end_exon_end, single_exon_start, single_exon_end = start_end(flair_talon)
    plot(start_exon_start, start_exon_end, end_exon_start, end_exon_end, single_exon_start, single_exon_end, snakemake.output.flair_talon)


main()

shell:
    '''
    NanoPlot \
    -t {threads} \
    --fastq {input.fq} \
    --plots dot kde \
    -p {params.prefix} \
    -o {params.outdir}
    '''

shell:
    '''
    samtools view -Sb {input.sam} | samtools sort -@ {threads} -o {output.bam}
    samtools index {output.bam}
    '''

shell:
    '''
    samtools view -q {params.min_mq} -F 2304 -b {input.bam} \
    | seqkit bam -j {threads} -x -T '{params.flt}' \
    | samtools sort -@ {threads} -o {output.bam_filtered}
    samtools index {output.bam_filtered}
    '''

shell:
    '''
    stringtie --rf -G {input.annotation} -L -p {threads} {params.opts} -o {output.gff} {input.bam_filtered}
    '''

shell:
    '''
    featureCounts \
        {input.bams} \
        -T {threads} \
        -a {input.annotation} \
        -L \
        -M \
        -S 1 \
        -o {output.counts}
    '''

shell:
    '''
    gffcompare \
        -r {input.annotation} \
        -e 100 \
        -d 100 \
        --no-merge \
        -T \
        -o {params.outdir} \
        {input.oxford} {input.flair} {input.talon}
    '''

shell:
    '''
    multiqc \
    {params.nanoplot_dir} \
    {params.samtools_subread_dir} \
    --outdir {params.outdir}
    '''