# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/src/nixtla_client.ipynb.

# %% auto 0
__all__ = ['ApiError', 'NixtlaClient']

# %% ../nbs/src/nixtla_client.ipynb 3
import logging
import math
import os
import warnings
from collections.abc import Sequence
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import (
    TYPE_CHECKING,
    Annotated,
    Any,
    Callable,
    Literal,
    Optional,
    TypeVar,
    Union,
)

import annotated_types
import httpcore
import httpx
import numpy as np
import orjson
import pandas as pd
import utilsforecast.processing as ufp
from tenacity import (
    RetryCallState,
    retry,
    retry_if_exception,
    stop_after_attempt,
    stop_after_delay,
    wait_fixed,
)
from utilsforecast.compat import DFType, DataFrame, pl_DataFrame
from utilsforecast.feature_engineering import _add_time_features, time_features
from utilsforecast.preprocessing import id_time_grid
from utilsforecast.validation import ensure_time_dtype, validate_format

if TYPE_CHECKING:
    try:
        from fugue import AnyDataFrame
    except ModuleNotFoundError:
        pass
    try:
        import matplotlib.pyplot as plt
    except ModuleNotFoundError:
        pass
    try:
        import plotly
    except ModuleNotFoundError:
        pass
    try:
        import triad
    except ModuleNotFoundError:
        pass
    try:
        from polars import DataFrame as PolarsDataFrame
    except ModuleNotFoundError:
        pass
    try:
        from dask.dataframe import DataFrame as DaskDataFrame
    except ModuleNotFoundError:
        pass
    try:
        from pyspark.sql import DataFrame as SparkDataFrame
    except ModuleNotFoundError:
        pass
    try:
        from ray.data import Dataset as RayDataset
    except ModuleNotFoundError:
        pass

# %% ../nbs/src/nixtla_client.ipynb 4
AnyDFType = TypeVar(
    "AnyDFType",
    "DaskDataFrame",
    pd.DataFrame,
    "PolarsDataFrame",
    "RayDataset",
    "SparkDataFrame",
)
DistributedDFType = TypeVar(
    "DistributedDFType",
    "DaskDataFrame",
    "RayDataset",
    "SparkDataFrame",
)
logging.basicConfig(level=logging.INFO)
logging.getLogger("httpx").setLevel(logging.ERROR)
logger = logging.getLogger(__name__)

# %% ../nbs/src/nixtla_client.ipynb 7
_PositiveInt = Annotated[int, annotated_types.Gt(0)]
_NonNegativeInt = Annotated[int, annotated_types.Ge(0)]
_Loss = Literal["default", "mae", "mse", "rmse", "mape", "smape"]
_Model = Literal["azureai", "timegpt-1", "timegpt-1-long-horizon"]
_Finetune_Depth = Literal[1, 2, 3, 4, 5]

_date_features_by_freq = {
    # Daily frequencies
    "B": ["year", "month", "day", "weekday"],
    "C": ["year", "month", "day", "weekday"],
    "D": ["year", "month", "day", "weekday"],
    # Weekly
    "W": ["year", "week", "weekday"],
    # Monthly
    "M": ["year", "month"],
    "SM": ["year", "month", "day"],
    "BM": ["year", "month"],
    "CBM": ["year", "month"],
    "MS": ["year", "month"],
    "SMS": ["year", "month", "day"],
    "BMS": ["year", "month"],
    "CBMS": ["year", "month"],
    # Quarterly
    "Q": ["year", "quarter"],
    "BQ": ["year", "quarter"],
    "QS": ["year", "quarter"],
    "BQS": ["year", "quarter"],
    # Yearly
    "A": ["year"],
    "Y": ["year"],
    "BA": ["year"],
    "BY": ["year"],
    "AS": ["year"],
    "YS": ["year"],
    "BAS": ["year"],
    "BYS": ["year"],
    # Hourly
    "BH": ["year", "month", "day", "hour", "weekday"],
    "H": ["year", "month", "day", "hour"],
    # Minutely
    "T": ["year", "month", "day", "hour", "minute"],
    "min": ["year", "month", "day", "hour", "minute"],
    # Secondly
    "S": ["year", "month", "day", "hour", "minute", "second"],
    # Milliseconds
    "L": ["year", "month", "day", "hour", "minute", "second", "millisecond"],
    "ms": ["year", "month", "day", "hour", "minute", "second", "millisecond"],
    # Microseconds
    "U": ["year", "month", "day", "hour", "minute", "second", "microsecond"],
    "us": ["year", "month", "day", "hour", "minute", "second", "microsecond"],
    # Nanoseconds
    "N": [],
}


def _retry_strategy(max_retries: int, retry_interval: int, max_wait_time: int):
    def should_retry(exc: Exception) -> bool:
        retriable_exceptions = (
            ConnectionResetError,
            httpcore.ConnectError,
            httpcore.RemoteProtocolError,
            httpx.ConnectTimeout,
            httpx.ReadError,
            httpx.RemoteProtocolError,
            httpx.ReadTimeout,
            httpx.PoolTimeout,
            httpx.WriteError,
            httpx.WriteTimeout,
        )
        retriable_codes = [408, 409, 429, 502, 503, 504]
        return isinstance(exc, retriable_exceptions) or (
            isinstance(exc, ApiError) and exc.status_code in retriable_codes
        )

    def after_retry(retry_state: RetryCallState) -> None:
        error = retry_state.outcome.exception()
        logger.error(f"Attempt {retry_state.attempt_number} failed with error: {error}")

    return retry(
        retry=retry_if_exception(should_retry),
        wait=wait_fixed(retry_interval),
        after=after_retry,
        stop=stop_after_attempt(max_retries) | stop_after_delay(max_wait_time),
        reraise=True,
    )


def _maybe_infer_freq(
    df: DataFrame,
    freq: Optional[str],
    id_col: str,
    time_col: str,
) -> str:
    if freq is not None and freq not in ["W", "M", "Q", "Y", "A"]:
        return freq
    if isinstance(df, pl_DataFrame):
        raise ValueError(
            "Cannot infer frequency for a polars DataFrame, please set the "
            "`freq` argument to a valid polars offset.\nYou can find them at "
            "https://pola-rs.github.io/polars/py-polars/html/reference/expressions/api/polars.Expr.dt.offset_by.html"
        )
    assert isinstance(df, pd.DataFrame)
    sizes = df[id_col].value_counts(sort=True)
    times = df.loc[df[id_col] == sizes.index[0], time_col].sort_values()
    if times.dt.tz is not None:
        times = times.dt.tz_convert("UTC").dt.tz_localize(None)
    inferred_freq = pd.infer_freq(times.values)
    if inferred_freq is None:
        raise RuntimeError(
            "Could not infer the frequency of the time column. This could be due "
            "to inconsistent intervals. Please check your data for missing, "
            "duplicated or irregular timestamps"
        )
    if freq is not None:
        # check we have the same base frequency
        # except when we have yearly frequency (A, and Y means the same)
        if (freq[0] != inferred_freq[0] and freq[0] not in ("A", "Y")) or (
            freq[0] in ("A", "Y") and inferred_freq[0] not in ("A", "Y")
        ):
            raise RuntimeError(
                f"Failed to infer special date, inferred freq {inferred_freq}"
            )
    logger.info(f"Inferred freq: {inferred_freq}")
    return inferred_freq


def _standardize_freq(freq: str) -> str:
    return freq.replace("mo", "MS")


def _array_tails(
    x: np.ndarray,
    indptr: np.ndarray,
    out_sizes: np.ndarray,
) -> np.ndarray:
    if (out_sizes > np.diff(indptr)).any():
        raise ValueError("out_sizes must be at most the original sizes.")
    idxs = np.hstack(
        [np.arange(end - size, end) for end, size in zip(indptr[1:], out_sizes)]
    )
    return x[idxs]


def _tail(proc: ufp.ProcessedDF, n: int) -> ufp.ProcessedDF:
    new_sizes = np.minimum(np.diff(proc.indptr), n)
    new_indptr = np.append(0, new_sizes.cumsum())
    new_data = _array_tails(proc.data, proc.indptr, new_sizes)
    return ufp.ProcessedDF(
        uids=proc.uids,
        last_times=proc.last_times,
        data=new_data,
        indptr=new_indptr,
        sort_idxs=None,
    )


