# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_downloader_and_summarizer.ipynb.

# %% auto 0
__all__ = ['NcbiDataSummarizer', 'get_files_from_data_folder', 'cli']

# %% ../nbs/03_downloader_and_summarizer.ipynb 3
# 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/
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 ncbi_datasets_util import (
    core,
    downloader,
    summarizer,
)

# Project specific libraries

import pandas as pd

# %% ../nbs/03_downloader_and_summarizer.ipynb 6
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 = 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


# def summarize_ncbi_dataset(ncbi_data_path):
#    for GCF_id in os.listdir(ncbi_data_path):


def get_files_from_data_folder(isolate_data_folder):
    file_paths = {
        "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

# %% ../nbs/03_downloader_and_summarizer.ipynb 10
from fastcore.script import call_parse


@call_parse
def cli(
    taxon: str = None,  # Taxonomy name
    output_folder: str = None,  # Output folder
    config=None,
) -> None:
    """
    This will print Hello World! with your name
    """
    # config = core.get_config(config_file)  # Set env vars and get config variables
    data_folder = os.path.join(output_folder, "ncbi_dataset", "data")
    stats_file = os.path.join(output_folder, "qc_stats.tsv")
    downloader.download_genomes_from_taxon(taxon, output_folder)
    stats = NcbiDataSummarizer.summarize_data_folder(data_folder)
    stats_df = pd.DataFrame(stats).T
    stats_df.to_csv(stats_file, sep="\t", index=False)
