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

__all__ = ["Client", "DataSource", "Model", "Prediction"]

# Internal Cell

""" Client module

    This module encapsulates all classes conncted to the API service.

"""

# Cell

from typing import *

# Internal Cell

import pandas as pd
import numpy as np
from time import sleep
from datetime import datetime, timedelta
from tqdm import tqdm
from fastcore.foundation import patch
import requests
import os
from datetime import timedelta

# Internal Cell

# for documentation purposes only


def add_spaces(s, num_add):
    white = " " * num_add
    return white + white.join(s.splitlines(1))


def num_spaces(s):
    return [len(line) - len(line.lstrip()) for line in s.splitlines()]


def add_last_cell_as_example_in_docstring(o: Any, *, _ih, _oh):
    i = len(_ih) - 2
    sp = " " * 4
    example = f"""\n\nExample
```python
{_ih[i]}
```

Output:
```
{_oh[i] if i in _oh else ''}
```
"""
    if o.__doc__ is None:
        o.__doc__ = example
    else:
        indentations = sorted(list(set(num_spaces(o.__doc__))))
        if len(indentations) == 1:
            indentation = indentations[0]
        else:
            indentation = indentations[1]
        example = add_spaces(example, indentation)
        o.__doc__ = o.__doc__ + example if o.__doc__ is not None else example

    print(f"Documentation for {o.__name__}:\n")
    print(o.__doc__)


# Internal Cell


def check_for_valid_status_codes(response):
    # Requests lib will return True if the status code was between 200 and 400, and False otherwise.
    # This Truth Value Test is made possible because __bool__() is an overloaded method on Response object
    if response:
        return response.json()
    else:
        #         raise ValueError(f"{response.status_code=} --- {response.request.url=} --- {response.text=}")
        raise ValueError(response.json()["detail"])


# Internal Cell

# Helper Funtion for POST request


def post_data(
    url: str,
    data_to_post: dict,
    token: Optional[str] = None,
    check_status_code: Optional[bool] = True,
):
    try:
        if token is not None:
            headers = {"Authorization": f"Bearer {token}"}
            response = requests.post(url, json=data_to_post, headers=headers)
        else:
            response = requests.post(url, data=data_to_post)
        return (
            check_for_valid_status_codes(response)
            if check_status_code
            else response.json()
        )
    except requests.exceptions.RequestException as e:
        raise ConnectionError("Error: {}".format(e))


# Internal Cell

# Helper Funtion for GET request


def get_data(url: str, token: str):
    try:
        headers = {"Authorization": f"Bearer {token}"}
        response = requests.get(url, headers=headers)
        return check_for_valid_status_codes(response)
    except requests.exceptions.RequestException as e:
        raise ConnectionError("Error: {}".format(e))


# Internal Cell


class StatusBase:
    """A base class for querying status of a remote operation"""

    def __init__(self, total_steps: int, url: str, token: str):
        """Initializer

        Args:
            total_steps: number of steps to perform.
            url: server URL to get the status.
            token: user token to pass in the GET headers.
        """
        self.total_steps = total_steps
        self.url = url
        self.token = token

    def is_ready(self):
        response = get_data(url=self.url, token=self.token)
        return True if response["completed_steps"] == response["total_steps"] else False

    def progress_bar(self, sleep_for: int = 5, timeout: int = 60):
        """Blocks execution while waiting for remote action to be completed and displays a progress bar

        Args:
            sleep_for: Interval between successive API calls.
            timeout: The time period to call the API. Once expired the while loop will be terminated.
        """
        with tqdm(total=self.total_steps) as pbar:
            prev_completed_steps = 0
            i = 0
            while True:
                if 0 < timeout <= i:
                    break
                response = get_data(url=self.url, token=self.token)
                completed_steps = response["completed_steps"]
                if completed_steps != prev_completed_steps:
                    pbar.update(completed_steps - prev_completed_steps)
                    prev_completed_steps = completed_steps
                if completed_steps == self.total_steps:
                    break
                sleep(sleep_for)
                i = i + sleep_for


