# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/09_rmlst.ipynb.

# %% auto 0
__all__ = ['process_rmlst_data', 'process_rmlst_data_from_cli']

# %% ../nbs/09_rmlst.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
from . import core
import pandas as pd

# %% ../nbs/09_rmlst.ipynb 6
def process_rmlst_data(
    input_path: str,
    output_path: str = "./output.tsv",
    parsed: bool = False,
    replace_header: str = None,
    filter_columns: str = None,
    add_header: str = None,
):
    """
    Command-line interface for processing rmlst data.

    This function sets up an argument parser to handle command-line arguments for processing rmlst data files.
    It supports specifying input and output file paths, replacing headers, filtering columns, and handling the presence or absence of headers in the input file.

    Arguments:
        input_path (str): Path to the input file.
        output_path (str): Path to the output file (default: './output.tsv').
        parsed (bool): Indicates if the input file is already parsed (default: False).
        replace_header (str): Header to replace the existing header (default: None).
        filter_columns (str): Columns to filter from the header (default: None).
        header_exists (int): Indicates if the header exists in the input file (default: 1).
        add_header (str): Header to add if the header does not exist in the input file (default: None).
    """

    df = core.DataFrame()

    if not os.path.exists(input_path):
        raise FileNotFoundError(f"The input file {input_path} does not exist.")

    if parsed:
        df.import_data(input_path, file_type="tsv")

        if filter_columns:
            df.filter_columns(filter_columns)

        if replace_header:
            df.rename_header(replace_header)

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

    else:
        if os.stat(input_path).st_size > 0:
            # df.import_data(input_path, file_type='tsv', add_header=add_header)
            with open(input_path) as rmlst_json:
                rmlst_dict = json.load(rmlst_json)
            rmlst_dict.pop(
                "exact_matches"
            )  # remove the big section with data per contig
            taxon_prediction_df = pd.json_normalize(
                rmlst_dict["taxon_prediction"]
            )  # this dict is a list and may contain several entries

            # print(taxon_prediction_df)
            def concatenate_vector(vector: list, sep=",") -> str:
                return ",".join([str(i) for i in vector])

            taxon_prediction_df = (
                taxon_prediction_df.apply(concatenate_vector, axis=0).to_frame().T
            )  # concatenate the list
            conc_df = taxon_prediction_df
            df.df = conc_df

            if filter_columns:
                df.filter_columns(filter_columns)

            if replace_header:
                df.rename_header(replace_header)

            # df.show()

            df.export_data(output_path, file_type="tsv")
        else:
            empty_df = pd.DataFrame(
                columns=[col.strip() for col in replace_header.split(",")]
            )
            empty_df.to_csv(output_path, index=False)


@call_parse
def process_rmlst_data_from_cli(
    input_path: str,
    output_path: str = "./output.tsv",
    parsed: bool = False,
    replace_header: str = None,
    filter_columns: str = None,
    add_header: str = None,
):
    process_rmlst_data(
        input_path, output_path, replace_header, filter_columns, add_header
    )
