import logging
import sys
from pathlib import Path

import pandas as pd

log_format = "%(levelname)-8s %(asctime)s   %(message)s"
date_format = "%d/%m %H:%M:%S"
logging.basicConfig(format=log_format, datefmt=date_format, level=logging.INFO)


def csv_to_parquet(project_folder, output_folder="."):
    """
    Given a list of csv files, convert into parquet files
    """
    project_folder = Path(project_folder)
    if output_folder == ".":
        output_folder = project_folder / "data_warehouse"
    else:
        pass

    logging.info(
        f"Grabbing all csv from folder {project_folder} and saving parquets in {output_folder}"
    )
    for i in project_folder.rglob("*.csv"):
        if "docs" in str(i):
            pass
        elif "dbt" in str(i):
            pass
        elif "assets" in str(i):
            pass
        elif "ipynb_checkpoints" in str(i):
            pass
        else:
            category = str(i).split("/")[3]
            output_subfolder = output_folder / category
            df = pd.read_csv(i)
            # df = df.fillna("")
            # df.columns = [c.replace(".", "_") for c in df.columns]
            output_parquet = output_subfolder / f"{i.stem}.parquet"
            logging.info(f"Converting {i} to {output_parquet}")
            output_subfolder.mkdir(parents=True, exist_ok=True)
            df.to_parquet(output_parquet)

    return


if __name__ == "__main__":
    csv_to_parquet(sys.argv[1])

import logging
import sys
from pathlib import Path

log_format = "%(levelname)-8s %(asctime)s   %(message)s"
date_format = "%d/%m %H:%M:%S"
logging.basicConfig(format=log_format, datefmt=date_format, level=logging.DEBUG)


def bigscape_copy(input_index, output_main):
    input = Path(input_index)
    output = Path(output_main)
    output.mkdir(parents=True, exist_ok=True)
    for item in input.parent.glob("*"):
        target = output / item.name
        logging.debug(f"Generating symlink for: {target} --> {item.resolve()}")
        try:
            target.symlink_to(item.resolve(), target_is_directory=True)
        except FileExistsError as e:
            logging.debug(f"Got error:\n{e}")
    return


if __name__ == "__main__":
    bigscape_copy(sys.argv[1], sys.argv[2])

import sys

import pandas as pd


def prep_bigslice(df_gtdb_output, df_bigslice_output):
    df = pd.read_csv(df_gtdb_output)
    # Format to BiG-Slice
    df_bigslice = df.rename(columns={"genome_id": "#Genome Folder"})
    df_bigslice["#Genome Folder"] = df_bigslice["#Genome Folder"].apply(
        lambda x: f"{x}/"
    )
    df_bigslice.to_csv(df_bigslice_output, index=False, sep="\t")
    return None


if __name__ == "__main__":
    prep_bigslice(sys.argv[1], sys.argv[2])

import json
import logging
import sys
from pathlib import Path

import pandas as pd
from alive_progress import alive_bar

log_format = "%(levelname)-8s %(asctime)s   %(message)s"
date_format = "%d/%m %H:%M:%S"
logging.basicConfig(format=log_format, datefmt=date_format, level=logging.DEBUG)


def combine_deeptfactor_prediction(input_json, filter_str="_deeptfactor"):
    container = {}
    logging.info("Reading json files...")
    with alive_bar(len(input_json), title="Merging DeepTFactor prediction:") as bar:
        for item in input_json:
            item = Path(item)
            genome_id = item.stem
            if filter_str in genome_id:
                genome_id = genome_id.replace(filter_str, "")
            logging.debug(f"Processing {genome_id}")
            with open(item, "r") as f:
                data = json.load(f)
                container.update(data)
            bar()
    return container


