"""Copyright 2020 Sarus SAS.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.

Sarus library to leverage sensitive data without revealing them

This lib contains classes and method to browse,
learn from & explore sensitive datasets.
It connects to a Sarus server, which acts as a gateway, to ensure no
results/analysis coming out of this lib are sensitive.
"""
from __future__ import annotations  # noqa: F407

import base64
import datetime
import decimal
import getpass
import hashlib
import io
import json
import logging
import os
import pickle as pkl
import re
import sys
import tarfile
import tempfile
import textwrap
import time
import typing as t
import warnings
import webbrowser
from dataclasses import dataclass
from io import BytesIO
from typing import Any, Dict, List, Optional

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq
import sarus_data_spec
from PIL import Image
from requests import Session

import sarus.typing as srt
from sarus.dataspec_wrapper import DataSpecWrapper
from sarus.manager.dataspec_api import compile_dataspec, pull_dataspec_graph
from sarus.manager.typing import SDKManager

with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    import sarus_data_spec.protobuf as sp
    import sarus_data_spec.typing as st
    from sarus_data_spec.constants import (
        ARROW_TASK,
        CACHE,
        CACHE_PATH,
        SCHEMA_TASK,
    )
    from sarus_data_spec.context import push_global_context
    from sarus_data_spec.dataspec_rewriter.simple_rules import attach_variant
    from sarus_data_spec.dataspec_validator.privacy_limit import (
        DeltaEpsilonLimit,
    )
    from sarus_data_spec.protobuf.utilities import (
        deserialize,
        dict_deserialize,
        serialize,
    )
    from sarus_data_spec.status import last_statuses, ready
    from sarus_data_spec.transform import extract, select_sql
    from sarus_data_spec.variant_constraint import (
        mock_constraint,
        pep_constraint,
    )

    from sarus.context.local_sdk import LocalSDKContext
    from sarus.pandas.dataframe import DataFrame

    try:
        import tensorflow as tf

        import sarus.legacy.tensorflow as sltf
    except ModuleNotFoundError:
        pass  # error message in sarus_data_spec.typing


@dataclass
class Synthetic:
    """Simple class to represent synthetic data."""

    data: Optional[pa.Table] = None
    rows: Optional[int] = None


class DatasetList(list):
    """Custom class to lazily fetch a Dataset when indexing."""

    def __init__(
        self, iterable: t.Iterable[ShallowDataset], client: Client
    ) -> None:
        super().__init__(iterable)
        self.client = client

    def __getitem__(self, i: int) -> Dataset:
        shallow_dataset = super().__getitem__(i)
        slugname = shallow_dataset.name
        return self.client.dataset(slugname=slugname)

    def __setitem__(self, i: int, value: t.Any) -> None:
        raise NotImplementedError("A DatasetList is immutable.")

    def append(self, value: t.Any) -> None:
        raise NotImplementedError("A DatasetList is immutable.")

    def insert(self, i: int, value: t.Any) -> None:
        raise NotImplementedError("A DatasetList is immutable.")

    def extend(self, other: DatasetList) -> None:
        raise NotImplementedError("A DatasetList is immutable.")


class ShallowDataset:
    """Shallow dataset that is merely used as a placeholder for slugname."""

    def __init__(self, id: int, slugname: str):
        self.id = id
        self.name = slugname

    def __repr__(self) -> str:
        return f"<Sarus Dataset slugname={self.name} id={self.id}>"

    @classmethod
    def _from_dict(cls, data: dict) -> ShallowDataset:
        # `data` is the JSON returned by calling GET /datasets/
        return cls(data.get("id"), data.get("name"))


