# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/02_seq_tools.ipynb.

# %% auto 0
__all__ = ['BioSeq', 'DnaSeq', 'ProteinSeq', 'Fastq', 'Fasta', 'NucleotideFasta', 'ProteinFasta', 'Assembly', 'SpadesAssembly',
           'translate', 'uniquify', 'filter_spades_batch']

# %% ../nbs/02_seq_tools.ipynb 3
# standard libs
import os
import sys
import re
import pandas

# Common to template
# add into settings.ini, requirements, package name is python-dotenv, for conda build ensure `conda config --add channels conda-forge`
import dotenv  # for loading config from .env files, https://pypi.org/project/python-dotenv/
import envyaml  # Allows to loads env vars into a yaml file, https://github.com/thesimj/envyaml
import fastcore  # To add functionality related to nbdev development, https://github.com/fastai/fastcore/
from fastcore import (
    test,
)
from fastcore.script import (
    call_parse,
)  # for @call_parse, https://fastcore.fast.ai/script
import json  # for nicely printing json and yaml

# import functions from core module (optional, but most likely needed).
from . import core

# Project specific libraries
from pathlib import Path
import copy
from math import floor

# %% ../nbs/02_seq_tools.ipynb 7
class BioSeq:
    def __init__(self, name: str, sequence: str) -> None:
        self.name = name
        self.sequence = sequence.upper()

    def __len__(self):
        return len(self.sequence)

    def __iter__(self):
        for nt in self.sequence:
            yield nt

    def __repr__(self):
        return f"< Biological sequence object {self.name} of length {len(self)} bp >"

    def __add__(self, other):
        return BioSeq(self.name, self.sequence + other.sequence)

    @classmethod
    def from_fasta(cls, input_file: Path):
        startflag = True
        with open(input_file) as f:
            for line in f:
                line = line.rstrip("\n")
                if line[0] == ">":
                    if startflag:
                        seq = ""
                        startflag = False
                    else:
                        print(
                            f"Fasta file {input_file} contains more than one entry. Only first entry is loaded into DnaSeq object.\nTo load all entries into Fasta object use Fasta.from_file({input_file})\n",
                            file=sys.stderr,
                        )
                        f.close()
                        return cls(sequence_name, seq)
                    sequence_name = line[1:]

                else:
                    seq += line
        f.close()
        return cls(sequence_name, seq)

    def print_fasta(self, file: Path, linelength: int = None) -> None:
        printlines = []
        if linelength is None:
            printlines += [">" + self.name, self.sequence]
        else:
            printlines += [">" + self.name] + [
                (self.sequence[i : i + linelength])
                for i in range(0, len(self.sequence), linelength)
            ]

        o = open(file, "w")
        o.write("\n".join(printlines) + "\n")
        o.close
        return None


