#!python

"""Matches barcodes to subassembled variants, makes count files.

Written by Danny Lawrence and Jesse Bloom."""

import time
import os
import re
import random
import dms_tools
import dms_tools.parsearguments
import dms_tools.file_io
import dms_tools.cutils


def parse_barcodes(r1files, r2files, minquality, bc_len, r1start, r2end, purgefrac):
    """Reads FASTQ files and returns barcodes passing quality threshold.

    Barcodes discarded if either read fails quality filter, if any site
    in barcode doesn't have read quality >= *minquality* for at least one 
    read, or if reads disagree at any site.

    `r1files` and `r2files` are the file names of FASTQ files, which
    are optionally gzipped. Both files must contain the same number of reads
    in the same order.

    `minquality` is an integer (0-50) indicating the Phred quality score 
    below which a base call will be ignored. Any reads failing the Illumina
    filter are also ignored.

    `bc_len` is the length of the barcode sequence. 

    `r1start` indicates the nucleotide in R1 at which the barcode starts,
    in 1, 2, ... indexing.

    `r2end` indicates how many nucleotides **before** the end of R2 the
    barcode ends.

    `purgefrac` indicates the fraction of read pairs to randomly purge.

    The return variable is the tuple *(bc_list, bc_stats)*:

        * *bc_list* is a list of all barcodes that pass the filters.

        * *bc_stats* is a dictionary with the counts of reads parsed. Keys:

            - *nreads* : total reads in files

            - *nfiltered* : number failed Illumina filter

            - *nlowquality* : number that failed quality filter

            - *nmismatched* : number purged due to mismatch
    """
    if all(map(lambda f: os.path.splitext(f)[-1] == '.gz', r1files + r2files)):
        gzipped = True
    elif all(map(lambda f: os.path.splitext(f)[-1] == '.fastq', r1files + r2files)):
        gzipped = False
    else:
        raise ValueError("FASTQ files must either all end in '.gz' or all end in '.fastq'. Instead we had:\n%s" '\n'.join(r1files + r2files))
    bc_list = []
    r1start -= 1 # decrement so we can use zero-based indexing
    assert r1start >= 0 and r2end >= 0
    minqchar = chr(minquality + 33) # minimum character for Q-score
    ntot = nfiltered = nlowquality = nmismatched = nrandomlypurged = 0
    for read_tup in dms_tools.file_io.IteratePairedFASTQ(r1files, r2files, gzipped, applyfilter=True):
        ntot += 1
        if purgefrac:
            if random.random() < purgefrac:
                nrandomlypurged += 1
                continue
        if not read_tup:
            nfiltered += 1
        else:
            (name, r1, r2, q1, q2) = read_tup
            assert len(r1) >= r1start + bc_len, "R1 read %s too short" % name
            assert len(r2) >= r2end + bc_len, "R2 read %s too short" % name
            r1 = r1[r1start : r1start + bc_len]
            q1 = q1[r1start : r1start + bc_len]
            if r2end:
                r2 = r2[-bc_len - r2end : -r2end]
                q2 = q2[-bc_len - r2end : -r2end]
            else:
                r2 = r2[-bc_len - r2end : ]
                q2 = q2[-bc_len - r2end : ]
            assert len(r1) == len(r2), "R1: %s, R2: %s" % (r1, r2)
            if any(map(lambda x: x[0] < minqchar and x[1] < minqchar, zip(q1, q2))):
                nlowquality += 1
                continue
            r2 = dms_tools.cutils.ReverseComplement(r2)
            if r1 != r2:
                nmismatched += 1
            else:
                bc_list.append(r1)
    assert ntot == len(bc_list) + nfiltered + nlowquality + nmismatched + nrandomlypurged
    bc_stats = {
            'nreads':ntot,
            'nfiltered':nfiltered,
            'nlowquality':nlowquality,
            'nmismatches':nmismatched,
            }
    if purgefrac:
        bc_stats['nrandomlypurged'] = nrandomlypurged
    return (bc_list, bc_stats)


