# AUTOGENERATED! DO NOT EDIT! File to edit: ../../notebooks/API_DataSource.ipynb.

# %% auto 0
__all__ = ['DataSource']

# %% ../../notebooks/API_DataSource.ipynb 4
from typing import *

# %% ../../notebooks/API_DataSource.ipynb 5
import os

import pandas as pd
from fastcore.foundation import patch
from datetime import datetime, timedelta
from pathlib import Path
import requests

from airt.logger import get_logger, set_level
from airt.helper import (
    get_data,
    post_data,
    delete_data,
    add_ready_column,
    get_values_from_item,
    generate_df,
    dict_to_df,
    add_example_to_docs,
)


from airt.components.client import Client
from airt.components.progress_status import ProgressStatus
from airt.components.model import Model

# %% ../../notebooks/API_DataSource.ipynb 7
logger = get_logger(__name__)

# %% ../../notebooks/API_DataSource.ipynb 10
class DataSource:
    """A class for managing datasources and training ML models on them.

    To instantiate the **DataSource** class, please call `DataBlob.to_datasource` method of the `DataBlob` class.

    The **DataSource** class has two categories of methods,

    * Methods for managing the datasources.
    * Method for training a model against a datasource.

    Methods such as `delete`, `ls`, `details`, `head`, etc., can be used to manage and obtain additional information from a datasource instance.

    And, the `train` method can be used to train a new model against a datasource instance.

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

    BASIC_DS_COLS = [
        "uuid",
        "datablob",
        "region",
        "cloud_provider",
        "tags",
        "pulled_on",
        "completed_steps",
        "total_steps",
        "no_of_rows",
        "folder_size",
    ]

    ALL_DS_COLS = BASIC_DS_COLS + ["user", "error", "disabled"]

    COLS_TO_RENAME = {
        "uuid": "datasource_uuid",
        "datablob": "datablob_uuid",
        "user": "user_uuid",
    }

    def __init__(
        self,
        uuid: str,
        datablob: Optional[str] = None,
        folder_size: Optional[int] = None,
        no_of_rows: Optional[int] = None,
        error: Optional[str] = None,
        disabled: Optional[bool] = None,
        created: Optional[str] = None,
        pulled_on: Optional[str] = None,
        user: Optional[str] = None,
        hash: Optional[str] = None,
        region: Optional[str] = None,
        cloud_provider: Optional[str] = None,
        tags: Optional[List[Dict[str, str]]] = None,
        total_steps: Optional[int] = None,
        completed_steps: Optional[int] = None,
    ):
        """Constructs a new `DataSource` instance.

        Warning:
            Do not construct this object directly by calling the constructor, please use `DataBlob.to_datasource` method instead.

        Args:
            uuid: DataSource uuid.
            datablob: Datablob uuid.
            folder_size: The uploaded datasource's size in bytes.
            no_of_rows: The number of records in the datasource.
            error: Contains the error message if the processing of the datasource fails.
            disabled: A flag that indicates the datasource's status. If the datasource is deleted, then **False** will be set.
            created: The datasource creation date.
            pulled_on: The most recent date the datasource was uploaded.
            user: The uuid of the user who created the datasource.
            hash: The datasource hash.
            region: The region name of the cloud provider where the datasource is stored
            cloud_provider: The name of the cloud storage provider where the datasource is stored.
            tags: Tag names associated with the datasource.
            total_steps: The number of steps required to upload the datasource to the server.
            completed_steps: The number of steps completed during the datasource's upload to the server.
        """
        self.uuid = uuid
        self.datablob = datablob
        self.folder_size = folder_size
        self.no_of_rows = no_of_rows
        self.error = error
        self.disabled = disabled
        self.created = created
        self.pulled_on = pulled_on
        self.user = user
        self.hash = hash
        self.region = region
        self.cloud_provider = cloud_provider
        self.tags = tags
        self.total_steps = total_steps
        self.completed_steps = completed_steps

    @property
    def dtypes(self) -> pd.DataFrame:
        """Return the dtypes of the datasource.

        Returns:
            A pandas DataFrame containing the data type of each column.

        Raises:
            ConnectionError: If the server address is invalid or not reachable.
        """
        dtypes = Client._get_data(relative_url=f"/datasource/{self.uuid}/dtypes")
        return pd.DataFrame([dtypes])

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

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

        Returns:
            A list of `DataSource` instances available in server.

        Raises:
            ConnectionError: If the server address is invalid or not reachable.
        """
        lists = Client._get_data(
            relative_url=f"/datasource/?disabled={disabled}&completed={completed}&offset={offset}&limit={limit}"
        )

        dsx = [
            DataSource(
                uuid=ds["uuid"],
                datablob=ds["datablob"],
                folder_size=ds["folder_size"],
                no_of_rows=ds["no_of_rows"],
                region=ds["region"],
                cloud_provider=ds["cloud_provider"],
                error=ds["error"],
                disabled=ds["disabled"],
                created=ds["created"],
                pulled_on=ds["pulled_on"],
                user=ds["user"],
                hash=ds["hash"],
                tags=ds["tags"],
                total_steps=ds["total_steps"],
                completed_steps=ds["completed_steps"],
            )
            for ds in lists
        ]

        return dsx

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

        Args:
            dsx: List of `DataSource` instances.

        Returns:
            Details of the datasources in a dataframe.

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

        ds_lists = [{i: getattr(ds, i) for i in DataSource.ALL_DS_COLS} for ds in dsx]

        for ds in ds_lists:
            ds["tags"] = get_values_from_item(ds["tags"], "name")

        lists_df = generate_df(ds_lists, DataSource.BASIC_DS_COLS)
        df = add_ready_column(lists_df)

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

        return df

    def is_ready(self):
        raise NotImplementedError()

    def progress_bar(self, sleep_for: Union[int, float] = 5, timeout: int = 0):
        raise NotImplementedError()

    def wait(self, sleep_for: Union[int, float] = 1, timeout: int = 0):
        raise NotImplementedError()

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

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

    def tag(self, name: str) -> pd.DataFrame:
        raise NotImplementedError()

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

    def train(
        self,
        *,
        client_column: str,
        timestamp_column: Optional[str] = None,
        target_column: str,
        target: str,
        predict_after: timedelta,
    ) -> "airt.client.Model":  # type: ignore
        raise NotImplementedError()

