# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/07_im2im_annotator.ipynb (unless otherwise specified).

__all__ = ['Im2ImAnnotator']

# Internal Cell
import io
import warnings
from pathlib import Path
from copy import deepcopy
from typing import Optional, Callable, Union, Iterable

from ipycanvas import Canvas
from ipywidgets import (AppLayout, VBox, HBox, Button, Layout, HTML, Output, Image)

from .base import BaseState, AppWidgetState, Annotator, AnnotatorStep
from .bbox_canvas import ImageRenderer
from .mltypes import OutputImageLabel, OutputLabel, InputImage
from .ipytyping.annotations import LabelStore, LabelStoreCaster
from .custom_widgets.grid_menu import GridMenu, Grid
from .navi_widget import Navi
from .storage import JsonLabelStorage
from IPython.display import display
from .docs.utils import is_building_docs
from PIL import Image as PILImage

# Internal Cell

class Im2ImState(BaseState):
    annotations: LabelStore = LabelStore()
    question_value: str = ''
    grid: Grid
    image_path: Optional[str]
    im_width: int = 300
    im_height: int = 300

# Cell
if is_building_docs():
    class ImCanvas(Image):
        def __init__(
            self,
            width: int = 150,
            height: int = 150,
            has_border: bool = False,
            fit_canvas: bool = False
        ):
            super().__init__(width=width, height=height)
            image = PILImage.new('RGB', (100, 100), (255, 255, 255))
            b = io.BytesIO()
            image.save(b, format='PNG')
            self.value = b.getvalue()

        def _draw_image(self, image_path: str):
            self.value = Image.from_file(image_path).value

        def _clear_image(self):
            pass

        def observe_client_ready(self, cb=None):
            pass
else:
    class ImCanvas(HBox):  # type: ignore
        def __init__(
            self,
            width: int = 150,
            height: int = 150,
            has_border: bool = False,
            fit_canvas: bool = False
        ):
            self.has_border = has_border
            self.fit_canvas = fit_canvas
            self._canvas = Canvas(width=width, height=height)
            super().__init__([self._canvas])

        def _draw_image(self, image_path: str):
            img_render_strategy = ImageRenderer(
                clear=True,
                has_border=self.has_border,
                fit_canvas=self.fit_canvas
            )

            self._image_scale = img_render_strategy.render(
                self._canvas,
                image_path
            )

        def _clear_image(self):
            self._canvas.clear()

        # needed to support voila
        # https://ipycanvas.readthedocs.io/en/latest/advanced.html#ipycanvas-in-voila
        def observe_client_ready(self, cb=None):
            self._canvas.on_client_ready(cb)

# Internal Cell