def _partition_series(
    payload: dict[str, Any], n_part: int, h: int
) -> list[dict[str, Any]]:
    parts = []
    series = payload.pop("series")
    n_series = len(series["sizes"])
    n_part = min(n_part, n_series)
    series_per_part = math.ceil(n_series / n_part)
    prev_size = 0
    for i in range(0, n_series, series_per_part):
        sizes = series["sizes"][i : i + series_per_part]
        curr_size = sum(sizes)
        part_idxs = slice(prev_size, prev_size + curr_size)
        prev_size += curr_size
        part_series = {
            "y": series["y"][part_idxs],
            "sizes": sizes,
        }
        if series["X"] is None:
            part_series["X"] = None
            if h > 0:
                part_series["X_future"] = None
        else:
            part_series["X"] = [x[part_idxs] for x in series["X"]]
            if h > 0:
                part_series["X_future"] = [
                    x[i * h : (i + series_per_part) * h] for x in series["X_future"]
                ]
        parts.append({"series": part_series, **payload})
    return parts


def _maybe_add_date_features(
    df: DFType,
    X_df: Optional[DFType],
    features: Union[bool, Sequence[Union[str, Callable]]],
    one_hot: Union[bool, list[str]],
    freq: str,
    h: int,
    id_col: str,
    time_col: str,
    target_col: str,
) -> tuple[DFType, Optional[DFType]]:
    if not features:
        return df, X_df
    if isinstance(features, list):
        date_features: Sequence[Union[str, Callable]] = features
    else:
        date_features = _date_features_by_freq.get(freq, [])
        if not date_features:
            warnings.warn(
                f"Non default date features for {freq} "
                "please provide a list of date features"
            )
    # add features
    if X_df is None:
        df, X_df = time_features(
            df=df,
            freq=freq,
            features=date_features,
            h=h,
            id_col=id_col,
            time_col=time_col,
        )
    else:
        df = _add_time_features(df, features=date_features, time_col=time_col)
        X_df = _add_time_features(X_df, features=date_features, time_col=time_col)
    # one hot
    if isinstance(one_hot, list):
        features_one_hot = one_hot
    elif one_hot:
        features_one_hot = [f for f in date_features if not callable(f)]
    else:
        features_one_hot = []
    if features_one_hot:
        X_df = ufp.assign_columns(X_df, target_col, 0)
        full_df = ufp.vertical_concat([df, X_df])
        if isinstance(full_df, pd.DataFrame):
            full_df = pd.get_dummies(full_df, columns=features_one_hot, dtype="float32")
        else:
            full_df = full_df.to_dummies(columns=features_one_hot)
        df = ufp.take_rows(full_df, slice(0, df.shape[0]))
        X_df = ufp.take_rows(full_df, slice(df.shape[0], full_df.shape[0]))
        X_df = ufp.drop_columns(X_df, target_col)
        X_df = ufp.drop_index_if_pandas(X_df)
    if h == 0:
        # time_features returns an empty df, we use it as None here
        X_df = None
    return df, X_df


def _validate_exog(
    df: DFType,
    X_df: Optional[DFType],
    id_col: str,
    time_col: str,
    target_col: str,
    hist_exog: Optional[list[str]],
) -> tuple[DFType, Optional[DFType]]:
    exogs = [c for c in df.columns if c not in (id_col, time_col, target_col)]
    if hist_exog is None:
        hist_exog = []
    if X_df is None:
        # all exogs must be historic
        ignored_exogs = [c for c in exogs if c not in hist_exog]
        if ignored_exogs:
            warnings.warn(
                f"`df` contains the following exogenous features: {ignored_exogs}, "
                "but `X_df` was not provided and they were not declared in `hist_exog_list`. "
                "They will be ignored."
            )
        exogs = [c for c in exogs if c in hist_exog]
        df = df[[id_col, time_col, target_col, *exogs]]
        return df, None

    # exogs in df that weren't declared as historic nor future
    futr_exog = [c for c in X_df.columns if c not in (id_col, time_col)]
    declared_exogs = {*hist_exog, *futr_exog}
    ignored_exogs = [c for c in exogs if c not in declared_exogs]
    if ignored_exogs:
        warnings.warn(
            f"`df` contains the following exogenous features: {ignored_exogs}, "
            "but they were not found in `X_df` nor declared in `hist_exog_list`. "
            "They will be ignored."
        )

    # future exogenous are provided in X_df that are not in df
    missing_futr = set(futr_exog) - set(exogs)
    if missing_futr:
        raise ValueError(
            "The following exogenous features are present in `X_df` "
            f"but not in `df`: {missing_futr}."
        )

    # Make sure df and X_df are in right order
    df = df[[id_col, time_col, target_col, *futr_exog, *hist_exog]]
    X_df = X_df[[id_col, time_col, *futr_exog]]

    return df, X_df


def _validate_input_size(
    df: DataFrame,
    id_col: str,
    model_input_size: int,
    model_horizon: int,
) -> None:
    min_size = ufp.counts_by_id(df, id_col)["counts"].min()
    if min_size < model_input_size + model_horizon:
        raise ValueError(
            "Some series are too short. "
            "Please make sure that each serie contains "
            f"at least {model_input_size + model_horizon} observations."
        )


def _prepare_level_and_quantiles(
    level: Optional[list[Union[int, float]]],
    quantiles: Optional[list[float]],
) -> tuple[Optional[list[Union[int, float]]], Optional[list[float]]]:
    if level is not None and quantiles is not None:
        raise ValueError("You should provide `level` or `quantiles`, but not both.")
    if quantiles is None:
        return level, quantiles
    # we recover level from quantiles
    if not all(0 < q < 1 for q in quantiles):
        raise ValueError("`quantiles` should be floats between 0 and 1.")
    level = [abs(int(100 - 200 * q)) for q in quantiles]
    return level, quantiles


def _maybe_convert_level_to_quantiles(
    df: DFType,
    quantiles: Optional[list[float]],
) -> DFType:
    if quantiles is None:
        return df
    out_cols = [c for c in df.columns if "-lo-" not in c and "-hi-" not in c]
    df = ufp.copy_if_pandas(df, deep=False)
    for q in sorted(quantiles):
        if q == 0.5:
            col = "TimeGPT"
        else:
            lv = int(100 - 200 * q)
            hi_or_lo = "lo" if lv > 0 else "hi"
            lv = abs(lv)
            col = f"TimeGPT-{hi_or_lo}-{lv}"
        q_col = f"TimeGPT-q-{int(q * 100)}"
        df = ufp.assign_columns(df, q_col, df[col])
        out_cols.append(q_col)
    return df[out_cols]


def _preprocess(
    df: DFType,
    X_df: Optional[DFType],
    h: int,
    freq: str,
    date_features: Union[bool, Sequence[Union[str, Callable]]],
    date_features_to_one_hot: Union[bool, list[str]],
    id_col: str,
    time_col: str,
    target_col: str,
) -> tuple[ufp.ProcessedDF, Optional[DFType], list[str], Optional[list[str]]]:
    df, X_df = _maybe_add_date_features(
        df=df,
        X_df=X_df,
        features=date_features,
        one_hot=date_features_to_one_hot,
        freq=freq,
        h=h,
        id_col=id_col,
        time_col=time_col,
        target_col=target_col,
    )
    processed = ufp.process_df(
        df=df, id_col=id_col, time_col=time_col, target_col=target_col
    )
    if X_df is not None:
        X_df = ensure_time_dtype(X_df, time_col=time_col)
        processed_X = ufp.process_df(
            df=X_df,
            id_col=id_col,
            time_col=time_col,
            target_col=None,
        )
        X_future = processed_X.data.T
        futr_cols = [c for c in X_df.columns if c not in (id_col, time_col)]
    else:
        X_future = None
        futr_cols = None
    x_cols = [c for c in df.columns if c not in (id_col, time_col, target_col)]
    return processed, X_future, x_cols, futr_cols


def _forecast_payload_to_in_sample(payload):
    in_sample_payload = {
        k: v
        for k, v in payload.items()
        if k not in ("h", "finetune_steps", "finetune_loss")
    }
    del in_sample_payload["series"]["X_future"]
    return in_sample_payload


def _maybe_add_intervals(
    df: DFType,
    intervals: Optional[dict[str, list[float]]],
) -> DFType:
    if intervals is None:
        return df
    intervals_df = type(df)(
        {f"TimeGPT-{k}": intervals[k] for k in sorted(intervals.keys())}
    )
    return ufp.horizontal_concat([df, intervals_df])


def _maybe_drop_id(df: DFType, id_col: str, drop: bool) -> DFType:
    if drop:
        df = ufp.drop_columns(df, id_col)
    return df


def _parse_in_sample_output(
    in_sample_output: dict[str, Union[list[float], dict[str, list[float]]]],
    df: DataFrame,
    processed: ufp.ProcessedDF,
    id_col: str,
    time_col: str,
    target_col: str,
) -> DataFrame:
    times = df[time_col].to_numpy()
    targets = df[target_col].to_numpy()
    if processed.sort_idxs is not None:
        times = times[processed.sort_idxs]
        targets = targets[processed.sort_idxs]
    times = _array_tails(times, processed.indptr, in_sample_output["sizes"])
    targets = _array_tails(targets, processed.indptr, in_sample_output["sizes"])
    uids = ufp.repeat(processed.uids, in_sample_output["sizes"])
    out = type(df)(
        {
            id_col: uids,
            time_col: times,
            target_col: targets,
            "TimeGPT": in_sample_output["mean"],
        }
    )
    return _maybe_add_intervals(out, in_sample_output["intervals"])  # type: ignore