# %% ../../notebooks/API_DataSource.ipynb 11
def _docstring_example():
    """
    Example:
        ```python
        # Importing necessary libraries
        from datetime import timedelta

        from  airt.client import Client, DataBlob, DataSource

        # Authenticate
        Client.get_token(username="{fill in username}", password="{fill in password}")

        # Create a datablob
        # In this example, the datablob will be stored in an AWS S3 bucket. The
        # access_key and the secret_key are set in the AWS_ACCESS_KEY_ID and
        # AWS_SECRET_ACCESS_KEY environment variables, and the region is set to
        # eu-west-3; feel free to change the cloud provider and the region to
        # suit your needs.
        db = DataBlob.from_s3(
            uri="{fill in uri}",
            cloud_provider="aws",
            region="eu-west-3"
        )

        # Display the status in a progress bar
        db.progress_bar()

        # Create a datasource
        ds = db.to_datasource(
            file_type="{fill in file_type}",
            index_column="{fill in index_column}",
            sort_by="{fill in sort_by}",
        )

        # Display the status in a progress bar
        # Call the wait method to wait for the progress to finish but
        # without displaying an interactive progress bar.
        ds.progress_bar()

        # Display the ready status
        print(ds.is_ready())

        # Display the data types of the datasource's columns.
        print(ds.dtypes)

        # Display the details of the datasource
        print(ds.details())

        # Display the details of all datasource created by the currently
        # logged-in user
        print(DataSource.as_df(DataSource.ls()))

        # Display the first few records of the datasource
        print(ds.head())

        # Train a model against the datasource.
        # This example predicts which users will perform a purchase
        # event ("*purchase") three hours before they actually do it.
        model = ds.train(
            client_column="{fill in client_column}",
            target_column="{fill in target_column}",
            target="*purchase",
            predict_after=timedelta(hours=3)
        )

        # Display the training status in a progress bar
        model.progress_bar()

        # Display the details of the newly created model
        print(model.details())

        # Tag the datasource
        print(ds.tag(name="{fill in tag_name}"))

        # Delete the datasource
        print(ds.delete())
        ```
    """
    pass

