# This file has been generated - DO NOT MODIFY
# API Version : {{api_version}}


{#

Take inspiration from:

- https://github.com/openapi-generators/openapi-python-client/blob/main/openapi_python_client/templates/endpoint_module.py.jinja

#}


{% macro to_arguments(parameters) %}
{% for parameter in parameters %}
    {{ parameter.to_argument() }}
{% endfor %}
{% endmacro %}


{% macro explode_form_data(operation) %}
    {% if operation.request_body.content_type == "multipart/form-data" %}
        {%- for parameter in operation.request_body.properties %}
            {{ parameter.to_argument() }}
        {% endfor %}
    {% else %}
        request: {{operation.request_body.model_name}},
    {% endif %}
{% endmacro %}

import logging


{% for model in models %}
from avatars.models import {{model}} # noqa: F401
{% endfor %}

{% for tag, ops in operations|groupby("tag") %}
from avatars.api_autogenerated import {{tag|to_pascal_case}} as _{{tag|to_pascal_case}}
{% endfor %}

from avatars.models import JobStatus
from avatars.models import ColumnDetail, ColumnType, FileType
from avatars.models import PrivacyMetricsParameters, SignalMetricsParameters
from avatars.models import AvatarizationPipelineCreate, AvatarizationPipelineResult
from avatars.arrow_utils import FileLike, ArrowDatasetBuilder, DataSourceItem, DataSourceItems, is_file_like
from avatars.client import ApiClient
import tempfile
import shutil
from pathlib import Path
from typing import Any, Dict, Optional, TypeVar, Union, List, Tuple
import warnings
from uuid import UUID
import pandas as pd
import numpy as np
from avatars.exceptions import Timeout
from avatars.constants import DEFAULT_TIMEOUT

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}'")



{# Use this to add/override custom methods to a class.#}
{% set custom_methods_by_tag = {
    "datasets" : ["create_dataset", "create_dataset_from_stream", "download_dataset_as_stream", "download_dataset"],
}

%}

{% set get_job_methods = ["get_privacy_metrics", "get_signal_metrics", "get_avatarization_job", "get_avatarization_batch_job", "get_privacy_metrics_batch_job", "get_signal_metrics_batch_job", "get_avatarization_time_series_job", "get_signal_metrics_time_series_job", "get_privacy_metrics_time_series_job", "get_avatarization_multi_table_job", "get_privacy_metrics_multi_table_job"] %}
{% for tag, ops in operations|groupby("tag") %}


class {{tag|to_pascal_case}}:
    def __init__(self, client: ApiClient) -> None:
        self.client = client

    {# Include template of customs methods if there are methods to override #}
    {%+ if tag in custom_methods_by_tag %}
        {% filter indent(width=4) %}
        {% include "custom_methods/" ~ tag ~ ".py" %}
        {% endfilter %}
    {% endif %}

    {% for operation in ops %}

    {% if operation.name in custom_methods_by_tag[tag] %}
        {# We skip inserting an autogenerated method if it's present in the mapping #}
        {% continue %}
    {% endif %}

    def {{ operation.name }}(
        self,
        {% if operation.request_body.model_name or operation.request_body.content_type == "multipart/form-data" %}
        {{explode_form_data(operation)}}
        {% endif %}
        {{ to_arguments(operation.path_parameters) | indent(8) }}
        {{ to_arguments(operation.query_parameters) | indent(8) }}
        *,
        {% if operation.name in get_job_methods %}
        per_request_timeout : Optional[int] = DEFAULT_TIMEOUT,
        {% endif %}
        timeout : Optional[int] = DEFAULT_TIMEOUT,
            ) -> {% if operation.response_model %} {{ operation.response_model.type_hint() }} {% else %} Any {%endif%}:
    {% if operation.description %}
        """{{ operation.description|indent(8) }}
        """
    {% endif %}

        kwargs: Dict[str, Any] = {
            {% if operation.name in get_job_methods %}
            "per_request_timeout" : per_request_timeout,
            {% endif %}
            "timeout": timeout,

        }

        args: List[Any] = [
            {% if operation.request_body.model_name or operation.request_body.content_type == "multipart/form-data" %}
            request,
            {%- if operation.request_body.properties|length > 1 %}
            {# We have additional form data fields alongside 'file', so we explode those. #}
            {% for parameter in operation.request_body.properties %}
                {% if parameter.name != "file" %}
                    {{parameter.name}},
                {% endif %}
            {% endfor %}
            {% endif %}

            {% endif %}
            {% for parameter in operation.path_parameters %}
                {{ parameter.name }},
            {% endfor %}
            {% for parameter in operation.query_parameters %}
                {{ parameter.name }},
            {% endfor %}
        ]


        {# This calls the method from api_autogenerated.py #}
        return _{{tag|to_pascal_case}}(self.client).{{ operation.name }}(*args, **kwargs)

    {% endfor %}
{% endfor %}

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
