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

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

# %% ../../notebooks/CLI_DataSource.ipynb 3
from typing import *

# %% ../../notebooks/CLI_DataSource.ipynb 4
import os

import typer
from typer import echo
from tabulate import tabulate
import datetime as dt
import pandas as pd

from airt.client import Client
from airt.cli import helper
from airt.logger import get_logger, set_level

# %% ../../notebooks/CLI_DataSource.ipynb 6
app = typer.Typer(
    help="A set of commands for managing datasources and training ML models on them."
)

# %% ../../notebooks/CLI_DataSource.ipynb 8
logger = get_logger(__name__)

# %% ../../notebooks/CLI_DataSource.ipynb 14
@app.command()
@helper.requires_auth_token
def dtypes(
    uuid: str = typer.Argument(
        ...,
        help="Datasource uuid.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
) -> None:
    """Return the dtypes of the datasource."""

    from airt.client import DataSource

    ds = DataSource(uuid=uuid)
    dtypes = ds.dtypes.T.rename(columns={0: "dtype"})
    typer.echo(dtypes)

# %% ../../notebooks/CLI_DataSource.ipynb 19
@app.command()
@helper.display_formated_table
@helper.requires_auth_token
def head(
    uuid: str = typer.Argument(
        ...,
        help="Datasource uuid.",
    ),
    format: Optional[str] = typer.Option(
        None,
        "--format",
        "-f",
        help="Format output and show only the given column(s) values.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
) -> Dict["str", Union[pd.DataFrame, str]]:
    """Return the first few rows of the datasource."""

    from airt.client import DataSource

    ds = DataSource(uuid=uuid)
    df = ds.head()

    return {"df": df}

# %% ../../notebooks/CLI_DataSource.ipynb 22
@app.command()
@helper.display_formated_table
@helper.requires_auth_token
def ls(
    offset: int = typer.Option(
        0,
        "--offset",
        "-o",
        help="The number of datasources to offset at the beginning. If **None**, then the default value **0** will be used.",
    ),
    limit: int = typer.Option(
        100,
        "--limit",
        "-l",
        help="The maximum number of datasources to return from the server. If **None**, then the default value **100** will be used.",
    ),
    disabled: bool = typer.Option(
        False,
        "--disabled",
        help="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: bool = typer.Option(
        False,
        "--completed",
        help="If set to **True**, then only the datasources that are successfully downloaded"
        "to the server will be returned. Else, the default value **False** will be used to"
        "return all the datasources.",
    ),
    format: Optional[str] = typer.Option(
        None,
        "--format",
        "-f",
        help="Format output and show only the given column(s) values.",
    ),
    quiet: bool = typer.Option(
        False,
        "--quiet",
        "-q",
        help="Output only datasource uuids separated by space",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
) -> Dict["str", Union[pd.DataFrame, str]]:
    """Return the list of datasources."""

    from airt.client import DataSource

    dsx = DataSource.ls(
        offset=offset, limit=limit, disabled=disabled, completed=completed
    )
    df = DataSource.as_df(dsx)

    df["pulled_on"] = helper.humanize_date(df["pulled_on"])
    df["no_of_rows"] = helper.humanize_number(df["no_of_rows"])
    df["folder_size"] = helper.humanize_size(df["folder_size"])

    return {"df": df, "quite_column_name": "datasource_uuid"}

# %% ../../notebooks/CLI_DataSource.ipynb 27
@app.command()
@helper.requires_auth_token
def train(
    uuid: str = typer.Option(
        ...,
        "--datasource_uuid",
        "-uuid",
        help="Datasource uuid.",
    ),
    client_column: str = typer.Option(
        ...,
        "--client_column",
        help="The column name that uniquely identifies the users/clients.",
    ),
    timestamp_column: Optional[str] = typer.Option(
        None,
        "--timestamp_column",
        help="The timestamp column indicating the time of an event. If not passed, then the default value **None** will be used. ",
    ),
    target_column: str = typer.Option(
        ...,
        "--target_column",
        help="Target column name that indicates the type of the event.",
    ),
    target: str = typer.Option(
        ...,
        "--target",
        help="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: str = typer.Option(
        ...,
        "--predict_after",
        help="Time delta in hours of the expected target event.",
    ),
    quiet: bool = typer.Option(
        False,
        "--quiet",
        "-q",
        help="Output model uuid only.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
):
    """Train a model against the datasource."""

    from airt.client import DataSource
    from datetime import timedelta

    ds = DataSource(uuid=uuid)

    model = ds.train(
        client_column=client_column,
        target_column=target_column,
        target=target,
        predict_after=timedelta(hours=int(predict_after)),
        timestamp_column=timestamp_column,
    )

    if quiet:
        model.wait()
        typer.echo(f"{model.uuid}")
    else:
        typer.echo(f"Training started for model uuid: {model.uuid}")
        model.progress_bar()

# %% ../../notebooks/CLI_DataSource.ipynb 32
@app.command()
@helper.display_formated_table
@helper.requires_auth_token
def details(
    uuid: str = typer.Argument(
        ...,
        help="Datasource uuid.",
    ),
    format: Optional[str] = typer.Option(
        None,
        "--format",
        "-f",
        help="Format output and show only the given column(s) values.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
) -> Dict["str", Union[pd.DataFrame, str]]:
    """Return details of a datasource."""

    from airt.client import DataSource

    ds = DataSource(uuid=uuid)
    df = ds.details()

    df["pulled_on"] = helper.humanize_date(df["pulled_on"])
    df["no_of_rows"] = helper.humanize_number(df["no_of_rows"])
    df["folder_size"] = helper.humanize_size(df["folder_size"])

    return {"df": df}

# %% ../../notebooks/CLI_DataSource.ipynb 36
@app.command()
@helper.display_formated_table
@helper.requires_auth_token
def rm(
    uuid: str = typer.Argument(
        ...,
        help="Datasource uuid.",
    ),
    format: Optional[str] = typer.Option(
        None,
        "--format",
        "-f",
        help="Format output and show only the given column(s) values.",
    ),
    quiet: bool = typer.Option(
        False,
        "--quiet",
        "-q",
        help="Output the deleted datasource uuid only.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
) -> Dict["str", Union[pd.DataFrame, str]]:
    """Delete a datasource from the server."""

    from airt.client import DataSource

    ds = DataSource(uuid=uuid)
    df = ds.delete()

    df["pulled_on"] = helper.humanize_date(df["pulled_on"])
    df["no_of_rows"] = helper.humanize_number(df["no_of_rows"])
    df["folder_size"] = helper.humanize_size(df["folder_size"])

    return {"df": df, "quite_column_name": "datasource_uuid"}

# %% ../../notebooks/CLI_DataSource.ipynb 40
@app.command()
@helper.display_formated_table
@helper.requires_auth_token
def tag(
    uuid: str = typer.Option(
        ...,
        "--datasource_uuid",
        "-uuid",
        help="Datasource uuid.",
    ),
    name: str = typer.Option(
        ...,
        "--name",
        "-n",
        help="A string to tag the datasource.",
    ),
    format: Optional[str] = typer.Option(
        None,
        "--format",
        "-f",
        help="Format output and show only the given column(s) values.",
    ),
    debug: bool = typer.Option(
        False,
        "--debug",
        "-d",
        help="Set logger level to DEBUG and output everything.",
    ),
) -> Dict["str", Union[pd.DataFrame, str]]:
    """Tag an existing datasource in server."""

    from airt.client import DataSource

    ds = DataSource(uuid=uuid)
    df = ds.tag(name=name)

    df["pulled_on"] = helper.humanize_date(df["pulled_on"])
    df["no_of_rows"] = helper.humanize_number(df["no_of_rows"])
    df["folder_size"] = helper.humanize_size(df["folder_size"])

    return {"df": df}