def _restrict_input_samples(level, input_size, model_horizon, h) -> int:
    if level is not None:
        # add sufficient info to compute
        # conformal interval
        # @AzulGarza
        #  this is an old opinionated decision
        #  about reducing the data sent to the api
        #  to reduce latency when
        #  a user passes level. since currently the model
        #  uses conformal prediction, we can change a minimum
        #  amount of data if the series are too large
        new_input_size = 3 * input_size + max(model_horizon, h)
    else:
        # we only want to forecast
        new_input_size = input_size
    return new_input_size

# %% ../nbs/src/nixtla_client.ipynb 8
class ApiError(Exception):
    status_code: Optional[int]
    body: Any

    def __init__(
        self, *, status_code: Optional[int] = None, body: Optional[Any] = None
    ):
        self.status_code = status_code
        self.body = body

    def __str__(self) -> str:
        return f"status_code: {self.status_code}, body: {self.body}"

# %% ../nbs/src/nixtla_client.ipynb 9
class NixtlaClient:

    def __init__(
        self,
        api_key: Optional[str] = None,
        base_url: Optional[str] = None,
        timeout: int = 60,
        max_retries: int = 6,
        retry_interval: int = 10,
        max_wait_time: int = 6 * 60,
    ):
        """
        Client to interact with the Nixtla API.

        Parameters
        ----------
        api_key : str, optional (default=None)
            The authorization api_key interacts with the Nixtla API.
            If not provided, will use the NIXTLA_API_KEY environment variable.
        base_url : str, optional (default=None)
            Custom base_url.
            If not provided, will use the NIXTLA_BASE_URL environment variable.
        timeout : int, optional (default=60)
            Request timeout in seconds. Set this to `None` to disable it.
        max_retries : int (default=6)
            The maximum number of attempts to make when calling the API before giving up.
            It defines how many times the client will retry the API call if it fails.
            Default value is 6, indicating the client will attempt the API call up to 6 times in total
        retry_interval : int (default=10)
            The interval in seconds between consecutive retry attempts.
            This is the waiting period before the client tries to call the API again after a failed attempt.
            Default value is 10 seconds, meaning the client waits for 10 seconds between retries.
        max_wait_time : int (default=360)
            The maximum total time in seconds that the client will spend on all retry attempts before giving up.
            This sets an upper limit on the cumulative waiting time for all retry attempts.
            If this time is exceeded, the client will stop retrying and raise an exception.
            Default value is 360 seconds, meaning the client will cease retrying if the total time
            spent on retries exceeds 360 seconds.
            The client throws a ReadTimeout error after 60 seconds of inactivity. If you want to
            catch these errors, use max_wait_time >> 60.
        """
        if api_key is None:
            api_key = os.environ["NIXTLA_API_KEY"]
        if base_url is None:
            base_url = os.getenv("NIXTLA_BASE_URL", "https://api.nixtla.io")
        self._client_kwargs = {
            "base_url": base_url,
            "headers": {
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json",
            },
            "timeout": timeout,
        }
        self._retry_strategy = _retry_strategy(
            max_retries=max_retries,
            retry_interval=retry_interval,
            max_wait_time=max_wait_time,
        )
        self._model_params: dict[tuple[str, str], tuple[int, int]] = {}
        self._is_azure = "ai.azure" in base_url
        if self._is_azure:
            self.supported_models = ["azureai"]
        else:
            self.supported_models = ["timegpt-1", "timegpt-1-long-horizon"]

    def _make_request(
        self, client: httpx.Client, endpoint: str, payload: dict[str, Any]
    ) -> dict[str, Any]:
        def ensure_contiguous_arrays(d: dict[str, Any]) -> None:
            for k, v in d.items():
                if isinstance(v, np.ndarray):
                    if np.issubdtype(v.dtype, np.floating):
                        v_cont = np.ascontiguousarray(v, dtype=np.float32)
                        d[k] = np.nan_to_num(
                            v_cont,
                            nan=np.nan,
                            posinf=np.finfo(np.float32).max,
                            neginf=np.finfo(np.float32).min,
                            copy=False,
                        )
                    else:
                        d[k] = np.ascontiguousarray(v)
                elif isinstance(v, dict):
                    ensure_contiguous_arrays(v)

        ensure_contiguous_arrays(payload)
        content = orjson.dumps(payload, option=orjson.OPT_SERIALIZE_NUMPY)
        content_size_mb = len(content) / (1024 * 1024)
        if content_size_mb > 200:
            raise ValueError(
                f"The payload is too large. Set num_partitions={math.ceil(content_size_mb / 200)}"
            )
        resp = client.post(url=endpoint, content=content)
        try:
            resp_body = orjson.loads(resp.content)
        except orjson.JSONDecodeError:
            raise ApiError(
                status_code=resp.status_code,
                body=f"Could not parse JSON: {resp.content}",
            )
        if resp.status_code != 200:
            raise ApiError(status_code=resp.status_code, body=resp_body)
        if "data" in resp_body:
            resp_body = resp_body["data"]
        return resp_body

    def _make_request_with_retries(
        self,
        client: httpx.Client,
        endpoint: str,
        payload: dict[str, Any],
    ) -> dict[str, Any]:
        return self._retry_strategy(self._make_request)(
            client=client,
            endpoint=endpoint,
            payload=payload,
        )

    def _make_partitioned_requests(
        self,
        client: httpx.Client,
        endpoint: str,
        payloads: list[dict[str, Any]],
    ) -> dict[str, Any]:
        from tqdm.auto import tqdm

        num_partitions = len(payloads)
        results: list[dict[str, Any]] = [{} for _ in range(num_partitions)]
        max_workers = min(10, num_partitions)
        with ThreadPoolExecutor(max_workers) as executor:
            future2pos = {
                executor.submit(
                    self._make_request_with_retries, client, endpoint, payload
                ): i
                for i, payload in enumerate(payloads)
            }
            for future in tqdm(as_completed(future2pos), total=len(future2pos)):
                pos = future2pos[future]
                results[pos] = future.result()
        resp = {"mean": np.hstack([res["mean"] for res in results])}
        first_res = results[0]
        for k in ("sizes", "anomaly"):
            if k in first_res:
                resp[k] = np.hstack([res[k] for res in results])
        if "idxs" in first_res:
            offsets = [0] + [sum(p["series"]["sizes"]) for p in payloads[:-1]]
            resp["idxs"] = np.hstack(
                [
                    np.array(res["idxs"], dtype=np.int64) + offset
                    for res, offset in zip(results, offsets)
                ]
            )
        if first_res["intervals"] is None:
            resp["intervals"] = None
        else:
            resp["intervals"] = {}
            for k in first_res["intervals"].keys():
                resp["intervals"][k] = np.hstack(
                    [res["intervals"][k] for res in results]
                )
        if "weights_x" not in first_res or first_res["weights_x"] is None:
            resp["weights_x"] = None
        else:
            resp["weights_x"] = [res["weights_x"] for res in results]
        if (
            "feature_contributions" not in first_res
            or first_res["feature_contributions"] is None
        ):
            resp["feature_contributions"] = None
        else:
            resp["feature_contributions"] = np.vstack(
                [np.stack(res["feature_contributions"], axis=1) for res in results]
            ).T
        return resp

    def _maybe_override_model(self, model: _Model) -> _Model:
        if self._is_azure and model != "azureai":
            warnings.warn("Azure endpoint detected, setting `model` to 'azureai'.")
            model = "azureai"
        return model

    def _get_model_params(self, model: _Model, freq: str) -> tuple[int, int]:
        key = (model, freq)
        if key not in self._model_params:
            logger.info("Querying model metadata...")
            payload = {"model": model, "freq": freq}
            with httpx.Client(**self._client_kwargs) as client:
                params = self._make_request_with_retries(
                    client, "model_params", payload
                )["detail"]
            self._model_params[key] = (params["input_size"], params["horizon"])
        return self._model_params[key]

    def _maybe_assign_weights(
        self,
        weights: Optional[Union[list[float], list[list[float]]]],
        df: DataFrame,
        x_cols: list[str],
    ) -> None:
        if weights is None:
            return
        if isinstance(weights[0], list):
            self.weights_x = [
                type(df)({"features": x_cols, "weights": w}) for w in weights
            ]
        else:
            self.weights_x = type(df)({"features": x_cols, "weights": weights})

    def _maybe_assign_feature_contributions(
        self,
        expected_contributions: bool,
        resp: dict[str, Any],
        x_cols: list[str],
        out_df: DataFrame,
        insample_feat_contributions: Optional[list[list[float]]],
    ) -> None:
        if not expected_contributions:
            return
        if "feature_contributions" not in resp:
            if self._is_azure:
                warnings.warn("feature_contributions aren't implemented in Azure yet.")
                return
            else:
                raise RuntimeError(
                    "feature_contributions expected in response but not found"
                )
        feature_contributions = resp["feature_contributions"]
        if feature_contributions is None:
            return
        shap_cols = x_cols + ["base_value"]
        shap_df = type(out_df)(dict(zip(shap_cols, feature_contributions)))
        if insample_feat_contributions is not None:
            insample_shap_df = type(out_df)(
                dict(zip(shap_cols, insample_feat_contributions))
            )
            shap_df = ufp.vertical_concat([insample_shap_df, shap_df])
        self.feature_contributions = ufp.horizontal_concat([out_df, shap_df])

    def _run_validations(
        self,
        df: DFType,
        X_df: Optional[DFType],
        id_col: str,
        time_col: str,
        target_col: str,
        model: _Model,
        validate_api_key: bool,
        freq: Optional[str],
    ) -> tuple[DFType, Optional[DFType], bool, str]:
        if validate_api_key and not self.validate_api_key(log=False):
            raise Exception("API Key not valid, please email ops@nixtla.io")
        if model not in self.supported_models:
            raise ValueError(
                f"unsupported model: {model}. supported models: {self.supported_models}"
            )
        drop_id = id_col not in df.columns
        if drop_id:
            df = ufp.copy_if_pandas(df, deep=False)
            df = ufp.assign_columns(df, id_col, 0)
            if X_df is not None:
                X_df = ufp.copy_if_pandas(X_df, deep=False)
                X_df = ufp.assign_columns(X_df, id_col, 0)
        if (
            isinstance(df, pd.DataFrame)
            and time_col not in df
            and pd.api.types.is_datetime64_any_dtype(df.index)
        ):
            df.index.name = time_col
            df = df.reset_index()
        df = ensure_time_dtype(df, time_col=time_col)
        validate_format(df=df, id_col=id_col, time_col=time_col, target_col=target_col)
        if ufp.is_nan_or_none(df[target_col]).any():
            raise ValueError(
                f"Target column ({target_col}) cannot contain missing values."
            )
        freq = _maybe_infer_freq(df, freq=freq, id_col=id_col, time_col=time_col)
        expected_ids_times = id_time_grid(
            df,
            freq=freq,
            start="per_serie",
            end="per_serie",
            id_col=id_col,
            time_col=time_col,
        )
        if len(df) != len(expected_ids_times):
            raise ValueError(
                "Series contain missing or duplicate timestamps, or the timestamps "
                "do not match the provided frequency.\n"
                "Please make sure that all series have a single observation from the first "
                "to the last timestamp and that the provided frequency matches the timestamps'.\n"
                "You can refer to https://docs.nixtla.io/docs/tutorials-missing_values "
                "for an end to end example."
            )
        return df, X_df, drop_id, freq

    def validate_api_key(self, log: bool = True) -> bool:
        """Returns True if your api_key is valid."""
        try:
            with httpx.Client(**self._client_kwargs) as client:
                validation = self._make_request_with_retries(
                    client, "validate_token", {}
                )
        except:
            validation = {}
        if "support" in validation and log:
            logger.info(f'Happy Forecasting! :), {validation["support"]}')
        return validation.get(
            "message", ""
        ) == "success" or "Forecasting! :)" in validation.get("detail", "")

    def _distributed_forecast(
        self,
        df: DistributedDFType,
        h: _PositiveInt,
        freq: Optional[str],
        id_col: str,
        time_col: str,
        target_col: str,
        X_df: Optional[DistributedDFType],
        level: Optional[list[Union[int, float]]],
        quantiles: Optional[list[float]],
        finetune_steps: _NonNegativeInt,
        finetune_depth: _Finetune_Depth,
        finetune_loss: _Loss,
        clean_ex_first: bool,
        validate_api_key: bool,
        add_history: bool,
        date_features: Union[bool, list[Union[str, Callable]]],
        date_features_to_one_hot: Union[bool, list[str]],
        model: _Model,
        num_partitions: Optional[int],
        feature_contributions: bool,
    ) -> DistributedDFType:
        import fugue.api as fa

        schema, partition_config = _distributed_setup(
            df=df,
            method="forecast",
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
            level=level,
            quantiles=quantiles,
            num_partitions=num_partitions,
        )
        if X_df is not None:

            def format_df(df: pd.DataFrame) -> pd.DataFrame:
                return df.assign(_in_sample=True)

            def format_X_df(
                X_df: pd.DataFrame,
                target_col: str,
                df_cols: list[str],
            ) -> pd.DataFrame:
                return X_df.assign(**{"_in_sample": False, target_col: 0.0})[df_cols]

            df = fa.transform(df, format_df, schema="*,_in_sample:bool")
            X_df = fa.transform(
                X_df,
                format_X_df,
                schema=fa.get_schema(df),
                params={"target_col": target_col, "df_cols": fa.get_column_names(df)},
            )
            df = fa.union(df, X_df)
        result_df = fa.transform(
            df,
            using=_forecast_wrapper,
            schema=schema,
            params=dict(
                client=self,
                h=h,
                freq=freq,
                id_col=id_col,
                time_col=time_col,
                target_col=target_col,
                level=level,
                quantiles=quantiles,
                finetune_steps=finetune_steps,
                finetune_depth=finetune_depth,
                finetune_loss=finetune_loss,
                clean_ex_first=clean_ex_first,
                validate_api_key=validate_api_key,
                add_history=add_history,
                date_features=date_features,
                date_features_to_one_hot=date_features_to_one_hot,
                model=model,
                num_partitions=None,
                feature_contributions=feature_contributions,
            ),
            partition=partition_config,
            as_fugue=True,
        )
        return fa.get_native_as_df(result_df)

    def forecast(
        self,
        df: AnyDFType,
        h: _PositiveInt,
        freq: Optional[str] = None,
        id_col: str = "unique_id",
        time_col: str = "ds",
        target_col: str = "y",
        X_df: Optional[AnyDFType] = None,
        level: Optional[list[Union[int, float]]] = None,
        quantiles: Optional[list[float]] = None,
        finetune_steps: _NonNegativeInt = 0,
        finetune_depth: _Finetune_Depth = 1,
        finetune_loss: _Loss = "default",
        clean_ex_first: bool = True,
        hist_exog_list: Optional[list[str]] = None,
        validate_api_key: bool = False,
        add_history: bool = False,
        date_features: Union[bool, list[Union[str, Callable]]] = False,
        date_features_to_one_hot: Union[bool, list[str]] = False,
        model: _Model = "timegpt-1",
        num_partitions: Optional[_PositiveInt] = None,
        feature_contributions: bool = False,
    ) -> AnyDFType:
        """Forecast your time series using TimeGPT.

        Parameters
        ----------
        df : pandas or polars DataFrame
            The DataFrame on which the function will operate. Expected to contain at least the following columns:
            - time_col:
                Column name in `df` that contains the time indices of the time series. This is typically a datetime
                column with regular intervals, e.g., hourly, daily, monthly data points.
            - target_col:
                Column name in `df` that contains the target variable of the time series, i.e., the variable we
                wish to predict or analyze.
            Additionally, you can pass multiple time series (stacked in the dataframe) considering an additional column:
            - id_col:
                Column name in `df` that identifies unique time series. Each unique value in this column
                corresponds to a unique time series.
        h : int
            Forecast horizon.
        freq : str
            Frequency of the data. By default, the freq will be inferred automatically.
            See [pandas' available frequencies](https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#offset-aliases).
        id_col : str (default='unique_id')
            Column that identifies each serie.
        time_col : str (default='ds')
            Column that identifies each timestep, its values can be timestamps or integers.
        target_col : str (default='y')
            Column that contains the target.
        X_df : pandas or polars DataFrame, optional (default=None)
            DataFrame with [`unique_id`, `ds`] columns and `df`'s future exogenous.
        level : list[float], optional (default=None)
            Confidence levels between 0 and 100 for prediction intervals.
        quantiles : list[float], optional (default=None)
            Quantiles to forecast, list between (0, 1).
            `level` and `quantiles` should not be used simultaneously.
            The output dataframe will have the quantile columns
            formatted as TimeGPT-q-(100 * q) for each q.
            100 * q represents percentiles but we choose this notation
            to avoid having dots in column names.
        finetune_steps : int (default=0)
            Number of steps used to finetune learning TimeGPT in the
            new data.
        finetune_depth: int (default=1)
            The depth of the finetuning. Uses a scale from 1 to 5, where 1 means little finetuning,
            and 5 means that the entire model is finetuned.
        finetune_loss : str (default='default')
            Loss function to use for finetuning. Options are: `default`, `mae`, `mse`, `rmse`, `mape`, and `smape`.
        clean_ex_first : bool (default=True)
            Clean exogenous signal before making forecasts using TimeGPT.
        hist_exog_list : list of str, optional (default=None)
            Column names of the historical exogenous features.
        validate_api_key : bool (default=False)
            If True, validates api_key before sending requests.
        add_history : bool (default=False)
            Return fitted values of the model.
        date_features : bool or list of str or callable, optional (default=False)
            Features computed from the dates.
            Can be pandas date attributes or functions that will take the dates as input.
            If True automatically adds most used date features for the
            frequency of `df`.
        date_features_to_one_hot : bool or list of str (default=False)
            Apply one-hot encoding to these date features.
            If `date_features=True`, then all date features are
            one-hot encoded by default.
        model : str (default='timegpt-1')
            Model to use as a string. Options are: `timegpt-1`, and `timegpt-1-long-horizon`.
            We recommend using `timegpt-1-long-horizon` for forecasting
            if you want to predict more than one seasonal
            period given the frequency of your data.
        num_partitions : int (default=None)
            Number of partitions to use.
            If None, the number of partitions will be equal
            to the available parallel resources in distributed environments.
        feature_contributions: bool (default=False)
            Compute SHAP values
            Gives access to computed SHAP values to explain the impact
            of features on the final predictions.

        Returns
        -------
        pandas, polars, dask or spark DataFrame or ray Dataset.
            DataFrame with TimeGPT forecasts for point predictions and probabilistic
            predictions (if level is not None).
        """
        if not isinstance(df, (pd.DataFrame, pl_DataFrame)):
            return self._distributed_forecast(
                df=df,
                h=h,
                freq=freq,
                id_col=id_col,
                time_col=time_col,
                target_col=target_col,
                X_df=X_df,
                level=level,
                quantiles=quantiles,
                finetune_steps=finetune_steps,
                finetune_depth=finetune_depth,
                finetune_loss=finetune_loss,
                clean_ex_first=clean_ex_first,
                validate_api_key=validate_api_key,
                add_history=add_history,
                date_features=date_features,
                date_features_to_one_hot=date_features_to_one_hot,
                model=model,
                num_partitions=num_partitions,
                feature_contributions=feature_contributions,
            )
        self.__dict__.pop("weights_x", None)
        self.__dict__.pop("feature_contributions", None)
        model = self._maybe_override_model(model)
        logger.info("Validating inputs...")
        df, X_df, drop_id, freq = self._run_validations(
            df=df,
            X_df=X_df,
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
            validate_api_key=validate_api_key,
            model=model,
            freq=freq,
        )
        df, X_df = _validate_exog(
            df=df,
            X_df=X_df,
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
            hist_exog=hist_exog_list,
        )
        level, quantiles = _prepare_level_and_quantiles(level, quantiles)
        standard_freq = _standardize_freq(freq)
        model_input_size, model_horizon = self._get_model_params(model, standard_freq)
        if finetune_steps > 0 or level is not None or add_history:
            _validate_input_size(df, id_col, model_input_size, model_horizon)
        if h > model_horizon:
            logger.warning(
                'The specified horizon "h" exceeds the model horizon. '
                "This may lead to less accurate forecasts. "
                "Please consider using a smaller horizon."
            )

        logger.info("Preprocessing dataframes...")
        processed, X_future, x_cols, futr_cols = _preprocess(
            df=df,
            X_df=X_df,
            h=h,
            freq=standard_freq,
            date_features=date_features,
            date_features_to_one_hot=date_features_to_one_hot,
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
        )
        restrict_input = finetune_steps == 0 and not x_cols and not add_history
        if restrict_input:
            logger.info("Restricting input...")
            new_input_size = _restrict_input_samples(
                level=level,
                input_size=model_input_size,
                model_horizon=model_horizon,
                h=h,
            )
            processed = _tail(processed, new_input_size)
        if processed.data.shape[1] > 1:
            X = processed.data[:, 1:].T
            if futr_cols is not None:
                logger.info(f"Using future exogenous features: {futr_cols}")
            if hist_exog_list is not None:
                logger.info(f"Using historical exogenous features: {hist_exog_list}")
        else:
            X = None

        logger.info("Calling Forecast Endpoint...")
        payload = {
            "series": {
                "y": processed.data[:, 0],
                "sizes": np.diff(processed.indptr),
                "X": X,
                "X_future": X_future,
            },
            "model": model,
            "h": h,
            "freq": standard_freq,
            "clean_ex_first": clean_ex_first,
            "level": level,
            "finetune_steps": finetune_steps,
            "finetune_depth": finetune_depth,
            "finetune_loss": finetune_loss,
            "feature_contributions": feature_contributions and X is not None,
        }
        with httpx.Client(**self._client_kwargs) as client:
            insample_feat_contributions = None
            if num_partitions is None:
                resp = self._make_request_with_retries(client, "v2/forecast", payload)
                if add_history:
                    in_sample_payload = _forecast_payload_to_in_sample(payload)
                    logger.info("Calling Historical Forecast Endpoint...")
                    in_sample_resp = self._make_request_with_retries(
                        client, "v2/historic_forecast", in_sample_payload
                    )
                    insample_feat_contributions = in_sample_resp.get(
                        "feature_contributions", None
                    )
            else:
                payloads = _partition_series(payload, num_partitions, h)
                resp = self._make_partitioned_requests(client, "v2/forecast", payloads)
                if add_history:
                    in_sample_payloads = [
                        _forecast_payload_to_in_sample(p) for p in payloads
                    ]
                    logger.info("Calling Historical Forecast Endpoint...")
                    in_sample_resp = self._make_partitioned_requests(
                        client, "v2/historic_forecast", in_sample_payloads
                    )
                    insample_feat_contributions = in_sample_resp.get(
                        "feature_contributions", None
                    )

        # assemble result
        out = ufp.make_future_dataframe(
            uids=processed.uids,
            last_times=type(processed.uids)(processed.last_times),
            freq=freq,
            h=h,
            id_col=id_col,
            time_col=time_col,
        )
        out = ufp.assign_columns(out, "TimeGPT", resp["mean"])
        out = _maybe_add_intervals(out, resp["intervals"])
        if add_history:
            in_sample_df = _parse_in_sample_output(
                in_sample_output=in_sample_resp,
                df=df,
                processed=processed,
                id_col=id_col,
                time_col=time_col,
                target_col=target_col,
            )
            in_sample_df = ufp.drop_columns(in_sample_df, target_col)
            out = ufp.vertical_concat([in_sample_df, out])
        out = _maybe_convert_level_to_quantiles(out, quantiles)
        self._maybe_assign_feature_contributions(
            expected_contributions=feature_contributions,
            resp=resp,
            x_cols=x_cols,
            out_df=out[[id_col, time_col, "TimeGPT"]],
            insample_feat_contributions=insample_feat_contributions,
        )
        if add_history:
            sort_idxs = ufp.maybe_compute_sort_indices(
                out, id_col=id_col, time_col=time_col
            )
            if sort_idxs is not None:
                out = ufp.take_rows(out, sort_idxs)
                out = ufp.drop_index_if_pandas(out)
                if hasattr(self, "feature_contributions"):
                    self.feature_contributions = ufp.take_rows(
                        self.feature_contributions, sort_idxs
                    )
                    self.feature_contributions = ufp.drop_index_if_pandas(
                        self.feature_contributions
                    )
        out = _maybe_drop_id(df=out, id_col=id_col, drop=drop_id)
        self._maybe_assign_weights(weights=resp["weights_x"], df=df, x_cols=x_cols)
        return out

    def _distributed_detect_anomalies(
        self,
        df: DistributedDFType,
        freq: Optional[str],
        id_col: str,
        time_col: str,
        target_col: str,
        level: Union[int, float],
        clean_ex_first: bool,
        validate_api_key: bool,
        date_features: Union[bool, list[str]],
        date_features_to_one_hot: Union[bool, list[str]],
        model: _Model,
        num_partitions: Optional[int],
    ) -> DistributedDFType:
        import fugue.api as fa

        schema, partition_config = _distributed_setup(
            df=df,
            method="detect_anomalies",
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
            level=level,
            quantiles=None,
            num_partitions=num_partitions,
        )
        result_df = fa.transform(
            df,
            using=_detect_anomalies_wrapper,
            schema=schema,
            params=dict(
                client=self,
                freq=freq,
                id_col=id_col,
                time_col=time_col,
                target_col=target_col,
                level=level,
                clean_ex_first=clean_ex_first,
                validate_api_key=validate_api_key,
                date_features=date_features,
                date_features_to_one_hot=date_features_to_one_hot,
                model=model,
                num_partitions=None,
            ),
            partition=partition_config,
            as_fugue=True,
        )
        return fa.get_native_as_df(result_df)

    def detect_anomalies(
        self,
        df: AnyDFType,
        freq: Optional[str] = None,
        id_col: str = "unique_id",
        time_col: str = "ds",
        target_col: str = "y",
        level: Union[int, float] = 99,
        clean_ex_first: bool = True,
        validate_api_key: bool = False,
        date_features: Union[bool, list[str]] = False,
        date_features_to_one_hot: Union[bool, list[str]] = False,
        model: _Model = "timegpt-1",
        num_partitions: Optional[_PositiveInt] = None,
    ) -> AnyDFType:
        """Detect anomalies in your time series using TimeGPT.

        Parameters
        ----------
        df : pandas or polars DataFrame
            The DataFrame on which the function will operate. Expected to contain at least the following columns:
            - time_col:
                Column name in `df` that contains the time indices of the time series. This is typically a datetime
                column with regular intervals, e.g., hourly, daily, monthly data points.
            - target_col:
                Column name in `df` that contains the target variable of the time series, i.e., the variable we
                wish to predict or analyze.
            Additionally, you can pass multiple time series (stacked in the dataframe) considering an additional column:
            - id_col:
                Column name in `df` that identifies unique time series. Each unique value in this column
                corresponds to a unique time series.
        freq : str
            Frequency of the data. By default, the freq will be inferred automatically.
            See [pandas' available frequencies](https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#offset-aliases).
        id_col : str (default='unique_id')
            Column that identifies each serie.
        time_col : str (default='ds')
            Column that identifies each timestep, its values can be timestamps or integers.
        target_col : str (default='y')
            Column that contains the target.
        level : float (default=99)
            Confidence level between 0 and 100 for detecting the anomalies.
        clean_ex_first : bool (default=True)
            Clean exogenous signal before making forecasts
            using TimeGPT.
        validate_api_key : bool (default=False)
            If True, validates api_key before sending requests.
        date_features : bool or list of str or callable, optional (default=False)
            Features computed from the dates.
            Can be pandas date attributes or functions that will take the dates as input.
            If True automatically adds most used date features for the
            frequency of `df`.
        date_features_to_one_hot : bool or list of str (default=False)
            Apply one-hot encoding to these date features.
            If `date_features=True`, then all date features are
            one-hot encoded by default.
        model : str (default='timegpt-1')
            Model to use as a string. Options are: `timegpt-1`, and `timegpt-1-long-horizon`.
            We recommend using `timegpt-1-long-horizon` for forecasting
            if you want to predict more than one seasonal
            period given the frequency of your data.
        num_partitions : int (default=None)
            Number of partitions to use.
            If None, the number of partitions will be equal
            to the available parallel resources in distributed environments.

        Returns
        -------
        pandas, polars, dask or spark DataFrame or ray Dataset.
            DataFrame with anomalies flagged by TimeGPT.
        """
        if not isinstance(df, (pd.DataFrame, pl_DataFrame)):
            return self._distributed_detect_anomalies(
                df=df,
                freq=freq,
                id_col=id_col,
                time_col=time_col,
                target_col=target_col,
                level=level,
                clean_ex_first=clean_ex_first,
                validate_api_key=validate_api_key,
                date_features=date_features,
                date_features_to_one_hot=date_features_to_one_hot,
                model=model,
                num_partitions=num_partitions,
            )
        self.__dict__.pop("weights_x", None)
        model = self._maybe_override_model(model)
        logger.info("Validating inputs...")
        df, _, drop_id, freq = self._run_validations(
            df=df,
            X_df=None,
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
            validate_api_key=validate_api_key,
            model=model,
            freq=freq,
        )
        standard_freq = _standardize_freq(freq)
        model_input_size, model_horizon = self._get_model_params(model, standard_freq)

        logger.info("Preprocessing dataframes...")
        processed, _, x_cols, _ = _preprocess(
            df=df,
            X_df=None,
            h=0,
            freq=standard_freq,
            date_features=date_features,
            date_features_to_one_hot=date_features_to_one_hot,
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
        )
        if processed.data.shape[1] > 1:
            X = processed.data[:, 1:].T
            logger.info(f"Using the following exogenous features: {x_cols}")
        else:
            X = None

        logger.info("Calling Anomaly Detector Endpoint...")
        payload = {
            "series": {
                "y": processed.data[:, 0],
                "sizes": np.diff(processed.indptr),
                "X": X,
            },
            "model": model,
            "freq": standard_freq,
            "clean_ex_first": clean_ex_first,
            "level": level,
        }
        with httpx.Client(**self._client_kwargs) as client:
            if num_partitions is None:
                resp = self._make_request_with_retries(
                    client, "v2/anomaly_detection", payload
                )
            else:
                payloads = _partition_series(payload, num_partitions, h=0)
                resp = self._make_partitioned_requests(
                    client, "v2/anomaly_detection", payloads
                )

        # assemble result
        out = _parse_in_sample_output(
            in_sample_output=resp,
            df=df,
            processed=processed,
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
        )
        out = ufp.assign_columns(out, "anomaly", resp["anomaly"])
        out = _maybe_drop_id(df=out, id_col=id_col, drop=drop_id)
        self._maybe_assign_weights(weights=resp["weights_x"], df=df, x_cols=x_cols)
        return out

    def _distributed_cross_validation(
        self,
        df: DistributedDFType,
        h: _PositiveInt,
        freq: Optional[str],
        id_col: str,
        time_col: str,
        target_col: str,
        level: Optional[list[Union[int, float]]],
        quantiles: Optional[list[float]],
        validate_api_key: bool,
        n_windows: _PositiveInt,
        step_size: Optional[_PositiveInt],
        finetune_steps: _NonNegativeInt,
        finetune_depth: _Finetune_Depth,
        finetune_loss: _Loss,
        clean_ex_first: bool,
        date_features: Union[bool, Sequence[Union[str, Callable]]],
        date_features_to_one_hot: Union[bool, list[str]],
        model: _Model,
        num_partitions: Optional[int],
    ) -> DistributedDFType:
        import fugue.api as fa

        schema, partition_config = _distributed_setup(
            df=df,
            method="forecast",
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
            level=level,
            quantiles=quantiles,
            num_partitions=num_partitions,
        )
        result_df = fa.transform(
            df,
            using=_cross_validation_wrapper,
            schema=schema,
            params=dict(
                client=self,
                h=h,
                freq=freq,
                id_col=id_col,
                time_col=time_col,
                target_col=target_col,
                level=level,
                quantiles=quantiles,
                validate_api_key=validate_api_key,
                n_windows=n_windows,
                step_size=step_size,
                finetune_steps=finetune_steps,
                finetune_depth=finetune_depth,
                finetune_loss=finetune_loss,
                clean_ex_first=clean_ex_first,
                date_features=date_features,
                date_features_to_one_hot=date_features_to_one_hot,
                model=model,
                num_partitions=None,
            ),
            partition=partition_config,
            as_fugue=True,
        )
        return fa.get_native_as_df(result_df)

    def cross_validation(
        self,
        df: AnyDFType,
        h: _PositiveInt,
        freq: Optional[str] = None,
        id_col: str = "unique_id",
        time_col: str = "ds",
        target_col: str = "y",
        level: Optional[list[Union[int, float]]] = None,
        quantiles: Optional[list[float]] = None,
        validate_api_key: bool = False,
        n_windows: _PositiveInt = 1,
        step_size: Optional[_PositiveInt] = None,
        finetune_steps: _NonNegativeInt = 0,
        finetune_depth: _Finetune_Depth = 1,
        finetune_loss: _Loss = "default",
        clean_ex_first: bool = True,
        date_features: Union[bool, list[str]] = False,
        date_features_to_one_hot: Union[bool, list[str]] = False,
        model: _Model = "timegpt-1",
        num_partitions: Optional[_PositiveInt] = None,
    ) -> AnyDFType:
        """Perform cross validation in your time series using TimeGPT.

        Parameters
        ----------
        df : pandas or polars DataFrame
            The DataFrame on which the function will operate. Expected to contain at least the following columns:
            - time_col:
                Column name in `df` that contains the time indices of the time series. This is typically a datetime
                column with regular intervals, e.g., hourly, daily, monthly data points.
            - target_col:
                Column name in `df` that contains the target variable of the time series, i.e., the variable we
                wish to predict or analyze.
            Additionally, you can pass multiple time series (stacked in the dataframe) considering an additional column:
            - id_col:
                Column name in `df` that identifies unique time series. Each unique value in this column
                corresponds to a unique time series.
        h : int
            Forecast horizon.
        freq : str
            Frequency of the data. By default, the freq will be inferred automatically.
            See [pandas' available frequencies](https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#offset-aliases).
        id_col : str (default='unique_id')
            Column that identifies each serie.
        time_col : str (default='ds')
            Column that identifies each timestep, its values can be timestamps or integers.
        target_col : str (default='y')
            Column that contains the target.
        level : float (default=99)
            Confidence level between 0 and 100 for prediction intervals.
        quantiles : list[float], optional (default=None)
            Quantiles to forecast, list between (0, 1).
            `level` and `quantiles` should not be used simultaneously.
            The output dataframe will have the quantile columns
            formatted as TimeGPT-q-(100 * q) for each q.
            100 * q represents percentiles but we choose this notation
            to avoid having dots in column names.
        validate_api_key : bool (default=False)
            If True, validates api_key before sending requests.
        n_windows : int (defaul=1)
            Number of windows to evaluate.
        step_size : int, optional (default=None)
            Step size between each cross validation window. If None it will be equal to `h`.
        finetune_steps : int (default=0)
            Number of steps used to finetune TimeGPT in the
            new data.
        finetune_depth: int (default=1)
            The depth of the finetuning. Uses a scale from 1 to 5, where 1 means little finetuning,
            and 5 means that the entire model is finetuned.
        finetune_loss : str (default='default')
            Loss function to use for finetuning. Options are: `default`, `mae`, `mse`, `rmse`, `mape`, and `smape`.
        clean_ex_first : bool (default=True)
            Clean exogenous signal before making forecasts
            using TimeGPT.
        date_features : bool or list of str or callable, optional (default=False)
            Features computed from the dates.
            Can be pandas date attributes or functions that will take the dates as input.
            If True automatically adds most used date features for the
            frequency of `df`.
        date_features_to_one_hot : bool or list of str (default=False)
            Apply one-hot encoding to these date features.
            If `date_features=True`, then all date features are
            one-hot encoded by default.
        model : str (default='timegpt-1')
            Model to use as a string. Options are: `timegpt-1`, and `timegpt-1-long-horizon`.
            We recommend using `timegpt-1-long-horizon` for forecasting
            if you want to predict more than one seasonal
            period given the frequency of your data.
        num_partitions : int (default=None)
            Number of partitions to use.
            If None, the number of partitions will be equal
            to the available parallel resources in distributed environments.

        Returns
        -------
        pandas, polars, dask or spark DataFrame or ray Dataset.
            DataFrame with cross validation forecasts.
        """
        if not isinstance(df, (pd.DataFrame, pl_DataFrame)):
            return self._distributed_cross_validation(
                df=df,
                h=h,
                freq=freq,
                id_col=id_col,
                time_col=time_col,
                target_col=target_col,
                level=level,
                quantiles=quantiles,
                n_windows=n_windows,
                step_size=step_size,
                validate_api_key=validate_api_key,
                finetune_steps=finetune_steps,
                finetune_depth=finetune_depth,
                finetune_loss=finetune_loss,
                clean_ex_first=clean_ex_first,
                date_features=date_features,
                date_features_to_one_hot=date_features_to_one_hot,
                model=model,
                num_partitions=num_partitions,
            )
        model = self._maybe_override_model(model)
        logger.info("Validating inputs...")
        df, _, drop_id, freq = self._run_validations(
            df=df,
            X_df=None,
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
            validate_api_key=validate_api_key,
            model=model,
            freq=freq,
        )
        standard_freq = _standardize_freq(freq)
        level, quantiles = _prepare_level_and_quantiles(level, quantiles)
        model_input_size, model_horizon = self._get_model_params(model, standard_freq)
        if step_size is None:
            step_size = h

        logger.info("Preprocessing dataframes...")
        processed, _, x_cols, _ = _preprocess(
            df=df,
            X_df=None,
            h=0,
            freq=standard_freq,
            date_features=date_features,
            date_features_to_one_hot=date_features_to_one_hot,
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
        )
        if isinstance(df, pd.DataFrame):
            # in pandas<2.2 to_numpy can lead to an object array if
            # the type is a pandas nullable type, e.g. pd.Float64Dtype
            # we thus use the dtype's type as the target dtype
            target_dtype = df.dtypes[target_col].type
            targets = df[target_col].to_numpy(dtype=target_dtype)
        else:
            targets = df[target_col].to_numpy()
        times = df[time_col].to_numpy()
        if processed.sort_idxs is not None:
            targets = targets[processed.sort_idxs]
            times = times[processed.sort_idxs]
        restrict_input = finetune_steps == 0 and not x_cols
        if restrict_input:
            logger.info("Restricting input...")
            new_input_size = _restrict_input_samples(
                level=level,
                input_size=model_input_size,
                model_horizon=model_horizon,
                h=h,
            )
            new_input_size += h + step_size * (n_windows - 1)
            orig_indptr = processed.indptr
            processed = _tail(processed, new_input_size)
            times = _array_tails(times, orig_indptr, np.diff(processed.indptr))
            targets = _array_tails(targets, orig_indptr, np.diff(processed.indptr))
        if processed.data.shape[1] > 1:
            X = processed.data[:, 1:].T
            logger.info(f"Using the following exogenous features: {x_cols}")
        else:
            X = None

        logger.info("Calling Cross Validation Endpoint...")
        payload = {
            "series": {
                "y": targets,
                "sizes": np.diff(processed.indptr),
                "X": X,
            },
            "model": model,
            "h": h,
            "n_windows": n_windows,
            "step_size": step_size,
            "freq": standard_freq,
            "clean_ex_first": clean_ex_first,
            "level": level,
            "finetune_steps": finetune_steps,
            "finetune_depth": finetune_depth,
            "finetune_loss": finetune_loss,
        }
        with httpx.Client(**self._client_kwargs) as client:
            if num_partitions is None:
                resp = self._make_request_with_retries(
                    client, "v2/cross_validation", payload
                )
            else:
                payloads = _partition_series(payload, num_partitions, h=0)
                resp = self._make_partitioned_requests(
                    client, "v2/cross_validation", payloads
                )

        # assemble result
        idxs = np.array(resp["idxs"], dtype=np.int64)
        sizes = np.array(resp["sizes"], dtype=np.int64)
        window_starts = np.arange(0, sizes.sum(), h)
        cutoff_idxs = np.repeat(idxs[window_starts] - 1, h)
        out = type(df)(
            {
                id_col: ufp.repeat(processed.uids, sizes),
                time_col: times[idxs],
                "cutoff": times[cutoff_idxs],
                target_col: targets[idxs],
            }
        )
        out = ufp.assign_columns(out, "TimeGPT", resp["mean"])
        out = _maybe_add_intervals(out, resp["intervals"])
        out = _maybe_drop_id(df=out, id_col=id_col, drop=drop_id)
        return _maybe_convert_level_to_quantiles(out, quantiles)

    def plot(
        self,
        df: Optional[DataFrame] = None,
        forecasts_df: Optional[DataFrame] = None,
        id_col: str = "unique_id",
        time_col: str = "ds",
        target_col: str = "y",
        unique_ids: Union[Optional[list[str]], np.ndarray] = None,
        plot_random: bool = True,
        max_ids: int = 8,
        models: Optional[list[str]] = None,
        level: Optional[list[Union[int, float]]] = None,
        max_insample_length: Optional[int] = None,
        plot_anomalies: bool = False,
        engine: Literal["matplotlib", "plotly", "plotly-resampler"] = "matplotlib",
        resampler_kwargs: Optional[dict] = None,
        ax: Optional[
            Union["plt.Axes", np.ndarray, "plotly.graph_objects.Figure"]
        ] = None,
    ):
        """Plot forecasts and insample values.

        Parameters
        ----------
        df : pandas or polars DataFrame, optional (default=None)
            The DataFrame on which the function will operate. Expected to contain at least the following columns:
            - time_col:
                Column name in `df` that contains the time indices of the time series. This is typically a datetime
                column with regular intervals, e.g., hourly, daily, monthly data points.
            - target_col:
                Column name in `df` that contains the target variable of the time series, i.e., the variable we
                wish to predict or analyze.
            Additionally, you can pass multiple time series (stacked in the dataframe) considering an additional column:
            - id_col:
                Column name in `df` that identifies unique time series. Each unique value in this column
                corresponds to a unique time series.
        forecasts_df : pandas or polars DataFrame, optional (default=None)
            DataFrame with columns [`unique_id`, `ds`] and models.
        id_col : str (default='unique_id')
            Column that identifies each serie.
        time_col : str (default='ds')
            Column that identifies each timestep, its values can be timestamps or integers.
        target_col : str (default='y')
            Column that contains the target.
        unique_ids : list[str], optional (default=None)
            Time Series to plot.
            If None, time series are selected randomly.
        plot_random : bool (default=True)
            Select time series to plot randomly.
        max_ids : int (default=8)
            Maximum number of ids to plot.
        models : list[str], optional (default=None)
            list of models to plot.
        level : list[float], optional (default=None)
            list of prediction intervals to plot if paseed.
        max_insample_length : int, optional (default=None)
            Max number of train/insample observations to be plotted.
        plot_anomalies : bool (default=False)
            Plot anomalies for each prediction interval.
        engine : str (default='matplotlib')
            Library used to plot. 'matplotlib', 'plotly' or 'plotly-resampler'.
        resampler_kwargs : dict
            Kwargs to be passed to plotly-resampler constructor.
            For further custumization ("show_dash") call the method,
            store the plotting object and add the extra arguments to
            its `show_dash` method.
        ax : matplotlib axes, array of matplotlib axes or plotly Figure, optional (default=None)
            Object where plots will be added.
        """
        try:
            from utilsforecast.plotting import plot_series
        except ModuleNotFoundError:
            raise Exception(
                "You have to install additional dependencies to use this method, "
                'please install them using `pip install "nixtla[plotting]"`'
            )
        if df is not None and id_col not in df.columns:
            df = ufp.copy_if_pandas(df, deep=False)
            df = ufp.assign_columns(df, id_col, "ts_0")
        df = ensure_time_dtype(df, time_col=time_col)
        if forecasts_df is not None:
            if id_col not in forecasts_df.columns:
                forecasts_df = ufp.copy_if_pandas(forecasts_df, deep=False)
                forecasts_df = ufp.assign_columns(forecasts_df, id_col, "ts_0")
            forecasts_df = ensure_time_dtype(forecasts_df, time_col=time_col)
            if "anomaly" in forecasts_df.columns:
                # special case to plot outputs
                # from detect_anomalies
                df = None
                forecasts_df = ufp.drop_columns(forecasts_df, "anomaly")
                cols = [
                    c.replace("TimeGPT-lo-", "")
                    for c in forecasts_df.columns
                    if "TimeGPT-lo-" in c
                ]
                level = [float(c) if "." in c else int(c) for c in cols]
                plot_anomalies = True
                models = ["TimeGPT"]
        return plot_series(
            df=df,
            forecasts_df=forecasts_df,
            ids=unique_ids,
            plot_random=plot_random,
            max_ids=max_ids,
            models=models,
            level=level,
            max_insample_length=max_insample_length,
            plot_anomalies=plot_anomalies,
            engine=engine,
            resampler_kwargs=resampler_kwargs,
            palette="tab20b",
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
            ax=ax,
        )