def main():
    """Main body of program."""
    random.seed(1)

    parser = dms_tools.parsearguments.MatchSubassembledBarcodesParser()
    args = vars(parser.parse_args())
    prog = parser.prog
    assert len(args['r1files']) == len(args['r2files']), "Different number of r1files and r2files."

    print("Beginning execution of %s in directory %s\n" % (prog, os.getcwd()))
    versionstring = dms_tools.file_io.Versions()
    print(versionstring)
    print("Parsed the following arguments:\n%s\n" % '\n'.join(['\t%s = %s' % tup for tup in args.items()]))
    print("Starting time is %s" % time.asctime())

    outfiles = {}
    for fbase in ['singlemutcounts', 'allmutcounts', 'variantcounts', 'stats']:
        fname = '%s_%s.txt' % (args['outprefix'], fbase)
        if os.path.isfile(fname):
            print("\nRemoving existing file %s" % fname)
            os.remove(fname)
        outfiles[fbase] = fname

    # Get subassembled variants for barcodes
    print("\nReading barcoded subassembled variants from %s" % args['subassembled'])
    (bc_dict, bc_length, wtseq, chartype) = dms_tools.file_io.ReadSubassembledVariants(args['subassembled'])
    print("Read %d barcoded subassembled variants. Each variant is a sequence of length %s." % (len(bc_dict), len(wtseq)))
    print("The mutations are of character type: %s" % chartype)
    print("The barcodes are of length %d nucleotides." % bc_length)

    # Get barcodes that pass quality filters.
    print("\nNow parsing barcodes from %s..." % ', '.join(args['r1files'] + args['r2files']))
    (bc_list, bc_stats) = parse_barcodes(args['r1files'], args['r2files'], args['minquality'], bc_length, args['r1start'], args['r2end'], args['purgefrac'])
    print("Parsed %d high-quality barcodes from %d total read pairs." % (len(bc_list), bc_stats['nreads']))

    # Now parse information to create output files
    # First, set up dictionaries to hold counts
    print("\nNow looping over all %d barcodes to count mutations..." % len(bc_list))
    if chartype  == 'DNA':
        sites = [str(r) for r in range(1, len(wtseq) + 1)]
        chars = dms_tools.nts
        wts = dict([(r, wtseq[int(r) - 1]) for r in sites])
        mutmatch = re.compile('^(?P<wt>[%s])(?P<r>\d+)(?P<mut>[%s])$' % (''.join(chars), ''.join(chars)))
    elif chartype  == 'aminoacid':
        sites = [str(r) for r in range(1, len(wtseq) + 1)]
        chars = dms_tools.aminoacids_withstop
        wts = dict([(r, wtseq[int(r) - 1]) for r in sites])
        mutmatch = re.compile('^(?P<wt>[%s])(?P<r>\d+)(?P<mut>[%s])$' % (''.join(chars), ''.join(chars)))
    elif chartype == 'codon':
        sites = [str(r) for r in range(1, len(wtseq) // 3 + 1)]
        chars = dms_tools.codons
        assert len(wtseq) % 3 == 0
        wts = dict([(r, wtseq[3 * (int(r) - 1) : 3 * int(r)]) for r in sites])
        mutmatch = re.compile('^(?P<wt>[%s]{3})(?P<r>\d+)(?P<mut>[%s]{3})$' % (''.join(dms_tools.nts), ''.join(dms_tools.nts)))
    else:
        raise ValueError("Invalid chartype %s" % chartype)
    allmutcounts = {}
    singlemutcounts = {}
    for r in sites:
        allmutcounts[r] = dict([(char, 0) for char in chars])
        singlemutcounts[r] = dict([(char, 0) for char in chars])
        allmutcounts[r]['WT'] = wts[r]
        singlemutcounts[r]['WT'] = wts[r]
    # now loop over barcodes
    bc_counts = {}
    bc_stats['nunrecognized'] = bc_stats['nretained'] = 0
    for bc in bc_list:
        if bc not in bc_dict:
            bc_stats['nunrecognized'] += 1
        else:
            if bc in bc_counts:
                bc_counts[bc] += 1
            else:
                bc_counts[bc] = 1
            bc_stats['nretained'] += 1
            muts = [mutmatch.search(m) for m in bc_dict[bc]]
            assert all(muts), "Failed to match mutation in %s" % bc_dict[bc]
            muts = dict([(m.group('r'), m.group('mut')) for m in muts])
            assert len(muts) == len(bc_dict[bc]), "Duplicate mutated site in %s" % bc_dict[bc]
            statskey = 'n%dmut' % len(muts)
            if statskey not in bc_stats:
                bc_stats[statskey] = 1
            else:
                bc_stats[statskey] += 1
            for r in sites:
                if r in muts:
                    allmutcounts[r][muts[r]] += 1
                    if len(muts) == 1:
                        singlemutcounts[r][muts[r]] += 1
                else:
                    allmutcounts[r][wts[r]] += 1
                    if not muts:
                        singlemutcounts[r][wts[r]] += 1
    print("Completed looping over all barcodes.")

    print("\nHere are the stats on the parsing of barcodes from the read pairs:")
    statkeys = bc_stats.keys()
    statkeys.sort()
    countkeymatch = re.compile('^n(?P<count>\d+)mut$')
    countkeys = [(int(countkeymatch.search(key).group('count')), key) for key in statkeys if countkeymatch.search(key)]
    statkeys = [key for key in statkeys if not countkeymatch.search(key)]
    countkeys.sort()
    statkeys += [tup[1] for tup in countkeys]
    print('\n'.join(["%s = %d" % (key, bc_stats[key]) for key in statkeys]))
    print("Writing these stats to %s" % outfiles['stats'])
    with open(outfiles['stats'], 'w') as f:
        f.write('\n'.join(["%s = %d" % (key, bc_stats[key]) for key in statkeys]))

    print("\nWriting single-mutant counts to %s" % outfiles['singlemutcounts'])
    print("This file contains mutation counts only from singly mutated variants, and wildtype counts only from unmutated variants.")
    dms_tools.file_io.WriteDMSCounts(outfiles['singlemutcounts'], singlemutcounts)

    print("\nWriting all mutant counts to %s" % outfiles['allmutcounts'])
    print("This file contains counts from all variants at all sites.")
    dms_tools.file_io.WriteDMSCounts(outfiles['allmutcounts'], allmutcounts)

    print("\nWriting counts of each barcode to %s" % outfiles['variantcounts'])
    lines = []
    for bc in bc_dict.keys():
        muts = bc_dict[bc]
        if muts:
            muts = ','.join(muts)
        else:
            muts = 'no_mutations'
        if bc in bc_counts:
            counts = bc_counts[bc]
        else:
            counts = 0
        lines.append((counts, '%s %s' % (bc, muts)))
    lines.sort()
    lines.reverse()
    lines = ['%d %s' % tup for tup in lines]
    with open(outfiles['variantcounts'], 'w') as f:
        f.write("#COUNTS BARCODE MUTATIONS\n")
        f.write('\n'.join(lines))

    print('\nSuccessfully completed %s at %s' % (prog, time.asctime()))


if __name__ == '__main__':
    main() #run the script
