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

# %% auto 0
__all__ = ['np_2_b64', 'b64_2_np', 'b64_2_file', 'bytes_2_file', 'file_2_bytes', 'csvs_2_cvat', 'cvat_2_csvs', 'df_2_yolo',
           'yolo_2_df']

# %% ../nbs/adapters.ipynb 2
import base64, cv2, numpy as np
from typing import Union
from .loader import *
from .paths import P, stems, stem, parent, makedir
from .markup import write_xml, read_xml, AttrDict

# %% ../nbs/adapters.ipynb 3
def np_2_b64(image: np.ndarray) -> str:
    """Convert a numpy image to base64 string"""
    img_str = cv2.imencode(".jpg", image)[1].tobytes()
    img_b64 = base64.b64encode(img_str).decode("utf-8")
    return img_b64


def b64_2_np(input: str) -> np.ndarray:
    """Converts a base64 encoded image to a NumPy array.

    Args:
        input (str): The base64 encoded image.

    Returns:
        np.ndarray: The NumPy array representation of the image in RGB format.
    """
    input = bytes(input, "utf-8")
    input = base64.b64decode(input)
    img_nparr = np.frombuffer(input, np.uint8)
    img_cv2 = cv2.imdecode(img_nparr, cv2.IMREAD_COLOR)
    img_rgb = cv2.cvtColor(img_cv2, cv2.COLOR_BGR2RGB)
    return img_rgb


def b64_2_file(
    input: str,  # base64 encoded string
    fpath: Union[str, P],  # Place where you want to save the file
) -> None:
    """Save a file encoded as a base64 `input` at given `fpath`"""
    with open(fpath, "wb") as pdf_file:
        input = bytes(input, "utf-8")
        input = base64.b64decode(input)
        pdf_file.write(input)
        logger.info(f"wrote pdf file to {fpath}")


def bytes_2_file(
    input: bytes,  # bytes
    fpath: Union[str, P],  # Place where you want to save the file
    silent: bool = False,
) -> None:
    """Save bytes `input` at given `fpath`"""
    with open(fpath, "wb") as file_to_write:
        file_to_write.write(input)
        if not silent:
            logger.info(f"wrote pdf file to {fpath}")


def file_2_bytes(fpath):
    with open(fpath, "rb") as f:
        output = f.read()
    return output

# %% ../nbs/adapters.ipynb 4
def _process(
    df: pd.DataFrame, label_column="readable_label", default_label="Background"
):
    """
    Process the given DataFrame and convert it into a list of records.

    Args:
        df (pd.DataFrame): The input DataFrame to be processed.
        label_column (str, optional): The column name for the label. Defaults to "readable_label".
        default_label (str, optional): The default label value. Defaults to "Background".

    Returns:
        list: A list of records, where each record is a dictionary representing a row in the DataFrame.
    """
    df["@xbr"] = df["X"]
    df["@xtl"] = df["x"]
    df["@ybr"] = df["Y"]
    df["@ytl"] = df["y"]
    df["@label"] = df[label_column] if label_column in df.columns else default_label
    df["@occluded"] = "0"
    df["@source"] = "manual"
    df["@z_order"] = "0"
    if "text" in df.columns:
        df["attribute"] = [
            [{"@name": "OCR", "#text": text if text == text else ""}]
            for text in df["text"]
        ]
    df.drop(
        [
            c
            for c in df.columns
            if c
            not in set(
                [
                    "@xbr",
                    "@xtl",
                    "@ybr",
                    "@ytl",
                    "@label",
                    "@occluded",
                    "@source",
                    "@z_order",
                    "attribute",
                ]
            )
        ],
        axis=1,
        inplace=True,
    )
    records = df.to_dict(orient="records")
    return records


def csvs_2_cvat(
    images_folder,
    csvs_folder,
    xml_output_file,
    items=None,
    parquet=False,
    relative_df=True,
    default_label="Background",
    extension="jpg",
):
    """
    Convert CSV annotations to CVAT XML format.

    Args:
        images_folder (str): Path to the folder containing the images.
        csvs_folder (str): Path to the folder containing the CSV annotations.
        xml_output_file (str): Path to the output XML file.
        items (list, optional): List of items to process. If None, all items will be processed. Defaults to None.
        parquet (bool, optional): Whether the annotations are stored in Parquet format. Defaults to False.
        relative_df (bool, optional): Whether the bounding box coordinates in the CSV are relative to the image size. Defaults to True.
        default_label (str, optional): Default label for the bounding boxes. Defaults to "Background".
        extension (str, optional): Image file extension. Defaults to "jpg".

    Returns:
        None
    """
    images_folder, csvs_folder = [P(_) for _ in [images_folder, csvs_folder]]
    data = AttrDict({"annotations": {"image": []}})
    if items is None:
        items = common(stems(images_folder), stems(csvs_folder))

    items = sorted(items)

    for ix, item in enumerate(track(items)):
        _ia = _image_annotation = AttrDict({})
        image = images_folder / f"{item}.{extension}"
        height, width = read(image).shape[:2]
        if parquet:
            df = pd.read_parquet(f"{csvs_folder}/{item}.parquet")
        else:
            df = pd.read_csv(f"{csvs_folder}/{item}.csv")
        if relative_df:
            df = to_absolute(df, height, width)
        _ia["@height"], _ia["@width"] = height, width
        _ia["@id"] = str(ix)
        _ia["@name"] = f"{item}.jpg"
        _ia["box"] = _process(df, default_label=default_label)
        data.annotations.image.append(_ia)
    write_xml(data, xml_output_file)


