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

__all__ = ["Model"]

# Cell

from typing import *

# Internal Cell

import os

import pandas as pd
from fastcore.foundation import patch

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

from .client import Client
from .progress_status import ProgressStatus
from .prediction import Prediction

# Internal Cell

logger = get_logger(__name__)

# Cell


class Model(ProgressStatus):
    """A class to query the status of the model training, evaluation, and prediction on the remote server.

    A `Model` class is instantiated automatically when the `DataSource.train` method of the `DataSource` class is called.
    Currently, it is the only way to instantiate the Model class.

    The model is trained on the connected data and tries 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 training status can either be viewed interactively by calling the `progress_bar` method or
    non-interactively by calling the `is_ready` method on the model instance which returns a boolean
    representing the training completion status.

    """

    def __init__(self, model_id: int):
        """Constructs a new `Model` instance

        Warning:
            Do not construct this object directly by calling the constructor, please use
            `DataSource.train()` function instead.

        Args:
            model_id: ID of the model in the airt service
        """
        self.model_id = model_id
        ProgressStatus.__init__(self, relative_url=f"/model/{self.model_id}")

    @staticmethod
    def _get_columns() -> list:
        """Returns the list of columns to be added to the DataFrame."""

        return ["id", "created", "total_steps", "completed_steps"]

    @staticmethod
    def ls(
        offset: int = 0,
        limit: int = 100,
        disabled: bool = False,
        completed: bool = False,
    ) -> pd.DataFrame:
        """Display the list of available models in airt service.

        Args:
            offset: The number of rows to offset at the beginning of the models
                list from the server.If **None**, then the default value **0** will be used.
            limit: The maximum number of rows to return from the server. If **None**,
                then the default value **100** will be used.
            disabled: If set to **True**, then only the deleted models will be displayed.
                Else, the default value **False** will be used to display only the list
                of active models.
            completed: If set to **True**, then only the models that are successfully trained
                in airt server will be displayed. Else, the default value **False** will be used
                to display all the models including the ones that are yet to finish training.

        Returns:
            A pandas dataframe with the list of available models.

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

        An example to list the available models:

        ```python
        Model.ls()
        ```
        """
        models = Client.get_data(
            relative_url=f"/model/?disabled={disabled}&completed={completed}&offset={offset}&limit={limit}"
        )

        columns = Model._get_columns()

        models_df = generate_df(models, columns)

        return add_ready_column(models_df)

    @staticmethod
    def details(id: int) -> pd.DataFrame:
        """Return details of a model

        Args:
            id: The id of the model in the airt service.

        Returns:
            A pandas dataframe with the details of the model.

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

        An example to get details of a model (id=1) from airt service:

        ```python
        Model.details(id=1)
        ```
        """

        details = Client.get_data(relative_url=f"/model/{id}")

        additional_cols = [
            "datasource_id",
            "user_id",
            "client_column",
            "target_column",
            "target",
            "predict_after",
            "timestamp_column",
            "error",
            "disabled",
        ]

        columns = Model._get_columns() + additional_cols

        details_df = pd.DataFrame(details, index=[0])[columns]

        return add_ready_column(details_df)

    @staticmethod
    def delete(id: int) -> pd.DataFrame:
        """Delete a model from airt service

        Args:
            id: The model id in airt service.

        Returns:
            A pandas DataFrame encapsulating the details of the deleted model id.

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

        An example to delete a model (id=1) from airt service:

        ```python
        Model.delete(id=1)
        ```
        """

        models = Client.delete_data(relative_url=f"/model/{id}")

        columns = Model._get_columns()

        models_df = pd.DataFrame(models, index=[0])[columns]

        return add_ready_column(models_df)

    def predict(self, data_id: Optional[int]) -> "airt.client.Prediction":  # type: ignore
        raise NotImplementedError()

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


# Cell


@patch
def evaluate(self: Model) -> pd.DataFrame:
    """A function to evaluate the performance of the trained model.

    This function 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 performance metrics 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 = Client.get_data(relative_url=f"/model/{self.model_id}/evaluate")
    return pd.DataFrame(dict(model_evaluate), index=[0]).T.rename(columns={0: "eval"})


# Cell


@patch
def predict(self: Model, data_id: Optional[int] = 0) -> Prediction:
    """A function that uses the trained model and makes predictions.

    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`

    Args:
        data_id: The data id for running the predictions. If the value is not passed then the data id used for
            training will be used for prediction aswell.
    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.
    """

    request_body = dict(data_id=data_id) if data_id else None

    response = Client.post_data(
        relative_url=f"/model/{self.model_id}/predict", data=request_body
    )

    return Prediction(
        prediction_id=response["id"], datasource_id=response["datasource_id"]
    )
