#!/usr/bin/env python3

import os
import argparse

from fqutil import Fastq, prefix_extension

def main():
    parser = argparse.ArgumentParser(
            description='Re-pair reads from discordant FASTQ files (for instance, after filtering reads by quality).')
    parser.add_argument('-u', '--keep-unique', default=False, const=True, action='store_const',
            help='Keep reads without a corresponding match. Unique reads are output to a separate file (filename_unique.fastq).')
    parser.add_argument('-o', '--output-dir', nargs=1, default=[None], type=str,
            help='Specify a custom output directory (defaults to directory of input files).')
    parser.add_argument('fastq1', nargs=1, type=str,
            help='First FASTQ file to be re-paired. May be gzipped.')
    parser.add_argument('fastq2', nargs=1, type=str,
            help='Second FASTQ file to be re-paired. May be gzipped.')
    argv = parser.parse_args()

    fq1 = Fastq(argv.fastq1[0])
    fq2 = Fastq(argv.fastq2[0])
    pair_reads(fq1, fq2, keep_unique=argv.keep_unique, output_dir=argv.output_dir[0])
    fq1.close()
    fq2.close()


def pair_reads(fastq1, fastq2, keep_unique=False, output_dir=None):
    '''
    Read-pairing algorithm
    '''

    # handle output directory properly
    if output_dir is None:
        output_dir = os.path.dirname(fastq1.filename)
    os.makedirs(output_dir, exist_ok=True)
    outfile1 = os.path.join(output_dir, os.path.basename(fastq1.filename))
    outfile2 = os.path.join(output_dir, os.path.basename(fastq2.filename))

    # index the second file for random-access
    idx = fastq2.index()

    # open output file handles
    fastq1_common = Fastq(prefix_extension(outfile1, '_common'), 'w')
    fastq2_common = Fastq(prefix_extension(outfile2, '_common'), 'w')
    if keep_unique:
        fastq1_unique = Fastq(prefix_extension(outfile1, '_unique'), 'w')

    # because stats are nice
    kept = 0

    # proceed through first file and look for matches against our index.
    while True:
        read1 = fastq1.get_read()
        if read1 is None:
            break
        readid = read1[0]
        if readid in idx.keys():
            # match found in index - dump both reads to disk and remove from index.
            fastq1_common.writelines(read1)
            fastq2.seek(idx.pop(readid))
            read2 = fastq2.get_read()
            fastq2_common.writelines(read2)
            kept += 1
        elif keep_unique:
            # write out to unique file for fastq1
            fastq1_unique.writelines(read1)
    fastq1_common.close()
    fastq2_common.close()

    if keep_unique:
        fastq1_unique.close()
        # all remaining keys in dictionary are the unique reads for fastq2
        with Fastq(prefix_extension(outfile2, '_unique'), 'w') as fastq2_unique:
            for remaining in idx:
                fastq2.seek(idx[remaining])
                read = fastq2.get_read()
                fastq2_unique.writelines(read)
    
    print_stats(fastq1, kept)
    print_stats(fastq2, kept)
    print('%s reads paired' % kept)


def print_stats(fastq, kept):
    n_uniq = fastq.lineno - kept
    if kept == 0:
        p_uniq = 100 - (kept / fastq.lineno * 100.)
    else:
        p_uniq = 100

    print('{}: {} unique reads identified out of a total of {} ({:.1f} unique)'
        .format(fastq.filename, n_uniq, fastq.lineno, p_uniq))


if __name__ == '__main__':
    main()