class DataConnectStatus(StatusBase):
    """A class for querying status of remote operations on data sources

    Args:
            total_steps: number of steps to perform.
            url: server URL to get the status.
            token: user token to pass in the GET headers.
    """

    def __init__(self, total_steps: int, url: str, token: str):
        """
        Initialise the class
        """
        StatusBase.__init__(self, total_steps, url, token)


# Cell


class Client:
    """A class for authenticating and accessing the airt service.

    Before you can use the service, you must acquire a username and password for your developer account. Please contact us by email info@airt.ai to get one.

    The username, password, and server address can be specified explicitly when initializing the `Client` object or it can be permanently stored in environment variables `AIRT_SERVICE_USERNAME`, `AIRT_SERVICE_PASSWORD`, and `AIRT_SERVER_URL`.

    Upon successful authentication, the airt services will be available to access.
    """

    def __init__(
        self, *, username: str = None, password: str = None, server: str = None
    ):
        """A function to authenticate and validate the developer token for accessing the airt services.

        Args:
            username: Username for your developer account. If not set (default value `None`), it will try to
                use the value from environment variable `AIRT_SERVICE_USERNAME`.
            password: Password for your developer account. If not set (default value `None`), it will try to
                use the value from environment variable `AIRT_SERVICE_PASSWORD`.
            server: Server address used to connect to. If not set (default value `None`), it will try to
                use the value from environment variable `AIRT_SERVER_URL`. If the variable is not set as well,
                then the default public server will be used. You should leave this to default value unless you
                are running your own server (please contact us for that possibility by email info@airt.ai).

        Returns:
            An instance of the `Client` class.

        Raises:
            ValueException: If the `username`/`password` pair does not match the one for sevice hosted at `server`.
            ConnectionError: If the server address is invalid or not reachable.
        """
        env_airt_service_username = "AIRT_SERVICE_USERNAME"
        env_airt_service_password = "AIRT_SERVICE_PASSWORD"
        env_airt_service_address = "AIRT_SERVER_URL"
        airt_prod_server_URI = "https://api.airt.ai"

        self.username = (
            self._check_env_variable(env_airt_service_username)
            if username is None
            else username
        )
        self.password = (
            self._check_env_variable(env_airt_service_password)
            if password is None
            else password
        )

        self.server = (
            self._check_env_variable(env_airt_service_address, raise_error=False)
            if server is None
            else server
        )
        if self.server is None:
            self.server = airt_prod_server_URI

        self.auth_api_response_json = post_data(
            url=f"{self.server}/token",
            data_to_post=dict(username=self.username, password=self.password),
        )
        self.token = self.auth_api_response_json["access_token"]

    def _check_env_variable(
        self, env_variable: str, raise_error: bool = True
    ) -> Optional[str]:
        """Checks if the given variable is present in the environment

        Args:
            env_variable: The name of the environment variable.
            raise_error: Flag whether to raise error or not.

        Returns:
            The value stored in the environment variable.

        Raises:
            ValueError: If no variable is set in the environment and raise_error is True.
        """
        value = os.environ.get(env_variable)
        if raise_error and value is None:
            raise ValueError(
                f"{env_variable} is either not set in environment variable nor passed in arguments"
            )
        return value


# Cell


