#!/usr/bin/env python
# -*- coding: utf-8; -*-# Copyright (c) 2021 Oracle and/or its affiliates.

# Licensed under the Universal Permissive License v 1.0 as shown at https://oss.oracle.com/licenses/upl/

from typing import Dict
from ads.common import auth as authutil
from ads.data_labeling.reader.dataset_reader import LabeledDatasetReader
from ads.data_labeling.visualizer import image_visualizer, text_visualizer
from IPython.core.display import HTML, Markdown, display

ROWS_TO_RENDER_LIMIT = 50


class DataLabelingAccessMixin:
    """Mixin class for labeled text data."""

    @staticmethod
    def read_labeled_data(
        path: str,
        auth: Dict = None,
        materialize: bool = False,
        encoding: str = "utf-8",
        include_unlabeled: bool = False,
    ):
        """Loads the dataset generated by data labeling service from the object storage.

        Parameters
        ----------
        path : str
            The metadata file path, can be either local or object storage path.
        auth: (dict, optional). Defaults to None.
            The default authetication is set using `ads.set_auth` API. If you need to override the
            default, use the `ads.common.auth.api_keys` or `ads.common.auth.resource_principal` to create appropriate
            authentication signer and kwargs required to instantiate IdentityClient object.
        materialize: (bool, optional). Defaults to False.
            Whether the content of the dataset file should be loaded or it should return the file path to the content.
            By default the content will not be loaded.
        encoding : str, optional
            encoding of files, by default 'utf-8'. Only used for "TEXT" dataset.
        include_unlabeled: (bool, optional). Default to False.
            Whether to load the unlabeled records or not.

        Returns
        -------
        DataFrame
            The labeled dataset in dataframe format.

        Examples
        --------
        >>> import pandas as pd
        >>> import ads
        >>> from ads.common import auth as authutil
        >>> df = pd.DataFrame.ads.read_labeled_data(path="path_to_your_metadata.jsonl",
        ...                                         auth=authutil.api_keys(),
        ...                                         materialize=False)

                                    Path       Content               Annotations
            --------------------------------------------------------------------
            0   path/to/the/content/file                                     yes
            1   path/to/the/content/file                                      no

        >>> df = pd.DataFrame.ads.read_labeled_data(path="path_to_your_metadata.jsonl",
        ...                                         auth=authutil.api_keys(),
        ...                                         materialize=True)

                                    Path         Content             Annotations
            --------------------------------------------------------------------
            0   path/to/the/content/file    file content                     yes
            1   path/to/the/content/file    file content                      no
        """
        auth = auth or authutil.default_signer()

        reader = LabeledDatasetReader.from_export(
            path=path,
            auth=auth,
            encoding=encoding,
            materialize=materialize,
            include_unlabeled=include_unlabeled,
        )

        result = reader.read()
        result.ads._info = reader.info()
        return result

    def render_bounding_box(
        self,
        options: Dict = None,
        content_column: str = "Content",
        annotations_column: str = "Annotations",
        limit: int = ROWS_TO_RENDER_LIMIT,
        path: str = None,
    ) -> None:
        """Renders bounding box dataset. Displays only first 50 rows.

        Parameters
        ----------
        options: dict
            The colors options specified for rendering.
        content_column: Optional[str]
            The column name with the content data.
        annotations_column: Optional[str]
            The column name for the annotations list.
        limit: Optional[int]. Defaults to 50.
            The maximum amount of records to display.
        path: Optional[str]
            Path to save the image with annotations to local directory.

        Returns
        -------
        None
            Nothing.

        Examples
        --------
        >>> import pandas as pd
        >>> import ads
        >>> from ads.common import auth as authutil
        >>> df = pd.DataFrame.ads.read_labeled_data(path="path_to_your_metadata.jsonl",
        ...                                         auth=authutil.api_keys(),
        ...                                         materialize=True)
        >>> df.ads.render_bounding_box(content_column="Content", annotations_column="Annotations")
        """
        items = image_visualizer._df_to_bbox_items(
            self._obj.iloc[0:limit] if self._obj.shape[0] > limit else self._obj,
            content_column=content_column,
            annotations_column=annotations_column,
        )
        image_visualizer.render(items, options=options, path=path)

    def render_ner(
        self,
        options: Dict = None,
        content_column: str = "Content",
        annotations_column: str = "Annotations",
        limit: int = ROWS_TO_RENDER_LIMIT,
    ) -> None:
        """Renders NER dataset. Displays only first 50 rows.

        Parameters
        ----------
        options: dict
            The colors options specified for rendering.
        content_column: Optional[str]
            The column name with the content data.
        annotations_column: Optional[str]
            The column name for the annotations list.
        limit: Optional[int]. Defaults to 50.
            The maximum amount of records to display.

        Returns
        -------
        None
            Nothing.

        Examples
        --------
        >>> import pandas as pd
        >>> import ads
        >>> from ads.common import auth as authutil
        >>> df = pd.DataFrame.ads.read_labeled_data(path="path_to_your_metadata.jsonl",
        ...                                         auth=authutil.api_keys(),
        ...                                         materialize=True)
        >>> df.ads.render_ner(content_column="Content", annotations_column="Annotations")
        """
        items = text_visualizer._df_to_ner_items(
            self._obj.iloc[0:limit] if self._obj.shape[0] > limit else self._obj,
            content_column=content_column,
            annotations_column=annotations_column,
        )
        result_html = text_visualizer.render(items=items, options=options)
        display(Markdown(result_html))
