"""I/O functions"""

# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_util.ipynb.

# %% auto 0
__all__ = ['read_gff', 'gff_attribute_selector', 'decorate', 'filter', 'syntenic_block_borders', 'read_aln', 'estimate_plot_size',
           'insert_gap', 'flip', 'insert_break']

# %% ../nbs/01_util.ipynb 4
import os
from collections.abc import Sequence
from pathlib import Path
from typing import Literal

import numpy as np
import pandas as pd

# %% ../nbs/01_util.ipynb 6
def read_gff(
    loc: str | Path,  # input filepath
    gff_columns: list = [
        "seqid",
        "source",
        "type",
        "start",
        "end",
        "score",
        "strand",
        "phase",
        "attributes",
    ],  # column names for the GFF3 file
    skiprows: int = 1,  # how many rows to skip in the beginning
    header: (
        int | Sequence[int] | None | Literal["infer"]
    ) = None,  # whether to expect a header or not
    sep: str = "\t",  # separator for the table
    **kwargs,  # various pd.read_csv() arguments
) -> pd.DataFrame:  # the GFF3 file in DataFrame form
    "A function to read a GFF3 file. Expects 9 tab-separated fields, and will try to name the columns according to https://github.com/The-Sequence-Ontology/Specifications/blob/master/gff3.md"
    gff = pd.read_csv(loc, sep=sep, header=header, skiprows=skiprows, **kwargs)
    gff.columns = gff_columns
    return gff

# %% ../nbs/01_util.ipynb 11
def gff_attribute_selector(
    line,  # the attributes field of a GFF3 line
    sep: str = ";",  # the field separator. Should be a semicolon for a GFF3 file.
    select="ID",  # the field ID. Should be one that is included in the GFF3 file. Refer to https://github.com/The-Sequence-Ontology/Specifications/blob/master/gff3.md for possibilities, or choose a manually defined tag that you know is present in the file.
) -> str:  # the value of field `select`
    "A function that extracts the value of a specified field from the attributes of a GFF3 line. Should be used on the `attributes` field of the corresponding pandas DataFrame."
    for field in line.split(sep):
        if field.startswith(select):
            value = field.split("=")[1]
            return value
    return None

# %% ../nbs/01_util.ipynb 16
def decorate(
    gff: pd.DataFrame,  # a GFF file in Pandas dataframe form
    attributes: dict = {
        "gene_id": "ID",
        "gene_name": "gene",
    },  # a dictionary of tags to extract from the attributes column. For each key/value pair, the value under the `key` tag will be saved in a new column named `value`.
):
    "A function that"
    for attr_add, attr_filter in attributes.items():
        gff[attr_add] = gff["attributes"].apply(
            lambda x: gff_attribute_selector(x, select=attr_filter)
        )


# | export


def filter(
    gff,
    filter_by_type=True,
    filter_type="gene",
    filter_by_field=True,
    field="gene_id",
    field_values=None,
):
    if filter_type and filter_type not in gff[filter_type]:
        raise ValueError(
            f'The value {filter_type} is not found in the "type" column of the input table.'
        )

    if filter_by_field:
        if field not in gff.columns:
            raise ValueError(
                f'The value {filter_type} is not found in the "type" column of the input table.'
            )

    keep = np.array([True] * gff.shape[0])
    if genes is not None:
        is_genes = gff["gene_id"].isin(genes)
        keep = keep & is_genes

    if is_type:
        is_gene = gff["type"] == "gene"

    slim = gff[is_gene & is_hox]
    return slim

# %% ../nbs/01_util.ipynb 17
def syntenic_block_borders(
    gff: pd.DataFrame,  # a GFF in Pandas dataframe form. Only includes the genes of the syntenic block in question.
    flank_length: int = None,  # the amount of space to be granted on both sides of the syntenic region, in basepairs. If unspecified, it will be set to 5% of the syntenic block length.
    start: str = "start",  # the GFF column with the start position of the gene ("start").
    end: str = "end",  # the GFF column with the end position of the gene ("end").
) -> (int, int):
    "A function to calculate the boundaries of a syntenic block. It automatically pads the boundaries by an additional 5% of total length on both ends."
    possible = [gff[start].min(), gff[start].max(), gff[end].min(), gff[end].max()]
    block_start = min(possible)
    block_end = max(possible)
    if flank_length is None:
        flank_length = (block_end - block_start) // 100 * 5
    if not flank_length.is_integer():
        raise ValueError(
            f'The parameter `flank_length` has to be an integer. You supplied the value "{flank_length}" which is {type(flank_length)}.'
        )
    return max(block_start - flank_length, 0), block_end + flank_length