class DnaSeq(BioSeq):
    def __init__(self, name: str, sequence: str, phred_scores: str = None) -> None:
        self.name = name
        self.sequence = sequence.upper()
        self.phred_scores = phred_scores

    def __repr__(self):
        return f"< DNA sequence object {self.name} of length {len(self)} bp >"

    def reverse_complement(self):
        complement = {
            "A": "T",
            "T": "A",
            "C": "G",
            "G": "C",
            "R": "Y",
            "Y": "R",
            "S": "W",
            "W": "S",
            "K": "M",
            "M": "K",
            "B": "V",
            "V": "B",
            "D": "H",
            "H": "D",
            "N": "N",
            "X": "X",
            "-": "-",
            ".": ".",
        }
        reverse_complement = "".join(
            complement.get(base, base) for base in self.sequence[::-1]
        )
        illegal_characters = set(self.sequence) - set(complement.keys())
        if len(illegal_characters) > 0:
            print(
                "Non standard UIPAC characters found in sequence. Replacing with N when reverse complementing.",
                file=sys.stderr,
            )
            for c in illegal_characters:
                reverse_complement = reverse_complement.replace(c, "N")
        if self.phred_scores is None:
            return DnaSeq(self.name, reverse_complement)
        else:
            return DnaSeq(self.name, reverse_complement, self.phred_scores[::-1])

    def translate(self):
        seq = self.sequence.replace("-", "")
        table = {
            "ATA": "I",
            "ATC": "I",
            "ATT": "I",
            "ATG": "M",
            "ACA": "T",
            "ACC": "T",
            "ACG": "T",
            "ACT": "T",
            "AAC": "N",
            "AAT": "N",
            "AAA": "K",
            "AAG": "K",
            "AGC": "S",
            "AGT": "S",
            "AGA": "R",
            "AGG": "R",
            "CTA": "L",
            "CTC": "L",
            "CTG": "L",
            "CTT": "L",
            "CCA": "P",
            "CCC": "P",
            "CCG": "P",
            "CCT": "P",
            "CAC": "H",
            "CAT": "H",
            "CAA": "Q",
            "CAG": "Q",
            "CGA": "R",
            "CGC": "R",
            "CGG": "R",
            "CGT": "R",
            "GTA": "V",
            "GTC": "V",
            "GTG": "V",
            "GTT": "V",
            "GCA": "A",
            "GCC": "A",
            "GCG": "A",
            "GCT": "A",
            "GAC": "D",
            "GAT": "D",
            "GAA": "E",
            "GAG": "E",
            "GGA": "G",
            "GGC": "G",
            "GGG": "G",
            "GGT": "G",
            "TCA": "S",
            "TCC": "S",
            "TCG": "S",
            "TCT": "S",
            "TTC": "F",
            "TTT": "F",
            "TTA": "L",
            "TTG": "L",
            "TAC": "Y",
            "TAT": "Y",
            "TAA": "*",
            "TAG": "*",
            "TGC": "C",
            "TGT": "C",
            "TGA": "*",
            "TGG": "W",
        }
        protein = ""
        mod = len(seq) % 3
        if not mod == 0:
            print(
                f"Warning: number of nucleotides in sequence is not divisible by 3. Ignoring {mod} nt at the end of sequence when translating to protein.",
                file=sys.stderr,
            )
        illegal_codons = 0
        for i in range(0, len(seq) - mod, 3):
            codon = seq[i : i + 3]
            if len(set(codon) - {"A", "C", "G", "T"}) > 0:
                illegal_codons += 1
                protein += "X"
            else:
                protein += table[codon]
        if illegal_codons > 0:
            print(
                f"WARNING: {illegal_codons} codons containing non A/C/G/T found in {self.name}. X has been inserted on those positions.",
                file=sys.stderr,
            )
        return ProteinSeq(self.name, protein)

    def __add__(self, other):
        if self.phred_scores is None or other.phred_scores is None:
            return DnaSeq(self.name, self.sequence + other.sequence)
        else:
            return DnaSeq(
                self.name,
                self.sequence + other.sequence,
                self.phred_scores,
                other.phred_scores,
            )

    @property
    def iupac_counts(self):
        base_types = ["A", "C", "G", "T"]
        amb_types = ["R", "Y", "S", "W", "K", "M", "B", "D", "H", "M"]
        gap_types = [".", "-"]
        nt_counts = {
            "A": 0,
            "C": 0,
            "G": 0,
            "T": 0,
            "N": 0,
            "R": 0,
            "Y": 0,
            "S": 0,
            "W": 0,
            "K": 0,
            "M": 0,
            "B": 0,
            "D": 0,
            "H": 0,
            "M": 0,
            ".": 0,
            "-": 0,
            "ACGT": 0,
            "amb": 0,
            "N": 0,
            "gap": 0,
            "nonIUPAC": 0,
        }
        for nt in self.sequence:
            try:
                nt_counts[nt] += 1
            except KeyError:
                nt_counts["nonIUPAC"] += 1
            if nt in base_types:
                nt_counts["ACGT"] += 1
            elif nt in amb_types:
                nt_counts["amb"] += 1
            elif nt in gap_types:
                nt_counts["gap"] += 1

        return nt_counts

    @staticmethod
    def phred_score_to_int(phred_scores: str) -> list:
        int_scores = []
        for score in phred_scores:
            int_scores.append(ord(score) - 33)
        return int_scores


