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

__all__ = ['simple_predict', 'TimeSeries']

# Cell
import concurrent.futures
import inspect
import os
import warnings
from collections import OrderedDict
from itertools import chain
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

import numpy as np
import pandas as pd
from numba import njit
from window_ops.shift import shift_array

from .utils import data_indptr_from_sorted_df

# Internal Cell
date_features_dtypes = {
    'year': np.uint16,
    'month': np.uint8,
    'day': np.uint8,
    'hour': np.uint8,
    'minute': np.uint8,
    'second': np.uint8,
    'dayofyear': np.uint16,
    'day_of_year': np.uint16,
    'weekofyear': np.uint8,
    'week': np.uint8,
    'dayofweek': np.uint8,
    'day_of_week': np.uint8,
    'weekday': np.uint8,
    'quarter': np.uint8,
    'daysinmonth': np.uint8,
    'is_month_start': np.uint8,
    'is_month_end': np.uint8,
    'is_quarter_start': np.uint8,
    'is_quarter_end': np.uint8,
    'is_year_start': np.uint8,
    'is_year_end': np.uint8,
}


@njit
def _append_new(data, indptr, new):
    """Append each value of new to each group in data formed by indptr."""
    n_series = len(indptr) - 1
    new_data = np.empty(data.size + new.size, dtype=data.dtype)
    new_indptr = indptr.copy()
    new_indptr[1:] += np.arange(1, n_series + 1)
    for i in range(n_series):
        prev_slice = slice(indptr[i], indptr[i + 1])
        new_slice = slice(new_indptr[i], new_indptr[i + 1] - 1)
        new_data[new_slice] = data[prev_slice]
        new_data[new_indptr[i + 1] - 1] = new[i]
    return new_data, new_indptr


# Internal Cell
class GroupedArray:
    """Array made up of different groups. Can be thought of (and iterated) as a list of arrays.

    All the data is stored in a single 1d array `data`.
    The indices for the group boundaries are stored in another 1d array `indptr`."""

    def __init__(self, data: np.ndarray, indptr: np.ndarray):
        self.data = data
        self.indptr = indptr
        self.ngroups = len(indptr) - 1

    def __len__(self) -> int:
        return self.ngroups

    def __getitem__(self, idx: int) -> np.ndarray:
        return self.data[self.indptr[idx] : self.indptr[idx + 1]]

    def take_from_groups(self, idx: Union[int, slice]) -> 'GroupedArray':
        """Takes `idx` from each group in the array."""
        ranges = [
            range(self.indptr[i], self.indptr[i + 1])[idx] for i in range(self.ngroups)
        ]
        items = [self.data[rng] for rng in ranges]
        sizes = np.array([item.size for item in items])
        data = np.hstack(items)
        indptr = np.append(0, sizes.cumsum())
        return GroupedArray(data, indptr)

    def append(self, new: np.ndarray) -> 'GroupedArray':
        """Appends each element of `new` to each existing group. Returns a copy."""
        if new.size != self.ngroups:
            raise ValueError(f'new must be of size {self.ngroups}')
        new_data, new_indptr = _append_new(self.data, self.indptr, new)
        return GroupedArray(new_data, new_indptr)

    def __repr__(self) -> str:
        return f'GroupedArray(ndata={self.data.size}, ngroups={self.ngroups})'


# Internal Cell
@njit
def _identity(x: np.ndarray) -> np.ndarray:
    """Do nothing to the input."""
    return x


def _as_tuple(x):
    """Return a tuple from the input."""
    if isinstance(x, tuple):
        return x
    return (x,)


@njit(nogil=True)
def _transform_series(data, indptr, updates_only, lag, func, *args) -> np.ndarray:
    """Shifts every group in `data` by `lag` and computes `func(shifted, *args)`.

    If `updates_only=True` only last value of the transformation for each group is returned,
    otherwise the full transformation is returned"""
    n_series = len(indptr) - 1
    if updates_only:
        out = np.empty_like(data[:n_series])
        for i in range(n_series):
            lagged = shift_array(data[indptr[i] : indptr[i + 1]], lag)
            out[i] = func(lagged, *args)[-1]
    else:
        out = np.empty_like(data)
        for i in range(n_series):
            lagged = shift_array(data[indptr[i] : indptr[i + 1]], lag)
            out[indptr[i] : indptr[i + 1]] = func(lagged, *args)
    return out


