# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['PACKAGE_NAME', 'DEV_MODE', 'PACKAGE_DIR', 'PROJECT_DIR', 'config', 'set_env_variables', 'get_config',
           'show_project_env_vars', 'get_samplesheet', 'DnaSeq', 'ProteinSeq', 'Fasta', 'ProteinFasta',
           'SpadesAssembly', 'Fastq', 'NcbiDataSummarizer']

# %% ../nbs/00_core.ipynb 4
# Need the ncbi_datasets_util for a few functions, this can be considered a static var

import importlib
import importlib.util
import os

PACKAGE_NAME: str = (
    "ncbi_datasets_util"  # Make sure to adjust this to your package name
)
DEV_MODE: bool = (
    False  # set below to override, as this is in an export block it'll be exported while the dev mode section is not
)

PACKAGE_DIR = None
try:
    spec = importlib.util.find_spec(PACKAGE_NAME)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    PACKAGE_DIR = os.path.dirname(module.__file__)
except ImportError:
    DEV_MODE = True
except AttributeError:
    DEV_MODE = True
PROJECT_DIR = os.getcwd()  # override value in dev mode
if PROJECT_DIR.endswith("nbs"):
    DEV_MODE = True
    PROJECT_DIR = os.path.split(PROJECT_DIR)[0]

# %% ../nbs/00_core.ipynb 10
# standard libs
import os
import re

# 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/
import pandas  # For sample sheet manipulation
from fastcore import (
    test,
)
from fastcore.script import (
    call_parse,
)  # for @call_parse, https://fastcore.fast.ai/script

# Project specific libraries

# %% ../nbs/00_core.ipynb 13
import importlib
import importlib.util


def set_env_variables(config_path: str, overide_env_vars: bool = True) -> bool:
    # Load dot env sets environmental values from a file, if the value already exists it will not be overwritten unless override is set to True.
    # If we have multiple .env files then we need to apply the one which we want to take precedence last with overide.

    # Order of precedence: .env file > environment variables > default values
    # When developing, making a change to the config will not be reflected until the environment is restarted

    # Set the env vars first, this is needed for the card.yaml to replace ENV variables
    # NOTE: You need to adjust PROJECT_NAME to your package name for this to work, the exception is only for dev purposes
    # This here checks if your package is installed, such as through pypi or through pip install -e  [.dev] for development. If it is then it'll go there and use the config files there as your default values.
    try:
        dotenv.load_dotenv(f"{PACKAGE_DIR}/config/config.default.env", override=False)
    except Exception as e:
        print(f"Error: {PACKAGE_DIR}/config/config.default.env does not exist")
        return False

    # 2. set values from file:
    if os.path.isfile(config_path):
        dotenv.load_dotenv(config_path, override=overide_env_vars)

    return True

# %% ../nbs/00_core.ipynb 15
import importlib
import importlib.util


def get_config(config_path: str = None, overide_env_vars: bool = True) -> dict:
    if config_path is None:
        config_path = ""
    # First sets environment with variables from config_path, then uses those variables to fill in appropriate values in the config.yaml file, the yaml file is then returned as a dict
    # If you want user env variables to take precedence over the config.yaml file then set overide_env_vars to False
    set_env_variables(config_path, overide_env_vars)

    config: dict = envyaml.EnvYAML(
        os.environ.get(
            "CORE_YAML_CONFIG_FILE", f"{PACKAGE_DIR}/config/config.default.yaml"
        ),
        strict=False,
    ).export()

    return config

# %% ../nbs/00_core.ipynb 17
# create a os.PathLike object
config = get_config(os.environ.get("CORE_CONFIG_FILE", ""))

# %% ../nbs/00_core.ipynb 19
def show_project_env_vars(config: dict) -> None:
    # Prints out all the project environment variables
    # This is useful for debugging and seeing what is being set
    for k, v in config.items():
        # If ENV var starts with PROJECTNAME_ then print
        if k.startswith(config["CORE_PROJECT_VARIABLE_PREFIX"]):
            print(f"{k}={v}")

# %% ../nbs/00_core.ipynb 22
import pandas as pd