class DataSource:
    """A class for encapsulating the data from sources like AWS S3 bucket or a database. The DataSource class is automatically instantiated by calling either the s3 or the db
        static methods of a DataSource class. Currently, it is the only way to instantiate this class.

    Currently, we support reading and pushing the data to a:

    - MySql database, and

    - AWS S3 bucket in the Parquet file format.

    We plan to add other databases and storage mediums in the future.

    For establishing the connection to the MySql database, parameters like host, port,
    database name, table name needs are required.

    And for establishing a connection to the S3 bucket, URI to the target parquet file
    is required.

    In case if access to the database requires authentication, the required username and
    password for the database are automatically read from either the environment variables
    `AIRT_CLIENT_DB_USERNAME` and `AIRT_CLIENT_DB_PASSWORD` or from the username and password
    parameters that are passed to the `DataSource.db` function.

    All the function calls to the library are asynchronous and they return immediately.
    To manage completion, all methods inside the returned object will return a status object
    and a method to display an interactive progress bar that can be called to check the progress.

    Below are code examples for accessing the above methods:

    An example to check for the status flag of s3 connection:

    ```python
    data_source_s3 = DataSource.s3(
        client,
        uri="s3://bucket/events.parquet"
    )
    status = data_source_s3.pull()
    status.is_ready()
    ```

    An example to display an interactive progress bar of s3 connection:

    ```python
    data_source_s3 = DataSource.s3(
        client,
        uri="s3://bucket/events.parquet"
    )
    status = data_source_s3.pull()
    status.progress_bar()
    ```

    """

    def __init__(
        self,
        total_steps: int,
        completed_steps: int,
        user_id: int,
        data_id: int,
        client: Client,
    ):
        """
        Constructs a new DataSource instance.
        """
        self.total_steps = total_steps
        self.completed_steps = completed_steps
        self.user_id = user_id
        self.id = data_id
        self.client = client

    @staticmethod
    def db(
        *,
        host: str,
        database: str,
        table: str,
        port: Optional[str] = None,
        engine: Optional[str] = None,
        username: Optional[str] = None,
        password: Optional[str] = None,
    ):
        """Create DB backed data source

        A static method that creates and returns an object that encapsulates the data
        from a database. In case if access to the database requires authentication,
        the username and password will be read either from the arguments or in the airt
        environment variables.

        The objects created by calling this method won't establish the connection yet.

        Args:
            host: The host name (subdomain) of the database as a string.
            database: The name of the database as a string.
            table: The name of the table as a string.
            port: The port number as a string. If the value is not passed then the default
                port number will be used (e.g. for MySQL we will use 3306)
            engine: The name of the database engine as a string. If the value is not passed
                then the default database engine for MySQL (MyISAM) will be used.
            username: The username to connect to the database as a string. If the value is
                not passed then the default username for MYSQL (root) will be used.
            password: The password to connect to the database as a string. If the value is
                not passed then the default password for MYSQL (no password) will be used.

        Returns:
            An instance of the `DataSource` class. For more information on the methods that
            are available in the returned object, please check the documentation of the
            `DataSource` class.

        An example function call to the DataSource.db:

        ```python
        data_source = DataSource.db(
            host="db.staging.airt.ai",
            database="test",
            table="events"
        )
        ```
        """
        return "DataSource"

    @staticmethod
    def s3(
        client: Client,
        uri: str = None,
        access_key: str = None,
        secret_key: str = None,
    ) -> "DataSource":
        """
        A static method that creates and returns an object that encapsulates the data from a AWS S3 bucket.

        Args:
            client: An instance of `Client` that encapsulates the parameters required
                for establishing the connection with the data source.
            uri: The path to the Parquet file located in the AWS S3 bucket as a string.
            access_key: The access key for the S3 bucket. If `None` (default value), then the value
                of environment variable `AWS_ACCESS_KEY_ID` is used.
            secret_key: The secret key for the S3 bucket. If `None` (default value), then the value
                of environment variable `AWS_SECRET_ACCESS_KEY` is used.

        Returns:
            An instance of the `DataSource` class. For more information on the methods that are available in
            the returned object, please check the documentation of the `DataSource` class.

        Raises:
            ValueError: If the parameters `client` and `URI` are empty or None.
            ValueError: If the input parameters to the API are invalid.
            ConnectionError: If the server address is invalid or not reachable.

        An example function call to the DataSource.s3:

        ```python
            data_source_s3 = DataSource.s3(
                client,
                uri="s3://bucket/events.parquet"
            )
        ```

        """

        if client is None or client == "":
            raise ValueError("client cannot be empty or None")

        if uri is None or uri == "":
            raise ValueError("uri cannot be empty or None")

        if access_key is None:
            access_key = os.environ["AWS_ACCESS_KEY_ID"]

        if secret_key is None:
            secret_key = os.environ["AWS_SECRET_ACCESS_KEY"]

        connect_api_response = post_data(
            url=f"{client.server}/data/s3",
            data_to_post=dict(uri=uri, access_key=access_key, secret_key=secret_key),
            token=client.token,
        )

        total_steps = connect_api_response["total_steps"]
        completed_steps = connect_api_response["completed_steps"]
        user_id = connect_api_response["user_id"]
        data_id = connect_api_response["id"]
        return DataSource(
            total_steps=total_steps,
            completed_steps=completed_steps,
            user_id=user_id,
            data_id=data_id,
            client=client,
        )

    @property
    def dtypes(self) -> Dict[str, str]:
        """
        This function parses the data source and returns the columns and their dtypes

        Returns:
            A dictionary that contains the columns and its type as string only key value pairs.

        Raises:
            ValueError: If the input parameters to the API are invalid.
            ConnectionError: If the server address is invalid or not reachable.

        e.g:
        ```python
        {
            "UserId": "int64",
            "Event": "str",
            "Time": "datetime64"
        }
        ```
        """

        column_response = get_data(
            url=f"{self.client.server}/data/{int(self.id)}/dtypes",
            token=self.client.token,
        )

        return column_response