def write_deeptf_table(input_json, deeptf_table):
    """
    Write df_deeptfactor.csv table in processed data
    """
    # Handle multiple json
    input_json = Path(input_json)
    logging.info(input_json)
    if input_json.is_file() and input_json.suffix == ".json":
        logging.info(f"Getting deepTFactor overview from a single file: {input_json}")
        input_json_files = input_json

    elif input_json.is_file() and input_json.suffix == ".txt":
        logging.info(f"Getting deepTFactor overview from a text file: {input_json}")
        with open(input_json, "r") as file:
            file_content = [i.strip("\n") for i in file.readlines()]
            if len(file_content) == 1:
                # Paths space-separated on a single line
                paths = file_content[0].split()
            else:
                # Paths written on separate lines
                paths = file_content
            input_json_files = [
                Path(path) for path in paths if Path(path).suffix == ".json"
            ]
    else:
        input_json_files = [
            Path(file)
            for file in str(input_json).split()
            if Path(file).suffix == ".json"
        ]
        logging.info(
            f"Getting deepTFactor overview from the given list of {len(input_json_files)} files..."
        )

    df = combine_deeptfactor_prediction(input_json_files)
    df = pd.DataFrame.from_dict(df).T
    df.index.name = "locus_tag"

    logging.debug(f"Writing file to: {deeptf_table}")

    # Save dataframes to csv tables
    df_table = Path(deeptf_table)
    df_table.parent.mkdir(parents=True, exist_ok=True)
    df.to_csv(deeptf_table, index=True)
    logging.info("Job done")
    return None


if __name__ == "__main__":
    write_deeptf_table(sys.argv[1], sys.argv[2])

import sys
from pathlib import Path

import pandas as pd


def deeptfactor_to_json(prediction, outfile, keep_false=True):
    genome_id = Path(prediction).parent.stem
    df = pd.read_csv(prediction, sep="\t")
    df.loc[:, "genome_id"] = genome_id
    # rename columns
    df = df.rename(
        columns={"prediction": "deeptfactor_prediction", "score": "deeptfactor_score"}
    )
    df = df.set_index("sequence_ID", drop=True)
    if keep_false:
        pass
    else:
        df = df[df.deeptfactor_prediction is True]
    df.T.to_json(outfile)
    return


if __name__ == "__main__":
    deeptfactor_to_json(sys.argv[1], sys.argv[2])

import logging
import sys
from pathlib import Path

import pandas as pd

log_format = "%(levelname)-8s %(asctime)s   %(message)s"
date_format = "%d/%m %H:%M:%S"
logging.basicConfig(format=log_format, datefmt=date_format, level=logging.DEBUG)


def extract_ncbi_information(assembly_report_path, outfile):
    """
    Extract NCBI assembly json reports generated by get_assembly_information.py in a given directory into a csv file

    Parameters
    ----------
    1. assembly_report_path : str / path
        Three different input types can be used:
        - A directory containing assembly information json files generated by get_assembly_information.py
        - A single json file generated by get_assembly_information.py
        - A list of files as string separated by spaces (put the string inside '' in bash expression)
        - A text file (.txt) containing the paths of the json files (one per line, or space-separated on a single line)
    2. outfile : str / path
        Location of the output csv file

    Returns
    -------
    1. outfile : .csv file
        A comma-separated table summarizing the NCBI assembly reports metadata
    """

    assembly_report_path = Path(assembly_report_path)
    outfile = Path(outfile)

    if assembly_report_path.is_dir():
        logging.info(
            f"Summarizing NCBI assembly information from {assembly_report_path}..."
        )
        ncbi_json_files = [
            file
            for file in assembly_report_path.iterdir()
            if file.is_file() and file.suffix == ".json"
        ]
    elif assembly_report_path.is_file() and assembly_report_path.suffix == ".json":
        logging.info(
            f"Getting NCBI assembly information from a single file: {assembly_report_path}"
        )
        ncbi_json_files = [assembly_report_path]
    elif assembly_report_path.is_file() and assembly_report_path.suffix == ".txt":
        logging.info(
            f"Reading NCBI assembly information from a text file: {assembly_report_path}"
        )
        with open(assembly_report_path, "r") as file:
            file_content = [i.strip("\n") for i in file.readlines()]
            if len(file_content) == 1:
                # Paths space-separated on a single line
                paths = file_content[0].split()
            else:
                # Paths written on separate lines
                paths = file_content
            ncbi_json_files = [
                Path(path) for path in paths if Path(path).suffix == ".json"
            ]
    else:
        ncbi_json_files = [
            Path(file)
            for file in str(assembly_report_path).split()
            if Path(file).suffix == ".json"
        ]
        logging.info(
            f"Summarizing NCBI assembly information from the given list of {len(ncbi_json_files)} files..."
        )

    dfs = []  # an empty list to store the data frames
    for file in ncbi_json_files:
        data = pd.read_json(file).T  # read data frame from json file
        dfs.append(data)  # append the data frame to the list

    df_ncbi = pd.concat(dfs)  # concatenate all the data frames in the list.
    df_ncbi.reset_index(inplace=True)
    df_ncbi = df_ncbi.rename(columns={"index": "genome_id"})

    logging.info(f"Summarized {len(df_ncbi.index)} NCBI assembly information.")
    df_ncbi.to_csv(outfile, index=False)
    return None


