# AUTOGENERATED! DO NOT EDIT! File to edit: notebooks/API_Prediction.ipynb (unless otherwise specified).

__all__ = ["Prediction"]

# Cell

from typing import *

# Internal Cell

import os
import requests

from pathlib import Path
import pandas as pd
from fastcore.foundation import patch
from tqdm import tqdm

from ..logger import get_logger, set_level
from ..helper import (
    get_data,
    post_data,
    delete_data,
    add_ready_column,
    generate_df,
    get_attributes_from_instances,
)

from .client import Client
from .progress_status import ProgressStatus
from ..constant import CLIENT_DB_USERNAME, CLIENT_DB_PASSWORD

# Internal Cell

logger = get_logger(__name__)

# Cell


class Prediction(ProgressStatus):
    """A class to manage and download the predictions.

    The `Prediction` class is automatically instantiated by calling the `Model.predict` method of a `Model` instance.
    Currently, it is the only way to instantiate this class.
    """

    BASIC_PRED_COLS = ["uuid", "created", "total_steps", "completed_steps"]
    ALL_PRED_COLS = BASIC_PRED_COLS + ["model", "datasource", "region", "error"]

    COLS_TO_RENAME = {
        "uuid": "prediction_uuid",
        "datasource": "datasource_uuid",
        "model": "model_uuid",
    }

    def __init__(
        self,
        uuid: str,
        datasource: Optional[str] = None,
        model: Optional[str] = None,
        created: Optional[pd.Timestamp] = None,
        total_steps: Optional[int] = None,
        completed_steps: Optional[int] = None,
        region: Optional[str] = None,
        error: Optional[str] = None,
        disabled: Optional[bool] = None,
    ):
        """Constructs a new `Prediction` instance

        Warning:
            Do not construct this object directly by calling the constructor, instead please use
            `Model.predict` method of the Model instance.

        Args:
            uuid: Prediction uuid in the server.
            datasource: DataSource uuid in the server.
            model: Model uuid in the server.
            created: Prediction creation date.
            total_steps: No of steps needed to complete the model prediction.
            completed_steps: No of steps completed so far in the model prediction.
            region: AWS bucket region.
            error: Error message while making the prediction.
            disabled: Flag to indicate the status of the prediction.
        """
        self.uuid = uuid
        self.datasource = datasource
        self.model = model
        self.created = created
        self.total_steps = total_steps
        self.completed_steps = completed_steps
        self.region = region
        self.error = error
        self.disabled = disabled
        ProgressStatus.__init__(self, relative_url=f"/prediction/{self.uuid}")

    @staticmethod
    def _download_prediction_file_to_local(
        file_name: str, url: str, path: Union[str, Path]
    ) -> None:
        """Download the file to local directory.

        Args:
            file_name: Name of the file
            url: Url of the file
            path: Local directory path

        Raises:
            HTTPError: If the **url** is invalid or not reachable.
        """
        try:
            response = requests.get(url)
            response.raise_for_status()

        except requests.exceptions.HTTPError as e:
            raise requests.exceptions.HTTPError(e)

        else:
            with open(Path(path) / file_name, "wb") as f:
                f.write(response.content)

    @staticmethod
    def ls(
        offset: int = 0,
        limit: int = 100,
        disabled: bool = False,
        completed: bool = False,
    ) -> List["Prediction"]:
        """Return the list of Prediction instances available in the server.

        Args:
            offset: The number of predictions to offset at the beginning. If None, then the default value **0** will be used.
            limit: The maximum number of predictions to return from the server. If None,
                then the default value **100** will be used.
            disabled: If set to **True**, then only the deleted predictions will be returned. Else, the default value
                **False** will be used to return only the list of active predictions.
            completed: If set to **True**, then only the predictions that are successfully processed in server will be returned.
                Else, the default value **False** will be used to return all the predictions.

        Returns:
            A list of Prediction instances available in the server.

        Raises:
            ConnectionError: If the server address is invalid or not reachable.

        An example to list the available predictions:

        ```python
        Prediction.ls()
        ```
        """
        lists = Client._get_data(
            relative_url=f"/prediction/?disabled={disabled}&completed={completed}&offset={offset}&limit={limit}"
        )

        predx = [
            Prediction(
                uuid=pred["uuid"],
                model=pred["model"],
                datasource=pred["datasource"],
                created=pred["created"],
                total_steps=pred["total_steps"],
                completed_steps=pred["completed_steps"],
                region=pred["region"],
                error=pred["error"],
                disabled=pred["disabled"],
            )
            for pred in lists
        ]

        return predx

    @staticmethod
    def as_df(predx: List["Prediction"]) -> pd.DataFrame:
        """Return the details of prediction instances as a pandas dataframe.

        Args:
            predx: List of prediction instances.

        Returns:
            Details of all the prediction in a dataframe.

        Raises:
            ConnectionError: If the server address is invalid or not reachable.

        An example get the details of available prediction:

        ```python
        predictionx = Prediction.ls()
        Prediction.as_df(predictionx)
        ```
        """
        response = get_attributes_from_instances(predx, Prediction.BASIC_PRED_COLS)  # type: ignore

        df = generate_df(response, Prediction.BASIC_PRED_COLS)

        df = df.rename(columns=Prediction.COLS_TO_RENAME)

        return add_ready_column(df)

    def details(self) -> pd.DataFrame:
        raise NotImplementedError()

    def to_pandas(self) -> pd.DataFrame:
        raise NotImplementedError()

    def delete(self) -> pd.DataFrame:
        raise NotImplementedError()

    def to_s3(
        self,
        uri: str,
        access_key: Optional[str] = None,
        secret_key: Optional[str] = None,
    ) -> ProgressStatus:
        raise NotImplementedError()

    def to_azure_blob_storage(
        self,
        uri: str,
        credential: Optional[str] = None,
    ) -> ProgressStatus:
        raise NotImplementedError()

    def to_local(
        self,
        path: Union[str, Path],
        show_progress: Optional[bool] = True,
    ) -> None:
        raise NotImplementedError()

    def to_mysql(
        self,
        *,
        host: str,
        database: str,
        table: str,
        port: int = 3306,
        username: Optional[str] = None,
        password: Optional[str] = None,
    ) -> ProgressStatus:
        raise NotImplementedError()

    def to_clickhouse(
        self,
        *,
        host: str,
        database: str,
        table: str,
        port: int = 0,
        protocol: str,
        username: Optional[str] = None,
        password: Optional[str] = None,
    ) -> ProgressStatus:
        raise NotImplementedError()


