# -*- coding: utf-8 -*-
"""
Pandas like dataset API
"""
import re
import math
import warnings
import pandas as pd
import numpy as np

from functools import lru_cache
from typing import Dict, List, Union, Callable, Optional

from relevanceai.api.client import BatchAPIClient
from relevanceai.warnings import warn_function_is_work_in_progress

from relevanceai.analytics_funcs import track


class Series(BatchAPIClient):
    """
    Dataset Series Object
    -----------------------------

    A wrapper class for being able to vectorize documents over field

    Parameters
    ----------
    project : str
        Project name on RelevanceAI
    api_key : str
        API key for RelevanceAI
    dataset_id : str
        Data type for the output Series. If not specified, this will be
        inferred from `data`.
        See the :ref:`user guide <basics.dtypes>` for more usages.
    field : str
        The name of the field with the Dataset.

    Examples
    --------
    Assuming the following code has been executed:

    .. code-block::

        from relevanceai import client
        relevanceai.datasets import get_dummy_ecommerce_dataset

        documents = get_dummy_ecommerce_dataset()
        client = Client()

        df = client.Dataset('ecommerce')
        df.create()
        df.insert_documents(documents)

    Retrieve a Series from your dataset

    .. code-block::

        product_images = df['product_image'] # A Series object of every every product image url in dataset

    """

    def __init__(
        self,
        project: str,
        api_key: str,
        firebase_uid: str,
        dataset_id: str,
        field: str,
        image_fields: Optional[List[str]] = None,
        audio_fields: Optional[List[str]] = None,
        highlight_fields: Optional[Dict[str, List]] = None,
        text_fields: Optional[List[str]] = None,
    ):
        super().__init__(project=project, api_key=api_key, firebase_uid=firebase_uid)
        self.project = project
        self.api_key = api_key
        self.firebase_uid = firebase_uid
        self.dataset_id = dataset_id
        self.field = field
        self.image_fields = [] if image_fields is None else image_fields
        self.audio_fields = [] if audio_fields is None else audio_fields
        self.highlight_fields = {} if highlight_fields is None else highlight_fields
        self.text_fields = [] if text_fields is None else text_fields

    def _repr_html_(self):
        if "_vector_" in self.field:
            include_vector = True
        else:
            include_vector = False

        documents = self._get_documents(
            dataset_id=self.dataset_id,
            select_fields=[self.field],
            include_vector=include_vector,
        )
        if include_vector:
            warnings.warn(
                "Displaying using pandas. To get image functionality please install RelevanceAI[notebook]. "
            )
            return pd.json_normalize(documents).set_index("_id")._repr_html_()

        try:
            return self._show_json(documents, return_html=True)
        except Exception as e:
            warnings.warn(
                "Displaying using pandas. To get image functionality please install RelevanceAI[notebook]. "
                + str(e)
            )
            return pd.json_normalize(documents).set_index("_id")._repr_html_()

    @track
    def list_aliases(self):
        fields = self._format_select_fields()
        fields = [field for field in fields if field != "_cluster_"]
        return pd.DataFrame(fields, columns=["_cluster_"])._repr_html_()

    def _format_select_fields(self):
        fields = [
            field
            for field in self.datasets.schema(self.dataset_id)
            if "_cluster_" in field
        ]
        return fields

    def _show_json(self, documents, **kw):
        from jsonshower import show_json

        if not self.text_fields:
            text_fields = pd.json_normalize(documents).columns.tolist()
        else:
            text_fields = self.text_fields

        return show_json(
            documents,
            image_fields=self.image_fields,
            audio_fields=self.audio_fields,
            highlight_fields=self.highlight_fields,
            text_fields=text_fields,
            **kw,
        )

    @track
    def sample(
        self,
        n: int = 1,
        frac: float = None,
        filters: Optional[list] = None,
        random_state: int = 0,
        include_vector: bool = True,
        output_format="pandas",
    ):
        """
        Return a random sample of items from a dataset.

        Parameters
        ----------
        n : int
            Number of items to return. Cannot be used with frac.
        frac: float
            Fraction of items to return. Cannot be used with n.
        filters: list
            Query for filtering the search results
        random_state: int
            Random Seed for retrieving random documents.

        Example
        -------
        .. code-block::

            from relevanceai import client

            client = Client()

            df = client.Dataset(dataset_id)
            df.sample(n=3)

        """
        filters = [] if filters is None else filters

        select_fields = [self.field]

        if frac and n:
            raise ValueError("Only one of n or frac can be provided")

        if frac:
            if frac > 1 or frac < 0:
                raise ValueError("Fraction must be between 0 and 1")
            n = math.ceil(
                self.get_number_of_documents(self.dataset_id, filters=filters) * frac
            )

        documents = self.datasets.documents.get_where(
            dataset_id=self.dataset_id,
            filters=filters,
            page_size=n,
            random_state=random_state,
            is_random=True,
            select_fields=select_fields,
            include_vector=include_vector,
        )["documents"]

        if output_format == "json":
            return documents
        elif output_format == "pandas":
            return pd.DataFrame.from_dict(documents, orient="records")
        else:
            raise ValueError("Incorrect output format")

    head = sample

    @track
    def all(
        self,
        chunksize: int = 1000,
        filters: Optional[List] = None,
        sort: Optional[List] = None,
        include_vector: bool = True,
        show_progress_bar: bool = True,
    ):
        filters = [] if filters is None else filters
        sort = [] if sort is None else sort

        select_fields = [self.field] if isinstance(self.field, str) else self.field
        return self._get_all_documents(
            dataset_id=self.dataset_id,
            chunksize=chunksize,
            filters=filters,
            sort=sort,
            select_fields=select_fields,
            include_vector=include_vector,
            show_progress_bar=show_progress_bar,
        )

    @track
    def apply(
        self,
        func: Callable,
        output_field: str,
        axis: int = 0,
    ):
        """
        Apply a function along an axis of the DataFrame.

        Objects passed to the function are Series objects whose index is either the DataFrame’s index (axis=0) or the DataFrame’s columns (axis=1). By default (result_type=None), the final return type is inferred from the return type of the applied function. Otherwise, it depends on the result_type argument.


        .. note::
            We recommend using the bulk_apply functionality
            if you are looking to have faster processing.

        Parameters
        --------------
        func: function
            Function to apply to each document

        axis: int
            Axis along which the function is applied.
            - 9 or 'index': apply function to each column
            - 1 or 'columns': apply function to each row

        output_field: str
            The field from which to output

        Example
        ---------------
        .. code-block::

            from relevanceai import Client

            client = Client()

            dataset_id = "sample_dataset_id"
            df = client.Dataset(dataset_id)

            df["sample_1_label"].apply(lambda x: x + 3, output_field="output_field")

        """
        if axis == 1:
            raise ValueError("We do not support column-wise operations!")

        def bulk_fn(documents):
            for d in documents:
                try:
                    if self.is_field(self.field, d):
                        self.set_field(
                            output_field, d, func(self.get_field(self.field, d))
                        )
                except Exception as e:
                    continue
            return documents

        return self.pull_update_push(
            self.dataset_id, bulk_fn, select_fields=[self.field]
        )

    @track
    def bulk_apply(
        self,
        bulk_func: Callable,
        retrieve_chunksize: int = 100,
        max_workers: int = 8,
        filters: Optional[list] = None,
        select_fields: Optional[list] = None,
        show_progress_bar: bool = True,
        use_json_encoder: bool = True,
    ):
        """
        Apply a bulk function along an axis of the DataFrame.

        Parameters
        ------------
        bulk_func: function
            Function to apply to a bunch of documents at a time
        retrieve_chunksize: int
            The number of documents that are received from the original collection with each loop iteration.
        max_workers: int
            The number of processors you want to parallelize with
        max_error: int
            How many failed uploads before the function breaks
        json_encoder : bool
            Whether to automatically convert documents to json encodable format
        axis: int
            Axis along which the function is applied.
            - 9 or 'index': apply function to each column
            - 1 or 'columns': apply function to each row

        Example
        ---------
        .. code-block::

            from relevanceai import Client
            client = Client()

            df = client.Dataset("sample_dataset_id")

            def update_documents(documents):
                for d in documents:
                    d["value"] = 10
                return documents

            df.bulk_apply(update_documents)
        """
        filters = [] if filters is None else filters
        select_fields = [] if select_fields is None else select_fields

        return self.pull_update_push(
            self.dataset_id,
            bulk_func,
            retrieve_chunk_size=retrieve_chunksize,
            max_workers=max_workers,
            filters=filters,
            select_fields=select_fields,
            show_progress_bar=show_progress_bar,
            use_json_encoder=use_json_encoder,
        )

    @track
    def numpy(self) -> np.ndarray:
        """
        Iterates over all documents in dataset and returns all numeric values in a numpy array.

        Parameters
        ---------
        None

        Returns
        -------
        vectors: np.ndarray
            an array/matrix of all numeric values selected

        Example
        ---------------
        .. code-block::

            from relevanceai import Client

            client = Client()

            dataset_id = "sample_dataset_id"
            df = client.Dataset(dataset_id)

            field = "sample_field"
            arr = df[field].numpy()
        """
        documents = self._get_all_documents(self.dataset_id, select_fields=[self.field])
        vectors = self.get_field_across_documents(self.field, documents)
        vectors = np.array(vectors)
        return vectors

    @track
    def value_counts(
        self,
        normalize: bool = False,
        ascending: bool = False,
        sort: bool = False,
        bins: Optional[int] = None,
    ):
        """
        Return a Series containing counts of unique values (or values with in a range if bins is set).

        Parameters
        ----------
        normalize : bool, default False
            If True then the object returned will contain the relative frequencies of the unique values.
        ascending : bool, default False
            Sort in ascending order.
        bins : int, optional
            Groups categories into 'bins'. These bins are good for representing groups within continuous series

        Returns
        ----------
        Series

        Example
        ---------------
        .. code-block::

            from relevanceai import Client

            client = Client()

            dataset_id = "sample_dataset_id"
            df = client.Dataset(dataset_id)

            field = "sample_field"
            value_counts_df = df[field].value_counts()
        """
        schema = self.datasets.schema(self.dataset_id)
        dtype = schema[self.field]

        if dtype == "numeric":
            agg_type = dtype
        else:
            agg_type = "category"

        groupby_query = [{"name": self.field, "field": self.field, "agg": agg_type}]
        aggregation = self.services.aggregate.aggregate(
            self.dataset_id,
            groupby=groupby_query,
            page_size=10000,
            asc=ascending,
        )

        total = self.get_number_of_documents(dataset_id=self.dataset_id)
        aggregation = pd.DataFrame(aggregation)

        if normalize:
            aggregation["frequency"] /= total

        if bins is not None:
            vals = []
            for agg in [[agg[0]] * int(agg[1]) for agg in aggregation.values]:
                vals += agg

            vals = pd.cut(vals, bins)

            categories = [
                "({}, {}]".format(interval.left, interval.right) for interval in vals
            ]
            unique_categories = list(set(categories))

            if sort:
                categories = sorted(
                    categories, key=lambda x: float(x.split(",")[0][1:])
                )

            aggregation = pd.DataFrame(
                [categories.count(cat) for cat in unique_categories],
                index=unique_categories,
            )
            aggregation.columns = ["Frequency"]

        return aggregation

    def __getitem__(self, loc: Union[int, str]):
        """
        Indexs a value with a series, usually to get a specific sample from a column in your dataset

        Parameters
        ----------
        loc : int or str, preferably a str
            if int, this operates exactly as indexing a regular python list
            if str, this will be a string corresponding to the _id of the document

        Returns
        ----------
        A single document

        Example
        ---------------
        .. code-block::

            from relevanceai import Client

            client = Client()

            dataset_id = "sample_dataset_id"
            df = client.Dataset(dataset_id)

            field = "sample_field"
            id = "sample_id"
            index = 56

            document = df[field][id]
            document = df[field][index]
        """
        if isinstance(loc, int):
            warnings.warn(
                "Integer selection of dataframe is not stable at the moment. Please use a string ID if possible to ensure exact selection."
            )
            return self.get_documents(loc + 1, select_fields=[self.field])[loc][
                self.field
            ]
        elif isinstance(loc, str):
            return self.datasets.documents.get(self.dataset_id, loc)[self.field]
        raise TypeError("Incorrect data type! Must be a string or an integer")

    @lru_cache(maxsize=1)
    def _get_pandas_series(self):
        documents = self._get_all_documents(
            dataset_id=self.dataset_id,
            select_fields=[self.field],
            include_vector=False,
            show_progress_bar=True,
        )

        try:
            df = pd.DataFrame(documents)
            df.set_index("_id", inplace=True)
            return df.squeeze()
        except KeyError:
            raise Exception("No documents found")

    def __getattr__(self, attr):
        series = self._get_pandas_series()
        try:
            return getattr(series, attr)
        except SyntaxError:
            raise AttributeError(f"'{attr}' is an invalid attribute")
