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

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

# %% ../nbs/timegpt.ipynb 5
import logging
import inspect
import json
import requests
from typing import Dict, List, Optional, Union

import pandas as pd

logger = logging.getLogger()

# %% ../nbs/timegpt.ipynb 7
class TimeGPT:
    """
    A class used to interact with the TimeGPT API.
    """

    def __init__(self, token: str):
        """
        Constructs all the necessary attributes for the TimeGPT object.

        Parameters
        ----------
        token : str
            The authorization token to interact with the TimeGPT API.
        """
        self.token = token
        self.api_url = "https://dashboard.nixtla.io/api"
        self.weights_x: pd.DataFrame = None

    @property
    def request_headers(self):
        headers = {
            "accept": "application/json",
            "content-type": "application/json",
            "authorization": f"Bearer {self.token}",
        }
        return headers

    def _parse_response(self, response) -> Dict:
        """Parses responde."""
        response.raise_for_status()
        try:
            resp = response.json()
        except Exception as e:
            raise Exception(response)
        return resp

    def validate_token(self) -> bool:
        """Returns True if your token is valid."""
        response = requests.post(
            f"{self.api_url}/validate_token_front",
            headers=self.request_headers,
        )
        valid = True
        try:
            response = self._parse_response(response)
        except:
            valid = False
        return valid

    def _input_size(self, freq: str):
        response_input_size = requests.post(
            f"{self.api_url}/timegpt_input_size",
            json={"freq": freq},
            headers=self.request_headers,
        )
        response_input_size = self._parse_response(response_input_size)
        return response_input_size["data"]

    def _validate_inputs(
        self,
        df: pd.DataFrame,
        X_df: pd.DataFrame,
        id_col: str,
        time_col: str,
        target_col: str,
    ):
        renamer = {id_col: "unique_id", time_col: "ds", target_col: "y"}
        df = df.rename(columns=renamer)
        if df.dtypes.ds != "object":
            df["ds"] = df["ds"].astype(str)
        drop_uid = False
        if "unique_id" not in df.columns:
            # Insert unique_id column
            df = df.assign(unique_id="ts_0")
            drop_uid = True
        if X_df is not None:
            X_df = X_df.rename(columns=renamer)
            if "unique_id" not in X_df.columns:
                X_df = X_df.assign(unique_id="ts_0")
            if X_df.dtypes.ds != "object":
                X_df["ds"] = X_df["ds"].astype(str)
        return df, X_df, drop_uid

    def _validate_outputs(
        self,
        fcst_df: pd.DataFrame,
        id_col: str,
        time_col: str,
        target_col: str,
        drop_uid: bool,
    ):
        renamer = {
            "unique_id": id_col,
            "ds": time_col,
            "target_col": target_col,
        }
        if drop_uid:
            fcst_df = fcst_df.drop(columns="unique_id")
        fcst_df = fcst_df.rename(columns=renamer)
        return fcst_df

    def _infer_freq(self, df: pd.DataFrame):
        unique_id = df.iloc[0]["unique_id"]
        df_id = df.query("unique_id == @unique_id")
        freq = pd.infer_freq(df_id["ds"])
        if freq is None:
            raise Exception(
                "Could not infer frequency of ds column. This could be due to "
                "inconsistent intervals. Please check your data for missing, "
                "duplicated or irregular timestamps"
            )
        return freq

    def _preprocess_inputs(
        self,
        df: pd.DataFrame,
        h: int,
        freq: str,
        X_df: Optional[pd.DataFrame] = None,
    ):
        input_size = self._input_size(freq)
        y_cols = ["unique_id", "ds", "y"]
        y = df[y_cols].groupby("unique_id").tail(input_size + h)
        to_dict_args = {"orient": "split"}
        if "index" in inspect.signature(pd.DataFrame.to_dict).parameters:
            to_dict_args["index"] = False
        if y["y"].isna().any():
            raise Exception("Your target variable contains NA, please check")
        y = y.to_dict(**to_dict_args)
        x_cols = []
        if X_df is None:
            x = None
        else:
            x_cols = X_df.drop(columns=["unique_id", "ds"]).columns.to_list()
            if not all(col in df.columns for col in x_cols):
                raise Exception(
                    "You must include the exogenous variables in the `df` object, "
                    f'exogenous variables {",".join(x_cols)}'
                )
            x = (
                df[["unique_id", "ds"] + x_cols]
                .groupby("unique_id")
                .tail(input_size + h)
            )
            x = pd.concat([x, X_df])
            if x[x_cols].isna().any().any():
                raise Exception(
                    "Some of your exogenous variables contain NA, please check"
                )
            x = x.sort_values(["unique_id", "ds"])
            x = x.to_dict(**to_dict_args)
        return y, x, x_cols

    def _multi_series(
        self,
        df: pd.DataFrame,
        h: int,
        freq: str,
        X_df: Optional[pd.DataFrame] = None,
        level: Optional[List[int]] = None,
        finetune_steps: int = 0,
        clean_ex_first: bool = True,
    ):
        if freq is None:
            freq = self._infer_freq(df)
        y, x, x_cols = self._preprocess_inputs(df=df, h=h, freq=freq, X_df=X_df)
        payload = dict(
            y=y,
            x=x,
            fh=h,
            freq=freq,
            level=level,
            finetune_steps=finetune_steps,
            clean_ex_first=clean_ex_first,
        )
        response_timegpt = requests.post(
            f"{self.api_url}/timegpt_multi_series",
            json=payload,
            headers=self.request_headers,
        )
        response_timegpt = self._parse_response(response_timegpt)
        if "weights_x" in response_timegpt["data"]:
            self.weights_x = pd.DataFrame(
                {
                    "features": x_cols,
                    "weights": response_timegpt["data"]["weights_x"],
                }
            )
        return pd.DataFrame(**response_timegpt["data"]["forecast"])

    def forecast(
        self,
        df: pd.DataFrame,
        h: int,
        freq: Optional[str] = None,
        id_col: str = "unique_id",
        time_col: str = "ds",
        target_col: str = "y",
        X_df: Optional[pd.DataFrame] = None,
        level: Optional[List[Union[int, float]]] = None,
        finetune_steps: int = 0,
        clean_ex_first: bool = True,
        validate_token: bool = False,
    ):
        """Forecast your time series using TimeGPT.

        Parameters
        ----------
        df : pandas.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.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.
        finetune_steps : int (default=0)
            Number of steps used to finetune TimeGPT in the
            new data.
        clean_ex_first : bool (default=True)
            Clean exogenous signal before making forecasts
            using TimeGPT.
        validate_token: bool (default=False)
            If True, validates token before
            sending requests.

        Returns
        -------
        fcsts_df : pandas.DataFrame
            DataFrame with TimeGPT forecasts for point predictions and probabilistic
            predictions (if level is not None).
        """
        if not self.validate_token():
            raise Exception(
                "Token not valid, please go to https://dashboard.nixtla.io/ to get yours"
            )

        df, X_df, drop_uid = self._validate_inputs(
            df=df,
            X_df=X_df,
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
        )
        fcst_df = self._multi_series(
            df=df,
            h=h,
            freq=freq,
            X_df=X_df,
            level=level,
            finetune_steps=finetune_steps,
            clean_ex_first=clean_ex_first,
        )
        fcst_df = self._validate_outputs(
            fcst_df=fcst_df,
            id_col=id_col,
            time_col=time_col,
            target_col=target_col,
            drop_uid=drop_uid,
        )
        return fcst_df