class ProteinSeq(BioSeq):
    def __init__(self, sequence_name: str, sequence: str) -> None:
        self.name = sequence_name
        self.sequence = sequence.upper()

    def __repr__(self):
        return f"< Protein sequence object {self.name} of length {len(self)} AA >"

    def __add__(self, other):
        return ProteinSeq(self.name, self.sequence + other.sequence)

    @classmethod
    def from_fasta(cls, input_file: Path):
        startflag = True
        with open(input_file) as f:
            for line in f:
                line = line.rstrip("\n")
                if line[0] == ">":
                    if startflag:
                        seq = ""
                        startflag = False
                    else:
                        f.close()
                        print(
                            f"Fasta file {input_file} contains more than one entry. Only first entry is loaded into DnaSeq object. To load all entries use Fasta.from_file()\n",
                            file=sys.stderr,
                        )
                        return cls(sequence_name, seq)
                    sequence_name = line[1:]

                else:
                    seq += line
        f.close()
        return cls(sequence_name, seq)

    def print_fasta(self, output_file: str, linelength: int = None) -> None:
        printlines = []
        if linelength is None:
            printlines += [">" + self.name, self.sequence]
        else:
            printlines += [">" + self.name] + [
                (self.sequence[i : i + linelength])
                for i in range(0, len(self.sequence), linelength)
            ]

        o = open(output_file, "w")
        o.write("\n".join(printlines) + "\n")
        o.close
        return None


import gzip


class Fastq:

    def __init__(self, dna_seqs: list, file: str = None) -> None:
        self.entries = dna_seqs
        self.file = file

    @classmethod
    def from_file(cls, input_file: str):
        entries = []
        if input_file.endswith(".gz"):
            with gzip.open(input_file, "rt") as f:
                for line_number, line in enumerate(f):
                    lines.append(line.rstrip("\n"))
                    if (line_number + 1) % 4 == 0:
                        entries.append(
                            DnaSeq(
                                sequence_name=line[0],
                                sequence=line[1],
                                phred_scores=line[3],
                            )
                        )
                        lines = []
            f.close()
        else:
            with open(input_file, "r") as f:
                for line_number, line in enumerate(f):
                    lines.append(line.rstrip("\n"))
                    if (line_number + 1) % 4 == 0:
                        entries.append(
                            DnaSeq(
                                sequence_name=line[0],
                                sequence=line[1],
                                phred_scores=line[3],
                            )
                        )
                        lines = []
            f.close()
        return cls(entries, file=input_file)

    def __iter__(self):
        for entry in self.entries:
            yield ((entry.name, entry.sequence, entry.phred_score))

    def __len__(self):
        length = 0
        for name, sequence in self:
            length += len(sequence)
        return length

    def __repr__(self):
        return f"< Fasta object containing {len(self.entries)} sequences, total length {len(self)} bp >"

    def write(self, output_file: str, gzip: bool = True) -> None:
        printlines = []
        for header, read, phred_score in self:
            printlines += ["@" + header, read, "+", phred_score]
        if gzip:
            o = open(output_file, "wb")
            o.write("\n".join(printlines) + "\n")
            o.close
        else:
            o = open(output_file, "w")
            o.write("\n".join(printlines) + "\n")
            o.close
        return None