# Cell


@patch
def details(self: Prediction) -> pd.DataFrame:
    """Return the details of a prediction.

    Returns:
        A pandas DataFrame encapsulating the details of the prediction.

    Raises:
        ConnectionError: If the server address is invalid or not reachable.

    An example to get details of a prediction:

    ```python
    prediction.details()
    ```
    """
    response = Client._get_data(relative_url=f"/prediction/{self.uuid}")

    df = pd.DataFrame(response, index=[0])[Prediction.ALL_PRED_COLS]

    df = df.rename(columns=Prediction.COLS_TO_RENAME)

    return add_ready_column(df)


# Cell


@patch
def delete(self: Prediction) -> pd.DataFrame:
    """Delete a prediction from the server.

    Returns:
        A pandas DataFrame encapsulating the details of the deleted prediction.

    Raises:
        ConnectionError: If the server address is invalid or not reachable.

    An example to delete a prediction:

    ```python
    prediction.delete()
    ```
    """
    response = Client._delete_data(relative_url=f"/prediction/{self.uuid}")

    df = pd.DataFrame(response, index=[0])[Prediction.BASIC_PRED_COLS]

    df = df.rename(columns=Prediction.COLS_TO_RENAME)

    return add_ready_column(df)


# Cell


@patch
def to_pandas(self: Prediction) -> pd.DataFrame:
    """Return the prediction results as a pandas DataFrame

    Returns:
        A pandas DataFrame encapsulating the results of the prediction.

    Raises:
        ConnectionError: If the server address is invalid or not reachable.

    An example to convert the prediction results into a pandas DataFrame:

    ```python
    predictions.to_pandas()
    ```
    """
    response = Client._get_data(relative_url=f"/prediction/{self.uuid}/pandas")
    keys = list(response.keys())
    keys.remove("Score")
    index_name = keys[0]
    return (
        pd.DataFrame(response)
        .set_index(index_name)
        .sort_values("Score", ascending=False)
    )


# Cell