def get_samplesheet(sample_sheet_config: dict) -> pd.DataFrame:
    # Load the sample sheet into a pandas dataframe
    # If columns is not None then it will only load those columns
    # If the sample sheet is a csv then it will load it as a csv, otherwise it will assume it's a tsv

    # Expected sample_sheet_config:
    # sample_sheet:
    #   path: path/to/sample_sheet.tsv
    #   delimiter: '\t' # Optional, will assume , for csv and \t otherwises
    #   header: 0 # Optional, 0 indicates first row is header, None indicates no header
    #   columns: ['column1', 'column2', 'column3'] # Optional, if not provided all columns will be used

    # Example sample sheet:
    # #sample_id	file_path	metadata1	metadata2
    # Sample1	/path/to/sample1.fasta	value1	option1
    # Sample2	/path/to/sample2.fasta	value2	option2
    # Sample3	/path/to/sample3.fasta	value3	option1
    # Sample4	/path/to/sample4.fasta	value1	option2
    # Sample5	/path/to/sample5.fasta	value2	option1

    # This function should also handle ensuring the sample sheet is in the correct format, such as ensuring the columns are correct and that the sample names are unique.
    if not os.path.isfile(sample_sheet_config["path"]):
        raise FileNotFoundError(f"File {sample_sheet_config['path']} does not exist")
    if "delimiter" in sample_sheet_config:
        delimiter = sample_sheet_config["delimiter"]
    else:
        # do a best guess based on file extension
        delimiter = "," if sample_sheet_config["path"].endswith(".csv") else "\t"
    header = 0
    # if "header" in sample_sheet_config:
    #     header = sample_sheet_config["header"]
    # else:
    #     # check if the first line starts with a #, if so lets assume it's a header otherwise assume there isn't one
    #     with open(sample_sheet_config["path"], "r") as f:
    #         first_line = f.readline()
    #         header = 0 if first_line.startswith("#") else None
    if "columns" in sample_sheet_config:
        columns = sample_sheet_config[
            "columns"
        ]  # note the # for the first item needs to be stripped to compare to the columns
    else:
        columns = None  # implies all columns
    try:
        # note when we have a header the first column may begin with a #, so we need to remove it
        df = pd.read_csv(
            sample_sheet_config["path"],
            delimiter=delimiter,
            header=header,
            comment=None,
        )
    except Exception as e:
        print(
            "Error: Could not load sample sheet into dataframe, you have a problem with your sample sheet or the configuration."
        )
        raise e

    # Check the first header has a # in it, if so remove it for only that item
    if df.columns[0].startswith("#"):
        df.columns = [col.lstrip("#") for col in df.columns]
    # Ensure the sample sheet has the correct columns
    if columns is not None and not all([col in df.columns for col in columns]):
        raise ValueError("Error: Sample sheet does not have the correct columns")
    # also drop columns which are not needed
    if columns is not None:
        df = df[columns]

    # Clean the df of any extra rows that can be caused by empty lines in the sample sheet
    df = df.dropna(how="all")
    return df

# %% ../nbs/00_core.ipynb 23
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,
            )

    def print_fasta(self, 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(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", "N"]
        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)

    def print_to_file(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


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"]
        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


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:

    @staticmethod
    def get_files_from_isolate_folder(isolate_data_folder):
        file_paths = {
            "accession": os.path.basename(isolate_data_folder),
            "genome": None,
            "cds": None,
            "gbff": None,
            "gff3": None,
            "protein": None,
            "seq-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["gbff"] = abs_path
            elif file == "genomic.gff":
                file_paths["gff3"] = abs_path
            elif file == "protein.faa":
                file_paths["protein"] = abs_path
            elif file == "cds_from_genomic.fna":
                file_paths["cds"] = abs_path
            elif file == "sequence_report.jsonl":
                file_paths["seq-report"] = abs_path
            elif file.endswith("_genomic.fna"):
                file_paths["genome"] = abs_path
        return file_paths

    @staticmethod
    def get_files_from_data_folder(data_folder):
        isolate_file_path_dict = {}
        isolate_names = os.listdir(data_folder)
        for isolate in isolate_names:
            isolate_data_folder = os.path.join(data_folder, isolate)
            if os.path.isdir(isolate_data_folder):
                isolate_file_path_dict[isolate] = (
                    NcbiDataSummarizer.get_files_from_isolate_folder(
                        isolate_data_folder
                    )
                )
        return isolate_file_path_dict

    @staticmethod
    def summarize_fasta(fasta_file):
        fasta = core.Fasta.from_file(fasta_file)
        summary_dict = {"length": len(fasta), "contig_count": len(fasta.entries)}
        summary_dict.update(fasta.iupac_counts)
        # for entry in fasta.entries:
        #    print(entry.iupac_counts)

        return summary_dict

    @staticmethod
    def summarize_data_folder(data_folder):
        file_paths = NcbiDataSummarizer.get_files_from_data_folder(data_folder)
        for isolate, files in file_paths.items():
            file_paths[isolate].update(
                NcbiDataSummarizer.summarize_fasta(files["genome"])
            )
        return file_paths