# Cell


@patch
def pull(self: DataSource) -> DataConnectStatus:
    """
    This function establishes the connection with the data source. The call to this function is asynchronous and the progress
    of the connection can be checked using the progress bar or the status flag. Please refer to `DataSource` class documentation for more information.

    Returns:
        An instance of `DataConnectStatus` class.

    Raises:
        ValueError: If the input parameters to the API are invalid.
        ConnectionError: If the server address is invalid or not reachable.

    Below example shows establishing a connection with the s3 bucket:

    ```python
    from datetime import timedelta

    client = Client()
    data_source_s3 = DataSource.s3(
        client,
        uri="s3://test-airt-service/ecommerce_behavior"
    )

    data_source_s3.pull()
    ```
    """
    connect_response = get_data(
        url=f"{self.client.server}/data/{int(self.id)}/pull", token=self.client.token
    )

    return DataConnectStatus(
        total_steps=self.total_steps,
        url=f"{self.client.server}/data/{int(self.id)}",
        token=self.client.token,
    )


# Cell


@patch
def head(self: DataSource) -> pd.DataFrame:
    """
    This function will return the first few records of the data source upon successful connection.

    Returns:
        A pandas dataframe that displays the first few records of the connected data source.

    Raises:
        ValueError: If the input parameters to the API are invalid.
        ConnectionError: If the server address is invalid or not reachable.

    An example to show the first few records of the data source

    ```python
    from datetime import timedelta

    client = Client()
    data_source_s3 = DataSource.s3(
        client,
        uri="s3://test-airt-service/ecommerce_behavior"
    )
    data_source_s3.pull()

    data_source_s3.head()
    ```
    """
    head_response = get_data(
        url=f"{self.client.server}/data/{int(self.id)}/head", token=self.client.token
    )

    return pd.DataFrame(head_response)


# Cell


class Model(StatusBase):
    """
    A `Model` class is automatically instantiated when the train() method of the `Client` instance is called. Currently, it is the only way to instantiate this class.

    The model is trained on the client's data to predict a specific event in the future. For the model training and prediction, we assume the input data includes the following:

    - a column identifying a client client_column (person, car, business, etc.),

    - a column specifying a type of event we will try to predict target_column (buy, checkout, etc.),

    - a timestamp column specifying the time of an occurred event.

    Along with the above mandatory fields, each row in the data might have additional columns of int, category, float, or datetime type and they will be used to make predictions more accurate.

    Finally, we need to know how much ahead we wish to make predictions for. That lead time varies widely from application to application and can be in minutes for a webshop or even several weeks for a banking product such as a loan.

    As always, the model training and prediction will happen asynchronously and can take a few hours based on the size of your dataset.

    The respective status can be viewed by calling the progress_bar() or the status flag available on the returned object. For more information, please check the documentation of `DataSource`

    """

    def __init__(self, model_id: int, data_id: int, token: str, server: Optional[str]):
        """
        Constructs a new Model instance
        """
        self.model_id = model_id
        self.data_id = data_id
        self.token = token
        self.server = server
        model_status_url = f"{self.server}/model/{self.model_id}"
        self.model_status = get_data(url=model_status_url, token=self.token)

        StatusBase.__init__(
            self,
            self.model_status["total_steps"],
            url=model_status_url,
            token=self.token,
        )