if __name__ == "__main__":
    extract_ncbi_information(sys.argv[1], sys.argv[2])

import sys

import pandas as pd


def extract_ncbi_information(
    ncbi_meta_path, patric_genome_summary, patric_genome_metadata, patric_meta_path
):
    """
    Extract PATRIC assembly json reports generated by get_assembly_information.py in a given directory into a csv file

    Parameters
    ----------
    1. assembly_report_path : str / path
        Three different input types can be used:
        - A directory containing assembly information json files generated by get_assembly_information.py
        - A single json file generated by get_assembly_information.py
        - A list of files as string separated by spaces (put the string inside '' in bash expression)
    2. outfile : str / path
        Location of the output csv file

    Returns
    -------
    1. outfile : .csv file
        A comma separated table summarizing the NCBI assembly reports metadata
    '''
    """

    df_ncbi = pd.read_csv(ncbi_meta_path, index_col="genome_id")
    df_patric_genome_summary = pd.read_csv(
        patric_genome_summary,
        sep="\t",
        dtype={"genome_id": str, "genome_name": str, "taxon_id": str, "strain": str},
    )
    df_patric_genome_summary.set_index("genome_id", inplace=True)

    df_patric_genome_metadata = pd.read_csv(
        patric_genome_metadata,
        sep="\t",
        dtype={"genome_id": str, "genome_name": str, "taxon_id": str},
    )
    df_patric_genome_metadata.set_index("genome_id", inplace=True)

    df_patric_genome_metadata = df_patric_genome_metadata[
        df_patric_genome_metadata["genome_status"] != "Plasmid"
    ]

    df_patric_meta = pd.DataFrame()

    for genome_id in df_ncbi.index:
        refseq = df_ncbi.loc[genome_id, "refseq"]
        genbank = df_ncbi.loc[genome_id, "genbank"]
        if refseq in df_patric_genome_metadata["assembly_accession"].tolist():
            patric_genome_ids_list = df_patric_genome_metadata[
                df_patric_genome_metadata["assembly_accession"] == refseq
            ].index.tolist()
        elif genbank in df_patric_genome_metadata["assembly_accession"].tolist():
            patric_genome_ids_list = df_patric_genome_metadata[
                df_patric_genome_metadata["assembly_accession"] == genbank
            ].index.tolist()

        for patric_genome_id in patric_genome_ids_list:
            df_patric_meta.loc[patric_genome_id, "bgcflow_genome_id"] = genome_id
            df_patric_meta.loc[
                patric_genome_id, df_patric_genome_metadata.columns
            ] = df_patric_genome_metadata.loc[
                patric_genome_id, df_patric_genome_metadata.columns
            ]
            if patric_genome_id in df_patric_genome_summary.index:
                for col in df_patric_genome_summary.columns:
                    if col not in df_patric_genome_metadata.columns:
                        df_patric_meta.loc[
                            patric_genome_id, col
                        ] = df_patric_genome_summary.loc[patric_genome_id, col]

    df_patric_meta.index.name = "patric_genome_id"
    df_patric_meta.to_csv(patric_meta_path)

    return None


if __name__ == "__main__":
    extract_ncbi_information(sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4])

import json
import logging
import sys
from pathlib import Path

import pandas as pd

log_format = "%(levelname)-8s %(asctime)s   %(message)s"
date_format = "%d/%m %H:%M:%S"
logging.basicConfig(format=log_format, datefmt=date_format, level=logging.DEBUG)