# %% ../../notebooks/API_DataSource.ipynb 13
add_example_to_docs(DataSource, _docstring_example.__doc__)  # type: ignore
add_example_to_docs(DataSource.dtypes, _docstring_example.__doc__)  # type: ignore
add_example_to_docs(DataSource.ls, _docstring_example.__doc__)  # type: ignore
add_example_to_docs(DataSource.as_df, _docstring_example.__doc__)  # type: ignore

# %% ../../notebooks/API_DataSource.ipynb 15
@patch
def is_ready(
    self: DataSource,
) -> bool:
    """Check if the method's progress is complete.

    Returns:
        **True** if the progress is completed, else **False**.
    """
    progress_status = ProgressStatus(relative_url=f"/datasource/{self.uuid}")

    return progress_status.is_ready()

# %% ../../notebooks/API_DataSource.ipynb 16
add_example_to_docs(DataSource.is_ready, _docstring_example.__doc__)  # type: ignore

# %% ../../notebooks/API_DataSource.ipynb 18
@patch
def progress_bar(self: DataSource, sleep_for: Union[int, float] = 5, timeout: int = 0):
    """Blocks the execution and displays a progress bar showing the remote action progress.

    Args:
        sleep_for: The time interval in seconds between successive API calls.
        timeout: The maximum time allowed in seconds for the asynchronous call to complete. If not the
            progressbar will be terminated.

    Raises:
        ConnectionError: If the server address is invalid or not reachable.
        TimeoutError: in case of connection timeout.
    """
    progress_status = ProgressStatus(
        relative_url=f"/datasource/{self.uuid}", sleep_for=sleep_for, timeout=timeout
    )

    progress_status.progress_bar()

# %% ../../notebooks/API_DataSource.ipynb 19
add_example_to_docs(DataSource.progress_bar, _docstring_example.__doc__)  # type: ignore

# %% ../../notebooks/API_DataSource.ipynb 20
@patch
def wait(self: DataSource, sleep_for: Union[int, float] = 1, timeout: int = 0):
    """Blocks execution while waiting for the remote action to complete.

    Args:
        sleep_for: The time interval in seconds between successive API calls.
        timeout: The maximum time allowed in seconds for the asynchronous call to complete. If not the
            progressbar will be terminated.

    Raises:
        ConnectionError: If the server address is invalid or not reachable.
        TimeoutError: in case of timeout.
    """

    progress_status = ProgressStatus(
        relative_url=f"/datasource/{self.uuid}", sleep_for=sleep_for, timeout=timeout
    )

    progress_status.wait()

# %% ../../notebooks/API_DataSource.ipynb 21
add_example_to_docs(DataSource.wait, _docstring_example.__doc__)  # type: ignore

