# This file has been generated - DO NOT MODIFY
# API Version : 10.2.0


import logging
import shutil
import tempfile
import warnings
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, TypeVar, Union
from uuid import UUID

import numpy as np
import pandas as pd

from avatars.api_autogenerated import Auth as _Auth
from avatars.api_autogenerated import Compatibility as _Compatibility
from avatars.api_autogenerated import Datasets as _Datasets
from avatars.api_autogenerated import Health as _Health
from avatars.api_autogenerated import Jobs as _Jobs
from avatars.api_autogenerated import Metrics as _Metrics
from avatars.api_autogenerated import Reports as _Reports
from avatars.api_autogenerated import Stats as _Stats
from avatars.api_autogenerated import Users as _Users
from avatars.arrow_utils import (
    ArrowDatasetBuilder,
    DataSourceItem,
    DataSourceItems,
    FileLike,
    is_file_like,
)
from avatars.client import ApiClient
from avatars.constants import DEFAULT_TIMEOUT
from avatars.exceptions import Timeout
from avatars.models import AdviceJob  # noqa: F401
from avatars.models import AdviceJobCreate  # noqa: F401
from avatars.models import AvatarizationJob  # noqa: F401
from avatars.models import AvatarizationJobCreate  # noqa: F401
from avatars.models import AvatarizationMultiTableJob  # noqa: F401
from avatars.models import AvatarizationMultiTableJobCreate  # noqa: F401
from avatars.models import AvatarizationWithTimeSeriesJob  # noqa: F401
from avatars.models import AvatarizationWithTimeSeriesJobCreate  # noqa: F401
from avatars.models import ClusterStats  # noqa: F401
from avatars.models import CompatibilityResponse  # noqa: F401
from avatars.models import Contributions  # noqa: F401
from avatars.models import CreateDataset  # noqa: F401
from avatars.models import CreateUser  # noqa: F401
from avatars.models import Dataset  # noqa: F401
from avatars.models import ExplainedVariance  # noqa: F401
from avatars.models import ForgottenPasswordRequest  # noqa: F401
from avatars.models import GenericJob  # noqa: F401
from avatars.models import Login  # noqa: F401
from avatars.models import LoginResponse  # noqa: F401
from avatars.models import PatchDataset  # noqa: F401
from avatars.models import PrivacyMetricsGeolocationJob  # noqa: F401
from avatars.models import PrivacyMetricsGeolocationJobCreate  # noqa: F401
from avatars.models import PrivacyMetricsJob  # noqa: F401
from avatars.models import PrivacyMetricsJobCreate  # noqa: F401
from avatars.models import PrivacyMetricsMultiTableJob  # noqa: F401
from avatars.models import PrivacyMetricsMultiTableJobCreate  # noqa: F401
from avatars.models import PrivacyMetricsWithTimeSeriesJob  # noqa: F401
from avatars.models import PrivacyMetricsWithTimeSeriesJobCreate  # noqa: F401
from avatars.models import Projections  # noqa: F401
from avatars.models import Report  # noqa: F401
from avatars.models import ReportCreate  # noqa: F401
from avatars.models import ReportFromDataCreate  # noqa: F401
from avatars.models import ReportGeolocationPrivacyCreate  # noqa: F401
from avatars.models import ResetPasswordRequest  # noqa: F401
from avatars.models import SignalMetricsJob  # noqa: F401
from avatars.models import SignalMetricsJobCreate  # noqa: F401
from avatars.models import SignalMetricsWithTimeSeriesJob  # noqa: F401
from avatars.models import SignalMetricsWithTimeSeriesJobCreate  # noqa: F401
from avatars.models import User  # noqa: F401
from avatars.models import (
    AvatarizationPipelineCreate,
    AvatarizationPipelineResult,
    ColumnDetail,
    ColumnType,
    FileType,
    JobStatus,
    PrivacyMetricsParameters,
    SignalMetricsParameters,
)

logger = logging.getLogger(__name__)

DEFAULT_RETRY_TIMEOUT = 60
MAX_ROWS_PER_FILE = 1_000_000
MAX_BYTES_PER_FILE = 100 * 1024 * 1024  # 100 MB