# %% ../nbs/01_util.ipynb 20
def read_aln(
    m8: str,  # the path to the MMseqs2 alignment table file
    id_sep: (
        str | None
    ) = None,  # (optional) a character that separates the species ID from the gene ID
    **kwargs,  # various arguments to be passed to the pd.read_csv() function
) -> pd.DataFrame:  # the tabulated form of the alignment results.
    "Reads"
    hits = pd.read_csv(m8, sep="\t", header=None)
    m8_columns = [
        "query",
        "target",
        "seq_id",
        "ali_len",
        "no_mism",
        "no_go",
        "q_start",
        "q_end",
        "t_start",
        "t_end",
        "eval",
        "bit",
    ]
    hits.columns = m8_columns
    # trim the query to just the ID
    if id_sep is not None:
        hits["query"] = hits["query"].str.split(id_sep).str[1].str.join(id_sep)
    return hits

# %% ../nbs/01_util.ipynb 21
def estimate_plot_size(
    gff,
    width_factor: int = 3,
    height: int = 2,
):
    x = len(gff) * width_factor
    return (x, height)

# %% ../nbs/01_util.ipynb 23
def insert_gap(
    gff: pd.DataFrame,
    locus1=None,
    locus2=None,
    identifier="gene_id",
    purge_columns=None,
    no_gaps=1,
) -> pd.DataFrame:
    "This function inserts a number of dummy entries between two loci (lines) in the GFF DataFrame."
    loc1_index = -1
    loc2_index = len(gff)
    if locus1 is not None:
        loc1_index = gff[gff[identifier] == locus1].index[0]
    if locus2 is not None:
        loc2_index = gff[gff[identifier] == locus2].index[0]
    if loc2_index - loc1_index != 1:
        raise ValueError(
            f"The two loci are not consecutive; their indices are {locus1}: {loc1_index} and {locus2}: {loc2_index}, respectively. Please refer to your input dataframe."
        )

    offset_step = 1 / (no_gaps + 1)

    copy = loc1_index
    if locus1 is None:
        copy = loc2_index

    for i in range(no_gaps):
        offset = (i + 1) * offset_step
        tmp = gff.loc[copy].copy()
        if locus1 is not None:
            tmp["start"] = tmp["end"] + 1 + (i * 2)
            tmp["end"] = tmp["start"] + 1
        else:
            tmp["start"] = tmp["start"] - 1 - (i * 2)
            tmp["end"] = tmp["start"] + 1

        tmp[identifier] = f"gap_{locus1}-{i}"
        for column in purge_columns:
            tmp[column] = ""
        gff.loc[loc1_index + offset] = tmp

    gff = gff.sort_index().reset_index(drop=True)
    return gff

# %% ../nbs/01_util.ipynb 29
def _flip_strand(s):
    if s == "+":
        return "-"
    elif s == "-":
        return "+"
    else:
        return s


# | export


def flip(gff):
    gff["start"] = (-1) * gff["start"]
    gff["end"] = (-1) * gff["end"]
    gff["strand"] = gff["strand"].apply(_flip_strand)
    gff.index = gff.index[::-1]
    return gff.sort_index()

# %% ../nbs/01_util.ipynb 31
def insert_break(
    gff: pd.DataFrame,
    locus1=None,
    locus2=None,
    identifier="gene_id",
) -> pd.DataFrame:
    "This function inserts a molecule break between two loci (lines) in the GFF DataFrame."
    loc1_index = gff[gff[identifier] == locus1].index[0]
    loc2_index = gff[gff[identifier] == locus2].index[0]
    if loc2_index - loc1_index != 1:
        raise ValueError(
            f"The two loci are not consecutive; their indices are {locus1}: {loc1_index} and {locus2}: {loc2_index}, respectively. Please refer to your input dataframe."
        )

    tmp = gff.loc[loc1_index].copy()
    for c in tmp.index:
        tmp[c] = ""
    tmp[identifier] = f"break"
    gff.loc[loc1_index + 0.5] = tmp

    gff = gff.sort_index().reset_index(drop=True)
    return gff