# %% ../nbs/02_seq_tools.ipynb 11
class Fasta:

    def __init__(self, dna_seqs: list[DnaSeq]) -> None:
        self.entries = dna_seqs

        ### Indexing for entries to allow for easy retrieval by sequence name, i.e. Fasta["sequence_name"] = [DnaSeq object]
        self.name_index = {}
        for i, dna_seq in enumerate(self.entries):
            if dna_seq.name in self.name_index:
                self.name_index[dna_seq.name].append(i)
            else:
                self.name_index[dna_seq.name] = [i]

    # Alternative constructor that initiates from a fasta file
    @classmethod
    def from_file(cls, input_file: str):
        entries = "Start"
        skipped_lines_count = 0
        with open(input_file) as f:
            for line in f:
                line = line.rstrip("\n")
                if not len(line) == 0:
                    if line[0] == ">":
                        if entries == "Start":
                            entries = []
                        else:
                            entries.append(DnaSeq(sequence_name, new_seq))
                        new_seq = ""
                        sequence_name = line[1:]

                    else:
                        new_seq += line
                else:
                    skipped_lines_count += 1
        entries.append(DnaSeq(sequence_name, new_seq))
        f.close()
        print(
            f"Warning: {skipped_lines_count} blank lines were skipped when loading fasta file"
        )
        return cls(entries)

    def check_sequence_type(self):

        return

    def __iter__(self) -> tuple[DnaSeq]:
        for entry in self.entries:
            yield entry

    def items(self) -> tuple[str, str]:
        for entry in self.entries:
            yield entry.name, entry.sequence

    def iterseqs(self) -> str:
        for entry in self.entries:
            yield entry.sequence

    def __getitem__(self, name) -> DnaSeq:
        return [self.entries[i] for i in self.name_index[name]]

    # len returns total number of nucleotides in fasta file
    def __len__(self):
        length = 0
        for name, sequence in self.items():
            length += len(sequence)
        return length

    def __repr__(self):
        return f"< Fasta object containing {len(self.entries)} sequences, total length {len(self)} bp >"

    def write(self, output_file: str = None, linelength: int = None) -> None:
        printlines = []
        if linelength is None:
            for name, sequence in self.items():
                printlines += [">" + name, sequence]
        else:
            for name, sequence in self.items():
                printlines += [">" + name] + [
                    (sequence[i : i + linelength])
                    for i in range(0, len(sequence), linelength)
                ]
        if output_file:
            o = open(output_file, "w")
            o.write("\n".join(printlines) + "\n")
            o.close
        else:
            for line in printlines:
                print(line)
        return None

    # Append a new DnaSeq object to exisisting Fasta object
    def append(self, entry: DnaSeq) -> None:
        self.entries.append(entry)
        if entry.name in self.name_index:
            self.name_index[entry.name].append(len(self.entries) - 1)
        else:
            self.name_index[entry.name] = [len(self.entries) - 1]

    ### Add all entries from another Fasta object to entries in self
    def add(self, other: DnaSeq):
        for entry in other:
            self.append(entry)
        return None

    def concat_seq(self):
        combined_seq = ""
        for name, sequence in self.items():
            combined_seq += sequence
        return combined_seq

    @property
    def sequence_names(self):
        seq_names = []
        for entry in self.entries:
            seq_names.append(entry.name)
        return seq_names

    @property
    def sequences(self) -> BioSeq:
        seqs = []
        for entry in self.entries:
            seqs.append(entry.sequence)
        return seqs

    def uniquify(self) -> BioSeq:
        seq_to_names = {}
        for name, sequence in self.items():
            if sequence in seq_to_names:
                seq_to_names[sequence].append(name)
            else:
                seq_to_names[sequence] = [name]
        ordered_seqs = [
            k for k, v in sorted(seq_to_names.items(), key=lambda item: len(item[1]))
        ]
        count_str_length = len(str(len(ordered_seqs)))
        seq_count = 0
        seq_to_unique_name = {}
        unique_fasta = Fasta([])  ### initialize empty Fasta object
        for sequence in ordered_seqs:
            seq_count += 1
            seq_count_str = "0" * (count_str_length - len(str(seq_count))) + str(
                seq_count
            )
            seq_unique_name = f"unique_sequence_{seq_count_str}"
            unique_fasta.append(BioSeq(name=seq_unique_name, sequence=sequence))
            seq_to_unique_name[sequence] = seq_unique_name

        seq_info = []
        for name, sequence in self.items():
            seq_info.append(
                (
                    name,
                    seq_to_unique_name[sequence],
                    len(seq_to_names[sequence]),
                    sequence,
                )
            )

        unique_fasta.uniquify_info = seq_info
        return unique_fasta


