import base64
from types import new_class
import ipywidgets as widgets
from pathlib import PurePath
from jupyter_bbox_widget import BBoxWidget
from ipywidgets import Layout
import numpy as np


class Labeler:
    def __init__(self, dataset=None):
        self.dataset = dataset

    def StartPyLaber(self, new_classes=None, image=None, yolo_model=None):
        """Display the bbox widget loaded with images and annotations from this dataset."""

        if "google.colab" in str(get_ipython()):
            from google.colab import output

            output.enable_custom_widget_manager()

        dataset = self.dataset
        widget_output = None

        files = dataset.df.img_filename.unique()
        files = files.tolist()

        global file_index

        if image == None:
            file_index = 0
            image = files[0]
        else:
            file_index = files.index(image)

        def GetBBOXs(image):
            # Make a dataframe with the annotations for a single image
            img_df = dataset.df.loc[dataset.df["img_filename"] == image]
            img_df_subset = img_df[
                [
                    "cat_name",
                    "ann_bbox_height",
                    "ann_bbox_width",
                    "ann_bbox_xmin",
                    "ann_bbox_ymin",
                ]
            ]
            # Rename the columns to match the format used by jupyter_bbox_widget
            img_df_subset.columns = ["label", "height", "width", "x", "y"]

            # Drop rows that have NaN, invalid bounding boxes
            img_df_subset = img_df_subset.dropna()
            bboxes_dict = img_df_subset.to_dict(orient="records")
            return bboxes_dict

        def GetImageLabel(image):
            """Returns the imagename or imagename + (not annotated)"""
            # Make a dataframe with the annotations for a single image
            img_df = dataset.df.loc[dataset.df["img_filename"] == image]

            if img_df.iloc[0].annotated == 1:
                return image
            else:
                return f"{image} (not annotated)"
            return

        bboxes_dict = GetBBOXs(image)

        img_folder = dataset.df.loc[dataset.df["img_filename"] == image].iloc[0][
            "img_folder"
        ]

        file_paths = [
            str(PurePath(dataset.path_to_annotations, img_folder, file))
            for file in files
        ]

        def encode_image(filepath):
            with open(filepath, "rb") as f:
                image_bytes = f.read()
            encoded = str(base64.b64encode(image_bytes), "utf-8")
            return "data:image/jpg;base64," + encoded

        def UpdateCategoryList(cat_dict, new_categories):
            from math import isnan

            # Remove invalid entries
            cat_dict.pop("", None)
            # cat_dict = {k: v for k, v in cat_dict.items() if not isnan(v)}

            for cat in new_categories:
                if len(cat_dict) == 0:
                    cat_dict[cat] = "0"
                elif cat in list(cat_dict.keys()):
                    continue
                else:
                    # Create a new cat id that 1+ the highest cat id value
                    new_cat_id = max([int(v) for v in cat_dict.values()]) + 1
                    cat_dict[cat] = str(new_cat_id)

            return cat_dict

        def on_submit(b):
            # save annotations for current image
            import pandas as pd

            global widget_output
            global file_index

            img_filename = files[file_index]

            widget_output = pd.DataFrame.from_dict(w_bbox.bboxes)

            # Check if there are any bounding boxes for the current image:
            if not widget_output.empty:
                # If there are bounding boxes then add them to the dataset
                widget_output = widget_output.rename(
                    columns={
                        "label": "cat_name",
                        "height": "ann_bbox_height",
                        "width": "ann_bbox_width",
                        "x": "ann_bbox_xmin",
                        "y": "ann_bbox_ymin",
                    }
                )

                widget_output["ann_area"] = (
                    widget_output["ann_bbox_height"] * widget_output["ann_bbox_width"]
                )
                widget_output["cat_name"] = widget_output["cat_name"].astype("string")
                categories = dict(zip(dataset.df.cat_name, dataset.df.cat_id))
                categories = UpdateCategoryList(
                    categories, list(widget_output.cat_name)
                )
                widget_output["cat_id"] = widget_output["cat_name"].map(categories)

            else:
                # Build a entry in the table with empty bounding box columns
                widget_output = pd.DataFrame(
                    [],
                    columns=[
                        "cat_name",
                        "cat_id",
                        "ann_bbox_height",
                        "ann_bbox_width",
                        "ann_bbox_xmin",
                        "ann_bbox_ymin",
                        "ann_area",
                    ],
                )

                widget_output.loc[0] = [
                    "",
                    np.NaN,
                    np.NaN,
                    np.NaN,
                    np.NaN,
                    np.NaN,
                    np.NaN,
                ]

            widget_output["img_filename"] = str(img_filename)
            widget_output["img_filename"] = widget_output["img_filename"].astype(
                "string"
            )
            widget_output.index.name = "id"

            img_df = dataset.df.loc[dataset.df["img_filename"] == img_filename]

            # Get the metadata associated with this image from the dataset
            metadata = img_df.iloc[0].to_frame().T

            metadata["img_filename"] = metadata["img_filename"].astype("string")

            # Drop the fields that are in the widget_output dataframe
            metadata.drop(
                [
                    "cat_name",
                    "cat_id",
                    "ann_area",
                    "ann_bbox_height",
                    "ann_bbox_width",
                    "ann_bbox_xmin",
                    "ann_bbox_ymin",
                ],
                axis=1,
                inplace=True,
            )

            widget_output = widget_output.merge(
                metadata, left_on="img_filename", right_on="img_filename"
            )
            widget_output = widget_output[dataset.df.columns]

            # Set annotated = 1, which means the annotates have been reviewed and accepted
            widget_output["annotated"] = 1

            # Now we have a dataframe with output of the bbox widget
            # Drop the current annotations for the image and add the the new ones
            dataset.df.drop(
                dataset.df[dataset.df["img_filename"] == img_filename].index,
                inplace=True,
            )
            dataset.df.reset_index(drop=True, inplace=True)

            dataset.df = dataset.df.append(widget_output).reset_index(drop=True)

            # move on to the next file
            on_next(b)

        def on_next(b):
            global file_index

            file_index += 1
            # open new image in the widget
            image_file = file_paths[file_index]
            w_bbox.image = encode_image(image_file)
            w_bbox.bboxes = GetBBOXs(files[file_index])
            progress_label.value = f"{file_index+1} / {len(files)}"
            current_img_name_label.value = GetImageLabel(files[file_index])

        def on_back(b):
            global file_index
            file_index -= 1
            # open new image in the widget
            image_file = file_paths[file_index]
            w_bbox.image = encode_image(image_file)
            w_bbox.bboxes = GetBBOXs(files[file_index])
            progress_label.value = f"{file_index+1} / {len(files)}"
            current_img_name_label.value = GetImageLabel(files[file_index])

        def on_add_class(b):
            if new_class_text.value.strip() != "":
                class_list = list(w_bbox.classes)
                class_list.append(new_class_text.value)
                w_bbox.classes = list(set(class_list))
                new_class_text.value = ""

        def on_predict(b):
            global file_index
            image_file = file_paths[file_index]

            result = yolo_model(image_file)
            result = result.pandas().xyxy[0]
            result["width"] = result.xmax - result.xmin
            result["height"] = result.ymax - result.ymin
            result.drop(["class", "confidence", "xmax", "ymax"], axis=1, inplace=True)
            result.columns = ["x", "y", "label", "width", "height"]
            result = result[["label", "height", "width", "x", "y"]]
            bboxes_dict = result.to_dict(orient="records")
            w_bbox.bboxes = bboxes_dict
            # Add the predicted classes to the widget so they can be selected by the user
            w_bbox.classes = list(set(w_bbox.classes + list(result["label"])))

        if new_classes:
            classes = dataset.analyze.classes + new_classes
        else:
            classes = dataset.analyze.classes

        # remove empty labels and duplicate labels
        classes = list(
            set([str(c).strip() for c in classes if len(str(c).strip()) > 0])
        )

        # Load BBoxWidget for first load on page
        w_bbox = BBoxWidget(
            image=encode_image(file_paths[file_index]),
            classes=classes,
            bboxes=bboxes_dict,
            hide_buttons=True,
        )

        progress_txt = f"{file_index+1} / {len(files)}"

        left_arrow_btn = widgets.Button(
            icon="fa-arrow-left", layout=Layout(width="35px")
        )
        progress_label = widgets.Label(value=progress_txt)
        current_img_name_label = widgets.Label(value=GetImageLabel(image))

        right_arrow_btn = widgets.Button(
            icon="fa-arrow-right", layout=Layout(width="35px")
        )
        save_btn = widgets.Button(
            icon="fa-check", description="Save", layout=Layout(width="70px")
        )
        predict_btn = widgets.Button(
            icon="fa-eye", description="Predict", layout=Layout(width="100px")
        )
        train_btn = widgets.Button(
            icon="fa-refresh", description="Train", layout=Layout(width="100px")
        )
        add_class_label = widgets.Label(value="Add class:")
        new_class_text = widgets.Text(layout=Layout(width="200px"))
        plus_btn = widgets.Button(icon="fa-plus", layout=Layout(width="35px"))

        button_row_list = [left_arrow_btn, progress_label, right_arrow_btn, save_btn]

        # If model arg is empty hide predict and train buttons
        if yolo_model != None:
            button_row_list = button_row_list + [predict_btn]  # train_btn

        button_row = widgets.HBox(button_row_list)
        current_img_details_row = widgets.HBox([current_img_name_label])

        bottom_row = widgets.HBox([add_class_label, new_class_text, plus_btn])

        pylabler = widgets.VBox(
            [current_img_details_row, button_row, w_bbox, bottom_row]
        )

        save_btn.on_click(on_submit)
        left_arrow_btn.on_click(on_back)
        right_arrow_btn.on_click(on_next)
        predict_btn.on_click(on_predict)
        plus_btn.on_click(on_add_class)
        new_class_text.on_submit(on_add_class)

        # Returning the container will show the widgets
        return pylabler
