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

__all__ = ['generate_daily_series', 'data_indptr_from_sorted_df', 'backtest_splits']

# Cell
import random
from itertools import chain
from math import ceil, log10
from typing import Tuple

import numpy as np
import pandas as pd
from numba import njit


# Cell
def generate_daily_series(
    n_series: int,
    min_length: int = 50,
    max_length: int = 500,
    n_static_features: int = 0,
    equal_ends: bool = False,
    seed: int = 0,
) -> pd.DataFrame:
    """Generates `n_series` of different lengths in the interval [`min_length`, `max_length`].

    If `n_static_features > 0`, then each serie gets static features with random values.
    If `equal_ends == True` then all series end at the same date."""
    rng = np.random.RandomState(seed)
    random.seed(seed)
    series_lengths = rng.randint(min_length, max_length + 1, n_series)
    total_length = series_lengths.sum()
    n_digits = ceil(log10(n_series))

    dates = pd.date_range('2000-01-01', periods=max_length, freq='D').values
    uids = [
        [f'id_{i:0{n_digits}}'] * serie_length
        for i, serie_length in enumerate(series_lengths)
    ]
    if equal_ends:
        ds = [dates[-serie_length:] for serie_length in series_lengths]
    else:
        ds = [dates[:serie_length] for serie_length in series_lengths]
    y = np.arange(total_length) % 7 + rng.rand(total_length) * 0.5
    series = pd.DataFrame(
        {
            'unique_id': list(chain.from_iterable(uids)),
            'ds': list(chain.from_iterable(ds)),
            'y': y,
        }
    )
    for i in range(n_static_features):
        static_values = [
            [random.randint(0, 100)] * serie_length for serie_length in series_lengths
        ]
        series[f'static_{i}'] = list(chain.from_iterable(static_values))
        series[f'static_{i}'] = series[f'static_{i}'].astype('category')
        if i == 0:
            series['y'] = series['y'] * (1 + series[f'static_{i}'].cat.codes)
    series['unique_id'] = series['unique_id'].astype('category')
    series['unique_id'] = series['unique_id'].cat.as_ordered()
    series = series.set_index('unique_id')
    return series


# Internal Cell
@njit
def _get_last_n_mask(x: np.ndarray, n: int) -> np.ndarray:
    n_samples = x.size
    mask = np.full(n_samples, True)
    n_first = max(0, n_samples - n)
    mask[:n_first] = False
    return mask


# Cell
def data_indptr_from_sorted_df(df: pd.DataFrame) -> Tuple[np.ndarray, np.ndarray]:
    grouped = df.groupby('unique_id')
    sizes = grouped.size().values
    indptr = np.append(0, sizes.cumsum())
    data = df['y'].values
    return data, indptr


# Internal Cell
@njit
def _get_mask(data: np.ndarray, indptr: np.ndarray, n: int) -> np.ndarray:
    mask = np.empty_like(data)
    for start, end in zip(indptr[:-1], indptr[1:]):
        mask[start:end] = _get_last_n_mask(data[start:end], n)
    return mask


def _get_dataframe_mask(df, n) -> pd.Series:
    data, indptr = data_indptr_from_sorted_df(df)
    mask = _get_mask(data, indptr, n)
    return mask.astype(bool)


def _split_frame(data, n_windows, window, valid_size):
    full_valid_size = (n_windows - window) * valid_size
    extra_valid_size = full_valid_size - valid_size
    if isinstance(data, pd.DataFrame):
        full_valid_mask = _get_dataframe_mask(data, full_valid_size)
        train_mask = ~full_valid_mask
        extra_valid_mask = _get_dataframe_mask(data, extra_valid_size)
    else:
        full_valid_mask = data.map_partitions(
            _get_dataframe_mask, full_valid_size, meta=bool
        )
        train_mask = ~full_valid_mask
        extra_valid_mask = data.map_partitions(
            _get_dataframe_mask, extra_valid_size, meta=bool
        )
    valid_mask = full_valid_mask & ~extra_valid_mask
    return data[train_mask], data[valid_mask]


# Cell
def backtest_splits(data, n_windows, window_size):
    """Returns a generator of `n_windows` for train, valid splits of
    `data` where each valid has `window_size` samples."""
    for window in range(n_windows):
        train, valid = _split_frame(data, n_windows, window, window_size)
        yield train, valid