class NucleotideFasta(Fasta):
    def __init__(self, dna_seqs: list[DnaSeq]) -> None:
        self.entries = dna_seqs

        ### Indexing for entries to allow for easy retrieval by sequence name, i.e. Fasta["sequence_name"] = [DnaSeq object]
        self.name_index = {}
        for i, dna_seq in enumerate(self.entries):
            if dna_seq.name in self.name_index:
                self.name_index[dna_seq.name].append(i)
            else:
                self.name_index[dna_seq.name] = [i]

    def uniquify(self) -> DnaSeq:
        seq_to_names = {}
        for name, sequence in self.items():
            if sequence in seq_to_names:
                seq_to_names[sequence].append(name)
            else:
                seq_to_names[sequence] = [name]
        ordered_seqs = [
            k for k, v in sorted(seq_to_names.items(), key=lambda item: len(item[1]))
        ]
        count_str_length = len(str(len(ordered_seqs)))
        seq_count = 0
        seq_to_unique_name = {}
        unique_fasta = Fasta([])  ### initialize empty Fasta object
        for sequence in ordered_seqs:
            seq_count += 1
            seq_count_str = "0" * (count_str_length - len(str(seq_count))) + str(
                seq_count
            )
            seq_unique_name = f"unique_sequence_{seq_count_str}"
            unique_fasta.append(DnaSeq(name=seq_unique_name, sequence=sequence))
            seq_to_unique_name[sequence] = seq_unique_name

        seq_info = []
        for name, sequence in self.items():
            seq_info.append(
                (
                    name,
                    seq_to_unique_name[sequence],
                    len(seq_to_names[sequence]),
                    sequence,
                )
            )

        unique_fasta.uniquify_info = seq_info
        return unique_fasta

    def translate(self):
        protein_fasta = ProteinFasta([])
        for entry in self:
            protein_fasta.append(entry.translate())
        return protein_fasta

    @property
    def iupac_counts(self):
        combined_seq = self.concat_seq()
        base_types = ["A", "C", "G", "T"]
        amb_types = ["R", "Y", "S", "W", "K", "M", "B", "D", "H", "M"]
        gap_types = [".", "-"]
        nt_counts = {
            "A": 0,
            "C": 0,
            "G": 0,
            "T": 0,
            "N": 0,
            "R": 0,
            "Y": 0,
            "S": 0,
            "W": 0,
            "K": 0,
            "M": 0,
            "B": 0,
            "D": 0,
            "H": 0,
            "M": 0,
            ".": 0,
            "-": 0,
            "ACGT": 0,
            "amb": 0,
            "N": 0,
            "gap": 0,
            "nonIUPAC": 0,
        }
        for nt in combined_seq:
            try:
                nt_counts[nt] += 1
            except KeyError:
                nt_counts["nonIUPAC"] += 1
            if nt in base_types:
                nt_counts["ACGT"] += 1
            elif nt in amb_types:
                nt_counts["amb"] += 1
            elif nt in gap_types:
                nt_counts["gap"] += 1

        return nt_counts

    def batch_iupac_counts(input_directory):
        input_directory = Path(input_directory)
        iupac_counts_dict = {}
        for p in input_directory.iterdir():
            if p.name.endswith((".fna", ".fa", ".fasta")):
                fasta = NucleotideFasta.from_file(p)
                iupac_counts_dict[p.name] = fasta.iupac_counts
        return iupac_counts_dict


class ProteinFasta(Fasta):

    def __init__(self, protein_seqs: list[ProteinSeq]) -> None:
        self.entries = protein_seqs

        ### Indexing for entries to allow for easy retrieval by sequence name, i.e. ProteinFasta["sequence_name"] = ProteinSeq object
        self.name_index = {}
        for i, dna_seq in enumerate(self.entries):
            self.name_index[dna_seq.name] = i

    # Alternative constructor that initiates from a fasta file
    @classmethod
    def from_file(cls, input_file: str):
        entries = "Start"
        with open(input_file) as f:
            for line in f:
                line = line.rstrip("\n")
                if line[0] == ">":
                    if entries == "Start":
                        entries = []
                    else:
                        entries.append(ProteinSeq(sequence_name, new_seq))
                    new_seq = ""
                    sequence_name = line[1:]

                else:
                    new_seq += line
        entries.append(ProteinSeq(sequence_name, new_seq))
        f.close()
        return cls(entries)

    def uniquify(self) -> ProteinSeq:
        seq_to_names = {}
        for name, sequence in self.items():
            if sequence in seq_to_names:
                seq_to_names[sequence].append(name)
            else:
                seq_to_names[sequence] = [name]
        ordered_seqs = [
            k for k, v in sorted(seq_to_names.items(), key=lambda item: len(item[1]))
        ]
        count_str_length = len(str(len(ordered_seqs)))
        seq_count = 0
        seq_to_unique_name = {}
        unique_fasta = ProteinFasta([])  ### initialize empty Fasta object
        for sequence in ordered_seqs:
            seq_count += 1
            seq_count_str = "0" * (count_str_length - len(str(seq_count))) + str(
                seq_count
            )
            seq_unique_name = f"unique_sequence_{seq_count_str}"
            unique_fasta.append(ProteinSeq(name=seq_unique_name, sequence=sequence))
            seq_to_unique_name[sequence] = seq_unique_name

        seq_info = []
        for name, sequence in self.items():
            seq_info.append(
                (
                    name,
                    seq_to_unique_name[sequence],
                    len(seq_to_names[sequence]),
                    sequence,
                )
            )

        unique_fasta.uniquify_info = seq_info
        return unique_fasta

    def __getitem__(self, name):
        return self.entries[self.name_index[name]]