# Internal Cell
def _build_transform_name(lag, tfm, *args) -> str:
    """Creates a name for a transformation based on `lag`, the name of the function and its arguments."""
    tfm_name = f'{tfm.__name__}_lag-{lag}'
    func_params = inspect.signature(tfm).parameters
    func_args = list(func_params.items())[1:]  # remove input array argument
    changed_params = [
        f'{name}-{value}'
        for value, (name, arg) in zip(args, func_args)
        if arg.default != value
    ]
    if changed_params:
        tfm_name += '_' + '_'.join(changed_params)
    return tfm_name


# Cell
def simple_predict(model, new_x: pd.DataFrame, *args) -> np.ndarray:
    """Drop the ds column from `new_x` and call `model.predict` on it."""
    # this avoids installing xgboost only to check if the model is instance of xgb.Booster
    model_is_xgb_booster = (
        type(model).__module__ == 'xgboost.core' and type(model).__name__ == 'Booster'
    )
    new_x = new_x.drop(columns='ds')
    if model_is_xgb_booster:
        import xgboost as xgb

        new_x = xgb.DMatrix(new_x)
    return model.predict(new_x)


# Cell
class TimeSeries:
    """Utility class for storing and transforming time series data."""

    def __init__(
        self,
        freq: str = 'D',
        lags: List[int] = [],
        lag_transforms: Dict[int, List[Tuple]] = {},
        date_features: List[str] = [],
        num_threads: Optional[int] = None,
    ):
        self.freq = pd.tseries.frequencies.to_offset(freq)
        self.num_threads = num_threads or os.cpu_count()
        self.date_features = date_features

        self.transforms: Dict[str, Tuple[Any, ...]] = OrderedDict()
        for lag in lags:
            self.transforms[f'lag-{lag}'] = (lag, _identity)
        for lag in lag_transforms.keys():
            for tfm_args in lag_transforms[lag]:
                tfm, *args = _as_tuple(tfm_args)
                tfm_name = _build_transform_name(lag, tfm, *args)
                self.transforms[tfm_name] = (lag, tfm, *args)

        self.ga: GroupedArray

    @property
    def features(self) -> List[str]:
        """Names of all computed features."""
        return list(self.transforms.keys()) + self.date_features

    def __repr__(self):
        return f'TimeSeries(freq={self.freq}, transforms={list(self.transforms.keys())}, date_features={self.date_features}, num_threads={self.num_threads})'

    def _apply_transforms(self, updates_only: bool = False) -> Dict[str, np.ndarray]:
        """Apply the transformations using the main process.

        If `updates_only` then only the updates are returned.
        """
        results = {}
        offset = 1 if updates_only else 0
        for tfm_name, (lag, tfm, *args) in self.transforms.items():
            results[tfm_name] = _transform_series(
                self.ga.data, self.ga.indptr, updates_only, lag - offset, tfm, *args
            )
        return results

    def _apply_multithreaded_transforms(
        self, updates_only: bool = False
    ) -> Dict[str, np.ndarray]:
        """Apply the transformations using multithreading.

        If `updates_only` then only the updates are returned.
        """
        future_to_result = {}
        results = {}
        offset = 1 if updates_only else 0
        with concurrent.futures.ThreadPoolExecutor(self.num_threads) as executor:
            for tfm_name, (lag, tfm, *args) in self.transforms.items():
                future = executor.submit(
                    _transform_series,
                    self.ga.data,
                    self.ga.indptr,
                    updates_only,
                    lag - offset,
                    tfm,
                    *args,
                )
                future_to_result[future] = tfm_name
            for future in concurrent.futures.as_completed(future_to_result):
                tfm_name = future_to_result[future]
                results[tfm_name] = future.result()
        return results

    def _compute_transforms(self) -> Dict[str, np.ndarray]:
        """Compute the transformations defined in the constructor.

        If `self.num_threads > 1` these are computed using multithreading."""
        if self.num_threads == 1 or len(self.transforms) == 1:
            return self._apply_transforms()
        return self._apply_multithreaded_transforms()

    def _update_y(self, new: np.ndarray) -> None:
        """Appends the elements of `new` to every time serie.

        These values are used to update the transformations and are stored as predictions."""
        if not hasattr(self, 'y_pred'):
            self.y_pred = []
        self.y_pred.append(new)
        new_arr = np.asarray(new)
        self.ga = self.ga.append(new_arr)

    def _update_features(self) -> pd.DataFrame:
        """Compute the current values of all the features using the latest values of the time series."""
        if not hasattr(self, 'curr_dates'):
            self.curr_dates = self.last_dates.copy()
            self.test_dates = []
        self.curr_dates += self.freq
        self.test_dates.append(self.curr_dates)

        if self.num_threads == 1 or len(self.transforms) == 1:
            features = self._apply_transforms(updates_only=True)
        else:
            features = self._apply_multithreaded_transforms(updates_only=True)

        for feature in self.date_features:
            feat_vals = getattr(self.curr_dates, feature).values
            features[feature] = feat_vals.astype(date_features_dtypes[feature])

        features_df = pd.DataFrame(features, columns=self.features, index=self.uids)
        nulls_in_cols = features_df.isnull().any()
        if any(nulls_in_cols):
            warnings.warn(
                f'Found null values in {", ".join(nulls_in_cols[nulls_in_cols].index)}.'
            )
        results_df = self.static_features.join(features_df)
        results_df['ds'] = self.curr_dates
        return results_df

    def _get_predictions(self) -> pd.DataFrame:
        """Get all the predicted values with their corresponding ids and datestamps."""
        n_preds = len(self.y_pred)
        idx = pd.Index(
            chain.from_iterable([uid] * n_preds for uid in self.uids),
            name='unique_id',
            dtype=self.uids.dtype,
        )
        df = pd.DataFrame(
            {
                'ds': np.array(self.test_dates).ravel('F'),
                'y_pred': np.array(self.y_pred).ravel('F'),
            },
            index=idx,
        )
        return df

    def _fit(
        self,
        df: pd.DataFrame,
        static_features: Optional[List[str]] = None,
    ) -> 'TimeSeries':
        """Save the series values, ids, last dates and static features."""
        data, indptr = data_indptr_from_sorted_df(df)
        if data.dtype not in (np.float32, np.float64):
            # since all transformations generate nulls, we need a float dtype
            data = data.astype(np.float32)
        self.ga = GroupedArray(data, indptr)
        self.uids = df.index.unique(level='unique_id')
        last_obs = df.iloc[indptr[1:] - 1]
        self.last_dates = last_obs.index.get_level_values('ds')
        self.static_features = last_obs.reset_index('ds').drop(columns=['ds', 'y'])
        if static_features is not None:
            self.static_features = self.static_features[static_features]
        return self

    def _transform(
        self,
        df: pd.DataFrame,
        dropna: bool = True,
        keep_last_n: Optional[int] = None,
    ) -> pd.DataFrame:
        """Add the features to `df`.

        if `dropna=True` then all the null rows are dropped.
        if `keep_last_n` is not None then that number of observations is kept across all series."""
        df = df.reset_index('ds')
        features = self._compute_transforms()
        for feat in self.transforms.keys():
            df[feat] = features[feat]

        if dropna:
            df.dropna(inplace=True)

        for feature in self.date_features:
            feat_vals = getattr(df.ds.dt, feature).values
            df[feature] = feat_vals.astype(date_features_dtypes[feature])

        if keep_last_n is not None:
            self.ga = self.ga.take_from_groups(slice(-keep_last_n, None))

        self._ga = GroupedArray(self.ga.data, self.ga.indptr)
        self.features_order_ = df.columns.drop(['ds', 'y'])
        return df

    def fit_transform(
        self,
        data: pd.DataFrame,
        static_features: Optional[List[str]] = None,
        dropna: bool = True,
        keep_last_n: Optional[int] = None,
    ) -> pd.DataFrame:
        """Add the features to `data` and save the required information for the predictions step.

        If not all features are static, specif which ones are in `static_features`.
        If you don't want to drop rows with null values after the transformations set `dropna=False`.
        If you want to keep only the last `n` values of each time serie set `keep_last_n=n`.
        """
        data = data.set_index('ds', append=True)
        if not data.index.is_monotonic_increasing:
            data = data.sort_index()
        self._fit(data, static_features)
        return self._transform(data, dropna, keep_last_n)

    def _predict_setup(self) -> None:
        self.curr_dates = self.last_dates.copy()
        self.test_dates = []
        self.y_pred = []
        self.ga = GroupedArray(self._ga.data, self._ga.indptr)

    def predict(
        self,
        model,
        horizon: int,
        predict_fn: Callable = simple_predict,
        **predict_fn_kwargs,
    ) -> pd.DataFrame:
        """Use `model` to predict the next `horizon` timesteps.

        `predict_fn(model, new_x, features_order, **predict_fn_kwargs)` is called on each timestep."""
        self._predict_setup()
        for _ in range(horizon):
            new_x = self._update_features()
            predictions = predict_fn(
                model, new_x, self.features_order_, **predict_fn_kwargs
            )
            self._update_y(predictions)
        return self._get_predictions()
