#!/usr/bin/env python
from more_itertools import chunked
import argparse
import gzip
import re

if __name__=='__main__':
    argparser = argparse.ArgumentParser(
        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
        description="""Trim vasa fastq files for homo-polymers
    """)


    argparser.add_argument('R2_fastq', metavar='R2_fastq', type=str)
    argparser.add_argument('R2_singletons_out', type=str)

    argparser.add_argument('-poly_length', type=int, default=10)
    argparser.add_argument('-min_read_len', type=int, default=20)
    args = argparser.parse_args()

    poly_A = args.poly_length*'A'
    poly_T = args.poly_length*'T'
    poly_G = args.poly_length*'G'

    poly_length = args.poly_length

    r2_trimmer = re.compile('[GA]*$')

    def trim_r2(header, sequence, comment, qualities, min_read_len ):

        start = sequence.find(poly_A)
        if start != -1:
            sequence, qualities = sequence[:start], qualities[:start]

        start = sequence.find(poly_G)
        if start != -1:
            sequence, qualities = sequence[:start], qualities[:start]

        # Trim any trailing A and G bases from the end and # Trim down 3 bases
        sequence = r2_trimmer.sub('',sequence)[:-3]
        qualities = qualities[:len(sequence)]


        return f'{header}{sequence.rstrip()}\n{comment}{qualities.rstrip()}\n', len(sequence)>=min_read_len

    def trim_r1(header, sequence, comment, qualities, min_read_len ):

        start = sequence.rfind(poly_T) # Trim poly T
        if start != -1:
            sequence, qualities = sequence[start+poly_length:], qualities[start+poly_length:]

        start = sequence.find(poly_G)
        if start != -1:
            sequence, qualities = sequence[:start], qualities[:start]


        return f'{header}{sequence.rstrip()}\n{comment}{qualities.rstrip()}\n', len(sequence)>=min_read_len

    with gzip.open(args.R2_fastq,'rt') as r2, \
         gzip.open(args.R2_singletons_out,'wt',compresslevel=1) as r2_single_out:

        for read2 in chunked(r2,4):
            (r2_o, valid_r2)  = trim_r2(*read2, args.min_read_len)
            if valid_r2:
                r2_single_out.write(r2_o)