class Assembly(NucleotideFasta):

    def assembly_stats(self, detailed=False):
        """
        Generate basic assembly stats for genome assembly
        Returns dictionary with L50, N50, L90, N90, AuN, AuN_percent and full list of Lx and Nx values
        """

        contig_lengths = []
        for seq in self.entries:
            contig_lengths.append(len(seq))
        total_length = sum(contig_lengths)
        contig_lengths.sort(reverse=True)
        nX_list = [0] * 100
        lX_list = [0] * 100
        previous_percent = 0
        AuN_normalized = 0
        for i, contig_length in enumerate(contig_lengths):
            cumsum_percent = floor(sum(contig_lengths[: (i + 1)]) / total_length * 100)
            for x in range(previous_percent, cumsum_percent):
                nX_list[x] = contig_length
                lX_list[x] = i + 1
                AuN_normalized += cumsum_percent
            previous_percent = cumsum_percent
        nX_list[0] = nX_list[1]
        l50 = lX_list[49]
        n50 = nX_list[49]
        l90 = lX_list[89]
        n90 = nX_list[89]
        if detailed:
            assembly_stats_dict = {
                "L50": l50,
                "N50": n50,
                "L90": l90,
                "N90": n90,
                "AuN": sum(nX_list) / 100,
                "AuN_percent": AuN_normalized / 100,
                "Lx_values": lX_list,
                "Nx_values": nX_list,
            }
        else:
            assembly_stats_dict = {
                "L50": l50,
                "N50": n50,
                "L90": l90,
                "N90": n90,
                "AuN": sum(nX_list) / 100,
                "AuN_percent": AuN_normalized / 100,
            }

        return assembly_stats_dict