class Dataset(DataSpecWrapper[t.Iterator[pa.RecordBatch]]):
    """A class representing a Sarus Dataset.

    This class is the interface to the protected data. It enables to inspect the
    Sarus dataset metadata, manipulate synthetic data locally, prepare
    processing steps and identify the dataset for executing remote private
    queries.

    Args:
        id: The dataset id.
        client: The Sarus client where the dataset is defined.
        type_metadata: A serialized json holding the dataset metadata.
        marginals: A serialized json holding the dataset marginals.
        human_description: A short human readable description.
        policy:
        transforms:
            A list of transformations to be applied to the original dataset.
    """

    def __init__(
        self,
        id: int,
        client: "Client",
        dataspec: st.DataSpec,
        is_bigdata: Optional[bool] = None,
        type_metadata: Optional[str] = None,
        human_description: Optional[str] = None,
        marginals: Optional[str] = None,
        policy: Optional[dict] = None,
        synthetic: Dict[str, Synthetic] = None,
        transforms: List = [],
    ):
        self.id: int = id
        self.client: "Client" = client
        self.is_bigdata: Optional[bool] = is_bigdata
        self.name = dataspec.properties().get("slugname")
        self.desc = f"<Sarus Dataset slugname={self.name} id={self.id}>"
        self._set_dataspec(dataspec)
        self.type_metadata: Optional[Dict[str, Any]] = None
        if type_metadata is not None:
            self.type_metadata = json.loads(type_metadata)
        self.human_description: Optional[str] = human_description
        if synthetic is None:
            self._synthetic = dict(original=Synthetic(), encoded=Synthetic())
        else:
            self._synthetic = synthetic
        self.marginals: Optional[Dict[str, Any]] = None
        if marginals is not None:
            self.marginals = json.loads(marginals)

        self.policy: Optional[dict] = policy

        # We set `unbatch` as the default first transform
        # for a better user experience.
        # TODO remove the default batching on the API side
        if len(transforms) == 0:
            transforms.append(("unbatch", {}))
        self.transforms = transforms

    @classmethod
    def _from_dict(cls, admin_ds: dict, client: Client) -> Dataset:
        """Get a dataset from the json data sent by the server.

        Deserializes the protobufs contained in the response and registers them
        in the storage.

        Args:
            admin_ds (dict): JSON data returned by calling GET /datasets/<id>
            client (Client): client used to get information from the server

        Returns:
            Dataset
        """
        dataspec_id = admin_ds.get("id")
        serialized_dataspec = admin_ds.get("dataspec")
        slugname = admin_ds.get("name")
        manager: SDKManager = client.context().manager()

        if serialized_dataspec is None:
            # TODO /datasets doesn't return the DataSpec def
            raise ValueError(f"{slugname} not returned by the server.")

        # Deserialize objects from the JSON response
        dataspec: st.Dataset = (
            client.context()
            .factory()
            .create(
                dict_deserialize(serialized_dataspec["dataset"]), store=False
            )
        )
        pull_dataspec_graph(client=client, uuid=dataspec.uuid())
        manager.set_remote(dataspec)

        syn_dataset_uuid, _ = compile_dataspec(
            client=client,
            uuid=dataspec.uuid(),
            constraint_kind=st.ConstraintKind.SYNTHETIC,
        )

        pull_dataspec_graph(client=client, uuid=syn_dataset_uuid)

        syn_dataset: st.Dataset = (
            client.context().storage().referrable(syn_dataset_uuid)
        )

        # Update statuses
        dataspec.status(task_names=[ARROW_TASK, SCHEMA_TASK])
        syn_dataset.status(task_names=[ARROW_TASK, SCHEMA_TASK])

        # Create mock
        mock_dataset: st.Dataset = extract(size=dataspec.manager()._mock_size)(
            syn_dataset
        )
        mock_constraint(mock_dataset)
        attach_variant(dataspec, mock_dataset, st.ConstraintKind.MOCK)

        # TODO fetch the PEP token from the API
        # (it is lazily computed on the server)
        h = hashlib.md5()
        h.update(dataspec.protobuf().SerializeToString())
        pep_token = h.hexdigest()
        pep_constraint(
            dataspec=dataspec,
            token=pep_token,
            required_context=[],
            privacy_limit=DeltaEpsilonLimit({0.0: 0.0}),
        )

        # Attach default delta to the source
        sources_ds = dataspec.sources(sp.type_name(sp.Dataset))
        source = sources_ds.pop()
        manager.set_admin_ds(source, admin_ds)

        return cls(
            id=dataspec_id,
            dataspec=dataspec,
            client=client,
            type_metadata=admin_ds.get("type_metadata"),
            marginals=admin_ds.get("marginals"),
            human_description=admin_ds.get("human_description"),
            is_bigdata=admin_ds.get("is_bigdata"),
        )

    def _split(
        self, split: t.Union[str, List[str]]
    ) -> t.Union[Dataset, List[Dataset]]:
        """Get subsets of the Sarus dataset.

        The `split` parameter is a string that specifies a slice of the dataset
        to select. The format is inspired by the format used by Tensorflow
        datasets slices.

        A slice consist of a `start` and an `end` separated by a semi-colon
        wrapped in square brackets. The `start` and `end` are integers,
        optionally followed by a percent sign. If the percent sign is omitted,
        the number indicates a row number. If the percent sign is present, the
        number indicates a proportion of a the dataset to select, the number of
        rows is rounded to the closest lower integer.

        Similar to python slices, an empty `start` value is considered to be the
        first row and an empty `end` value is considered to be the last row.


        >>> sarus_ds.split("[10:40]")  # select rows 10 to 39
        >>> sarus_ds.split("[10:70%]")  # select the first 70% rows but skip the first 10 rows
        >>> sarus_ds.split("[:-100]")  # select all but the last 100 rows

        Args:
            split (str): specifies the portion of the dataset to select.

        Returns:
            Dataset: The selected split as a `Sarus Dataset`.

        """
        split_list = [split] if isinstance(split, str) else split
        split_pattern = r"^\[(?P<start>|-?\d+%?):(?P<end>|-?\d+%?)\]$"
        split_re = re.compile(split_pattern)

        out_datasets = []
        for spl in split_list:
            match = split_re.match(spl)
            if match is None:
                raise ValueError(
                    f"Could not match the split pattern {spl}. "
                    "Expected pattern is `[start(%):end(%)]`."
                )

            # Get start and end rows
            start_str, end_str = match.group("start"), match.group("end")
            if self.marginals is None:
                raise ValueError("Cannot split dataset as marginals is None.")
            ds_size = int(self.marginals["rows"])
            if start_str == "":
                start = 0
            else:
                num = (
                    int(start_str[:-1]) * ds_size // 100
                    if start_str.endswith("%")
                    else int(start_str)
                )
                start = ds_size + num if num < 0 else num

            if end_str == "":
                end = ds_size
            else:
                num = (
                    int(end_str[:-1]) * ds_size // 100
                    if end_str.endswith("%")
                    else int(end_str)
                )
                end = ds_size + num if num < 0 else num

            # Check consistency
            if start >= end:
                raise ValueError(
                    f"Start ({start}) is greater that end ({end}) "
                    f"in split {spl}"
                )
            if start < 0:
                raise ValueError(
                    f"Start ({start}) is lower that 0 " f"in split {spl}"
                )
            if end > ds_size:
                raise ValueError(
                    f"End ({end}) is greater that the dataset size ({ds_size}) "
                    f"in split {spl}"
                )

            split_transform = ("split", {"start": start, "end": end})
            out_datasets.append(self._add_transform(split_transform))

        if isinstance(split, str):
            return out_datasets[0]
        else:
            return out_datasets

    @property
    def features(self) -> Optional[Dict[str, Dict]]:
        """Features of the Sarus dataset and associated metadata.

        Returns:
            Dict[str, Dict]: A dictionary holding metadata where each key
            is a table name and each value is a dict with features.
        """
        if self.type_metadata is None:
            return None
        if len(self.type_metadata) > 1:
            features: Dict[str, Dict] = {
                table["name"]: table["features"]
                for table in self.type_metadata
            }
        else:
            features: Dict[str, Dict] = self.type_metadata[0]["features"]
        return features

    def __repr__(self) -> str:
        return self.desc

    def _ipython_display_(self) -> None:
        """Custom ipython display."""
        print(repr(self))

    def as_tensorflow(
        self,
        max_download_size: Optional[int] = None,
        original: bool = False,
    ) -> sltf.Dataset:
        """Return the corresponding `sarus.tensorflow.Dataset` object.

        This allows to manipulate the Sarus `Dataset` as a Tensorflow dataset.

        Args:
            max_download_size (int, optional): Max number of synthetic data rows
                to download locally. Indicates the number of synthetic data
                rows to download from the Sarus server. It will not download
                more than the maximum number of available synthetic data. If
                `None`, it will download all the synthetic data. If different
                from `None`, all local computations will be done on the local
                synthetic sample so local results will differ from remote
                results.

            original (bool): Returns categories original values.
                If True will return categories as original values. If False,
                will encode categories as integers.

        Returns:
            A sarus_tensorflow.Dataset.
        """
        logging.warning(
            "'as_tensorflow()' uses the old API, not preprocessing."
        )
        if self.is_bigdata:
            raise Exception(
                "`as_tensorflow()` is not supported for this dataset"
            )
        return sltf.Dataset(
            self,
            max_download_size=max_download_size,
            original=original,
        )

    def as_pandas(self) -> DataFrame:
        if self.is_bigdata:
            raise Exception("`as_pandas()` is not supported for this dataset")
        return DataFrame.from_dataspec(self._dataspec)

    def __getitem__(self, item: t.List[str]) -> Dataset:
        return self.table(item)

    def table(self, table_name: t.List[str]) -> Dataset:
        """Get a table from the dataset.

        Args:
            table_name (List[str]): Name of a form ['namespace_1','namespace_2'...,'table_name'].
                One can omit namespaces if table names are not ambiguous.
        Returns:
            Table: Table fitting the given name.
        """

        dataspec_tables = [
            el.to_strings_list()[0] for el in self._dataspec.schema().tables()
        ]
        # remove protection name and add schema_name
        dataspec_tables = [
            [self._dataspec.schema().name(), *(val for val in el[1:])]
            for el in dataspec_tables
        ]

        possible_tables_indices = []
        # iterate over tables and select_ones that have same path
        n_elements = len(table_name)
        for i, table_path in enumerate(dataspec_tables):
            if table_path[-n_elements:] == table_name:
                possible_tables_indices.append(i)

        try:
            assert len(possible_tables_indices) == 1
        except AssertionError:
            if len(possible_tables_indices) == 0:
                raise ValueError(f"Table {table_name} not found")
            else:
                raise ValueError(f"Ambiguous table name {table_name}")

        else:
            table_path = self._dataspec.schema().tables()[
                possible_tables_indices[0]
            ]

        ds_filter = self._dataspec.schema().data_type().get(table_path)
        filtered_ds = sarus_data_spec.transform.filter(filter=ds_filter)(
            self._dataspec
        )
        selected_ds = sarus_data_spec.transform.get_item(path=table_path)(
            filtered_ds
        )

        # create dataset
        table_ds = Dataset.from_dataspec(selected_ds)
        table_ds.desc = f"<Sarus Table name={table_name}>"
        table_ds.is_bigdata = self.is_bigdata

        return table_ds

    def tables(self) -> t.List[t.List[str]]:
        """For given parameters of a Sarus Dataset return a list
        of Sarus Tables."""
        dataspec_tables = [
            el.to_strings_list()[0] for el in self._dataspec.schema().tables()
        ]
        # remove protection name and add schema_name
        dataspec_tables = [
            [self._dataspec.schema().name(), *(val for val in el[1:])]
            for el in dataspec_tables
        ]
        return dataspec_tables

    def sql(self, query: str) -> Dataset:
        """Apply an SQL query to the dataset

        Args:
            query (str): SQL query

        Returns:
            an instance of Sarus Dataset.
        """
        ds = select_sql(query)(self._dataspec)
        ds_wrapper = self.from_dataspec(ds)
        ds_wrapper.name = self.name
        ds_wrapper.id = self.id
        ds_wrapper.is_bigdata = False
        ds_wrapper.desc = f"<Sarus Dataset query='{query}'>"
        return ds_wrapper

    @property
    def epsilon(self) -> float:
        """Retrieve the remaining global privacy budget (epsilon) of the
        current access rule.

        Returns:
            float: The remaining privacy budget (global epsilon) of the access
            rule.
        """
        resp = self.client.session().get(
            f"{self.client.base_url}/datasets/{self.id}",
        )
        if resp.status_code > 200:
            raise Exception(
                f"Error while retrieving the current value of epsilon. "
                f"Gateway answer was: \n{resp}"
            )
        return float(resp.json()["accesses"][0]["current_epsilon"])

    @property
    def max_epsilon(self) -> float:
        """Retrieve the maximum global privacy budget (epsilon) granted by the
        Data preparator, for the current access rule.

        Returns:
            float: The maximum privacy budget (global epsilon) of the access
            rule.
        """
        resp = self.client.session().get(
            f"{self.client.base_url}/datasets/{self.id}",
        )
        if resp.status_code > 200:
            raise Exception(
                f"Error while retrieving the current value of epsilon. "
                f"Gateway answer was: \n{resp}"
            )
        return float(resp.json()["accesses"][0]["max_epsilon"])

    def _fetch_synthetic(
        self,
        rows_number: Optional[int] = None,
        force_refresh: bool = False,
        original: bool = True,
        table_name: Optional[str] = None,
    ) -> pa.Table:
        """Fetch synthetic data as a pyarrow.Table.

        Downloads them if they are not in memory yet or if more rows are
        required. Do nothing if enough data have already been downloaded.

        Args:
            rows_number (int, optional): number of rows to return
            force_refresh (bool): if True, always fetch from server
            original (bool): if False, get categorical values as integers
        """
        if table_name is None and len(self.type_metadata) > 1:
            raise ValueError("Table name needs to be specified in this case.")

        dataset_size = int(self.marginals["rows"])
        if rows_number is None:
            rows_number = dataset_size

        if rows_number < dataset_size:
            logging.warning(
                "Requested `max_download_size` is lower than "
                "total number of remote synth data, so"
                "local operations' results will differ from remote results. \n"
                "Use `max_download_size`=None to locally download all "
                "available synth data"
            )

        if rows_number > dataset_size:
            logging.warning(
                "Cannot satisfy `max_download_size` as there is not enough "
                "remote synth data. \nDownloading max available synth data..."
            )

        rows_number = min(dataset_size, rows_number)

        synthetic_type = "original" if original else "encoded"
        has_enough_synthetic = (
            self._synthetic[synthetic_type].rows is not None
            and rows_number <= self._synthetic[synthetic_type].rows
        )

        if not has_enough_synthetic or force_refresh:
            # Fetch synthetic data
            uri = f"{self.client.base_url}/synthetic_data/{self.id}"
            if table_name is not None:
                uri = f"{uri}/{table_name}"

            resp = self.client.session().get(
                uri,
                stream=True,
                params={
                    "textual_categories": original,
                    "rows_number": rows_number,
                },
            )
            if resp.status_code > 200:
                raise Exception(
                    f"Error while retrieving synthetic data. "
                    f"Gateway answer was: \n{resp}"
                )

            synthetic_table = pq.ParquetFile(io.BytesIO(resp.content)).read()
            self._synthetic[synthetic_type].data = synthetic_table
            self._synthetic[synthetic_type].rows = rows_number

        return self._synthetic[synthetic_type].data.slice(length=rows_number)

    def _synthetic_as_pd_dataframe(
        self,
        rows_number: Optional[int] = None,
        force_refresh: bool = False,
        original: bool = True,
        table_name: Optional[str] = None,
    ):
        """Return synthetic data as a pandas.DataFrame.

         Args:
            rows_number (int, optional): number of rows to return
            force_refresh (bool): if True, always fetch from server
            original (bool): if False, get categorical values as integers

        Returns:
            pandas.DataFrame: synthetic data
        """
        if table_name is None and len(self.type_metadata) > 1:
            raise ValueError("Table name needs to be specified in this case.")

        synthetic_data = self._fetch_synthetic(
            rows_number=rows_number,
            force_refresh=force_refresh,
            original=original,
        )

        synthetic_df = synthetic_data.to_pandas()

        # Convert images encoded as bytes to 3d arrays
        if table_name is None:
            features = self.features
        else:
            features = self.features[table_name]

        for feature in features:
            dtype = list(feature["type"].keys())[0]
            if dtype == "image":
                synthetic_df[feature["name"]] = synthetic_df[
                    feature["name"]
                ].map(Dataset.decode_image)
        # replace min values = -9e+18 by pd.NA and change the data types
        for column in synthetic_df.columns:
            if synthetic_df[column].dtype in [
                np.dtype("int64"),
                np.dtype("int32"),
                np.dtype("int16"),
            ]:
                synthetic_df[column] = synthetic_df[column].astype("Int64")
                synthetic_df.loc[
                    synthetic_df[column] <= np.iinfo(np.int64).min, column
                ] = pd.NA

        return synthetic_df

    @staticmethod
    def _sarus_features_to_tf_spec(
        features: Dict[str, Any],
        original: bool,
    ) -> Dict[str, tf.TensorSpec]:
        """Convert Sarus features to Tensorflow spec."""
        mapping = {
            "categorical": tf.string if original else tf.int16,
            "boolean": tf.bool,
            "integer": tf.int64,
            "real": tf.float32,
            "text": tf.string,
            "datetime": tf.int64,
            "image": tf.uint8,
        }

        def get_tensorspec(feature) -> tf.TensorSpec:
            """Return the tensorflow spec of a Sarus feature."""
            feature_type = list(feature["type"].keys())[0]
            dtype = mapping[feature_type]
            if feature_type == "image":
                width = feature["type"]["image"]["shape"]["width"]
                height = feature["type"]["image"]["shape"]["height"]
                channels = feature["type"]["image"]["shape"]["channel"]
                shape = (None, width, height, channels)
            else:
                shape = (None,)
            return tf.TensorSpec(dtype=dtype, shape=shape)

        return {
            feature["name"]: get_tensorspec(feature) for feature in features
        }

    @staticmethod
    def decode_image(serialized_image: bytes) -> np.ndarray:
        # Need to be careful when decoding images with PIL while encoding with
        # OpenCV. See https://stackoverflow.com/questions/58861577/differences-between-pil-image-open-and-cv2-imdecode
        return np.load(BytesIO(serialized_image))

    @staticmethod
    def _adapt_for_tf(
        batch: Dict[str, Any],
        features: Dict[str, Any],
        original: bool,
    ) -> Dict[str, Any]:
        """Convert python data to tensorflow compatible data.

        Convert datetime.datetime to nanoseconds.
        Convert images serialized as bytes by petastorm to 3d arrays.
        """
        adapted_batch = dict()
        for feature in features:
            dtype = list(feature["type"].keys())[0]
            if dtype == "datetime":

                def decode_datetime(datetime_value: datetime.datetime) -> int:
                    if datetime_value is None:
                        return np.iinfo(np.int64).min
                    return int(datetime_value.timestamp()) * int(1e9)

                adapted_batch[feature["name"]] = tf.constant(
                    list(map(decode_datetime, batch[feature["name"]])),
                    dtype=tf.int64,
                )
            elif dtype == "image":
                adapted_batch[feature["name"]] = tf.stack(
                    list(map(Dataset.decode_image, batch[feature["name"]])),
                    axis=0,
                )

            elif dtype == "text":
                adapted_batch[feature["name"]] = tf.constant(
                    list(map(str, batch[feature["name"]])), dtype=tf.string
                )

            elif dtype == "categorical":
                dtype = tf.string if original else tf.int16
                adapted_batch[feature["name"]] = tf.constant(
                    batch[feature["name"]], dtype=dtype
                )

            elif dtype == "integer":
                # replace None by min_int_32 to avoid crash
                cleaned_batch = [
                    x if x is not None else np.iinfo(np.int64).min
                    for x in batch[feature["name"]]
                ]
                adapted_batch[feature["name"]] = tf.constant(
                    cleaned_batch, dtype=tf.int64
                )

            elif dtype == "real":
                # replace None by NaN to avoid crash
                cleaned_batch = [
                    x if x is not None else float("NaN")
                    for x in batch[feature["name"]]
                ]
                adapted_batch[feature["name"]] = tf.constant(
                    cleaned_batch, dtype=tf.float32
                )

            elif dtype == "boolean":
                adapted_batch[feature["name"]] = tf.constant(
                    batch[feature["name"]], dtype=tf.bool
                )

        return adapted_batch

    def _synthetic_as_tf_dataset(
        self,
        batch_size: int,
        rows_number: Optional[int] = None,
        original: bool = False,
        force_refresh: bool = False,
        table_name: Optional[str] = None,
    ) -> tf.data.Dataset:
        """Return synthetic data as a tensorflow.data.Dataset.

        Args:
            batch_size (int): size of the batches in the dataset
            rows_number (int, optional): number of rows in the dataset
            original (bool): if False, return categories as integers
            force_refresh (bool): if True, always fetch from server

        Returns:
            tensorflow.data.Dataset: synthetic data
        """
        if table_name is None and len(self.type_metadata) > 1:
            raise ValueError("Table name needs to be specified in this case.")

        if force_refresh:
            self._fetch_synthetic(
                force_refresh=force_refresh,
                original=original,
                rows_number=rows_number,
            )

        if table_name is None:
            features = self.features
        else:
            features = self.features[table_name]

        # Generator function iterating pyarrow RecordBatches
        def generator() -> Dict[str, List[Any]]:
            synthetic_table = self._fetch_synthetic(
                original=original, rows_number=rows_number
            )
            for batch in synthetic_table.to_batches(max_chunksize=batch_size):
                yield Dataset._adapt_for_tf(
                    batch=batch.to_pydict(),
                    features=features,
                    original=original,
                )

        tf_signature = Dataset._sarus_features_to_tf_spec(features, original)

        return tf.data.Dataset.from_generator(
            generator, output_signature=tf_signature
        )

    def _plot_marginal_feature(
        self,
        marginal_feature: Dict[str, Any],
        width: float = 1.5,
        heigth: float = 1.5,
    ) -> Optional[str]:
        if "statistics" not in marginal_feature:
            return None

        # text-based representations
        # count for categories
        distrib = marginal_feature["statistics"].get("distribution")
        if distrib:
            html_response = ""
            distrib_s = sorted(distrib, key=lambda x: -x["probability"])
            if len(distrib_s) > 5:
                others_count = len(distrib_s) - 5
                others_sum = sum([x["probability"] for x in distrib_s[5:]])
                distrib_s = distrib_s[0:5]
                distrib_s.append(
                    {
                        "name": f"Other ({others_count})",
                        "probability": others_sum,
                        "class_other": "True",
                    }
                )
            for item in distrib_s:
                html_response += "<div><div class='category "
                if "class_other" in item:
                    html_response += "other"
                html_response += f"''>\
                    {item['name']}\
                  </div>\
                  <div class='number'> {round(100*item['probability'],2)}%\
                  </div>\
                 </div>"
            return html_response

        # Graph-based representation
        _ = plt.figure(figsize=(width, heigth))
        # cumulDistribution for real
        cumul = marginal_feature["statistics"].get("cumulativeDistribution")
        if cumul:
            try:
                plt.fill_between(
                    [vp["value"] for vp in cumul],
                    [vp["probability"] for vp in cumul],
                )
            except Exception:
                pass
        fi = io.BytesIO()
        plt.tight_layout()
        plt.savefig(fi, format="svg")
        plt.clf()
        svg_dta = fi.getvalue()  # this is svg data
        return svg_dta.decode()

    def _marginals_to_html(self, display_type: bool = False) -> str:
        """Return a HTML representation of this dataset.

        To be displayedin a Notebook for example.
        We'd like to render something like: https://www.kaggle.com/fmejia21/
        demographics-of-academy-awards-oscars-winners

        Args:
            display_type (bool): if True, display each column type in the html

        Returns:
            str: HTML representation of the dataset
        """
        css = """<style>
        @import url('https://rsms.me/inter/inter.css');
        @supports (font-variation-settings: normal) {
            html {font-family: 'Inter var', sans-serif; }
            table {font-size: 12px; border-collapse: collapse;}
            td, th {
                border: 1px solid rgb(222, 223, 224);
                font-weight: 500;
                color: rgba(0,0,0,0.7);
                padding: 8px;
                vertical-align:top;
                }
            tr.desc>td>div {
                display: flex; width: 140px;
                flex-wrap: wrap;
                justify-content: space-between;
            }
            tr.desc>td {
                border-bottom-width: 4px;
                }
            div.category {
                width: 70px;
                padding: 4px;
                margin-bottom: 4px;
                color: black;
                }
            div.number {
                padding: 4px;
                margin-bottom: 4px;
                color: rgb(0, 138, 188);
                text-align: right;
            }
            td>div:hover {
                background-color: rgba(0,0,0,0.03);
              }
            tr.synthetic {
                font-family: 'Roboto Mono', Monaco, Consolas, monospace;
              }
            tr.synthetic>td {
                padding: 8px 4px;
                color: rgba(0, 0, 0, 0.7);
              }
            div.other {color: rgba(0, 0, 0, 0.4);!important}
         </style>"""

        table = "<table>\
                <thead><tr>\n"
        if self.type_metadata is None:
            raise ValueError(
                "The dataset has no type_metadata, may be pending..."
            )
        try:
            columns = self.type_metadata["features"]
        except Exception:
            raise Exception(
                "Dataset does not have proper typing yet. Maybe pending..."
            )
        for c in columns:
            table += f"<th>{c['name']}</th>\n"

        table += """</thead></tr>
                <tbody>
                <tr class='desc'>"""
        if self.marginals is None:
            raise ValueError(
                "The dataset has no marginals yet, may be pending..."
            )
        try:
            for c in self.marginals["features"]:
                table += f"<td>{self._plot_marginal_feature(c)}</td>\n"
        except Exception:
            raise Exception(
                "Dataset does not have proper marginals yet. Maybe pending..."
            )
        table += "</tr><tr>"
        if display_type:
            for c in columns:
                table += f"<td>{c['type']}</td>\n"

        table += "</tr></tbody></table>"

        return f"<html>{css}<body>\n \
                   {table}\
                 </body></html>"

    def to_html(
        self, rows_number: int = None, force_refresh: bool = False
    ) -> str:
        """Return synthetic data as html.

        Args:
            rows_number (int): number of rows to display

            force_refresh (bool): if True, does not use cached synthetic data

        Returns:
            str: HTML representation of the synthetic data
        """
        synthetic: pd.DataFrame = self._synthetic_as_pd_dataframe(
            rows_number=rows_number, force_refresh=force_refresh
        )
        image_cols = [
            f["name"]
            for f in self.type_metadata["features"]
            if "image" in f["type"]
        ]

        def _image_formatter(img_array):
            img_array = np.squeeze(img_array)
            image_from_array = Image.fromarray(img_array)
            buffered = BytesIO()
            image_from_array.save(buffered, format="png")
            return (
                f'<img src="data:image/png;base64,'
                f'{base64.b64encode(buffered.getvalue()).decode()}">'
            )

        return synthetic.to_html(
            formatters={c: _image_formatter for c in image_cols},
            escape=False,
        )