class Im2ImAnnotatorGUI(AppLayout):
    debug_output = Output(layout={'border': '1px solid black'})

    def __init__(
        self,
        app_state: AppWidgetState,
        im2im_state: Im2ImState,
        state_to_widget: LabelStoreCaster,
        label_autosize=False,
        on_save_btn_clicked: Callable = None,
        on_grid_box_clicked: Callable = None,
        on_navi_clicked: Callable = None,
        has_border: bool = False,
        fit_canvas: bool = False
    ):
        self._app_state = app_state
        self._im2im_state = im2im_state
        self._on_save_btn_clicked = on_save_btn_clicked
        self._on_navi_clicked = on_navi_clicked
        self._on_grid_box_clicked = on_grid_box_clicked
        self.state_to_widget = state_to_widget

        if label_autosize:
            if self._im2im_state.im_width < 100 or self._im2im_state.im_height < 100:
                self._im2im_state.grid.width = 10
                self._im2im_state.grid.height = 10
            elif self._im2im_state.im_width > 1000 or self._im2im_state.im_height > 1000:
                self._im2im_state.grid.width = 50
                self._im2im_state.grid.height = 10
            else:
                label_width = min(self._im2im_state.im_width, self._im2im_state.im_height) // 10
                self._im2im_state.grid.width = label_width
                self._im2im_state.grid.height = label_width

        self._image = ImCanvas(
            width=self._im2im_state.im_width,
            height=self._im2im_state.im_height,
            has_border=has_border,
            fit_canvas=fit_canvas
        )
        self._image.layout = Layout(
            display='block',
            margin='auto'
        )

        self._navi = Navi()
        self._navi.on_navi_clicked = self.on_navi_clicked
        self._save_btn = Button(description="Save",
                                layout=Layout(width='auto'))

        self._controls_box = HBox(
            [self._navi, self._save_btn],
            layout=Layout(
                display='flex',
                justify_content='center',
                flex_flow='wrap',
                align_items='center'
            )
        )

        self._grid_box = GridMenu(self._im2im_state.grid)

        self._grid_label = HTML(value="<b>LABEL</b>",)
        self._labels_box = VBox(
            children=[self._grid_label, self._grid_box],
            layout=Layout(
                display='flex',
                justify_content='center',
                align_items='center')
        )

        self._save_btn.on_click(self._on_btn_clicked)
        self._grid_box.on_click(self.on_grid_clicked)

        if self._app_state.max_im_number:
            self._set_navi_max_im_number(self._app_state.max_im_number)

        if self._im2im_state.annotations:
            self._grid_box.load(
                self.state_to_widget(self._im2im_state.annotations)
            )

        if self._im2im_state.question_value:
            self._set_label(self._im2im_state.question_value)

        self._im2im_state.subscribe(self._set_label, 'question_value')
        self._im2im_state.subscribe(self._image._draw_image, 'image_path')
        self._im2im_state.subscribe(self.load_menu, 'annotations')

        layout = Layout(
            display='flex',
            justify_content='center',
            align_items='center'
        )

        im2im_display = HBox([
            VBox([self._image, self._controls_box]),
            self._labels_box
        ], layout=layout)

        super().__init__(
            header=None,
            left_sidebar=None,
            center=im2im_display,
            right_sidebar=None,
            footer=None,
            pane_widths=(6, 4, 0),
            pane_heights=(1, 1, 1))

    @debug_output.capture(clear_output=False)
    def load_menu(self, annotations: LabelStore):
        self._grid_box.load(
            self.state_to_widget(annotations)
        )

    @debug_output.capture(clear_output=False)
    def on_navi_clicked(self, index: int):
        if self._on_navi_clicked:
            self._on_navi_clicked(index)

    def _set_navi_max_im_number(self, max_im_number: int):
        self._navi.max_im_num = max_im_number

    def _set_label(self, question_value: str):
        self._grid_label.value = question_value

    def _on_btn_clicked(self, *args):
        if self._on_save_btn_clicked:
            self._on_save_btn_clicked(*args)
        else:
            warnings.warn("Save button click didn't triggered any event.")

    @debug_output.capture(clear_output=False)
    def on_grid_clicked(self, event, value=None):
        if self._on_grid_box_clicked:
            self._on_grid_box_clicked(event, value)
        else:
            warnings.warn("Grid box click didn't triggered any event.")

    def on_client_ready(self, callback):
        self._image.observe_client_ready(callback)

# Internal Cell
def _label_state_to_storage_format(label_state):
    return [Path(k).name for k, v in label_state.items() if v['answer']]

# Internal Cell
def _storage_format_to_label_state(
    storage_format,
    label_names,
    label_dir: str
):
    try:
        path = Path(label_dir)
        return {str(path / label): {
            'answer': label in storage_format} for label in label_names}
    except Exception:
        return {label: {'answer': label in storage_format} for label in label_names}

# Internal Cell

class Im2ImAnnotatorController:
    debug_output = Output(layout={'border': '1px solid black'})

    def __init__(
        self,
        app_state: AppWidgetState,
        im2im_state: Im2ImState,
        storage: JsonLabelStorage,
        input_item=None,
        output_item=None,
        question=None,
    ):
        self._app_state = app_state
        self._im2im_state = im2im_state
        self._storage = storage

        self.input_item = input_item
        self.output_item = output_item

        self.images = self._storage.get_im_names(None)
        self._app_state.max_im_number = len(self.images)

        self.labels = self._storage.get_labels()
        self.labels_num = len(self.labels)

        # Tracks the app_state.index history
        self._last_index = 0

        if question:
            self._im2im_state.question_value = (f'<center><p style="font-size:20px;">'
                                                f'{question}</p></center>')

    def _update_im(self):
        # print('_update_im')
        index = self._app_state.index
        self._im2im_state.image_path = str(self.images[index])

    def _update_state(self, change=None):  # from annotations
        # print('_update_state')
        image_path = self._im2im_state.image_path

        if not image_path:
            return
        tmp_annotations = LabelStore()
        if image_path in self._storage:
            current_annotation = self._storage.get(str(image_path)) or {}
            tmp_annotations.update(
                _storage_format_to_label_state(
                    storage_format=current_annotation or [],
                    label_names=self.labels,
                    label_dir=self._storage.label_dir
                )
            )
            self._im2im_state.annotations = tmp_annotations

    def _update_annotations(self, index: int):  # from screen
        # print('_update_annotations')
        image_path = self._im2im_state.image_path
        if image_path:
            self._storage[image_path] = _label_state_to_storage_format(
                self._im2im_state.annotations
            )

    def save_annotations(self, index: int):  # to disk
        # print('_save_annotations')
        self._update_annotations(index)
        self._storage.save()

    def idx_changed(self, index: int):
        """ On index change save old state
            and update current state for visualisation
        """
        # print('_idx_changed')
        self._app_state.set_quietly('index', index)
        self.save_annotations(self._last_index)
        # update new screen
        self._update_im()
        self._update_state()
        self._last_index = index

    @debug_output.capture(clear_output=False)
    def handle_grid_click(self, event, name):
        # print('_handle_grid_click')
        label_changed = name

        # check if the im2im is using the label as path
        # otherwise it uses the iterable of labels
        if isinstance(self._storage.label_dir, Path):
            label_changed = self._storage.label_dir / name

            if label_changed.is_dir():
                # button without image - invalid
                return

            label_changed = str(label_changed)
        current_label_state = deepcopy(self._im2im_state.annotations)

        # inverse state
        current_label_state[label_changed] = {
            'answer': not self._im2im_state.annotations[label_changed].get('answer', False)
        }

        # change traitlets.Dict entirely to have change events issued
        self._im2im_state.annotations = current_label_state

    def handle_client_ready(self):
        self._update_im()
        self._update_state()

    def to_dict(self, only_annotated: bool) -> dict:
        return self._storage.to_dict(only_annotated)