def summarize_gtdb_json(accession_list, df_gtdb_output):
    """
    Given a string of json filepaths generated by gtdb_prep.py,
    merge them together into one csv table
    """

    # Reading input
    logging.info("Reading GTDB metadata .json files...")

    input_json = Path(accession_list)

    if input_json.is_file() and input_json.suffix == ".json":
        logging.info(f"Getting GTDB taxonomy from a single file: {input_json}")
        input_json_files = input_json

    elif input_json.is_file() and input_json.suffix == ".txt":
        logging.info(f"Getting GTDB taxonomy from a text file: {input_json}")
        with open(input_json, "r") as file:
            file_content = [i.strip("\n") for i in file.readlines()]
            if len(file_content) == 1:
                # Paths space-separated on a single line
                paths = file_content[0].split()
            else:
                # Paths written on separate lines
                paths = file_content
            input_json_files = [
                Path(path) for path in paths if Path(path).suffix == ".json"
            ]
    else:
        input_json_files = [
            Path(file)
            for file in str(input_json).split()
            if Path(file).suffix == ".json"
        ]
        logging.info(
            f"Getting GTDB_taxonomy from the given list of {len(input_json_files)} files..."
        )

    accession = input_json_files
    out = []
    for a in accession:
        logging.info(f"Reading {a}...")
        with open(a, "r") as f:
            out.append(json.load(f))
    df = pd.DataFrame(out).set_index("genome_id", drop=False)

    # Getting taxonomic information
    logging.info("Getting taxonomic information...")
    df_taxonomy = pd.DataFrame.from_dict(
        {i: df.loc[i, "gtdb_taxonomy"] for i in df.index}
    ).T
    # Split between species (specific epithet) and organism
    df_taxonomy["organism"] = df_taxonomy["species"]
    for idx in df_taxonomy.index:
        try:
            df_taxonomy.loc[idx, "species"] = df_taxonomy.loc[idx, "species"].split(
                " "
            )[1]
        except IndexError:  # leave blank for empty taxonomy
            pass
    # Tidying up
    df_taxonomy.columns = [c.title() for c in df_taxonomy.columns]

    # Getting other metadata
    try:
        logging.info("Getting metadata into table...")
        if "metadata" not in df.columns:
            logging.warning(
                "metadata is not in the column information. Adding default values..."
            )
            df["metadata"] = [{"genome_id": genome_id} for genome_id in df.index]
        if "gtdb_release" not in df.columns:
            logging.warning(
                "gtdb_release is not in the column information. Adding default values..."
            )
            df["gtdb_release"] = "unknown"
        metadata = pd.DataFrame.from_dict(
            {i: df.loc[i, "metadata"] for i in df.index}
        ).T

        metadata_container = []
        for c in metadata.columns:
            value = {i: metadata.loc[i, c] for i in metadata.index}
            try:
                metadata_container.append(pd.DataFrame.from_dict(value).T)
            except ValueError:
                metadata_container.append(
                    pd.DataFrame([value]).T.rename(columns={0: c})
                )
        df_metadata = pd.concat(metadata_container, axis=1)

        logging.info("Finalizing table...")
        df_final = pd.concat(
            [df.loc[:, ["genome_id", "gtdb_release"]], df_taxonomy, df_metadata], axis=1
        )

    except KeyError:
        logging.info("No additional metadata found.")
        logging.info("Finalizing table...")
        df_final = pd.concat(
            [df.loc[:, ["genome_id", "gtdb_release"]], df_taxonomy], axis=1
        )

    # save to file
    logging.info(f"Writing to file: {df_gtdb_output}")
    df_final.to_csv(df_gtdb_output, index=False)
    return None


if __name__ == "__main__":
    summarize_gtdb_json(sys.argv[1], sys.argv[2])

import json
import logging
import re
import subprocess
import sys
from datetime import datetime
from pathlib import Path

import pandas as pd
from Bio import SeqIO

log_format = "%(levelname)-8s %(asctime)s   %(message)s"
date_format = "%d/%m %H:%M:%S"
logging.basicConfig(format=log_format, datefmt=date_format, level=logging.DEBUG)