# Cell


class Prediction(StatusBase):
    """
    The `Predict` class is automatically instantiated by calling the `train` method of a `Client` instance. Currently, it is the only way to instantiate this class. The returned object will have utility methods like converting the prediction results into a pandas dataframe and pushing the prediction results into one of the supported data sources.

    For more information on the supported data sources, please refer to the documentation on `DataSource` class

    """

    def __init__(
        self,
        prediction_id: int,
        datasource_id: int,
        model_id: int,
        token: str,
        server: Optional[str],
    ):
        """
        Constructs a new Prediction instance
        """
        self.prediction_id = prediction_id
        self.datasource_id = datasource_id
        self.model_id = model_id
        self.token = token
        self.server = server
        pred_status_url = f"{self.server}/prediction/{self.prediction_id}"
        self.pred_status = get_data(url=pred_status_url, token=self.token)

        StatusBase.__init__(
            self, self.pred_status["total_steps"], url=pred_status_url, token=self.token
        )


# Cell


@patch
def train(
    self: Client,
    data_source: "DataSource",
    *,
    client_column: str,
    timestamp_column: Optional[str] = None,
    target_column: str,
    target: str,
    predict_after: timedelta,
) -> Model:
    """This function trains the model for predicting which clients are most likely to have a specified
    event in the future. For more information on the model, please check the documentation of `Model` class.

    Args:
        data_source: An instance of `DataSource` class
        client_column: Name of the client column as string
        timestamp_column: Optional; Name of the timestamp_column specifying the time of an
            occurred event as a string. If the value is not passed then the timestamp_column will
            not be used for model training.
        target_column: Name of the target column that captures the type of event as string. This will
            be used for training the model as well as for making predictions for our target event.
        target: Name of the target event for as string. You can pass regular expressions as well to
            this parameter for making predictions for more than one event. For example, the
            passing `*checkout` will train a model to predict which users will do any kind of a
            checkout event.
        predict_after: Time delta in hours of the expected target event mentioned as timedelta.

    Returns:
        An instance of the `Model` class.

    Raises:
        ValueError: If any of the required parameters are empty or None.
        ValueError: If the input parameters to the API are invalid.
        ConnectionError: If the server address is invalid or not reachable.

    Below is an example for training a model  to predict which users will do any kind of a checkout event (*checkout) 3 hours before they do it:

    ```python
    from datetime import timedelta

    model = client.train(
        data_source = data_source_s3,
        client_column="user_id",
        target_column="event_type",
        target="*purchase",
        predict_after=timedelta(hours=3)
    )

    model.progress_bar()
    ```
    """

    if not isinstance(client_column, str) or len(client_column) == 0:
        raise ValueError("client_column should be a non-empty string")

    if not isinstance(target_column, str) or len(target_column) == 0:
        raise ValueError("target_column should be a non-empty string")

    if not isinstance(target, str) or len(target) == 0:
        raise ValueError("target should be a non-empty string")

    if not isinstance(predict_after, timedelta):
        raise ValueError("predict_after should be a datetime object")

    train_api_response = post_data(
        url=f"{self.server}/model/train",
        data_to_post=dict(
            data_id=int(data_source.id),
            client_column=client_column,
            target_column=str(target_column),
            target=str(target),
            predict_after=int(predict_after.total_seconds()),
        ),
        token=self.token,
    )

    model_id = train_api_response["id"]
    data_id = train_api_response["datasource_id"]
    return Model(
        model_id=model_id, data_id=data_id, token=self.token, server=self.server
    )


# Cell