# Cell

class Im2ImAnnotator(Annotator):
    """
    Represents image-to-image annotator.

    Gives an ability to itarate through image dataset,
    map images with labels for classification,
    export final annotations in json format

    """

    def __init__(
        self,
        project_path: Path,
        input_item: InputImage,
        output_item: Union[OutputImageLabel, OutputLabel],
        annotation_file_path,
        n_rows=None,
        n_cols=None,
        label_autosize=False,
        question=None,
        has_border=False
    ):
        assert input_item, "WARNING: Provide valid Input"
        assert output_item, "WARNING: Provide valid Output"

        self.project_path = project_path
        self.input_item = input_item
        self.output_item = output_item
        app_state = AppWidgetState(uuid=str(id(self)))

        super().__init__(app_state)

        grid = Grid(
            width=output_item.width,
            height=output_item.height,
            n_rows=n_rows,
            n_cols=n_cols
        )

        self.im2im_state = Im2ImState(
            uuid=str(id(self)),
            grid=grid,
            annotations=LabelStore(),
            im_height=input_item.height,
            im_width=input_item.width
        )

        self.storage = JsonLabelStorage(
            im_dir=project_path / input_item.dir,
            label_dir=self._get_label_dir(),
            annotation_file_path=annotation_file_path
        )

        self.controller = Im2ImAnnotatorController(
            app_state=self.app_state,
            im2im_state=self.im2im_state,
            storage=self.storage,
            input_item=input_item,
            output_item=output_item,
            question=question,
        )

        self.state_to_widget = LabelStoreCaster(output_item)

        self.view = Im2ImAnnotatorGUI(
            app_state=self.app_state,
            im2im_state=self.im2im_state,
            state_to_widget=self.state_to_widget,
            label_autosize=label_autosize,
            on_navi_clicked=self.controller.idx_changed,
            on_save_btn_clicked=self.controller.save_annotations,
            on_grid_box_clicked=self.controller.handle_grid_click,
            has_border=has_border,
            fit_canvas=input_item.fit_canvas
        )

        self.app_state.subscribe(self._on_annotation_step_change, 'annotation_step')

        # draw current image and bbox only when client is ready
        self.view.on_client_ready(self.controller.handle_client_ready)

    def _on_annotation_step_change(self, annotation_step: AnnotatorStep):
        disabled_buttons = annotation_step == AnnotatorStep.EXPLORE
        if disabled_buttons:
            self.view._grid_box.clear()
        self.state_to_widget.widgets_disabled = disabled_buttons
        self.view._save_btn.disabled = disabled_buttons

        # forces annotator to have img loaded
        self.controller._update_im()
        self.controller._update_state()
        self.view.load_menu(self.im2im_state.annotations)

    def _get_label_dir(self) -> Union[Iterable[str], Path]:
        if isinstance(self.output_item, OutputImageLabel):
            return self.project_path / self.output_item.dir
        elif isinstance(self.output_item, OutputLabel):
            return self.output_item.class_labels
        else:
            raise ValueError(
                "output_item should have type OutputLabel or OutputImageLabel"
            )

    def __repr__(self):
        display(self.view)
        return ""

    def to_dict(self, only_annotated=True):
        return self.controller.to_dict(only_annotated)