# %% ../nbs/src/nixtla_client.ipynb 10
def _forecast_wrapper(
    df: pd.DataFrame,
    client: NixtlaClient,
    h: _PositiveInt,
    freq: Optional[str],
    id_col: str,
    time_col: str,
    target_col: str,
    level: Optional[list[Union[int, float]]],
    quantiles: Optional[list[float]],
    finetune_steps: _NonNegativeInt,
    finetune_depth: _Finetune_Depth,
    finetune_loss: _Loss,
    clean_ex_first: bool,
    validate_api_key: bool,
    add_history: bool,
    date_features: Union[bool, list[Union[str, Callable]]],
    date_features_to_one_hot: Union[bool, list[str]],
    model: _Model,
    num_partitions: Optional[_PositiveInt],
    feature_contributions: bool,
) -> pd.DataFrame:
    if "_in_sample" in df:
        in_sample_mask = df["_in_sample"]
        X_df = df.loc[~in_sample_mask].drop(columns=["_in_sample", target_col])
        df = df.loc[in_sample_mask].drop(columns="_in_sample")
    else:
        X_df = None
    return client.forecast(
        df=df,
        h=h,
        freq=freq,
        id_col=id_col,
        time_col=time_col,
        target_col=target_col,
        X_df=X_df,
        level=level,
        quantiles=quantiles,
        finetune_steps=finetune_steps,
        finetune_depth=finetune_depth,
        finetune_loss=finetune_loss,
        clean_ex_first=clean_ex_first,
        validate_api_key=validate_api_key,
        add_history=add_history,
        date_features=date_features,
        date_features_to_one_hot=date_features_to_one_hot,
        model=model,
        num_partitions=num_partitions,
        feature_contributions=feature_contributions,
    )