# %% ../../notebooks/API_DataSource.ipynb 27
@patch
def delete(self: DataSource) -> pd.DataFrame:
    """Delete a datasource from the server.

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

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

    response = Client._delete_data(relative_url=f"/datasource/{self.uuid}")

    response["tags"] = get_values_from_item(response["tags"], "name")

    df = pd.DataFrame([response])[DataSource.BASIC_DS_COLS]

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

    return add_ready_column(df)

# %% ../../notebooks/API_DataSource.ipynb 28
add_example_to_docs(DataSource.delete, _docstring_example.__doc__)  # type: ignore

# %% ../../notebooks/API_DataSource.ipynb 31
@patch
def details(self: DataSource) -> pd.DataFrame:
    """Return details of a datasource.

    Returns:
        The datasource details as a pandas dataframe.

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

    response = Client._get_data(relative_url=f"/datasource/{self.uuid}")

    response["tags"] = get_values_from_item(response["tags"], "name")

    df = pd.DataFrame([response])[DataSource.ALL_DS_COLS]

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

    return add_ready_column(df)

# %% ../../notebooks/API_DataSource.ipynb 32
add_example_to_docs(DataSource.details, _docstring_example.__doc__)  # type: ignore

# %% ../../notebooks/API_DataSource.ipynb 34
@patch
def tag(self: DataSource, name: str) -> pd.DataFrame:
    """Tag an existing datasource in server.

    Args:
        name: A string to tag the datasource.

    Returns:
        A pandas dataframe with the details of the tagged datasource.

    Raises:
        ConnectionError: If the server address is invalid or not reachable.
    """
    response = Client._post_data(
        relative_url=f"/datasource/{self.uuid}/tag", json=dict(name=name)
    )

    response["tags"] = get_values_from_item(response["tags"], "name")

    df = pd.DataFrame([response])[DataSource.BASIC_DS_COLS]
    df = df.rename(columns=DataSource.COLS_TO_RENAME)

    return add_ready_column(df)

# %% ../../notebooks/API_DataSource.ipynb 35
add_example_to_docs(DataSource.tag, _docstring_example.__doc__)  # type: ignore

# %% ../../notebooks/API_DataSource.ipynb 37
@patch
def head(self: DataSource) -> pd.DataFrame:
    """Return the first few rows of the datasource.

    Returns:
        The first few rows of the datasource as a pandas dataframe.

    Raises:
        ConnectionError: If the server address is invalid or not reachable.
    """
    response = Client._get_data(relative_url=f"/datasource/{self.uuid}/head")
    df = dict_to_df(response)

    return df

# %% ../../notebooks/API_DataSource.ipynb 38
add_example_to_docs(DataSource.head, _docstring_example.__doc__)  # type: ignore

# %% ../../notebooks/API_DataSource.ipynb 40
@patch
def train(
    self: DataSource,
    *,
    client_column: str,
    timestamp_column: Optional[str] = None,
    target_column: str,
    target: str,
    predict_after: timedelta,
) -> Model:
    """Train a model against the datasource.

    This method trains the model for predicting which clients are most likely to have a specified
    event in the future.

    The call to this method is asynchronous and the progress can be checked using the progress bar method
    or the status flag attribute available in the `DataSource` class.

    For more model specific information, please check the documentation of `Model` class.

    Args:
        client_column: The column name that uniquely identifies the users/clients.
        timestamp_column: The timestamp column indicating the time of an event. If not passed,
            then the default value **None** will be used.
        target_column: Target column name that indicates the type of the event.
        target: Target event name to train and make predictions. You can pass the target event as a string or as a
            regular expression for predicting more than one event. For example, passing ***checkout** will
            train a model to predict any checkout event.
        predict_after: Time delta in hours of the expected target event.

    Returns:
        An instance of the `Model` class.

    Raises:
        ValueError: If the input parameters to the API are invalid.
        ConnectionError: If the server address is invalid or not reachable.
    """
    response = Client._post_data(
        relative_url=f"/model/train",
        json=dict(
            data_uuid=self.uuid,
            client_column=client_column,
            target_column=target_column,
            target=target,
            predict_after=int(predict_after.total_seconds()),
        ),
    )

    return Model(uuid=response["uuid"])

# %% ../../notebooks/API_DataSource.ipynb 41
add_example_to_docs(DataSource.train, _docstring_example.__doc__)  # type: ignore