PARQUET_MAGIC_BYTES = b"PAR1"

T = TypeVar("T")


def to_column_type(s: str) -> ColumnType:
    if "float" in s:
        return ColumnType.float
    if "int" in s:
        return ColumnType.int
    if "bool" in s:
        return ColumnType.bool
    if "datetime" in s:
        return ColumnType.datetime
    if "object" in s or s == "category" or s == "str":
        return ColumnType.category
    raise TypeError(f"Unknown column type: '{s}'")


def from_column_type(s: ColumnType) -> str:
    # We don't want to use 'category' here, as we store it as 'object' in the API
    if s == ColumnType.category:
        return "object"
    else:
        return str(s.value)

    raise TypeError(f"Unknown column type: '{s}'")


class Auth:
    def __init__(self, client: ApiClient) -> None:
        self.client = client

    def login(
        self,
        request: Login,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> LoginResponse:
        """Login the user."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            request,
        ]

        return _Auth(self.client).login(*args, **kwargs)

    def refresh(
        self,
        token: str,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> LoginResponse:

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            token,
        ]

        return _Auth(self.client).refresh(*args, **kwargs)

    def forgotten_password(
        self,
        request: ForgottenPasswordRequest,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> Any:

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            request,
        ]

        return _Auth(self.client).forgotten_password(*args, **kwargs)

    def reset_password(
        self,
        request: ResetPasswordRequest,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> Any:

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            request,
        ]

        return _Auth(self.client).reset_password(*args, **kwargs)


class Compatibility:
    def __init__(self, client: ApiClient) -> None:
        self.client = client

    def is_client_compatible(
        self,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> CompatibilityResponse:
        """Verify if the client is compatible with the API."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = []

        return _Compatibility(self.client).is_client_compatible(*args, **kwargs)


class Datasets:
    def __init__(self, client: ApiClient) -> None:
        self.client = client

    def create_dataset_from_stream(
        self,
        request: Optional[FileLike] = None,
        name: Optional[str] = None,
        source: Optional[
            DataSourceItem
        ] = None,  # optional because we still have to support the old way # TODO: Remove once deprecated
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> Dataset:
        """Create a dataset by streaming chunks of the dataset.

        DEPRECATED: Please use create_dataset instead.
        """

        warnings.warn(
            DeprecationWarning(
                "create_dataset_from_stream is deprecated. Use create_dataset instead."
            )
        )

        _source: Optional[Union[str, FileLike]] = request or source
        return self.create_dataset(source=_source, name=name, timeout=timeout)

    def create_dataset(
        self,
        request: Optional[FileLike] = None,  # TODO: Remove once deprecated
        name: Optional[str] = None,
        source: Optional[
            DataSourceItems
        ] = None,  # optional because we still have to support the old way
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> Dataset:
        """Create a dataset from file upload."""

        if request:
            warnings.warn(
                DeprecationWarning("request is deprecated. Use file instead.")
            )

        if request is not None and source is not None:
            raise ValueError("You cannot pass both request and source.")

        if request is None and source is None:
            raise ValueError("You need to pass in a source.")

        _source = request or source

        if _source is None:
            raise ValueError("You need to pass in a source.")

        ds, inferred_type = ArrowDatasetBuilder().to_dataset(_source)

        filetype: Optional[FileType] = (
            FileType[inferred_type] if inferred_type else None
        )

        kwargs = {
            "method": "post",
            "url": "/datasets/stream",
            "timeout": timeout,
            "dataset": ds,
            "params": dict(name=name, filetype=filetype),
        }

        result = self.client.request(**kwargs)

        return Dataset(**result)

    def download_dataset_as_stream(
        self,
        id: str,
        destination: Optional[Union[str, FileLike]] = None,
        *,
        filetype: Optional[FileType] = None,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> Any:
        """Download a dataset by streaming chunks of it.

        DEPRECATED: Please use download_dataset instead.
        """
        warnings.warn(
            DeprecationWarning(
                "download_dataset_as_stream is deprecated. Use download_dataset instead."
            )
        )

        # Ignoring return type because download_dataset's logic makes sure
        # that the return type will be BytesIO
        # No point in going through the hassle of using typing.overload for something
        # that will be removed soon
        return self.download_dataset(
            id,
            destination=destination,
            timeout=timeout,
            filetype=filetype,
            from_download_as_stream=True,
        )

    def download_dataset(
        self,
        id: str,
        destination: Optional[Union[str, FileLike]] = None,
        filetype: Optional[FileType] = None,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
        from_download_as_stream: bool = False,  # TODO: Remove once deprecated
    ) -> Any:
        """Download a dataset."""
        # Download dataset metadata
        dataset_info = self.client.datasets.get_dataset(id, timeout=timeout)

        if destination is None:
            warnings.warn(
                DeprecationWarning(
                    "Please specify the destination argument. The return type will change in the future."
                )
            )

            if not from_download_as_stream:
                # Download as CSV was the old return value when using download_dataset
                filetype = filetype or FileType.csv
            else:
                # Download as the filetype of the dataset on the server was the old
                # return value when using download_dataset_as_stream
                filetype = filetype or dataset_info.filetype
        else:
            if not (isinstance(destination, str) or is_file_like(destination)):
                raise TypeError(
                    f"Expected destination to be a string or a buffer, got {type(destination)} instead"
                )

        return self.client.request(
            method="get",
            url=f"/datasets/{id}/download/stream",
            params={"filetype": filetype},
            should_stream=True,
            want_content=destination is None and not from_download_as_stream,
            timeout=timeout,
            destination=destination,
        )

    def find_all_datasets_by_user(
        self,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> List[Dataset]:
        """List all datasets of the current_user."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = []

        return _Datasets(self.client).find_all_datasets_by_user(*args, **kwargs)

    def get_dataset(
        self,
        id: str,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> Dataset:
        """Get a dataset."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            id,
        ]

        return _Datasets(self.client).get_dataset(*args, **kwargs)

    def patch_dataset(
        self,
        request: PatchDataset,
        id: str,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> Dataset:
        """Modify a dataset."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            request,
            id,
        ]

        return _Datasets(self.client).patch_dataset(*args, **kwargs)

    def analyze_dataset(
        self,
        id: str,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> Dataset:
        """Start the analysis of a dataset."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            id,
        ]

        return _Datasets(self.client).analyze_dataset(*args, **kwargs)

    def get_dataset_correlations(
        self,
        id: str,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> Any:
        """Get a dataset's correlations."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            id,
        ]

        return _Datasets(self.client).get_dataset_correlations(*args, **kwargs)


class Health:
    def __init__(self, client: ApiClient) -> None:
        self.client = client

    def get_root(
        self,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> Any:
        """Verify server health."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = []

        return _Health(self.client).get_root(*args, **kwargs)

    def get_health(
        self,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> Any:
        """Verify server health."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = []

        return _Health(self.client).get_health(*args, **kwargs)

    def get_health_db(
        self,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> Any:
        """Verify connection to the db health."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = []

        return _Health(self.client).get_health_db(*args, **kwargs)


class Jobs:
    def __init__(self, client: ApiClient) -> None:
        self.client = client

    def find_all_jobs_by_user(
        self,
        nb_days: Optional[int] = None,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> List[GenericJob]:
        """Retrieve all jobs executed by the current user.

        Jobs are filtered by execution date, by default only the last 5 days are displayed,
        a parameter can be provided to go further back in time.
        """

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            nb_days,
        ]

        return _Jobs(self.client).find_all_jobs_by_user(*args, **kwargs)

    def create_full_avatarization_job(
        self,
        request: AvatarizationJobCreate,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> AvatarizationJob:
        """Create an avatarization job, then calculate metrics."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            request,
        ]

        return _Jobs(self.client).create_full_avatarization_job(*args, **kwargs)

    def cancel_job(
        self,
        id: str,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> GenericJob:
        """Cancel any kind of job.

        If the job hadn't been started yet, revoke it.
        If the job is ongoing, gently kill it.
        If the job is done, do nothing.
        """

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            id,
        ]

        return _Jobs(self.client).cancel_job(*args, **kwargs)

    def create_avatarization_job(
        self,
        request: AvatarizationJobCreate,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> AvatarizationJob:
        """Create an avatarization job."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            request,
        ]

        return _Jobs(self.client).create_avatarization_job(*args, **kwargs)

    def create_avatarization_with_time_series_job(
        self,
        request: AvatarizationWithTimeSeriesJobCreate,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> AvatarizationWithTimeSeriesJob:
        """Create an avatarization with time series job."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            request,
        ]

        return _Jobs(self.client).create_avatarization_with_time_series_job(
            *args, **kwargs
        )

    def create_avatarization_multi_table_job(
        self,
        request: AvatarizationMultiTableJobCreate,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> AvatarizationMultiTableJob:
        """Create an avatarization for relational data."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            request,
        ]

        return _Jobs(self.client).create_avatarization_multi_table_job(*args, **kwargs)

    def create_signal_metrics_job(
        self,
        request: SignalMetricsJobCreate,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> SignalMetricsJob:
        """Create a signal metrics job."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            request,
        ]

        return _Jobs(self.client).create_signal_metrics_job(*args, **kwargs)

    def create_privacy_metrics_job(
        self,
        request: PrivacyMetricsJobCreate,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> PrivacyMetricsJob:
        """Create a privacy metrics job."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            request,
        ]

        return _Jobs(self.client).create_privacy_metrics_job(*args, **kwargs)

    def create_privacy_metrics_time_series_job(
        self,
        request: PrivacyMetricsWithTimeSeriesJobCreate,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> PrivacyMetricsWithTimeSeriesJob:
        """Create a privacy metrics with time series job."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            request,
        ]

        return _Jobs(self.client).create_privacy_metrics_time_series_job(
            *args, **kwargs
        )

    def create_signal_metrics_time_series_job(
        self,
        request: SignalMetricsWithTimeSeriesJobCreate,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> SignalMetricsWithTimeSeriesJob:
        """Create a signal metrics with time series job."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            request,
        ]

        return _Jobs(self.client).create_signal_metrics_time_series_job(*args, **kwargs)

    def create_privacy_metrics_multi_table_job(
        self,
        request: PrivacyMetricsMultiTableJobCreate,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> PrivacyMetricsMultiTableJob:
        """Create a privacy metrics job."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            request,
        ]

        return _Jobs(self.client).create_privacy_metrics_multi_table_job(
            *args, **kwargs
        )

    def create_privacy_metrics_geolocation_job(
        self,
        request: PrivacyMetricsGeolocationJobCreate,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> PrivacyMetricsGeolocationJob:
        """Create a geolocation privacy metrics job."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            request,
        ]

        return _Jobs(self.client).create_privacy_metrics_geolocation_job(
            *args, **kwargs
        )

    def get_privacy_metrics_geolocation_job(
        self,
        id: str,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> PrivacyMetricsGeolocationJob:
        """Get a geolocation privacy metrics job."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            id,
        ]

        return _Jobs(self.client).get_privacy_metrics_geolocation_job(*args, **kwargs)

    def get_avatarization_job(
        self,
        id: str,
        *,
        per_request_timeout: Optional[int] = DEFAULT_TIMEOUT,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> AvatarizationJob:
        """Get an avatarization job."""

        kwargs: Dict[str, Any] = {
            "per_request_timeout": per_request_timeout,
            "timeout": timeout,
        }

        args: List[Any] = [
            id,
        ]

        return _Jobs(self.client).get_avatarization_job(*args, **kwargs)

    def get_avatarization_time_series_job(
        self,
        id: str,
        *,
        per_request_timeout: Optional[int] = DEFAULT_TIMEOUT,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> AvatarizationWithTimeSeriesJob:
        """Get an avatarization time series job."""

        kwargs: Dict[str, Any] = {
            "per_request_timeout": per_request_timeout,
            "timeout": timeout,
        }

        args: List[Any] = [
            id,
        ]

        return _Jobs(self.client).get_avatarization_time_series_job(*args, **kwargs)

    def get_avatarization_multi_table_job(
        self,
        id: str,
        *,
        per_request_timeout: Optional[int] = DEFAULT_TIMEOUT,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> AvatarizationMultiTableJob:
        """Get a multi table avatarization job."""

        kwargs: Dict[str, Any] = {
            "per_request_timeout": per_request_timeout,
            "timeout": timeout,
        }

        args: List[Any] = [
            id,
        ]

        return _Jobs(self.client).get_avatarization_multi_table_job(*args, **kwargs)

    def get_signal_metrics(
        self,
        id: str,
        *,
        per_request_timeout: Optional[int] = DEFAULT_TIMEOUT,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> SignalMetricsJob:
        """Get a signal metrics job."""

        kwargs: Dict[str, Any] = {
            "per_request_timeout": per_request_timeout,
            "timeout": timeout,
        }

        args: List[Any] = [
            id,
        ]

        return _Jobs(self.client).get_signal_metrics(*args, **kwargs)

    def get_signal_metrics_time_series_job(
        self,
        id: str,
        *,
        per_request_timeout: Optional[int] = DEFAULT_TIMEOUT,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> SignalMetricsWithTimeSeriesJob:
        """Get a signal metrics time series job."""

        kwargs: Dict[str, Any] = {
            "per_request_timeout": per_request_timeout,
            "timeout": timeout,
        }

        args: List[Any] = [
            id,
        ]

        return _Jobs(self.client).get_signal_metrics_time_series_job(*args, **kwargs)

    def get_privacy_metrics(
        self,
        id: str,
        *,
        per_request_timeout: Optional[int] = DEFAULT_TIMEOUT,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> PrivacyMetricsJob:
        """Get a privacy metrics job."""

        kwargs: Dict[str, Any] = {
            "per_request_timeout": per_request_timeout,
            "timeout": timeout,
        }

        args: List[Any] = [
            id,
        ]

        return _Jobs(self.client).get_privacy_metrics(*args, **kwargs)

    def get_privacy_metrics_time_series_job(
        self,
        id: str,
        *,
        per_request_timeout: Optional[int] = DEFAULT_TIMEOUT,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> PrivacyMetricsWithTimeSeriesJob:
        """Get a privacy metrics time series job."""

        kwargs: Dict[str, Any] = {
            "per_request_timeout": per_request_timeout,
            "timeout": timeout,
        }

        args: List[Any] = [
            id,
        ]

        return _Jobs(self.client).get_privacy_metrics_time_series_job(*args, **kwargs)

    def get_privacy_metrics_multi_table_job(
        self,
        id: str,
        *,
        per_request_timeout: Optional[int] = DEFAULT_TIMEOUT,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> PrivacyMetricsMultiTableJob:
        """Get a privacy metrics multi table job."""

        kwargs: Dict[str, Any] = {
            "per_request_timeout": per_request_timeout,
            "timeout": timeout,
        }

        args: List[Any] = [
            id,
        ]

        return _Jobs(self.client).get_privacy_metrics_multi_table_job(*args, **kwargs)

    def create_advice(
        self,
        request: AdviceJobCreate,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> AdviceJob:
        """Create advice on anonymization parameters."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            request,
        ]

        return _Jobs(self.client).create_advice(*args, **kwargs)

    def get_advice(
        self,
        id: str,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> AdviceJob:
        """Get advice result."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            id,
        ]

        return _Jobs(self.client).get_advice(*args, **kwargs)


class Metrics:
    def __init__(self, client: ApiClient) -> None:
        self.client = client

    def get_job_projections(
        self,
        job_id: str,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> Projections:
        """Get the projections of records and avatars in 3D.

        See https://saiph.readthedocs.io/en/latest/ for more information.

        Arguments
        ---------
            job_id:
                avatarization or privacy job id used to fit the model
        """

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            job_id,
        ]

        return _Metrics(self.client).get_job_projections(*args, **kwargs)

    def get_variable_contributions(
        self,
        job_id: str,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> Contributions:
        """Get the contributions of the dataset variables within the fitted space.

        See https://saiph.readthedocs.io/en/latest for more information.

        Arguments
        ---------
            job_id:
                avatarization or privacy job id used to fit the model
        """

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            job_id,
        ]

        return _Metrics(self.client).get_variable_contributions(*args, **kwargs)

    def get_explained_variance(
        self,
        job_id: str,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> ExplainedVariance:
        """Get the explained variance of records.

        See https://saiph.readthedocs.io/en/latest/ for more information.

        Arguments
        ---------
            job_id:
                avatarization or privacy job id used to fit the model
        """

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            job_id,
        ]

        return _Metrics(self.client).get_explained_variance(*args, **kwargs)


class Reports:
    def __init__(self, client: ApiClient) -> None:
        self.client = client

    def create_report(
        self,
        request: ReportCreate,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> Report:
        """Create an anonymization report."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            request,
        ]

        return _Reports(self.client).create_report(*args, **kwargs)

    def get_report(
        self,
        id: str,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> Report:

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            id,
        ]

        return _Reports(self.client).get_report(*args, **kwargs)

    def download_report(
        self,
        id: str,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> Any:
        """Download a report."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            id,
        ]

        return _Reports(self.client).download_report(*args, **kwargs)

    def create_report_from_data(
        self,
        request: ReportFromDataCreate,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> Report:
        """Create an anonymization report without avatarization job."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            request,
        ]

        return _Reports(self.client).create_report_from_data(*args, **kwargs)

    def create_geolocation_privacy_report(
        self,
        request: ReportGeolocationPrivacyCreate,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> Report:
        """Create an anonymization report without avatarization job."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            request,
        ]

        return _Reports(self.client).create_geolocation_privacy_report(*args, **kwargs)


class Stats:
    def __init__(self, client: ApiClient) -> None:
        self.client = client

    def get_cluster_stats(
        self,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> ClusterStats:
        """Get insights into the cluster's usage."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = []

        return _Stats(self.client).get_cluster_stats(*args, **kwargs)


class Users:
    def __init__(self, client: ApiClient) -> None:
        self.client = client

    def find_users(
        self,
        email: Optional[str] = None,
        username: Optional[str] = None,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> List[User]:
        """Get users, optionally filtering them by username or email.

        This endpoint is protected with rate limiting.
        """

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            email,
            username,
        ]

        return _Users(self.client).find_users(*args, **kwargs)

    def create_user(
        self,
        request: CreateUser,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> User:
        """Create a user.

        This endpoint is protected with rate limiting.
        """

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            request,
        ]

        return _Users(self.client).create_user(*args, **kwargs)

    def get_me(
        self,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> User:
        """Get my own user."""

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = []

        return _Users(self.client).get_me(*args, **kwargs)

    def get_user(
        self,
        id: str,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> User:
        """Get a user by id.

        This endpoint is protected with rate limiting.
        """

        kwargs: Dict[str, Any] = {
            "timeout": timeout,
        }

        args: List[Any] = [
            id,
        ]

        return _Users(self.client).get_user(*args, **kwargs)


class PandasIntegration:
    def __init__(self, client: ApiClient) -> None:
        self.client = client

    def upload_dataframe(
        self,
        request: "pd.DataFrame",
        name: Optional[str] = None,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
        identifier_variables: List[str] = [],
        **kwargs: Dict[str, Any],  # to collect should_stream
    ) -> Dataset:

        # Reset the index of the dataframe
        request = request.reset_index(drop=True)

        if "should_stream" in kwargs:
            warnings.warn(
                "The `should_stream` parameter is deprecated and will be removed in a future version. "
                "All uploads are now streamed by default.",
                DeprecationWarning,
            )

        if not set(identifier_variables).issubset(set(request.columns)):
            missing_columns = set(identifier_variables) - set(request.columns)
            raise ValueError(
                "Expected all identifier variables to be present in the dataframe."
                f" Got {missing_columns}"
            )

        for col in request.columns:
            if pd.api.types.infer_dtype(request[col], skipna=True) in (
                "mixed-integer",
                "mixed",
            ):
                raise ValueError(
                    f"Expected column '{col}' to have either str or numeric values."
                    " Consider harmonizing columns prior to upload."
                )

        df_types = request.dtypes

        dataset = self.client.datasets.create_dataset(
            source=request,
            name=name,
            timeout=timeout,
        )

        columns = []
        for index, dtype in zip(df_types.index, df_types):
            column_detail = ColumnDetail(
                type=to_column_type(str(dtype)),
                label=index,
                is_identifier=index in identifier_variables,
            )
            columns.append(column_detail)

        dataset = self.client.datasets.patch_dataset(
            id=str(dataset.id),
            request=PatchDataset(columns=columns),
        )
        return dataset

    def download_dataframe(
        self,
        id: str,
        *,
        filetype: Optional[FileType] = None,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
        **kwargs: Any,  # to collect should_stream
    ) -> pd.DataFrame:
        if "should_stream" in kwargs:
            warnings.warn(
                "The `should_stream` parameter is deprecated and will be removed in a future version. "
                "All uploads are now streamed by default.",
                DeprecationWarning,
            )

        dataset_info = self.client.datasets.get_dataset(id, timeout=timeout)

        _filetype = filetype or dataset_info.filetype
        with tempfile.TemporaryDirectory() as download_dir:
            path = Path(download_dir) / "file"

            try:
                Datasets(client=self.client).download_dataset(
                    id, str(path), timeout=timeout, filetype=_filetype
                )
                if _filetype is FileType.parquet:
                    df = pd.read_parquet(path, engine="pyarrow")
                elif _filetype is FileType.csv:
                    df = pd.read_csv(path)
                else:
                    raise ValueError(f"Unsupported filetype: {_filetype}")
            except Exception as e:
                # Save file somewhere to allow for investigation
                tdir = tempfile.mkdtemp()
                tfile = Path(tdir) / f"dataframe-{id}.fail"
                shutil.copyfile(path, tfile)
                warnings.warn(f"download failed: file is here: {str(tfile)}")
                raise e

        # We apply datetime columns separately as 'datetime' is not a valid pandas dtype
        dtypes = {c.label: from_column_type(c.type) for c in dataset_info.columns or {}}
        datetime_columns = [
            label for label, type in dtypes.items() if type == ColumnType.datetime.value
        ]

        # Remove datetime columns
        for label in list(dtypes.keys()):
            if label in datetime_columns:
                dtypes.pop(label, None)

        for name, dtype in dtypes.items():
            df[name] = df[name].astype(dtype)

        df[datetime_columns] = df[datetime_columns].astype("datetime64[ns]")

        return df


class Pipelines:
    def __init__(self, client: ApiClient) -> None:
        self.client = client

    def avatarization_pipeline_with_processors(
        self,
        request: AvatarizationPipelineCreate,
        *,
        timeout: Optional[int] = DEFAULT_TIMEOUT,
        per_request_timeout: Optional[int] = DEFAULT_TIMEOUT,
    ) -> AvatarizationPipelineResult:
        """Create a Pipeline of avatarization with processors."""

        # Upload the dataset if needed
        # User can either specify the dataset he already uploaded or let this pipeline upload it
        df = request.df.copy()
        original_dataset_id = (
            request.avatarization_job_create.parameters.dataset_id
            or self.client.pandas_integration.upload_dataframe(df, timeout=timeout).id
        )

        # Pre process the dataframe and upload it
        processors = request.processors
        for p in processors:
            df = p.preprocess(df)
        dataset = self.client.pandas_integration.upload_dataframe(
            df, timeout=per_request_timeout
        )

        # Avatarize the uploaded dataframe
        request.avatarization_job_create.parameters.dataset_id = dataset.id
        avatarization_job = self.client.jobs.create_avatarization_job(
            request.avatarization_job_create
        )
        logger.info(f"launching avatarization job with id={avatarization_job.id}")

        avatarization_job = self.client.jobs.get_avatarization_job(
            str(avatarization_job.id),
            timeout=timeout,
            per_request_timeout=per_request_timeout,
        )
        if avatarization_job.status == JobStatus.failure:
            raise Exception(
                f"Got error during the avatarization job: {avatarization_job.error_message}"
            )

        if (
            avatarization_job.status == JobStatus.pending
            or not avatarization_job.result
        ):
            raise Timeout(
                f"The avatarization job '{avatarization_job.id}' timed out."
                ""
                """Try increasing the timeout with the `timeout` parameter."""
            )

        # Download the dataframe, postprocess it and upload the new dataframe
        sensitive_unshuffled_avatars = (
            self.client.pandas_integration.download_dataframe(
                str(avatarization_job.result.sensitive_unshuffled_avatars_datasets.id),
                timeout=timeout,
            )
        )
        for p in reversed(processors):
            sensitive_unshuffled_avatars = p.postprocess(
                request.df, sensitive_unshuffled_avatars
            )
        unshuffled_dataset = self.client.pandas_integration.upload_dataframe(
            sensitive_unshuffled_avatars, timeout=timeout
        )

        # Calculate privacy metrics on the post processed dataset vs the original one
        privacy_job = self.client.jobs.create_privacy_metrics_job(
            PrivacyMetricsJobCreate(
                parameters=PrivacyMetricsParameters(
                    original_id=original_dataset_id,
                    unshuffled_avatars_id=unshuffled_dataset.id,
                )
            ),
            timeout=per_request_timeout,
        )
        logger.info(f"launching privacy metrics job with id={privacy_job.id}")

        # Calculate signal metrics
        signal_job = self.client.jobs.create_signal_metrics_job(
            SignalMetricsJobCreate(
                parameters=SignalMetricsParameters(
                    original_id=original_dataset_id, avatars_id=unshuffled_dataset.id
                )
            ),
            timeout=per_request_timeout,
        )
        logger.info(f"launching signal metrics job with id={signal_job.id}")

        # Get the job results
        signal_job = self.client.jobs.get_signal_metrics(
            str(signal_job.id), timeout=timeout, per_request_timeout=per_request_timeout
        )
        privacy_job = self.client.jobs.get_privacy_metrics(
            str(privacy_job.id),
            timeout=timeout,
            per_request_timeout=per_request_timeout,
        )
        if signal_job.status == JobStatus.failure:
            raise Exception(
                f"Got error during the signal metrics job: {signal_job.error_message}"
            )
        if privacy_job.status == JobStatus.failure:
            raise Exception(
                f"Got error during the privacy metrics job: {privacy_job.error_message}"
            )

        if signal_job.status == JobStatus.pending or not signal_job.result:
            raise Timeout(
                f"The signal metrics job '{signal_job.id}' timed out."
                ""
                """Try increasing the timeout with the `timeout` parameter."""
            )

        if privacy_job.status == JobStatus.pending or not privacy_job.result:
            raise Timeout(
                f"The privacy metrics job '{privacy_job.id}' timed out."
                ""
                """Try increasing the timeout with the `timeout` parameter."""
            )

        # Shuffle sensitive_unshuffled_avatars for security reasons
        random_gen = np.random.default_rng()
        map = random_gen.permutation(sensitive_unshuffled_avatars.index.values).tolist()
        post_processed_avatars = sensitive_unshuffled_avatars.iloc[map].reset_index(
            drop=True
        )

        return AvatarizationPipelineResult(
            privacy_metrics=privacy_job.result,
            signal_metrics=signal_job.result,
            post_processed_avatars=post_processed_avatars,
            avatarization_job_id=avatarization_job.id,
            signal_job_id=signal_job.id,
            privacy_job_id=privacy_job.id,
        )


def upload_batch_and_get_order(
    client: ApiClient,
    training: pd.DataFrame,
    splits: List[pd.DataFrame],
    timeout: int = DEFAULT_TIMEOUT,
) -> Tuple[UUID, List[UUID], Dict[UUID, pd.Index]]:
    """Upload batches to the server
    Arguments
    ---------
        client:
            Api client
        training:
            Dataframe used to train the anonymization model. This dataframe should contain all modalities of the categorical variables.
        splits:
            All other batches
    Returns
    -------
        training_dataset_id:
            The dataset id of the training dataset
        datasets_split_ids:
            The dataset id of all other batches
        batch_mapping:
            The index mapping for each dataset batch
    """
    training_dataset = client.pandas_integration.upload_dataframe(
        training, timeout=timeout
    )

    datasets_split_ids = [
        client.pandas_integration.upload_dataframe(split, timeout=timeout).id
        for split in splits
    ]
    batch_mapping: Dict[UUID, pd.Index] = {training_dataset.id: training.index}
    for dataset, dataframe in zip(datasets_split_ids, splits):
        batch_mapping[dataset] = dataframe.index

    return training_dataset.id, datasets_split_ids, batch_mapping


# This file has been generated - DO NOT MODIFY