@patch
def to_s3(
    self: Prediction,
    uri: str,
    access_key: Optional[str] = None,
    secret_key: Optional[str] = None,
) -> ProgressStatus:
    """Push the prediction results to the target AWS S3 bucket.

    Args:
        uri: Target S3 bucket uri.
        access_key: Access key for the target S3 bucket. If **None** (default value), then the value
            from **AWS_ACCESS_KEY_ID** environment variable is used.
        secret_key: Secret key for the target S3 bucket. If **None** (default value), then the value
            from **AWS_SECRET_ACCESS_KEY** environment variable is used.

    Returns:
        An instance of `ProgressStatus` class.

    Raises:
        ConnectionError: If the server address is invalid or not reachable.

    An example to push the prediction results to an AWS s3 bucket:

    ```python
        status = prediction.to_s3(
            uri="s3://target-bucket/"
        )
        status.progress_bar()
    ```
    """
    access_key = (
        access_key if access_key is not None else os.environ["AWS_ACCESS_KEY_ID"]
    )
    secret_key = (
        secret_key if secret_key is not None else os.environ["AWS_SECRET_ACCESS_KEY"]
    )

    response = Client._post_data(
        relative_url=f"/prediction/{self.uuid}/to_s3",
        data=dict(uri=uri, access_key=access_key, secret_key=secret_key),
    )

    return ProgressStatus(relative_url=f"/prediction/push/{response['uuid']}")


# Cell


@patch
def to_azure_blob_storage(
    self: Prediction,
    uri: str,
    credential: Optional[str] = None,
) -> ProgressStatus:
    """Push the prediction results to the target Azure Blob Storage.

    Args:
        uri: Target Azure Blob Storage uri.
        credential: Azure Blob Storage credential. If **None** (default value), then the value
            from **AZURE_STORAGE_CREDENTIAL** environment variable will be used.

    Returns:
        An instance of `ProgressStatus` class.

    Raises:
        ConnectionError: If the server address is invalid or not reachable.


    To push the prediction results to an Azure Blob Storage, you must have a valid Azure Blob Storage credential.

    If you already have one, please set the credential in the **AZURE_STORAGE_CREDENTIAL** environment variable or pass it in the parameter.

    If you don't know how to get the Azure Blob Storage credential, you can follow the below python example. It's one of the ways to get the Azure
    Blob Storage credential.

    - If you don't already have it, please install the Azure Storage Management (azure-mgmt-storage) and Azure Resource Management (azure-mgmt-resource)
    python client libraries using pip.

    - Import the **DefaultAzureCredential** and **StorageManagementClient** class from the respective package

    - Ensure the following four environment variables are set into your current working environment with appropriate values.

        - AZURE_TENANT_ID

        - AZURE_CLIENT_ID

        - AZURE_CLIENT_SECRET

        - AZURE_SUBSCRIPTION_ID

    - Assign the resource group name in the **GROUP_NAME** variable and the storage account name in the **STORAGE_ACCOUNT_NAME** variable.

    - Finally, use the StorageManagementClient and DefaultAzureCredential classes to get the credential.

    - Below is a sample code. Please copy it and replace the placeholders with appropriate values to get the Azure Blob Storage credential:

        ```python
            # Importing necessary libraries
            import os
            from azure.identity import DefaultAzureCredential
            from azure.mgmt.storage import StorageManagementClient

            # Setting the required environment variables
            os.environ["AZURE_SUBSCRIPTION_ID"] = "{fill in here}"
            os.environ["AZURE_CLIENT_ID"] = "{fill in here}"
            os.environ["AZURE_CLIENT_SECRET"] = "{fill in here}"
            os.environ["AZURE_TENANT_ID"]= "{fill in here}"

            # Setting the resource group name and storage account name
            GROUP_NAME = "{fill in here}"
            STORAGE_ACCOUNT_NAME = "{fill in here}"

            # Retrieving the credential
            storage_client = StorageManagementClient(
                DefaultAzureCredential(), os.environ["AZURE_SUBSCRIPTION_ID"]
            )
            storage_keys = storage_client.storage_accounts.list_keys(GROUP_NAME, STORAGE_ACCOUNT_NAME)
            storage_keys = {v.key_name: v.value for v in storage_keys.keys}
            credential = storage_keys['key1']
        ```

    - Once you have the credential, please set it in the **AZURE_STORAGE_CREDENTIAL** environment variable or pass it in the parameter
    to the `to_azure_blob_storage` method to push the prediction results to an Azure Blob Storage.

        ```python
            status = prediction.to_azure_blob_storage(
                uri="{fill in here}",
                credential=credential
            )
            status.progress_bar()
        ```
    """
    credential = (
        credential if credential is not None else os.environ["AZURE_STORAGE_CREDENTIAL"]
    )
    response = Client._post_data(
        relative_url=f"/prediction/{self.uuid}/to_azure_blob_storage",
        data=dict(uri=uri, credential=credential),
    )

    return ProgressStatus(relative_url=f"/prediction/push/{response['uuid']}")


# Cell


