"""Count occurrence of bases in a set of sequences.
"""

from Bio import SeqIO
from itertools import combinations
import Bio.Seq
import Bio.SeqRecord
import sys
import csv

def register_arguments(parser):
    parser.add_argument('--sequences', '-s', required=True, help="sequences in fasta format")
    parser.add_argument('--output', '-o', help="tab-delimited file containing the number of bases per sequence in the given file. Output columns include strain, length, and counts for A, C, G, T, N, other valid IUPAC characters, ambiguous characters ('?' and '-'), and other invalid characters.", required=True)
    parser.add_argument('--verbose', '-v', action="store_true", help="print index statistics to stdout")

def index_sequence(sequence, values):
    """Count the number of nucleotides for a given sequence record.

    Parameters
    ----------
    sequence : Bio.SeqRecord.SeqRecord
        sequence record to index.

    values : list of sets of str
        values to count; sets must be non-overlapping and contain only
        single-character, lowercase strings

    Returns
    -------
    list :
        summary of the given sequence's strain name, length, nucleotide counts
        for the given values, and a final column with the number of characters
        that didn't match any of those in the given values.

    >>> other_IUPAC = {'r', 'y', 's', 'w', 'k', 'm', 'd', 'h', 'b', 'v'}
    >>> values = [{'a'},{'c'},{'g'},{'t'},{'n'}, other_IUPAC, {'-'}, {'?'}]
    >>> sequence_a = Bio.SeqRecord.SeqRecord(seq=Bio.Seq.Seq("ACTGN-?XWN"), id="seq_A")
    >>> index_sequence(sequence_a, values)
    ['seq_A', 10, 1, 1, 1, 1, 2, 1, 1, 1, 1]

    >>> sequence_b = Bio.SeqRecord.SeqRecord(seq=Bio.Seq.Seq("ACTGACTG"), id="seq_B")
    >>> index_sequence(sequence_b, values)
    ['seq_B', 8, 2, 2, 2, 2, 0, 0, 0, 0, 0]

    Characters in the given sequence that are not in the given list of values to
    count get counted in the final column.

    >>> sequence_c = Bio.SeqRecord.SeqRecord(seq=Bio.Seq.Seq("ACTG%@!!!NN"), id="seq_C")
    >>> index_sequence(sequence_c, values)
    ['seq_C', 11, 1, 1, 1, 1, 2, 0, 0, 0, 5]

    The list of value sets must not overlap.

    >>> sequence_d = Bio.SeqRecord.SeqRecord(seq=Bio.Seq.Seq("A!C!TGXN"), id="seq_D")
    >>> index_sequence(sequence_d, [set('actg'), set('xn'), set('n')]) # doctest: +ELLIPSIS
    Traceback (most recent call last):
      ...
    ValueError: character sets ... and {'n'} overlap: {'n'}

    Value sets must contain only single-character, lowercase strings.

    >>> index_sequence(sequence_d, [{'a'}, {'c'}, {'T'}, {'g'}])
    Traceback (most recent call last):
      ...
    ValueError: character set {'T'} contains a non-lowercase character: 'T'

    >>> index_sequence(sequence_d, [{'actg'}])
    Traceback (most recent call last):
      ...
    ValueError: character set {'actg'} contains a multi-character (or maybe zero-length) string: 'actg'

    >>> index_sequence(sequence_d, [{'a', 'c'}, {0, 1}])
    Traceback (most recent call last):
      ...
    ValueError: character set {0, 1} contains a non-string element: 0
    """
    # Sets must be non-overlapping and contain only single-character, lowercase
    # strings, otherwise our assumptions are broken.  This is not much hardship
    # on the caller.
    for v in values:
        for c in v:
            if not isinstance(c, str):
                raise ValueError(f"character set {v!r} contains a non-string element: {c!r}")
            if len(c) != 1:
                raise ValueError(f"character set {v!r} contains a multi-character (or maybe zero-length) string: {c!r}")
            if c != c.lower():
                raise ValueError(f"character set {v!r} contains a non-lowercase character: {c!r}")

    for a, b in combinations(values, 2):
        if not a.isdisjoint(b):
            raise ValueError(f"character sets {a!r} and {b!r} overlap: {a & b!r}")

    counts = []
    seq = sequence.seq.lower()
    l = len(seq)

    for v in values:
        counts.append(sum(map(lambda x: seq.count(x), v)))

    invalid_nucleotides = l-sum(counts)
    row = [sequence.id, l]+counts+[invalid_nucleotides]

    return row


def index_sequences(sequences_path, sequence_index_path):
    """Count the number of A, C, T, G, N, other IUPAC nucleotides, ambiguous bases
    ("?" and "-"), and other invalid characters in a set of sequences and write
    the composition as a data frame to the given sequence index path.

    Parameters
    ----------
    sequences_path : str or Path-like
        path to a sequence file to index.

    sequence_index_path : str or Path-like
        path to a tab-delimited file containing the composition details for each
        sequence in the given input file.

    Returns
    -------
    int :
        number of sequences indexed

    int :
        total length of sequences indexed

    """
    #read in files
    try:
        seqs = SeqIO.parse(sequences_path, 'fasta')
    except ValueError as error:
        print("ERROR: Problem reading in {}:".format(sequences_path), file=sys.stderr)
        print(error, file=sys.stderr)
        return 1

    other_IUPAC = {'r', 'y', 's', 'w', 'k', 'm', 'd', 'h', 'b', 'v'}
    values = [{'a'},{'c'},{'g'},{'t'},{'n'},other_IUPAC,{'-'},{'?'}]
    labels = ['A','C','G','T','N','other_IUPAC','-','?']

    tot_length = 0
    num_of_seqs = 0

    with open(sequence_index_path, 'wt') as out_file:
        tsv_writer = csv.writer(out_file, delimiter = '\t')

        #write header i output file
        header = ['strain', 'length']+labels+['invalid_nucleotides']
        tsv_writer.writerow(header)

        for record in seqs:
            #index the sequence and write row in output file
            row = index_sequence(record, values)
            tsv_writer.writerow(row)

            tot_length += row[1]
            num_of_seqs += 1

    seqs.close()

    return num_of_seqs, tot_length


def run(args):
    '''
    runs index_sequences which counts the number of A, C, T, G, N, other IUPAC nucleotides, ambiguous bases
    ("?" and "-"), and other invalid characters in a set of sequences and write
    the composition as a data frame to the given sequence index path.
    '''
    num_of_seqs, tot_length = index_sequences(args.sequences, args.output)

    if args.verbose:
        print("Analysed %i sequences with an average length of %i nucleotides." % (num_of_seqs, int(tot_length / num_of_seqs)))