def get_git_version():
    """
    Get the sha1 of the current git version
    """

    git_version = ""
    try:
        version_cmd = subprocess.run(
            ["git", "rev-parse", "--short", "HEAD"],
            universal_newlines=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )
        status_cmd = subprocess.run(
            ["git", "status", "--porcelain"],
            universal_newlines=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
        )

        git_version = str(version_cmd.stdout.strip())
        changes = str(status_cmd.stdout).strip()
        if changes:
            git_version += "(changed)"
    except OSError:
        pass

    return git_version


def get_version(version):
    """
    Get the current version string
    """

    git_version = get_git_version()
    if git_version:
        version += "-%s" % git_version
    return version


def add_bgcflow_comments(gbk_in_path, version, json_path, genome_id, gbk_out_path):
    """
    Add bgcflow meta-annotation to genbank output
    """

    version = get_version(version)

    logging.info(f"Formatting genbank file: {gbk_in_path}")

    temp_gbk = Path(gbk_in_path).parent / f"{genome_id}-change_log.gbk"
    try:
        [record for record in SeqIO.parse(gbk_in_path, "genbank")]
    except ValueError as e:
        logging.warning(f"Parsing fail: {e}")
        logging.info("Attempting to fix genbank file...")
        change_log_path = Path(gbk_in_path).parent / f"{genome_id}-change_log.json"
        change_log = correct_collided_headers(
            gbk_in_path, genome_id, temp_gbk, json_dump=change_log_path
        )
        logging.info("Retry parsing with Bio.SeqIO...")

    if temp_gbk.is_file():
        records = SeqIO.parse(temp_gbk, "genbank")
        temp_gbk.unlink()
    else:
        records = SeqIO.parse(gbk_in_path, "genbank")

    df_gtdb_meta = pd.read_json(json_path, orient="index").T
    df_gtdb_meta.fillna("Unclassified", inplace=True)
    df_tax = pd.DataFrame([i for i in df_gtdb_meta.gtdb_taxonomy])
    for col in df_tax.columns:
        df_tax[col] = df_tax[col].apply(lambda x: x.split("__")[1])
    df_tax.columns = [c.title() for c in df_tax.columns]
    tax_levels = ["Domain", "Phylum", "Class", "Order", "Family", "Genus", "Species"]
    taxonomy_str = df_tax.loc[0, tax_levels].tolist()
    logging.debug(f"Taxonomy found: {taxonomy_str}")

    bgcflow_comment = (
        "##BGCflow-Data-START##\n"
        "Version      :: {version}\n"
        "Run date     :: {date}\n"
    ).format(
        version=version,
        date=str(datetime.now().strftime("%Y-%m-%d %H:%M:%S")),
    )

    new_records = []
    ctr = 0
    for record in records:
        logging.info(f"Formatting record {ctr}...")
        comment = bgcflow_comment
        if "comment" in record.annotations:
            record.annotations["comment"] += "\n" + comment
        else:
            record.annotations["comment"] = comment

        try:
            if not change_log[ctr]["accept_format"]:
                record.annotations[
                    "comment"
                ] += f"Original ID  :: {change_log[ctr]['original_id']}\n"
        except UnboundLocalError:
            pass

        record.annotations["comment"] += "##BGCflow-Data-END##"

        if "organism" in record.annotations:
            organism = record.annotations["organism"]
            if "Unclassified" in organism:
                record.annotations["organism"] = organism.split(" Unclassified")[
                    0
                ].strip()

        record.annotations["taxonomy"] = taxonomy_str

        new_records.append(record)

        ctr = ctr + 1

    logging.info(f"Writing final output: {gbk_out_path}")
    with open(gbk_out_path, "w") as output_handle:
        SeqIO.write(new_records, output_handle, "genbank")


# -- Correcting IDs --#
# 1. Traverse text, make a library of record header, length, and ids
# 2. Check the library for collided locus name and length
# 3. Propose changes
# 4. Make sure changes is unique
# 5. Write a new genbank file and recorded changes as json file