class Client:
    """Entry point for the Sarus API client."""

    def _url_validator(self, url):
        """URL validator.

        From https://stackoverflow.com/questions/7160737/
        python-how-to-validate-a-url-in-python-malformed-or-not
        """
        regex = re.compile(
            r"^(?:http|ftp)s?://"  # http:// or https://
            r"(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+(?:[A-Z]{2,6}\.?|"
            r"[A-Z0-9-]{2,}\.?)|"  # domain...
            r"localhost|"  # localhost...
            r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})"  # ...or ip
            r"(?::\d+)?"  # optional port
            r"(?:/?|[/?]\S+)$",
            re.IGNORECASE,
        )
        return re.match(regex, url) is not None

    def __init__(
        self,
        url="http://0.0.0.0:5000",
        google_login=False,
        email=None,
        password=None,
    ):
        # TODO : progress bar self.progress_bar = Progbar(100,
        # stateful_metrics=None)

        if self._url_validator(url):
            self.base_url = url
        else:
            raise Exception("Bad url")
        self._session = Session()
        if google_login:
            self._oidc_login()
        else:
            self._credentials_login(email, password)

        self._context = LocalSDKContext(self)
        self._context.set_sync_policy(srt.SyncPolicy.SEND_ON_VALUE)
        push_global_context(self._context)

    def context(self) -> LocalSDKContext:
        return self._context

    def url(self):
        return self.base_url

    def session(self):
        return self._session

    def save(self, wrapper: srt.DataSpecWrapper, path: str) -> None:
        """Save the representation of a sarus object to a file.

        Args:
            wrapper (DataSpecWrapper): The sarus object to save.
            path (str): The path to save the object to.

        Returns:
            None
        """
        manager: SDKManager = self.context().manager()

        user_ds = wrapper.dataspec(srt.DataSpecVariant.USER_DEFINED)
        synth_ds = wrapper.dataspec(srt.DataSpecVariant.SYNTHETIC)
        mock_ds = wrapper.dataspec(srt.DataSpecVariant.MOCK)

        user_graph = manager.computation_graph(user_ds)
        mock_graph = manager.computation_graph(mock_ds)

        user_refs = (
            set(user_graph["dataspecs"])
            | set(user_graph["transforms"])
            | set(user_graph["attributes"])
            | set(user_graph["variant_constraints"])
        )

        synth_refs = {synth_ds}

        mock_refs = (
            set(mock_graph["dataspecs"])
            | set(mock_graph["transforms"])
            | set(mock_graph["attributes"])
            | set(mock_graph["variant_constraints"])
        )

        refs = user_refs | synth_refs | mock_refs
        serialized_refs = [serialize(ref.protobuf()) for ref in refs]

        file_obj = tarfile.open(path, "w")

        with open("data", "wb") as f:
            pkl.dump(
                {
                    "serialized_obj": serialized_refs,
                    "user_defined": user_ds.uuid(),
                },
                f,
            )
        file_obj.add("data")
        os.remove("data")

        # saving the cache is there is one for the mock
        wrapper._mock_value()
        for last_status in last_statuses(mock_ds, task=CACHE):
            task_stage = last_status.task(CACHE)
            if task_stage and task_stage.ready():
                path_cache = task_stage[CACHE_PATH]
                cache_name = os.path.basename(path_cache)
                file_obj.add(path_cache, arcname=cache_name)
        file_obj.close()

    def load(self, path: str) -> srt.DataSpecWrapper:
        """Load a sarus object from a path.

        Args:
            path (str): The file path to load the sarus object from.

        Returns:
            DataSpecWrapper: The sarus object (eg: DataFrame,model..).
        """
        with tarfile.open(path, mode="r") as tar:
            tar.extract("data")
            with open("data", "rb") as f:
                data = pkl.load(f)
            os.remove("data")
        tar.close()

        refs = data["serialized_obj"]
        user_defined = data["user_defined"]

        factory = self.context().factory()
        storage = self.context().storage()
        # Two passes to have all the correct referrings
        refs = [factory.create(deserialize(ref)) for ref in refs]
        _ = [storage.store(ref) for ref in refs]

        user_ds = self.context().storage().referrable(user_defined)
        mock_ds = user_ds.variant(st.ConstraintKind.MOCK)

        # loading the cache of the mock (only for the dataspec we load)
        with tarfile.open(path, mode="r") as tar:
            tar_filenames = tar.getnames()
            for filename in tar_filenames:
                if f"{mock_ds.uuid()}" in filename:
                    cache_dir = self.context().manager().parquet_dir()
                    path_cache = os.path.join(cache_dir, filename)
                    tar.extract(filename, cache_dir)
                    ready(
                        mock_ds,
                        CACHE,
                        properties={CACHE_PATH: path_cache},
                    )
        tar.close()
        return self.context().wrapper_factory().create(user_ds)

    def _oidc_login(self):
        oidc_login_url = f"{self.base_url}/oidc_login?headless=true"
        try:
            from IPython.display import Javascript, clear_output

            display(  # noqa: F821
                Javascript(f'window.open("{oidc_login_url}");')
            )
            display(clear_output())  # noqa: F821
        except Exception:
            webbrowser.open(oidc_login_url)
        token = getpass.getpass(
            "Logging in via google.\nYou will be redirected to a login page "
            "where you will obtain a token to paste below.\nIf you are not "
            f"redirected automatically, you can visit {oidc_login_url}\n"
        )
        self.session().cookies.set(
            "session", base64.b64decode(token).decode("ascii")
        )
        # just to check that the login is successful
        try:
            self._available_datasets()
        except Exception:
            raise Exception("Error during login: incorrect token")

    def _credentials_login(self, email=None, password=None):
        if email is None:
            raise ValueError("Please enter your email")

        credentials = {}
        credentials["email"] = email

        if password is not None:
            credentials["password"] = password
        else:
            credentials["password"] = getpass.getpass(
                prompt="Password: ", stream=None
            )

        response = self.session().post(
            f"{self.base_url}/login",
            json=credentials,
            headers={"Content-Type": "application/json"},
        )

        if response.status_code == 401:
            raise ValueError("Error during login: incorrect credentials")

        # Let `requests` handle unexpected HTTP status codes.
        response.raise_for_status()

    def _available_datasets(self) -> List[str]:
        """List available datasets.

        Returns:
            List[str]: list of available dataset names.
        """
        request = self.session().get(f"{self.base_url}/datasets")
        return [ds["name"] for ds in request.json()]

    def list_datasets(self) -> DatasetList:
        """List available Sarus datasets.

        Returns:
            DatasetList: List of available Sarus datasets.
        """
        request = self.session().get(f"{self.base_url}/datasets")
        unsorted_datasets = [
            ShallowDataset._from_dict(ds_json)
            for ds_json in request.json()
            if ds_json["status"] == "ready"
        ]
        sorted_datasets = sorted(
            unsorted_datasets,
            key=lambda x: x.name,
        )
        return DatasetList(sorted_datasets, self)

    def dataset(self, slugname: str = None, id: int = None) -> Dataset:
        """Select a dataset from the Sarus Gateway.

        Either `slugname` or `id` should be provided. If both are provided, then
        only the `slugname` is considered.

        Args:
            slugname (str): the slugname of the Dataset to select.
            id (int): the id of the Dataset to select.

        Returns:
            Dataset: The selected Sarus dataset.

        """
        if slugname:
            return self._fetch_dataset_by_name(slugname)
        else:
            return self._fetch_dataset_by_id(id)

    def _fetch_dataset_by_id(self, id: int) -> Dataset:
        """Fetch a dataset from the Sarus Gateway.

        Args:
            id (int): id of the dataset to be fetched

        Returns:
            an instance of Dataset
        """
        request = self.session().get(f"{self.base_url}/datasets/{id}")
        dataset = Dataset._from_dict(request.json(), self)
        return dataset

    def _fetch_dataset_by_name(self, name: str) -> Dataset:
        """Fetch a dataset from the Sarus Gateway.

        Args:
            name (string): name of the dataset to be fetched

        Returns:
            Dataset: an instance of Dataset
        """
        request = self.session().get(
            f"{self.base_url}/datasets/name/{name}",
        )
        dataset = Dataset._from_dict(request.json(), self)
        return dataset

    def query(
        self,
        query: str,
        target_epsilon: Optional[float] = None,
        verbose: bool = True,
        use_old_query=False,
    ) -> int:
        """Execute a SQL query.

        Args:
            query (String): SQL query

            target_epsilon (Optional[float]): Maximum privacy budget (epsilon) to assign to the query.
                If 0, runs on the synthetic data.
                If >0, result is a combination of a query on the synthetic data and a
                Differentially-Private query on the real data.
                If None, a default target epsilon specific to the current user and access rule is used.
                Default target epsilon is 0 if the access is a Differentially-Private access with
                per-user or per-group limit; default value equals to per-query limit if the access is
                a Differentially-Private access with a per-query limit only. Meaning Sarus maximizes result
                accuracy in the given privacy constraints. See user documentation to know more.

            use_old_query (bool): Whether to use the v1 version of the query task.

        Returns:
            int: Id of the task.
        """
        endpoint = "query"
        if use_old_query:
            endpoint = "query_v1"

        payload = {
            "query": query,
        }
        if target_epsilon is not None:
            payload["target_epsilon"] = target_epsilon

        request = self.session().post(
            f"{self.base_url}/{endpoint}",
            json=payload,
        )
        if request.status_code > 200:
            if request.status_code == 403:
                raise ValueError(
                    "Query failed with the following error: Privacy budget "
                    "limit exceeded"
                )
            if request.status_code == 404:
                raise ValueError(
                    "Couldn't run the query: wrong table name or unauthorised "
                    "access"
                )
            raise Exception(
                f"Error while sending a query.\
                                                     Full Gateway answer was:{request}"
            )

        task_id = request.json()["task"]
        dataset_id = request.json()["dataset"]["id"]
        start_eps = request.json()["dataset"]["accesses"][0]["current_epsilon"]
        status = self._poll_query_status(task_id)
        error_message = status.get("error_message", None)
        if error_message is not None:
            raise RuntimeError(
                f"Query failed with the following error:\n"
                f"{textwrap.indent(error_message, '  |')}"
            )
        if verbose:
            ds = self._fetch_dataset_by_id(dataset_id)
            logging.info(
                f"Actual privacy consumption (epsilon): "
                f"{ds.epsilon-start_eps:.03f}"
            )

        if status["status"] == "SUCCESS":
            status_with_pyobjs = self._convert_to_result_with_pyobjs(status)

            return status_with_pyobjs
        else:
            return status

    @staticmethod
    def _convert_to_result_with_pyobjs(
        status: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Convert the query result values to native Python objects using the
           description field.

        Args:
            status (dict): the status response for a given completed query task returned by the server.

        Returns:
            dict: a new status dict with the initialized Python objects.
        """
        result = []
        for row in status["result"]:
            for idx, col_desc in enumerate(status["description"]):
                _, class_name, _, _, _, _, _ = col_desc
                if class_name == str(decimal.Decimal):
                    row[idx] = decimal.Decimal(row[idx])
            result.append(row)
        return {
            **status,
            "result": result,
        }

    def _abort_training(self, id: int):
        """Abort a training on the Sarus Gateway.

        Args:
            id (int): id of the task to abort (provided by the fit method).
        """
        resp = self.session().delete(
            f"{self.base_url}/training_tasks/{id}/abort",
        )
        if resp.status_code != 204:
            raise Exception(
                f"Error while trying to abort task:\n{resp.content}"
            )

    def _training_status(self, id: int) -> dict:
        """Fetch a dataset from the Sarus Gateway.

        Args:
            id (int): id of the task to be queried. It was provided by the fit
            method

        Returns:
            dict: a dict with the status of a training tasks
        """
        request = self.session().get(
            f"{self.base_url}/training_tasks/{id}",
        )
        return request.json()

    def _poll_training_status(self, id: int, timeout: int = 1000) -> dict:
        """Poll & display the status of a training task.

        Args:
            id (int): id of the task to be queried. It was provided by the fit
            method

            timeout (int): in seconds

        Returns:
            dict: The training status at the end of the task

        Raises:
            TimeoutError: if timeout is reached before the training finishes
        """
        offset = 0
        elapsed_time = 0.0
        while elapsed_time < timeout:
            elapsed_time += 0.5
            request = self.session().get(
                f"{self.base_url}/training_tasks/{id}",
                params=dict(offset=offset),
            )
            response_dict = request.json()
            offset = response_dict.get("next_offset", 0)
            if "progress" in response_dict:
                progress = base64.b64decode(
                    response_dict["progress"].encode("ascii")
                ).decode("ascii")
                if progress:
                    sys.stdout.write(progress)
            else:
                # this is the end of the training
                sys.stdout.write("\n")
                return response_dict
            sys.stdout.flush()
            time.sleep(0.5)
        raise TimeoutError(
            "Timeout reached while waiting for the model training to finish."
        )

    def _poll_query_status(self, id: int, timeout: int = 1000) -> dict:
        """Poll & display the status of a query task.

        Args:
            id (int): id of the task to be queried. It was provided by the fit
            method

            timeout (int): in seconds

        Returns:
            dict: The query status at the end of the task

        Raises:
            TimeoutError: if timeout is reached before the query finishes
        """
        offset = 0
        elapsed_time = 0.0
        while elapsed_time < timeout:
            elapsed_time += 0.5
            request = self.session().get(
                f"{self.base_url}/query_tasks/{id}",
                params=dict(offset=offset),
            )
            response_dict = request.json()
            status = response_dict["status"]
            if status != "PENDING":
                return response_dict
            time.sleep(0.5)
        raise TimeoutError(
            "Timeout reached while waiting for the model training to finish."
        )

    def _fetch_model(self, id: int) -> tf.keras.Model:
        """Fetch a trained model from the Sarus Gateway.

        Args:
            id (int): id of the task to be queried. It was provided by the fit
                method

        Returns:
            tf.keras.Model: a Keras model
        """
        response = self.session().get(
            f"{self.base_url}/models/{id}",
        )
        # apparently we need to save to a temp file
        # https://github.com/keras-team/keras/issues/9343
        with tempfile.TemporaryDirectory() as _dir:
            f = tarfile.open(fileobj=io.BytesIO(response.content), mode="r:gz")
            f.extractall(_dir, members=None)

            return tf.keras.models.load_model(_dir)


def _save_model(model: tf.keras.Model) -> bytes:
    """Convert a keras Model to compressed archive format."""
    with tempfile.TemporaryDirectory() as _dir:
        model.save(_dir)
        with tempfile.TemporaryDirectory() as _second_dir:
            path = os.path.join(_second_dir, "tmpzip")
            with tarfile.open(path, mode="w:gz") as archive:
                archive.add(_dir, recursive=True, arcname="")
            with open(path, "rb") as f:
                ret = f.read()
                return ret