def _get_attribute_columns(column):
    """
    Get attribute columns from a given column.

    Args:
        column (pandas.Series): The input column.

    Returns:
        set: A set of attribute columns extracted from the input column.
    """

    def _get_columns_from_row(item):
        if item != item:
            return []
        if isinstance(item, dict):
            return [item["@name"]]
        else:
            return [_item["@name"] for _item in item]

    output = set(flatten(column.map(_get_columns_from_row)))
    return output


def _get_attribute_data(item, column_name):
    """
    Retrieves the attribute data for a given item and column name.

    Parameters:
    - item: The item to retrieve the attribute data from.
    - column_name: The name of the column to retrieve the attribute data for.

    Returns:
    - The attribute data for the given item and column name, or np.nan if not found.
    """
    if item != item:
        return item
    if isinstance(item, dict):
        if item["@name"] == column_name:
            return item.get("#text", np.nan)
        else:
            return np.nan
    elif isinstance(item, list):
        item = [_item for _item in item if _item["@name"] == column_name]
        if item:
            return item[0].get("#text", np.nan)
        else:
            return np.nan


def _cvat_ann_2_csv(ann):
    """
    Convert CVAT annotation to a pandas DataFrame in CSV format.

    Args:
        ann (dict): CVAT annotation dictionary.

    Returns:
        pandas.DataFrame: DataFrame containing the converted annotation data in CSV format.
    """
    if "box" not in ann:
        return pd.DataFrame()
    if isinstance(ann.box, AttrDict):
        rows = [a.to_dict() for a in [ann.box]]
    else:
        rows = [a.to_dict() for a in ann.box]
    df = pd.DataFrame(rows)
    df.rename(
        {
            "@xtl": "x",
            "@ytl": "y",
            "@xbr": "X",
            "@ybr": "Y",
            "@label": "readable_label",
        },
        inplace=1,
        axis=1,
    )
    for col in "xyXY":
        df[col] = df[col].map(lambda x: int(float(x)))
    df.drop(["@z_order", "@source", "@occluded"], axis=1, inplace=True)
    height, width = int(ann["@height"]), int(ann["@width"])
    df = to_relative(df, height, width)

    if "attribute" in df.columns:
        new_columns = _get_attribute_columns(df.attribute)
        for column in new_columns:
            df[column] = df["attribute"].map(lambda x: _get_attribute_data(x, column))
        df.drop(["attribute"], axis=1, inplace=True)
    return df


def cvat_2_csvs(xmlfile, csvs_folder):
    """
    Convert CVAT XML annotations to CSV files.

    Args:
        xmlfile (str): Path to the CVAT XML file.
        csvs_folder (str): Path to the folder where the CSV files will be saved.

    Returns:
        None
    """
    data = read_xml(xmlfile)
    items = data.annotations.image
    if not isinstance(items, list):
        items = [items]
    for item in items:
        try:
            df = _cvat_ann_2_csv(item)
            save_at = f'{csvs_folder}/{stem(item["@name"])}.csv'
            makedir(parent(save_at))
            df.to_csv(save_at, index=False)
        except Exception as e:
            Warn(f'{e} @ {item["@name"]}')

# %% ../nbs/adapters.ipynb 5
def df_2_yolo(df, h, w, class2id, class_column):
    yolo_data = []
    for index, row in df.iterrows():
        cls = row[class_column]
        if cls not in class2id:
            continue
        class_id = class2id[cls]
        x_center = (row["x"] + row["X"]) / 2 / w
        y_center = (row["y"] + row["Y"]) / 2 / h
        width = (row["X"] - row["x"]) / w
        height = (row["Y"] - row["y"]) / h
        yolo_data.append(
            f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}"
        )
    return yolo_data


def yolo_2_df(yolo, h, w, id2class, class_column):
    data = []

    for line in yolo:
        class_id, x_center, y_center, width, height = map(float, line.split())
        class_name = id2class[int(class_id)]

        x_center *= w
        y_center *= h
        width *= w
        height *= h

        x_min = int(x_center - (width / 2))
        y_min = int(y_center - (height / 2))
        x_max = int(x_center + (width / 2))
        y_max = int(y_center + (height / 2))

        data.append([x_min, y_min, x_max, y_max, class_name])

    df = pd.DataFrame(data, columns=["x", "y", "X", "Y", class_column])
    return df
