#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@author: kcurry
"""

import os
import argparse
import pathlib
import subprocess
from sys import stdout
from operator import add, mul
from pathlib import Path

import math
import pysam
import numpy as np
import pandas as pd
from flatten_dict import unflatten
from Bio import SeqIO
from Bio.Seq import Seq
from Bio.SeqRecord import SeqRecord
from pandarallel import pandarallel

# static global variables
CIGAR_OPS = [1, 2, 4, 10]
CIGAR_OPS_ALL = [0, 1, 2, 4]
TAXONOMY_RANKS = ['species', 'genus', 'family', 'order', 'class', 'phylum', 'superkingdom']


def validate_input(path):
    """Validate input file is either: fasta, fastq, or sam alignement file.

        path(str): path to input file
    """
    # pass if input is sam file
    sam_pysam = None
    try:
        # pylint: disable=maybe-no-member
        sam_pysam = pysam.AlignmentFile(path)
    except (ValueError, OSError):
        pass
    if sam_pysam:
        return

    # fail if input is not fasta/q
    fasta_rd, fastq_rd = None, None
    try:
        fasta_rd = SeqIO.to_dict(SeqIO.parse(path, "fasta"))
        fastq_rd = SeqIO.to_dict(SeqIO.parse(path, "fastq"))
    except (UnicodeDecodeError, ValueError):
        pass
    if not (fasta_rd or fastq_rd):
        raise TypeError("Input must be in desired format: fasta or fastq")


def get_align_stats(alignment):
    """Retrieve list of inquired cigar stats (I,D,S,X) for alignment

        alignment (pysam.AlignmentFile): align of interest
        return (list(int)): list of counts for each cigar operation defined in (I,D,S,X)
    """
    cigar_stats = alignment.get_cigar_stats()[0]
    n_mismatch = cigar_stats[10] - cigar_stats[1] - cigar_stats[2]
    return [cigar_stats[1], cigar_stats[2], cigar_stats[4], n_mismatch]


def get_align_len(alignment):
    """Retrieve number of columns in alignment

        alignment (pysam.AlignmentFile): align of interest
        return (int): number of columns in alignment
    """
    return sum([alignment.get_cigar_stats()[0][cigar_op] for cigar_op in CIGAR_OPS_ALL])

def output_unclassified(sam_path, seq_output_path, input_type):
    """ Gather list of input sequences that are deemed unclassified by minimap2,
        and output as seq_output_path

        sam_path (str): path to sam file generated by minimap2
        fa_output_path (str): output path for fasta of unclassified sequences
    """
    output_records = []
    # pylint: disable=maybe-no-member
    sam_pysam = pysam.AlignmentFile(sam_path)
    if input_type == "fasta":
        for alignment in sam_pysam.fetch():
            if alignment.reference_name is None:
                record = SeqRecord(
                    Seq(alignment.query_sequence),
                    id=alignment.query_name,
                    name=alignment.query_name,
                )
                output_records.append(record)
        SeqIO.write(output_records, f"{seq_output_path}.fa", "fasta")
    elif input_type == "fastq":
        for alignment in sam_pysam.fetch():
            if alignment.reference_name is None:
                record = SeqRecord(
                    Seq(alignment.query_sequence),
                    id=alignment.query_name,
                    name=alignment.query_name,
                    letter_annotations={"phred_quality": alignment.query_qualities}
                )
                output_records.append(record)
        SeqIO.write(output_records, f"{seq_output_path}.fq", "fastq")

def get_cigar_op_log_probabilities(sam_path):
    """P(align_type) for each type in CIGAR_OPS by counting how often the corresponding
            operations occur in the primary alignments and by normalizing over the total
            sum of operations.

        sam_path(str): path to sam file of interest
        return: log probabilities (list(float)) for each cigar operation defined in CIGAR_OPS,
                    where p > 0
                zero_locs (list(int)): list of indices (int) where probability == 0
                dict_longest_align (dict[str]:(int)): dict of max alignment length
                    for each query read
    """
    cigar_stats_primary = [0] * len(CIGAR_OPS)
    dict_longest_align = {}
    # pylint: disable=maybe-no-member
    sam_pysam = pysam.AlignmentFile(sam_path)
    for alignment in sam_pysam.fetch():
        align_len = get_align_len(alignment)
        if align_len not in dict_longest_align.keys():
            dict_longest_align[alignment.query_name] = align_len
        if not alignment.is_secondary and not alignment.is_supplementary \
                and alignment.reference_name:
            cigar_stats_primary = list(map(add, cigar_stats_primary, get_align_stats(alignment)))
            if dict_longest_align[alignment.query_name] < align_len:
                dict_longest_align[alignment.query_name] = align_len
    zero_locs = [i for i, e in enumerate(cigar_stats_primary) if e == 0]
    if zero_locs:
        for i in sorted(zero_locs, reverse=True):
            del cigar_stats_primary[i]
    n_char = sum(cigar_stats_primary)
    return [math.log(x) for x in np.array(cigar_stats_primary)/n_char], zero_locs, \
           dict_longest_align


def compute_log_prob_rgs(alignment, cigar_stats, log_p_cigar_op, dict_longest_align, align_len):
    """ log(L(r|s)) = log(P(cigar_op)) × n_cigar_op for CIGAR_OPS

        alignment(pysam.AlignmentFile): pysam alignment to score
        cigar_stats(list(int)): list of cigar stats to compute
        log_p_cigar_op(list(float)): list of cigar_op probabilities corresponding to cigar_stats;
                                        computed from primary alignments
        dict_longest_align (dict[str]:(int)): dict of max alignment length for each query read
        align_len (int): number of columns in the alignment
        return: log_score (float): log(L(r|s))
                query_name (str): query name in alignment
                species_tid (int): species-level taxonomy id corresponding to ref_name
    """

    ref_name, query_name = alignment.reference_name, alignment.query_name
    log_score = sum(list(map(mul, log_p_cigar_op, cigar_stats))) * \
                (dict_longest_align[query_name]/align_len)
    species_tid = int(ref_name.split(":")[0])
    return log_score, query_name, species_tid

def log_prob_rgs_dict(sam_path, log_p_cigar_op, dict_longest_align, p_cigar_op_zero_locs=None):
    """dict containing log(L(read|seq)) for all pairwise alignments in sam file

        sam_path(str): path to sam file
        log_p_cigar_op(list(float)): probability for each cigar operation defined in CIGAR_OPS,
                                        where p > 0
        dict_longest_align (dict[str]:(int)): dict of max alignment length for each query read
        zero_locs(list(int)): list of indices (int) where probability == 0
        return ({[str,int]:float}): dict[(query_name,ref_tax_id)]=log(L(query_name|ref_tax_id))
                int: unassigned read count
                int: assigned read count
    """
    # calculate log(L(read|seq)) for all alignments
    log_p_rgs, unassigned_set = {}, set()
    # pylint: disable=maybe-no-member
    sam_filename = pysam.AlignmentFile(sam_path, 'rb')

    if not p_cigar_op_zero_locs:
        for alignment in sam_filename.fetch():
            align_len = get_align_len(alignment)
            if alignment.reference_name and align_len:
                cigar_stats = get_align_stats(alignment)
                log_score, query_name, species_tid = \
                    compute_log_prob_rgs(alignment, cigar_stats, log_p_cigar_op,
                                         dict_longest_align, align_len)
                if (((query_name, species_tid) not in log_p_rgs or
                     log_p_rgs[(query_name, species_tid)] < log_score)):
                    log_p_rgs[(query_name, species_tid)] = log_score
            else:
                unassigned_set.add(alignment.query_name)
    else:
        for alignment in sam_filename.fetch():
            align_len = get_align_len(alignment)
            if alignment.reference_name and align_len:
                cigar_stats = get_align_stats(alignment)
                if sum([cigar_stats[x] for x in p_cigar_op_zero_locs]) == 0:
                    for i in sorted(p_cigar_op_zero_locs, reverse=True):
                        del cigar_stats[i]
                    log_score, query_name, species_tid = \
                        compute_log_prob_rgs(alignment, cigar_stats, log_p_cigar_op,
                                             dict_longest_align, align_len)
                    if (((query_name, species_tid) not in log_p_rgs or
                         log_p_rgs[(query_name, species_tid)] < log_score)):
                        log_p_rgs[(query_name, species_tid)] = log_score
            else:
                unassigned_set.add(alignment.query_name)

    assigned_reads = set(unflatten(log_p_rgs).keys())
    unassigned_reads = unassigned_set - assigned_reads
    unassigned_count = len(unassigned_reads)
    stdout.write(f"Unassigned read count: {unassigned_count}\n")

    ## remove if p(r|s) < 0.01
    #min_p_thresh = math.log(0.01)
    #log_p_rgs = {r_map: val for r_map, val in log_p_rgs.items() if val > min_p_thresh}
    return log_p_rgs, unassigned_count, len(assigned_reads)


def expectation_maximization(log_p_rgs, freq):
    """One iteration of the EM algorithm. Updates the relative abundance estimation in f based on
            probabilities in log_L_rgs.

        log_p_rgs({[str,int]:float}): dict[(query_name,ref_tax_id)]=log(L(query_name|ref_tax_id))
        freq{int:float}: dict[species_tax_id]:likelihood species is present in sample
        returns: f {int:float}: dict[species_tax_id]:updated likelihood species is present in sample
                total_log_likelihood (float): log likelihood updated f is accurate
    """
    # log_p_rns = log(L(r|s))+log(f(s)) for each alignment
    log_p_rns = {}
    for (read, seq) in log_p_rgs:
        if seq in freq and freq[seq] != 0:
            log_p_rns[(read, seq)] = (log_p_rgs[(read, seq)] + math.log(freq[seq]))
    log_c = {r: -max(smap.values()) for r, smap in unflatten(log_p_rns).items()}
    p_rns_c = {(r, s): (math.e ** (v + log_c[r])) for (r, s), v in log_p_rns.items()}
    p_r_c = {r: sum(s.values()) for r, s in unflatten(p_rns_c).items()}

    # P(s|r), likelihood read r emanates db seq s
    p_sgr = unflatten({(seq, read): v / p_r_c[read] for (read, seq), v in p_rns_c.items()})

    # update total likelihood and f vector
    log_p_r = {read: math.log(v) - log_c[read] for read, v in p_r_c.items()}
    n_reads = len(log_p_r)
    return {seq: (sum(r_map.values()) / n_reads)
            for seq, r_map in p_sgr.items()}, \
           sum(log_p_r.values()), p_sgr


def expectation_maximization_iterations(log_p_rgs, db_ids, lli_thresh, input_threshold):
    """Full expectation maximization algorithm for alignments in log_L_rgs dict

        log_p_rgs{[str,int]:float}: dict[(query_name,ref_tax_id)]=log(L(query_name|ref_tax_id))
        db_ids(list(int)): list of each unique species taxonomy id present in database
        lli_thresh(float): log likelihood increase minimum to continue EM iterations
        input_threshold(float): minimum relative abundance in output
        return: {int:float}: dict[species_tax_id]:estimated likelihood species is present in sample
                float: min abundance threshold
    """
    n_db = len(db_ids)
    n_reads = len(unflatten(log_p_rgs))
    stdout.write(f"Assigned read count: {n_reads}\n")
    if n_reads == 0:
        raise ValueError("0 reads assigned")
    freq, counter = dict.fromkeys(db_ids, 1 / n_db), 1

    # set output abundance threshold
    freq_thresh = 1/n_reads
    if n_reads > 1000:
        freq_thresh = 10/n_reads

    total_log_likelihood = -math.inf
    while True:
        freq, updated_log_likelihood, _ = expectation_maximization(log_p_rgs, freq)

        # check f vector sums to 1
        freq_sum = sum(freq.values())
        if not .9 <= freq_sum <= 1.1:
            raise ValueError(f"f sums to {freq_sum}, rather than 1")

        # confirm log likelihood increase
        log_likelihood_diff = updated_log_likelihood - total_log_likelihood
        total_log_likelihood = updated_log_likelihood
        if log_likelihood_diff < 0:
            raise ValueError("total_log_likelihood decreased from prior iteration")

        # exit loop if log likelihood increase less than threshold
        if log_likelihood_diff < lli_thresh:
            stdout.write(f"Number of EM iterations: {counter}\n")
            freq = {k: v for k, v in freq.items() if v >= freq_thresh}
            freq_full, updated_log_likelihood, p_sgr = expectation_maximization(log_p_rgs, freq)
            freq_set_thresh = None
            if freq_thresh < input_threshold:
                freq = {k: v for k, v in freq_full.items() if v >= input_threshold}
                freq_set_thresh, updated_log_likelihood, p_sgr = \
                    expectation_maximization(log_p_rgs, freq)
            return freq_full, freq_set_thresh, p_sgr

        #output current estimation
        #freq_to_lineage_df(freq, f"{out_file}_{counter}", df_nodes, df_names)
        counter += 1


def lineage_dict_from_tid(taxid, nodes_dict, names_dict):
    """Retrieve dict of lineage for given taxid

        tid(int): tax id to retrieve lineage dict
        nodes_df(df): pandas df of nodes.dmp with columns ['tax_id', 'parent_tax_id', 'rank'];
                            tax_id as index
        names_df(df): pandas df of names.dmp with columns ['tax_id', 'name_txt']; tax_id as index
        return {str:str}: dict[rank]:taxonomy name at rank
    """

    lineage_dict = {}
    while names_dict[taxid] != "root":
        tup = nodes_dict[taxid]
        lineage_dict[tup[1]] = names_dict[taxid]
        taxid = tup[0]
    return lineage_dict


def freq_to_lineage_df(freq, tsv_output_path, taxonomy_df, assigned_count,
                       unassigned_count, counts=False):
    """Converts freq to a pandas df where each row contains abundance and tax lineage for
                classified species in f.keys(). Stores df as .tsv file in tsv_output_path.

        freq{int:float}: dict[species_tax_id]:estimated likelihood species is present in sample
        tsv_output_path(str): path to output .tsv file
        taxonomy_df(df): pandas df of all db sequence taxonomy with index 'tax_id'
        assigned_count(int): number of assigned reads
        unassigned_count(int): number of unassigned reads
        counts(boolean): True if include estimated counts in output .tsv file
        returns(df): pandas df with lineage and abundances for values in f
    """

    results_df = pd.DataFrame(zip(list(freq.keys()) + ['unassigned'],
                                  list(freq.values()) + ["-"]),
                              columns=["tax_id", "abundance"]).set_index('tax_id')
    results_df = results_df.join(taxonomy_df, how='left').reset_index()
    if counts:
        counts_series = (results_df["abundance"] * assigned_count)[:-1].append(pd.Series(unassigned_count)).reset_index(drop=True)
        results_df["estimated counts"] = counts_series
    results_df.to_csv(f"{tsv_output_path}.tsv", sep='\t', index=False)
    return results_df

def generate_alignments(in_file_list, out_basename, database):
    """ Generate .sam alignment file

        in_file_list(list(str)): list of path(s) to input sequences
        out_basename(str): path and basename for output files
    """
    input_file = " ".join(in_file_list)
    filetype = pathlib.PurePath(args.input_file[0]).suffix
    if filetype == '.sam':
        args.keep_files = True
        return input_file
    sam_align_file = f"{out_basename}_emu_alignments.sam"
    db_sequence_file = os.path.join(database, 'species_taxid.fasta')

    subprocess.check_output(
            f"minimap2 -ax {args.type} -t {args.threads} -N {args.N} -p .9 "
            f"-K {args.K} {db_sequence_file} {input_file} -o {sam_align_file}",
            shell=True)
    return sam_align_file

def output_read_assignments(p_sgr, tsv_output_path):
    """ Output file of read assignment distributions for all

        p_sgr({tid:{read_id:probability}}): P(s|r), likelihood read r emanates db seq s
        tsv_output_path(str): path to output .tsv file
        returns(df): pandas df of read assignment distributions
    """

    dist_df = pd.DataFrame(p_sgr)
    dist_df.to_csv(f"{tsv_output_path}.tsv", sep='\t')
    return dist_df

def create_nodes_dict(nodes_path):
    """convert nodes.dmp file into pandas dataframe

        nodes_path(str): path to nodes.dmp file
        returns(df): pandas df containing 'tax_id', 'parent_tax_id', and 'rank'
    """
    node_headers = ['tax_id', 'parent_tax_id', 'rank']
    nodes_df = pd.read_csv(nodes_path, sep='\t', header=None, dtype=str)[[0, 2, 4]]
    nodes_df.columns = node_headers
    return dict(zip(nodes_df['tax_id'], tuple(zip(nodes_df['parent_tax_id'], nodes_df['rank']))))


def create_names_dict(names_path):
    """convert names.dmp file into pandas dataframe

        names_path(str): path to names.dmp file
        returns(df): pandas df containing ['tax_id','name_txt'] for each row with
                    'name_class' = 'scientific name'
    """
    name_headers = ['tax_id', 'name_txt', 'unique_name', 'name_class']
    names_df = pd.read_csv(names_path, sep='\t', index_col=False, header=None, dtype=str)\
        .drop([1, 3, 5, 7], axis=1)
    names_df.columns = name_headers
    names_df = names_df[names_df["name_class"] == "scientific name"]
    names_df = names_df.drop(columns=['unique_name', 'name_class'])
    return dict(zip(names_df['tax_id'], names_df['name_txt']))

def build_taxonomy(unique_tids, nodes_dict, names_dict, threads):
    """Converts freq to a pandas df where each row contains abundance and tax lineage for
                classified species in f.keys(). Stores df as .tsv file in tsv_output_path.

        unique_tids(list): list of ints of each unique taxid in list of database sequences
        nodes_dict{int:[int, str]}: dict of nodes.dmp with 'tax_id' as keys,
                                        tuple ('parent_taxid', 'rank') as values
        names_dict{int:str}: pandas dict of names.dmp with 'tax_id' as keys,
                                        'name_txt' as values
        threads(int): number of threads to use in parallelization
        returns(df): pandas df with lineage for all taxids in unique_tids
    """
    pandarallel.initialize(nb_workers=threads)
    results_df = pd.DataFrame(list(unique_tids), columns=["tax_id"])
    lineages = results_df["tax_id"].parallel_apply\
        (lambda x: lineage_dict_from_tid(x, nodes_dict, names_dict))
    results_df = pd.concat([results_df, pd.json_normalize(lineages)], axis=1)
    header_order = ["tax_id", "species", "genus", "family", "order", "class",
                    "phylum", "clade", "superkingdom", "subspecies",
                    "species subgroup", "species group"]
    for col in header_order:
        if col not in results_df.columns:
            results_df[col] = ""
    results_df = results_df.sort_values(header_order[8:0:-1]).reset_index(drop=True)
    results_df = results_df.reindex(header_order, axis=1)
    return results_df


def get_species_tid(tid, nodes_dict):
    """ Get lowest taxid down to species-level in lineage for taxid [tid]

        tid(int): taxid for species level or more specific
        nodes_dict{int:[int, str]}: dict of nodes.dmp with 'tax_id' as keys,
                                        tuple ('parent_taxid', 'rank') as values
        return(int): species taxid in lineage
    """
    if str(tid) not in nodes_dict.keys():
        raise ValueError(f"Taxid:{tid} not found in nodes file")
    while nodes_dict[str(tid)][1] not in TAXONOMY_RANKS:
        tid = nodes_dict[str(tid)][0]
    return tid


def create_seq2tax_dict(seq2tax_path, nodes_dict):
    """Convert seqid-taxid mapping in seq2tax_path to dict mapping seqid to species level taxid

        seq2tax_path(str): path to seqid-taxid mapping file
        nodes_dict{int:[int, str]}: dict of nodes.dmp with 'tax_id' as keys,
                                        tuple ('parent_taxid', 'rank') as values
        returns {str:int}: dict[seqid] = species taxid
    """
    seq2tax_dict, species_id_dict = {}, {}
    with open(seq2tax_path) as file:
        for line in file:
            (seqid, tid) = line.rstrip().split("\t")
            if tid in species_id_dict.keys():
                species_tid = species_id_dict[tid]
            else:
                species_tid = get_species_tid(tid, nodes_dict)
                species_id_dict[tid] = species_tid
            seq2tax_dict[seqid] = species_tid
    return seq2tax_dict


def create_unique_seq_dict(db_fasta_path, seq2tax_dict):
    """ Creates dict of unique sequences to species taxids connected with the sequence

        db_fasta_path(str): path to fasta file of database sequences
        seq2tax_dict{str:int}: dict[seqid] = species taxid
        returns {str:{int:[str]}}: dict[seq] = {species_taxid: [list of sequence ids]}
    """
    fasta_dict = {}
    for record in SeqIO.parse(db_fasta_path, "fasta"):
        tid = seq2tax_dict[record.id]
        if tid:
            if record.seq in fasta_dict.keys():
                if tid in fasta_dict[record.seq].keys():
                    fasta_dict[record.seq][tid] += [record.description]
                else:
                    fasta_dict[record.seq][tid] = [record.description]
            elif record.seq.reverse_complement() in fasta_dict.keys():
                if tid in fasta_dict[record.seq.reverse_complement()].keys():
                    fasta_dict[record.seq.reverse_complement()][tid] += [record.description]
                else:
                    fasta_dict[record.seq.reverse_complement()][tid] = [record.description]
            else:
                fasta_dict[record.seq] = {tid: [record.description]}
    return fasta_dict


def create_reduced_fasta(fasta_dict, db_name):
    """ Creates fasta file of taxid for each sequences in fasta_dict with id
            'species_taxid:db_name:sequence_id'

        fasta_dict{str:{int:[str]}}: dict[seq] = {species_taxid: [list of sequence ids]}
        db_name(str): name to represent database represented in fasta_dict
        returns (list[Bio.SeqRecord]): list of sequences for output fasta file
    """
    records, count = [], 1
    for seq, tid_dict in fasta_dict.items():
        for taxid, descriptions in tid_dict.items():
            records += [SeqRecord(seq,
                                  id=f"{taxid}:{db_name}:{count}", description=f"{descriptions}")]
            count += 1
    return records

def collapse_rank(path, rank):
    """ Stores a version of emu-output (path) collapsed at the specified taxonomic rank in same
            folder as input.

        path(str): path to emu output
        rank(str): taxonomic rank for collapsed abundance: ["species", "genus", "family",
            "order", "class", "phylum", "clade", "superkingdom"]
    """
    df_emu = pd.read_csv(path, sep='\t')
    tax_order = list(df_emu.columns[1:9])
    if rank not in tax_order:
        raise ValueError(f"Specified rank must be in list: {TAXONOMY_RANKS}")
    taxa_remaining = tax_order[tax_order.index(rank):]
    # pylint: disable=unsubscriptable-object
    if "estimated counts" in df_emu.columns:
        df_emu_copy = df_emu[['abundance', 'estimated counts'] + taxa_remaining]
    else:
        df_emu_copy = df_emu[['abundance'] + taxa_remaining]
    df_emu_copy = df_emu_copy.groupby(taxa_remaining, dropna=False).sum()
    output_path = f"{os.path.splitext(path)[0]}-{rank}.tsv"
    df_emu_copy.to_csv(output_path, sep='\t')
    stdout.write(f"File generated: {output_path}\n")


if __name__ == "__main__":
    __version__ = "3.3.0"
    parser = argparse.ArgumentParser()
    parser.add_argument('--version', '-v', action='version', version='%(prog)s v' + __version__)
    subparsers = parser.add_subparsers(dest="subparser_name", help='sub-commands')
    abundance_parser = subparsers.\
        add_parser("abundance", help="Generate relative abundance estimates")
    abundance_parser.add_argument(
        'input_file', type=str, nargs='+',
        help='filepath to input nt sequence file')
    abundance_parser.add_argument(
        '--type', '-x', choices=['map-ont', 'map-pb', 'sr'], default='map-ont',
        help='short-read: sr, Pac-Bio:map-pb, ONT:map-ont [map-ont]')
    abundance_parser.add_argument(
        '--min-abundance', '-a', type=float, default=0.0001,
        help='min species abundance in results [0.0001]')
    abundance_parser.add_argument(
        '--db', type=str, default=os.environ.get("EMU_DATABASE_DIR"),
        help='path to emu database containing: names_df.tsv, '
             'nodes_df.tsv, species_taxid.fasta, unqiue_taxids.tsv [$EMU_DATABASE_DIR]')
    abundance_parser.add_argument(
        '--N', '-N', type=int, default=50,
        help='minimap max number of secondary alignments per read [50]')
    abundance_parser.add_argument(
        '--K', '-K', type=int, default=500000000,
        help='minibatch size for minimap2 mapping [500M]')
    abundance_parser.add_argument(
        '--output-dir', type=str, default="./results",
        help='output directory name [./results]')
    abundance_parser.add_argument(
        '--output-basename', type=str,
        help='basename for all emu output files [{input_file}]')
    abundance_parser.add_argument(
        '--keep-files', action="store_true",
        help='keep working files in output-dir')
    abundance_parser.add_argument(
        '--keep-counts', action="store_true",
        help='include estimated read counts in output')
    abundance_parser.add_argument(
        '--keep-read-assignments', action="store_true",
        help='output file of read assignment distribution')
    abundance_parser.add_argument(
        '--output-unclassified', action="store_true",
        help='output unclassified sequences')
    abundance_parser.add_argument(
        '--threads', type=int, default=3,
        help='threads utilized by minimap [3]')

    build_db_parser = subparsers.add_parser("build-database", help="Build custom Emu database")
    build_db_parser.add_argument(
        'db_name', type=str,
        help='custom database name')
    build_db_parser.add_argument(
        '--names', type=str,
        help='path to names.dmp file')
    build_db_parser.add_argument(
        '--nodes', type=str,
        help='path to nodes.dmp file')
    build_db_parser.add_argument(
        '--sequences', type=str,
        help='path to fasta of database sequences')
    build_db_parser.add_argument(
        '--seq2tax', type=str,
        help='path to tsv mapping species tax id to fasta sequence headers')
    build_db_parser.add_argument(
        '--threads', type=int, default=3,
        help='threads [3]')

    collapse_parser = subparsers.add_parser("collapse-taxonomy",
                                            help="Collapse emu output at specified taxonomic rank")
    collapse_parser.add_argument(
        'input_path', type=str,
        help='emu output filepath')
    collapse_parser.add_argument(
        'rank', type=str,
        help='collapsed taxonomic rank')
    args = parser.parse_args()


    if args.subparser_name == "abundance":
        # check input file is fasta/q or sam alignment file
        #validate_input(args.input_file[0])

        # convert taxonomy files to dataframes
        if not args.db:
            raise ValueError("Database not specified. "
                             "Either 'export EMU_DATABASE_DIR=<path_to_database>' or "
                             "utilize '--db' parameter.")
        df_taxonomy = pd.read_csv(os.path.join(args.db, "taxonomy.tsv"), sep='\t',
                                  index_col='tax_id', dtype=str)
        db_species_tids = df_taxonomy.index

        # set up output paths
        if not os.path.exists(args.output_dir):
            os.makedirs(args.output_dir)
        out_file = os.path.join(args.output_dir, "-".join([Path(v).stem for v in args.input_file]))
        if args.output_basename:
            out_file = os.path.join(args.output_dir, args.output_basename)

        # perform EM algorithm & generate output
        sam_file = generate_alignments(args.input_file, out_file, args.db)
        log_prob_cigar_op, locs_p_cigar_zero, longest_align_dict = \
            get_cigar_op_log_probabilities(sam_file)
        log_prob_rgs, counts_unassigned, counts_assigned = log_prob_rgs_dict(
            sam_file, log_prob_cigar_op, longest_align_dict, locs_p_cigar_zero)
        f_full, f_set_thresh, read_dist = expectation_maximization_iterations(log_prob_rgs,
                                                                   db_species_tids,
                                                                   .01, args.min_abundance)
        freq_to_lineage_df(f_full, f"{out_file}_rel-abundance", df_taxonomy,
                           counts_assigned, counts_unassigned, args.keep_counts)

        if args.keep_read_assignments:
            output_read_assignments(read_dist, f"{out_file}_read-assignment-distributions")

        if f_set_thresh:
            freq_to_lineage_df(
                f_set_thresh,
                f"{out_file}_rel-abundance-threshold-{args.min_abundance}",
                df_taxonomy, counts_assigned, counts_unassigned, args.keep_counts)

        if args.output_unclassified:
            ext = os.path.splitext(args.input_file[0])[-1]
            INPUT_FILETYPE = "fasta"
            if ext in [".fastq", ".fq"]:
                INPUT_FILETYPE = "fastq"
            output_unclassified(sam_file, f"{out_file}_unclassified", INPUT_FILETYPE)

        # clean up extra file
        if not args.keep_files:
            if os.path.exists(sam_file):
                os.remove(sam_file)

    if args.subparser_name == "build-database":
        emu_db_path = os.getcwd()
        custom_db_path = os.path.join(emu_db_path, args.db_name)
        if not os.path.exists(custom_db_path):
            os.makedirs(custom_db_path)
        stdout.write(f"Emu custom database generating at path: {custom_db_path} ...\n")

        dict_names = create_names_dict(args.names)
        dict_nodes = create_nodes_dict(args.nodes)
        seq2tax = create_seq2tax_dict(args.seq2tax, dict_nodes)
        db_unique_ids = set(seq2tax.values())

        dict_fasta = create_unique_seq_dict(args.sequences, seq2tax)
        fasta_records = create_reduced_fasta(dict_fasta, args.db_name)
        SeqIO.write(fasta_records, os.path.join(custom_db_path, 'species_taxid.fasta'), "fasta")

        df_taxonomy = build_taxonomy(db_unique_ids, dict_nodes, dict_names, args.threads)
        df_taxonomy.to_csv(os.path.join(custom_db_path, 'taxonomy.tsv'), index=False, sep='\t')
        stdout.write("Database creation successful\n")

    if args.subparser_name == "collapse-taxonomy":
        collapse_rank(args.input_path, args.rank)
