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

# %% auto 0
__all__ = ['DnaSeq', 'ProteinSeq', 'Fastq', 'NcbiDataSummarizer', 'Fasta', 'ProteinFasta', 'SpadesAssembly']

# %% ../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

# %% ../nbs/02_seq_tools.ipynb 7
class DnaSeq:
    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 __len__(self):
        return len(self.sequence)

    def __iter__(self):
        if self.phred_scores is None:
            for nt in self.sequence:
                yield nt
        else:
            for i in range(len(self)):
                yield (self.sequence[i], self.phred_scores[i])

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

    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,
            )

    @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. To load all entries use Fasta.from_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

    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)

    @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 = {"ACGT": 0, "amb": 0, "N": 0, "gap": 0, "nonIUPAC": 0}
        for nt in base_types:
            count = self.sequence.count(nt)
            nt_counts[nt] = count
            nt_counts["ACGT"] += count
        for nt in amb_types:
            count = self.sequence.count(nt)
            nt_counts[nt] = count
            nt_counts["amb"] += count
        nt_counts["N"] = self.sequence.count("N")
        for nt in gap_types:
            count = self.sequence.count(nt)
            nt_counts[nt] = count
            nt_counts["gap"] += count
        nonIUPAC_characters = set(self.sequence) - set(
            base_types + amb_types + gap_types + ["N"]
        )
        for c in nonIUPAC_characters:
            nt_counts["nonIUPAC"] += self.sequence.count(c)

        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:
    def __init__(self, sequence_name: str, sequence: str) -> None:
        self.name = sequence_name
        self.sequence = sequence.upper()

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

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

    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


class NcbiDataSummarizer:

    def get_files(isolate_data_folder):
        file_paths = {
            "genomic_fna": None,
            "cds_genomic_fna": None,
            "genomic_gbff": None,
            "genomic_gff": None,
            "protein_faa": None,
            "sequence_report": None,
        }
        files = os.listdir(isolate_data_folder)
        for file in files:
            abs_path = os.path.abspath(os.path.join(isolate_data_folder, file))
            if file == "genomic.gbff":
                file_paths["genomic_gbff"] = abs_path
            elif file == "genomic.gff":
                file_paths["genomic_gff"] = abs_path
            elif file == "protein.faa":
                file_paths["protein_faa"] = abs_path
            elif file == "cds_from_genomic.fna":
                file_paths["cds_genomic_fna"] = abs_path
            elif file == "sequence_report.jsonl":
                file_paths["sequence_report"] = abs_path
            elif file.endswith("_genomic.fna"):
                file_paths["genomic_fna"] = abs_path
        return file_paths

    def summarize_genomic_fna(genomic_fna_file):
        fasta = Fasta.from_file(genomic_fna_file)
        for entry in fasta.entries:
            print(entry.iupac_counts)

        return fasta.iupac_counts

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

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

    # 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(DnaSeq(sequence_name, new_seq))
                    new_seq = ""
                    sequence_name = line[1:]

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

    # Iterator returns tuple containing header and sequence of each entry
    # for header, sequence in self:
    #     do something
    def __iter__(self):
        for entry in self.entries:
            yield ((entry.name, entry.sequence))

    # len returns total number of nucleotides in fasta file
    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.sequences)} sequences, total length {len(self)} bp >"

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

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

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

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

        return nt_counts


class ProteinFasta:

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

    # 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)


class SpadesAssembly(Fasta):

    # 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 add_name is provided, this name will be added to the start of each contig. F.ex. add_name = "Ecoli-1013__"
    def filter_contigs(
        self,
        length_threshold: int = 200,
        coverage_threshold: int = 10,
        add_name: str = None,
    ):
        passed_entries = []
        for name, sequence in self:
            contig_len = len(sequence)
            contig_cov = float(name.split("_")[-1])
            if contig_len >= length_threshold and contig_cov >= coverage_threshold:
                if add_name is None:
                    passed_entries.append(DnaSeq(name, sequence))
                else:
                    passed_entries.append(DnaSeq(add_name + 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

    def contig_stats(
        self,
        length_threshold: int = 200,
        coverage_threshold: int = 10,
        add_name: str = None,
    ):
        passed_entries = []
        for name, sequence in self:
            contig_len = len(sequence)
            contig_cov = float(name.split("_")[-1])
            if contig_len >= length_threshold and contig_cov >= coverage_threshold:
                if add_name is None:
                    passed_entries.append(DnaSeq(name, sequence))
                else:
                    passed_entries.append(DnaSeq(add_name + 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