def correct_collided_headers(
    gbk_in_path, accession_id, outfile, json_dump="header.json"
):
    record_headers = get_record_headers(gbk_in_path)
    record_headers = shorten_record_headers(record_headers, accession_id)

    with open(json_dump, "w", encoding="utf8") as json_file:
        json.dump(record_headers, json_file, indent=4)

    logging.info(f"Writing result to: {outfile}")
    with open(gbk_in_path) as f:
        s = f.read()

    with open(outfile, "w") as f:
        for k in record_headers.keys():
            if not record_headers[k]["accept_format"]:
                s = s.replace(
                    record_headers[k]["original_header"],
                    record_headers[k]["new_header"],
                )
        f.write(s)
    return record_headers


def get_record_headers(gbk_in_path):
    """
    Get the header information for each records in a genbank file.
    Returns a json-like python dictionary.
    """
    record_headers = {}

    # Find all headers
    with open(gbk_in_path, "r") as file:
        logging.info(f"Reading file as text: {gbk_in_path}")
        ctr = 0
        for line in file:
            if line.startswith("LOCUS"):
                record_headers[ctr] = {"original_header": line}
            if "source" in line:
                length = line.split()[-1].split("..")[-1]
                record_headers[ctr]["length"] = length
                ctr = ctr + 1
    logging.debug(f"Found {len(record_headers)} records.")

    # Check for collided headers
    logging.info("Checking header format...")
    for k in record_headers.keys():
        query = record_headers[k]["original_header"].split()
        if query != 7 and query[3] != "bp":
            logging.warning("Record name and length collide in the LOCUS line")
            query_length = record_headers[k]["length"]
            if query_length in query[1]:
                original_id = query[1].replace(query_length, "")
                record_headers[k]["original_id"] = original_id
                record_headers[k]["accept_format"] = False
        else:
            record_headers[k]["original_id"] = query[1]
            record_headers[k]["accept_format"] = True
    return record_headers


def modify_id(accession_id, original_id, locus_index, record_counts, id_collection):
    """
    Attempt to fix collided name and length in the locus name by shortening id
    """
    logging.info(f"{original_id}: Attempting to change locus name...")

    # For refseq record, remove the accession prefix (first two digits) from string
    if accession_id.startswith("GCF"):
        logging.debug(
            f"{original_id}: Assuming locus came from Refseq. Removing refseq accession prefix."
        )
        refseq_type, refseq_number, refseq_version = re.split("_|[.]", original_id)
        new_id = f"{refseq_number}.{refseq_version}"

    elif accession_id.startswith("GCA"):
        logging.info(
            f"{original_id}: Assuming locus came from Genbank. Removing version from locus name..."
        )
        new_id, genbank_version = re.split("[.]", original_id)

    # For unknown source remove last 4 digit, add the contig index in front.
    # example: NZAJABAQG010000001.1 --> c1|NZAJABAQG010000
    else:
        logging.info(
            f"{original_id}: Cannot determine source. Shortening locus name..."
        )
        digit = len(str(record_counts))
        contig_number = str(locus_index + 1)

        new_id = f"c{contig_number.zfill(digit)}_{original_id[:-(digit + 4)]}"

    # make sure new_id is unique
    logging.info(f"{original_id}: Making sure id is unique...")
    if new_id in id_collection:
        raise
    else:
        logging.debug(f"{original_id}: shortened to {new_id}")
    return new_id


def shorten_record_headers(record_headers, accession_id):
    """
    Shorten record headers
    """
    record_counts = len(record_headers.keys())
    id_collection = [record_headers[k]["original_id"] for k in record_headers.keys()]

    for k in record_headers.keys():
        if record_headers[k]["accept_format"]:
            pass
        else:
            original_id = record_headers[k]["original_id"]
            old_header = record_headers[k]["original_header"]

            # Correct header name
            new_id = modify_id(
                accession_id, original_id, k, record_counts, id_collection
            )
            new_id_replacement_value = (
                f"{new_id}{(len(original_id) - len(new_id)) * ' '}"
            )
            new_header = old_header.replace(original_id, new_id_replacement_value)

            # Add value to dict
            record_headers[k]["new_id"] = new_id
            record_headers[k]["new_header"] = new_header

    return record_headers


if __name__ == "__main__":
    add_bgcflow_comments(
        sys.argv[1], sys.argv[2], sys.argv[3], sys.argv[4], sys.argv[5]
    )