#!/usr/bin/env python3

# Filter FASTA files based on genbank information

import argparse
import logging
from pathlib import Path
import sys

from Bio import SeqIO
from Bio.SeqFeature import FeatureLocation, CompoundLocation
from Bio.SeqRecord import SeqRecord


def parse_args(args=None):
    parser = argparse.ArgumentParser()
    parser.add_argument("fastafile", type=argparse.FileType('r'),
                        metavar="FASTA_FILE")
    parser.add_argument("gbdir", metavar="GENBANK_DIR")
    parser.add_argument('-d', '--debug', action='store_true')
    parser.add_argument('-o', '--outfile', type=argparse.FileType('a'),
                        default=sys.stdout, metavar="OUTPUT_FILE")
    parser.add_argument('-p', '--padding', type=int, default=0,
                        help="extra padding around genbank regions")
    args = parser.parse_args(args)
    return args


def init_logging(debug=False):
    msgFormat = '%(asctime)s %(levelname)s %(message)s'
    dateFormat = '%m/%d/%Y %H:%M:%S'
    level = logging.INFO
    if debug:
        level = logging.DEBUG
    logging.basicConfig(format=msgFormat, datefmt=dateFormat, level=level)


def fragment_accession(accession_version, fragment):
    accession, version = accession_version.split('.')
    return '{}-fragment{}.{}'.format(accession, fragment, version)


def write_sequence(out_fp, sequence, desc, accession):
    header_fmt = 'gi|NONE|gb|{}|Generated by gbfilter'
    header = header_fmt.format(accession)
    record = SeqRecord(sequence, id=header, description=desc)
    SeqIO.write(record, out_fp, 'fasta')


def feature_description(feature, sub_features=None):
    type_string = f'{feature.type}'
    if sub_features:
        subs = '-'.join(sub_features)
        type_string = f'{type_string} {subs}'
    return f'{type_string} {feature.location}'


def pad_compound_location(location, seq_len, padding):
    logging.debug('CompoundLocation: %r', location)
    if not (len(location.parts) == 2 and location.start == 0 and
            location.end == seq_len):
        logging.warning('compound location does not wrap from end to start: %r',
                        location)
        # raise AssertionError("Compound location does not wrap from end to start")
    new_parts = [pad_location(loc, seq_len, padding) for loc in location.parts]
    new_location = CompoundLocation(new_parts, location.operator)
    return new_location


def pad_location(location, seq_len, padding):
    """Given a FeatureLocation, pad it and return the new location.
    """
    # logging.debug('seq_len = %r', seq_len)
    # logging.debug('padding = %r', padding)
    # logging.debug('location.start = %r', location.start)
    # logging.debug('location.end = %r', location.end)
    new_start = max(0, location.start - padding)
    new_end = min(seq_len, location.end + padding)
    # logging.debug('new_start = %r', new_start)
    # logging.debug('new_end = %r', new_end)
    return FeatureLocation(new_start, new_end, strand=location.strand)


def padded_location(location, seq_len, padding):
    if isinstance(location, CompoundLocation):
        return pad_compound_location(location, seq_len, padding)
    return pad_location(location, seq_len, padding)


def get_filtered_sequences(gb_file, padding):
    # This is a generator, yielding sequences instead of returning a
    # list of all sequences. This way we don't keep a long list in
    # memory when we're just writing them out to the contrast fasta
    # file
    filtered_features = ['tRNA', 'rRNA', 'tmRNA']
    filtered_cds_genes = frozenset(['rpoA', 'RpoA',
                                    'rpoB', 'RpoB',
                                    'rpoC', 'RpoC',
                                    'rpoD', 'RpoD',
                                    'dnaA', 'DnaA',
                                    'dnaB', 'DnaB',
                                    'gyrB', 'GyrB'])
    for record in SeqIO.parse(str(gb_file), 'genbank'):
        for feature in record.features:
            if feature.type in filtered_features:
                feature.location = padded_location(feature.location,
                                                   len(record.seq),
                                                   padding)
                # logging.debug('location.start = %r', feature.location.start)
                # logging.debug('location.end = %r', feature.location.end)
                contrast_seq = feature.extract(record.seq)
                yield contrast_seq, feature_description(feature)
            elif feature.type == 'CDS':
                # Coding Sequence
                gene = set(feature.qualifiers.get('gene', []))
                gene_names = gene.intersection(filtered_cds_genes)
                if not gene_names:
                    continue
                # The gene is one we want to filter
                # (code reuse is copying code so it exists in two places, right?
                feature.location = padded_location(feature.location,
                                                   len(record.seq),
                                                   padding)
                contrast_seq = feature.extract(record.seq)
                yield contrast_seq, feature_description(feature, gene_names)


def append_sequences(gb_file, fasta_fp, accession_version, padding):
    # Given a genbank file (gb_file) and an open fasta file for
    # output, write sequences that should be given to contrast based
    # on rules encoded in get_filtered_sequences. The
    # accession_version of the original sequence is used to generate
    # an accession version for each fragment output to the fasta file.
    i = 0
    for seq, desc in get_filtered_sequences(gb_file, padding):
        i += 1
        accession = fragment_accession(accession_version, i)
        write_sequence(fasta_fp, seq, desc, accession)


def main(argv=None):
    args = parse_args(argv)
    init_logging(args.debug)

    # Convert to a pathlib.Path for easier appending
    gb_path = Path(args.gbdir)

    # read the fasta file
    # for each accession:
    #    if there is a matching gb file
    #        parse the gb file
    #
    for rec in SeqIO.parse(args.fastafile, 'fasta'):
        dfields = rec.description.split('|')
        accession = None
        if dfields[0] == 'gi' and dfields[2] == 'gb':
            accession = dfields[3]
        # Parse other FASTA header types here
        gb_file = gb_path / (accession + '.gb')
        if gb_file.exists():
            logging.debug(gb_file)
            append_sequences(gb_file, args.outfile, accession, args.padding)


if __name__ == '__main__':
    main()