@patch
def evaluate(self: Model) -> pd.DataFrame:
    """
    This function evaluates the performance of the trained model and returns the performance metrics like accuracy, precision, and recall. Currently, the function returns only the above-mentioned metrics and we plan to add more performance metrics in the future.

    Returns:
        A pandas Series that has the accuracy, recall, and precision of the trained model.

    Raises:
        ValueError: If the input parameters to the API are invalid.
        ConnectionError: If the server address is invalid or not reachable.
    """
    model_evaluate = get_data(
        url=f"{self.server}/model/{self.model_id}/evaluate", token=self.token
    )

    return pd.DataFrame(dict(model_evaluate), index=[0]).T.rename(columns={0: "eval"})


# Cell


@patch
def predict(self: Model) -> "Prediction":
    """
    This function leverages the trained model and makes predictions for a specified event in the future. As always, this function is asynchronous and can take a few hours based on the size of your dataset.

    The status of the model prediction can be viewed interactively by calling the `progress_bar` method available on the returned object. For more information, please check the documentation of `DataSource`

    Returns:
        An instance of the `Prediction` class. For more information on the methods that are available in the returned object, please check the documentation of the `Prediction` class

    Raises:
        ValueError: If the input parameters to the API are invalid.
        ConnectionError: If the server address is invalid or not reachable.
    """

    predict_res_json = post_data(
        url=f"{self.server}/model/{self.model_id}/predict",
        data_to_post=dict(data_id=self.data_id),
        token=self.token,
    )

    return Prediction(
        prediction_id=predict_res_json["id"],
        datasource_id=predict_res_json["datasource_id"],
        model_id=predict_res_json["model_id"],
        token=self.token,
        server=self.server,
    )


# Cell


@patch
def to_pandas(self: Prediction) -> pd.DataFrame:
    """
    This function converts the predicted results and returns them as a Pandas DataFrame object.
    This allows the end users to download the prediction dataset locally.

    Returns:
        The result of the model predictions in a Pandas DataFrame.

    Raises:
        ValueError: If the input parameters to the API are invalid.
        ConnectionError: If the server address is invalid or not reachable.

    The below example illustrates the usage of to_pandas function:

    ```python
    from datetime import timedelta

    client = Client()
    data_source_s3 = DataSource.s3(
        client,
        uri="s3://test-airt-service/ecommerce_behavior"
    )
    data_source_s3.pull()
    model = client.train(
        data_source_s3,
        client_column="user_id",
        target_column="event_type",
        target="*purchase",
        predict_after=timedelta(hours=3),
    )

    predictions = model.predict()
    predictions.to_pandas()
    ```
    """
    pandas_res_json = get_data(
        url=f"{self.server}/prediction/{self.prediction_id}/pandas", token=self.token
    )
    keys = list(pandas_res_json.keys())
    keys.remove("Score")
    index_name = keys[0]
    return (
        pd.DataFrame(pandas_res_json)
        .set_index(index_name)
        .sort_values("Score", ascending=False)
    )


# Cell


@patch
def push(self: Prediction, data_source: DataSource):
    """
    This function pushes the prediction results into the target data source.

    For more information on the supported data sources, please refer to the documentation on `DataSource` class

    Args:
        data_source: An instance of the `DataSource` class that encapsulates the data.

    Raises:
        ValueError: If the input parameters to the API are invalid.
        ConnectionError: If the server address is invalid or not reachable.


    The below example illustrates pushing the prediction results to a database:

    ```python
    from datetime import timedelta

    client = Client()
    data_source_s3 = DataSource.s3(
        client,
        uri="s3://test-airt-service/ecommerce_behavior"
    )
    data_source_s3.pull()
    model = client.train(
        data_source_s3,
        client_column="user_id",
        target_column="event_type",
        target="*purchase",
        predict_after=timedelta(hours=3),
    )
    data_source_pred = DataSource.s3(
        client,
        uri="s3://target-bucket"
    )
    predictions = model.predict()
    predictions.push(data_source_pred)
    ```

    """

    push_to_res_json = post_data(
        url=f"{self.server}/prediction/{self.prediction_id}/push",
        data_to_post=dict(data_id=data_source.id),
        token=self.token,
    )