def _detect_anomalies_wrapper(
    df: pd.DataFrame,
    client: NixtlaClient,
    freq: Optional[str],
    id_col: str,
    time_col: str,
    target_col: str,
    level: Union[int, float],
    clean_ex_first: bool,
    validate_api_key: bool,
    date_features: Union[bool, list[str]],
    date_features_to_one_hot: Union[bool, list[str]],
    model: _Model,
    num_partitions: Optional[_PositiveInt],
) -> pd.DataFrame:
    return client.detect_anomalies(
        df=df,
        freq=freq,
        id_col=id_col,
        time_col=time_col,
        target_col=target_col,
        level=level,
        clean_ex_first=clean_ex_first,
        validate_api_key=validate_api_key,
        date_features=date_features,
        date_features_to_one_hot=date_features_to_one_hot,
        model=model,
        num_partitions=num_partitions,
    )


def _cross_validation_wrapper(
    df: pd.DataFrame,
    client: NixtlaClient,
    h: _PositiveInt,
    freq: Optional[str],
    id_col: str,
    time_col: str,
    target_col: str,
    level: Optional[list[Union[int, float]]],
    quantiles: Optional[list[float]],
    validate_api_key: bool,
    n_windows: _PositiveInt,
    step_size: Optional[_PositiveInt],
    finetune_steps: _NonNegativeInt,
    finetune_depth: _Finetune_Depth,
    finetune_loss: _Loss,
    clean_ex_first: bool,
    date_features: Union[bool, list[str]],
    date_features_to_one_hot: Union[bool, list[str]],
    model: _Model,
    num_partitions: Optional[_PositiveInt],
) -> pd.DataFrame:
    return client.cross_validation(
        df=df,
        h=h,
        freq=freq,
        id_col=id_col,
        time_col=time_col,
        target_col=target_col,
        level=level,
        quantiles=quantiles,
        validate_api_key=validate_api_key,
        n_windows=n_windows,
        step_size=step_size,
        finetune_steps=finetune_steps,
        finetune_depth=finetune_depth,
        finetune_loss=finetune_loss,
        clean_ex_first=clean_ex_first,
        date_features=date_features,
        date_features_to_one_hot=date_features_to_one_hot,
        model=model,
        num_partitions=num_partitions,
    )