@patch
def to_local(
    self: Prediction,
    path: Union[str, Path],
    show_progress: Optional[bool] = True,
) -> None:
    """Download the prediction results to a local directory.

    Args:
        path: Local directory path.
        show_progress: Flag to set the progressbar visibility. If not passed, then the default value **True** will be used.

    Raises:
        FileNotFoundError: If the **path** is invalid.
        HTTPError: If the presigned AWS s3 uri to download the prediction results are invalid or not reachable.

    An example to download the predictions to local:

    ```python
        prediction.to_local(
            path=Path('path-to-local-directory')
        )
    ```
    """
    response = Client._get_data(relative_url=f"/prediction/{self.uuid}/to_local")

    # Initiate progress bar
    t = tqdm(total=len(response), disable=not show_progress)

    for file_name, url in response.items():
        Prediction._download_prediction_file_to_local(file_name, url, Path(path))
        t.update()

    t.close()


# Cell


@patch
def to_mysql(
    self: Prediction,
    *,
    host: str,
    database: str,
    table: str,
    port: int = 3306,
    username: Optional[str] = None,
    password: Optional[str] = None,
) -> ProgressStatus:
    """Push the prediction results to a mysql database.

    If the database requires authentication, pass the username/password as parameters or store it in
    the **AIRT_CLIENT_DB_USERNAME** and **AIRT_CLIENT_DB_PASSWORD** environment variables.

    Args:
        host: Database host name.
        database: Database name.
        table: Table name.
        port: Host port number. If not passed, then the default value **3306** will be used.
        username: Database username. If not passed, then the value set in the environment variable
            **AIRT_CLIENT_DB_USERNAME** will be used else the default value "root" will be used.
        password: Database password. If not passed, then the value set in the environment variable
            **AIRT_CLIENT_DB_PASSWORD** will be used else the default value "" will be used.

    Returns:
        An instance of `ProgressStatus` class.

    Raises:
        ConnectionError: If the server address is invalid or not reachable.

    An example to push the prediction results to a mysql database:

    ```python
        status = prediction.to_mysql(
            host="host_name",
            database="database_name",
            table="table_name"
        )
        status.progress_bar()
    ```
    """
    username = (
        username if username is not None else os.environ.get(CLIENT_DB_USERNAME, "root")
    )

    password = (
        password if password is not None else os.environ.get(CLIENT_DB_PASSWORD, "")
    )

    _body = dict(
        host=host,
        port=port,
        username=username,
        password=password,
        database=database,
        table=table,
    )

    response = Client._post_data(
        relative_url=f"/prediction/{self.uuid}/to_mysql", data=_body
    )

    return ProgressStatus(relative_url=f"/prediction/push/{response['uuid']}")


# Cell


@patch
def to_clickhouse(
    self: Prediction,
    *,
    host: str,
    database: str,
    table: str,
    protocol: str,
    port: int = 0,
    username: Optional[str] = None,
    password: Optional[str] = None,
) -> ProgressStatus:
    """Push the prediction results to a clickhouse database.

    If the database requires authentication, pass the username/password as parameters or store it in
    the **CLICKHOUSE_USERNAME** and **CLICKHOUSE_PASSWORD** environment variables.

    Args:
        host: Remote database host name.
        database: Database name.
        table: Table name.
        protocol: Protocol to use (native/http).
        port: Host port number. If not passed, then the default value **0** will be used.
        username: Database username. If not passed, then the value set in the environment variable
            **CLICKHOUSE_USERNAME** will be used else the default value "root" will be used.
        password: Database password. If not passed, then the value set in the environment variable
            **CLICKHOUSE_PASSWORD** will be used else the default value "" will be used.

    Returns:
        An instance of `ProgressStatus` class.

    Raises:
        ConnectionError: If the server address is invalid or not reachable.

    An example to push the prediction results to a clickhouse database:

    ```python
        status = prediction.to_clickhouse(
            host="host_name",
            database="database_name",
            table="table_name",
            protocol="native",
        )
        status.progress_bar()
    ```
    """
    username = (
        username
        if username is not None
        else os.environ.get("CLICKHOUSE_USERNAME", "root")
    )

    password = (
        password if password is not None else os.environ.get("CLICKHOUSE_PASSWORD", "")
    )

    _body = dict(
        host=host,
        database=database,
        table=table,
        protocol=protocol,
        port=port,
        username=username,
        password=password,
    )

    response = Client._post_data(
        relative_url=f"/prediction/{self.uuid}/to_clickhouse", data=_body
    )

    return ProgressStatus(relative_url=f"/prediction/push/{response['uuid']}")