class SpadesAssembly(Assembly):

    def filter_contigs(
        self,
        length_threshold: int = 200,
        coverage_threshold: int = 10,
        name_prefix: str = None,
    ):
        """
        Create new SpadesAssembly object where contigs less than length_threshold and kmer coverage less than coverage_threshold have been removed
        Length is based on length of sequence while kmer coverage is read from contig header (assumes Spades standard naming)
        if name_prefix is provided, this name will be added to the start of each contig. F.ex. name_prefix = "Ecoli-1013__"
        """
        passed_entries = []
        for name, sequence in self.items():
            contig_len = len(sequence)
            contig_cov = float(name.split("_")[-1])
            if contig_len >= length_threshold and contig_cov >= coverage_threshold:
                if name_prefix is None:
                    passed_entries.append(DnaSeq(name, sequence))
                else:
                    passed_entries.append(DnaSeq(name_prefix + name, sequence))
        trimmed_spades = SpadesAssembly(passed_entries)
        print(
            f"Removed contigs from spades assembly with contig length < {length_threshold} or kmer coverage < {coverage_threshold}\nInput sequence count: {len(self.entries)}, total length: {len(self)} bp\nSequence count after trimming: {len(trimmed_spades.entries)}, total length: {len(trimmed_spades)} bp",
            file=sys.stderr,
        )
        return trimmed_spades

    @staticmethod
    def batch_filter_contigs(
        input_directory: Path,
        output_directory: Path,
        length_threshold: int = 200,
        coverage_threshold: int = 10,
        linelength=None,
    ):
        """
        Static method for filtering outputs given a folder containing spades output folders from multiple isolates.
        """
        input_directory = Path(input_directory)
        output_directory = Path(output_directory)
        output_directory.mkdir(parents=True, exist_ok=True)
        stats_list = []
        for p in input_directory.iterdir():
            contigs_path = p.joinpath("contigs.fasta")
            if contigs_path.exists():
                sample_name = p.name
                output_fasta = output_directory.joinpath(f"{sample_name}.fasta")
                assembly = SpadesAssembly.from_file(contigs_path)
                assembly_filtered = assembly.filter_contigs(
                    length_threshold=length_threshold,
                    coverage_threshold=coverage_threshold,
                    name_prefix=f"{sample_name}__",
                )
                stats_list.append(
                    (
                        sample_name,
                        len(assembly_filtered.entries),
                        len(assembly.entries),
                        len(assembly.entries) - len(assembly_filtered.entries),
                        len(assembly_filtered),
                        len(assembly),
                        len(assembly) - len(assembly_filtered),
                    )
                )
                assembly_filtered.write(output_file=output_fasta, linelength=linelength)

        df = pandas.DataFrame(
            stats_list,
            columns=[
                "sample_name",
                "contigs_after_filtering",
                "contigs_before_filtering",
                "contigs_filtered",
                "length_after_filtering",
                "length_before_filtering",
                "length_filtered",
            ],
        )
        df.to_csv(output_directory.joinpath("filter_stats.tsv"), sep="\t", index=False)
        return None

# %% ../nbs/02_seq_tools.ipynb 16
from fastcore.script import call_parse


@call_parse
def translate(
    input: Path,  # Path to nucleotide fasta file
    output: Path = None,  # Path to protein fasta file
    line_length: int = None,  # Number of characters per sequence line in fasta output
) -> None:
    f"""
    Translate all sequences in a nucleotide fasta file. 
    Output will be printed to stdout if no --output is provided.
    """
    input_nt_fasta = NucleotideFasta.from_file(input)
    output_prot_fasta = input_nt_fasta.translate()
    output_prot_fasta.write(output_file=output, linelength=line_length)


@call_parse
def uniquify(
    input_fasta: Path,  # Path to nucleotide or protein fasta file
    output_fasta: Path = None,  # Path to output fasta file containing only unqiue sequences from input file
    output_table: Path = None,  # Path to output table containing mapping of input sequence names to unique sequence names (tsv-format)
    line_length: int = None,  # Number of characters per sequence line in fasta output
) -> None:
    """
    Identify all unique sequences in fasta file and print them in fasta format.
    Fasta output will be printed to stdout if no --output_fasta is provided.

    """

    input_fasta = Fasta.from_file(input_file=input_fasta)
    uniq_fasta = input_fasta.uniquify()
    uniq_fasta.write(output_file=output_fasta, linelength=line_length)
    if output_table:
        df = pandas.DataFrame(
            uniq_fasta.uniquify_info,
            columns=["name", "unique_name", "unique_count", "sequence"],
        )
        df.to_csv(output_table, sep="\t", index=False)


@call_parse
def filter_spades_batch(
    input_folder: Path,  # Path to folder containing spades output folders
    output_folder: Path = Path(
        "."
    ),  # Path to folder to print filtered assemblies and stats to.
    length_threshold: int = 200,  # length threshold that contigs must pass to be included after filtering
    coverage_trehold: int = 10,  # kmer coverage threshold that contigs must pass to be included after filtering
    line_length: int = None,  # Number of characters per sequence line in fasta output
) -> None:
    """
    Provided a folder containing spades output folders from individual samples, filter contigs.fasta from each sample based on provided thresholds and print filtered fasta to output folder. Filename will also be changed to <sample_name>.fasta and <sample_name>__ will be added to the start of each contig header line

    """

    SpadesAssembly.batch_filter_contigs(
        input_directory=input_folder,
        output_directory=output_folder,
        length_threshold=length_threshold,
        coverage_threshold=coverage_trehold,
        linelength=line_length,
    )