def _get_schema(
    df: "AnyDataFrame",
    method: str,
    id_col: str,
    time_col: str,
    target_col: str,
    level: Optional[Union[int, float, list[Union[int, float]]]],
    quantiles: Optional[list[float]],
) -> "triad.Schema":
    import fugue.api as fa

    base_cols = [id_col, time_col]
    if method != "forecast":
        base_cols.append(target_col)
    schema = fa.get_schema(df).extract(base_cols).copy()
    schema.append("TimeGPT:double")
    if method == "detect_anomalies":
        schema.append("anomaly:bool")
    elif method == "cross_validation":
        schema.append(("cutoff", schema[time_col].type))
    if level is not None and quantiles is not None:
        raise ValueError("You should provide `level` or `quantiles` but not both.")
    if level is not None:
        if not isinstance(level, list):
            level = [level]
        level = sorted(level)
        schema.append(",".join(f"TimeGPT-lo-{lv}:double" for lv in reversed(level)))
        schema.append(",".join(f"TimeGPT-hi-{lv}:double" for lv in level))
    if quantiles is not None:
        quantiles = sorted(quantiles)
        q_cols = [f"TimeGPT-q-{int(q * 100)}:double" for q in quantiles]
        schema.append(",".join(q_cols))
    return schema


def _distributed_setup(
    df: "AnyDataFrame",
    method: str,
    id_col: str,
    time_col: str,
    target_col: str,
    level: Optional[Union[int, float, list[Union[int, float]]]],
    quantiles: Optional[list[float]],
    num_partitions: Optional[int],
) -> tuple["triad.Schema", dict[str, Any]]:
    from fugue.execution import infer_execution_engine

    if infer_execution_engine([df]) is None:
        raise ValueError(
            f"Could not infer execution engine for type {type(df).__name__}. "
            "Expected a spark or dask DataFrame or a ray Dataset."
        )
    schema = _get_schema(
        df=df,
        method=method,
        id_col=id_col,
        time_col=time_col,
        target_col=target_col,
        level=level,
        quantiles=quantiles,
    )
    partition_config: dict[str, Any] = dict(by=id_col, algo="coarse")
    if num_partitions is not None:
        partition_config["num"] = num_partitions
    return schema, partition_config
