# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/07_bracken.ipynb.

# %% auto 0
__all__ = ['process_bracken_data', 'process_bracken_data_from_cli']

# %% ../nbs/07_bracken.ipynb 2
# That export there, it makes sure this code goes into the module.

# 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
from fastcore import test
import pandas as pd

#!export
from . import core

# %% ../nbs/07_bracken.ipynb 5
def process_bracken_data(
    input_path: str,
    output_path: str = "./output.tsv",
    add_header: str = "%ofreads, reads, notsure, rank, taxid, name",
    replace_header: str = None,
    filter_columns: str = None,
    transpose: bool = True,
):
    """
    Process Bracken data.

    This function processes Bracken data files by importing the data, optionally replacing the header,
    transposing the data, filtering columns, and exporting the processed data to an output file.

    Arguments:
        input_path (str): Path to the input file.
        output_path (str): Path to the output file (default: './output.tsv').
        replace_header (str): Header to replace the existing header (default: None).
        filter_columns (str): Columns to filter from the header (default: None).
        transpose (bool): Whether to transpose the data (default: True).
    """

    df = core.DataFrame()

    if not os.path.exists(input_path):
        raise FileNotFoundError(f"The input file {input_path} does not exist.")
    df.import_data(input_path, file_type="tsv", add_header=add_header)

    df.df = pd.concat(
        [
            df.df[df.df["rank"] == "U"],
            df.df[df.df["rank"] == "S"].sort_values(by="%ofreads", ascending=False),
        ]
    )
    df.df = df.df.head(3)

    df_output = core.DataFrame()
    df_output.df = pd.DataFrame(
        {
            "species1_unclassified_name": [df.df.iloc[1]["name"] + " + unclassified"],
            "species1_unclassified_pct": [
                df.df.iloc[1]["%ofreads"] + df.df.iloc[0]["%ofreads"]
            ],
            "species1_name": [df.df.iloc[1]["name"]],
            "species1_pct": [df.df.iloc[1]["%ofreads"]],
            "species2_name": [df.df.iloc[2]["name"]],
            "species2_pct": [df.df.iloc[2]["%ofreads"]],
            "unclassified_name": [df.df.iloc[0]["name"]],
            "unclassified_pct": [df.df.iloc[0]["%ofreads"]],
        }
    )

    if replace_header:
        df_output.rename_header(replace_header)

    if filter_columns:
        df_output.filter_columns(filter_columns)

    df_output.export_data(output_path, file_type="tsv")


@call_parse
def process_bracken_data_from_cli(
    input_path: str,
    output_path: str = "./output.tsv",
    add_header: str = "%ofreads, reads, notsure, rank, taxid, name",
    replace_header: str = None,
    filter_columns: str = None,
):
    process_bracken_data(input_path, output_path, replace_header, filter_columns)
