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

__all__ = ['StatsForecast']

# Cell
import inspect
import logging
from functools import partial
from os import cpu_count

import numpy as np
import pandas as pd

# Internal Cell
logging.basicConfig(
    format='%(asctime)s %(name)s %(levelname)s: %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S',
)
logger = logging.getLogger(__name__)

# Internal Cell
class GroupedArray:

    def __init__(self, data, indptr):
        self.data = data
        self.indptr = indptr
        self.n_groups = self.indptr.size - 1

    def __getitem__(self, idx):
        if isinstance(idx, int):
            return self.data[self.indptr[idx] : self.indptr[idx + 1]]
        elif isinstance(idx, slice):
            idx = slice(idx.start, idx.stop + 1, idx.step)
            new_indptr = self.indptr[idx].copy()
            new_data = self.data[new_indptr[0] : new_indptr[-1]].copy()
            new_indptr -= new_indptr[0]
            return GroupedArray(new_data, new_indptr)
        raise ValueError(f'idx must be either int or slice, got {type(idx)}')

    def __len__(self):
        return self.n_groups

    def __repr__(self):
        return f'GroupedArray(n_data={self.data.size:,}, n_groups={self.n_groups:,})'

    def __eq__(self, other):
        if not hasattr(other, 'data') or not hasattr(other, 'indptr'):
            return False
        return np.allclose(self.data, other.data) and np.array_equal(self.indptr, other.indptr)

    def compute_forecasts(self, h, func, xreg=None, fitted=False, level=None, *args):
        has_level = 'level' in inspect.signature(func).parameters and level is not None
        if has_level:
            out = np.full((h * self.n_groups, 2 * len(level) + 1), np.nan, dtype=np.float32)
            func = partial(func, level=level)
        else:
            out = np.full(h * self.n_groups, np.nan, dtype=np.float32)
        if fitted:
            fitted_vals = np.full(self.data.shape[0], np.nan, dtype=np.float32)
        xr = None
        keys = None
        for i, grp in enumerate(self):
            if xreg is not None:
                xr = xreg[i]
            res_fn = func(grp, h, xr, fitted, *args)
            if has_level:
                if keys is None:
                    keys = [key for key in res_fn.keys() if key not in ['fitted']]
                for j, key in enumerate(keys):
                    out[h * i : h * (i + 1), j] = res_fn[key]
            else:
                out[h * i : h * (i + 1)] = res_fn['mean']
            if fitted:
                fitted_vals[self.indptr[i] : self.indptr[i + 1]] = res_fn['fitted']
        result = {'forecasts': out, 'keys': keys}
        if fitted:
            result['fitted'] = {'values': fitted_vals}
        return result

    def compute_cv(self, h, test_size, func, step_size=1, input_size=None, fitted=False, *args):
        # output of size: (ts, window, h)
        if (test_size - h) % step_size:
            raise Exception('`test_size - h` should be module `step_size`')
        n_windows = int((test_size - h) / step_size) + 1
        out = np.full((self.n_groups, n_windows, h), np.nan, dtype=np.float32)
        out_test = np.full((self.n_groups, n_windows, h), np.nan, dtype=np.float32)
        if fitted:
            fitted_vals = np.full((self.data.shape[0], n_windows), np.nan, dtype=np.float32)
            fitted_idxs = np.full_like(fitted_vals, False, dtype=bool)
            last_fitted_idxs = np.full_like(fitted_vals, False, dtype=bool)
        for i_ts, grp in enumerate(self):
            for i_window, cutoff in enumerate(range(-test_size, -h + 1, step_size), start=0):
                end_cutoff = cutoff + h
                in_size_disp = cutoff if input_size is None else input_size
                y_train = grp[(cutoff - in_size_disp):cutoff]
                y_test = grp[cutoff:] if end_cutoff == 0 else grp[cutoff:end_cutoff]
                future_xreg = y_test[:, 1:] if (y_test.ndim == 2 and y_test.shape[1] > 1) else None
                res_fn = func(y_train, h, future_xreg, fitted, *args)
                out[i_ts, i_window] = res_fn['mean']
                out_test[i_ts, i_window] = y_test[:, 0] if y_test.ndim == 2 else y_test
                if fitted:
                    fitted_vals[self.indptr[i_ts] : self.indptr[i_ts + 1], i_window][
                        (cutoff - in_size_disp):cutoff
                    ] = res_fn['fitted']
                    fitted_idxs[self.indptr[i_ts] : self.indptr[i_ts + 1], i_window][
                        (cutoff - in_size_disp):cutoff
                    ] = True
                    last_fitted_idxs[self.indptr[i_ts] : self.indptr[i_ts + 1], i_window][cutoff-1] = True
        result = {'forecasts': out, 'y': out_test}
        if fitted:
            result['fitted'] = {'values': fitted_vals, 'idxs': fitted_idxs, 'last_idxs': last_fitted_idxs}
        return result

    def split(self, n_chunks):
        return [self[x[0] : x[-1] + 1] for x in np.array_split(range(self.n_groups), n_chunks) if x.size]

# Internal Cell
def _grouped_array_from_df(df, sort_df):
    df = df.set_index('ds', append=True)
    if not df.index.is_monotonic_increasing and sort_df:
        df = df.sort_index()
    data = df.values.astype(np.float32)
    indices_sizes = df.index.get_level_values('unique_id').value_counts(sort=False)
    indices = indices_sizes.index
    sizes = indices_sizes.values
    cum_sizes = sizes.cumsum()
    dates = df.index.get_level_values('ds')[cum_sizes - 1]
    indptr = np.append(0, cum_sizes).astype(np.int32)
    return GroupedArray(data, indptr), indices, dates, df.index

# Internal Cell
def _cv_dates(last_dates, freq, h, test_size, step_size=1):
    #assuming step_size = 1
    if (test_size - h) % step_size:
        raise Exception('`test_size - h` should be module `step_size`')
    n_windows = int((test_size - h) / step_size) + 1
    if len(np.unique(last_dates)) == 1:
        if issubclass(last_dates.dtype.type, np.integer):
            total_dates = np.arange(last_dates[0] - test_size + 1, last_dates[0] + 1)
            out = np.empty((h * n_windows, 2), dtype=last_dates.dtype)
            freq = 1
        else:
            total_dates = pd.date_range(end=last_dates[0], periods=test_size, freq=freq)
            out = np.empty((h * n_windows, 2), dtype='datetime64[s]')
        for i_window, cutoff in enumerate(range(-test_size, -h + 1, step_size), start=0):
            end_cutoff = cutoff + h
            out[h * i_window : h * (i_window + 1), 0] = total_dates[cutoff:] if end_cutoff == 0 else total_dates[cutoff:end_cutoff]
            out[h * i_window : h * (i_window + 1), 1] = np.tile(total_dates[cutoff] - freq, h)
        dates = pd.DataFrame(np.tile(out, (len(last_dates), 1)), columns=['ds', 'cutoff'])
    else:
        dates = pd.concat([_cv_dates(np.array([ld]), freq, h, test_size, step_size) for ld in last_dates])
        dates = dates.reset_index(drop=True)
    return dates

# Internal Cell
def _build_forecast_name(model, *args, idx_remove=4) -> str:
    model_name = f'{model.__name__}'
    func_params = inspect.signature(model).parameters
    func_args = list(func_params.items())[idx_remove:]  # remove input array, horizon and xreg
    changed_params = [
        f'{name}-{value}'
        for value, (name, arg) in zip(args, func_args)
        if arg.default != value
    ]
    if changed_params:
        model_name += '_' + '_'.join(changed_params)
    return model_name

# Internal Cell
def _as_tuple(x):
    if isinstance(x, tuple):
        return x
    return (x,)

# Internal Cell
def _get_n_jobs(n_groups, n_jobs, ray_address):
    if ray_address is not None:
        logger.info(
            'Using ray address,'
            'using available resources insted of `n_jobs`'
        )
        try:
            import ray
        except ModuleNotFoundError as e:
            msg = (
                '{e}. To use a ray cluster you have to install '
                'ray. Please run `pip install ray`. '
            )
            raise ModuleNotFoundError(msg) from e
        if not ray.is_initialized():
            ray.init(ray_address, ignore_reinit_error=True)
        actual_n_jobs = int(ray.available_resources()['CPU'])
    else:
        if n_jobs == -1 or (n_jobs is None):
            actual_n_jobs = cpu_count()
        else:
            actual_n_jobs = n_jobs
    return min(n_groups, actual_n_jobs)

# Cell
class StatsForecast:

    def __init__(self, df, models, freq, n_jobs=1, ray_address=None, sort_df=True):
        # needed for residuals, think about it later
        self.ga, self.uids, self.last_dates, self.ds = _grouped_array_from_df(df, sort_df)
        self.models = models
        self.freq = pd.tseries.frequencies.to_offset(freq)
        self.n_jobs = _get_n_jobs(len(self.ga), n_jobs, ray_address)
        self.ray_address = ray_address
        self.sort_df = sort_df

    def forecast(self, h, xreg=None, fitted=False, level=None):
        if xreg is not None:
            expected_shape = (h * len(self.ga), self.ga.data.shape[1])
            if xreg.shape != expected_shape:
                raise ValueError(f'Expected xreg to have shape {expected_shape}, but got {xreg.shape}')
            xreg, _, _, _ = _grouped_array_from_df(xreg, sort_df=self.sort_df)
        forecast_kwargs = dict(
            h=h, test_size=None, step_size=None,
            input_size=None,
            xreg=xreg, fitted=fitted,
            level=level, mode='forecast',
        )
        if self.n_jobs == 1:
            res_fcsts = self._sequential(**forecast_kwargs)
        else:
            res_fcsts = self._data_parallel(**forecast_kwargs)
        if fitted:
            self.fcst_fitted_ = res_fcsts['fitted']
        fcsts = res_fcsts['fcsts']
        if issubclass(self.last_dates.dtype.type, np.integer):
            last_date_f = lambda x: np.arange(x + 1, x + 1 + h, dtype=self.last_dates.dtype)
        else:
            last_date_f = lambda x: pd.date_range(x + self.freq, periods=h, freq=self.freq)
        if len(np.unique(self.last_dates)) == 1:
            dates = np.tile(last_date_f(self.last_dates[0]), len(self.ga))
        else:
            dates = np.hstack([
                last_date_f(last_date)
                for last_date in self.last_dates
            ])
        idx = pd.Index(np.repeat(self.uids, h), name='unique_id')
        return pd.DataFrame({'ds': dates, **fcsts}, index=idx)

    def forecast_fitted_values(self):
        if not hasattr(self, 'fcst_fitted_'):
            raise Exception('Please run `forecast` mehtod using `fitted=True`')
        fcst_fitted = {key: val['values'] for key, val in self.fcst_fitted_.items()}
        fcst_fitted['y'] = self.ga.data[:, 0]
        return pd.DataFrame({**fcst_fitted}, index=self.ds).reset_index(level=1)

    def cross_validation(self, h, n_windows=1, step_size=1, test_size=None, input_size=None, fitted=False):
        if test_size is None:
            test_size = h + step_size * (n_windows - 1)
        elif n_windows is None:
            if (test_size - h) % step_size:
                raise Exception('`test_size - h` should be module `step_size`')
            n_windows = int((test_size - h) / step_size) + 1
        elif (n_windows is None) and (test_size is None):
            raise Exception('you must define `n_windows` or `test_size`')
        else:
            raise Exception('you must define `n_windows` or `test_size` but not both')

        cv_kwargs = dict(
            h=h, test_size=test_size, step_size=step_size, input_size=input_size,
            xreg=None, fitted=fitted, level=None, mode='cv',
        )
        if self.n_jobs == 1:
            res_fcsts = self._sequential(**cv_kwargs)
        else:
            res_fcsts = self._data_parallel(**cv_kwargs)
        if fitted:
            self.cv_fitted_ = res_fcsts['fitted']
            self.n_cv_ = n_windows
        fcsts = res_fcsts['fcsts']
        dates = _cv_dates(last_dates=self.last_dates, freq=self.freq, h=h, test_size=test_size, step_size=step_size)
        dates = {'ds': dates['ds'].values, 'cutoff': dates['cutoff'].values}
        idx = pd.Index(np.repeat(self.uids, h * n_windows), name='unique_id')
        return pd.DataFrame({**dates, **fcsts}, index=idx)

    def cross_validation_fitted_values(self):
        if not hasattr(self, 'cv_fitted_'):
            raise Exception('Please run `cross_validation` mehtod using `fitted=True`')
        index = pd.MultiIndex.from_tuples(np.tile(self.ds, self.n_cv_), names=['unique_id', 'ds'])
        res = pd.DataFrame(index=index, columns=['cutoff', 'y'] + list(self.cv_fitted_.keys()))
        for model, res_ in self.cv_fitted_.items():
            res[model] = res_['values'].flatten('F')
        res['cutoff'] = res_['last_idxs'].flatten('F')
        res['y'] = np.tile(self.ga.data.flatten(), self.n_cv_)
        idxs = res_['idxs'].flatten('F')
        res = res.iloc[idxs].reset_index(level=1)
        res['cutoff'] = res['ds'].where(res['cutoff']).bfill()
        return res

    def _sequential(self, h, test_size, step_size, input_size, xreg, fitted, level, mode='forecast'):
        result = {'fcsts': {}, 'fitted': {}}
        logger.info('Computing forecasts')
        for model_args in self.models:
            model, *args = _as_tuple(model_args)
            model_name = _build_forecast_name(model, *args)
            if mode == 'forecast':
                res_fcsts = self.ga.compute_forecasts(h, model, xreg, fitted, level, *args)
                values = res_fcsts['forecasts']
                keys = res_fcsts['keys']
            elif mode == 'cv':
                res_fcsts = self.ga.compute_cv(h, test_size, model, step_size, input_size, fitted, *args)
                values = res_fcsts['forecasts']
                test_values = res_fcsts['y']
                keys = None
            if keys is not None:
                for j, key in enumerate(keys):
                    result['fcsts'][f'{model_name}_{key}'] = values[:, j]
            else:
                result['fcsts'][model_name] = values.flatten()
            if fitted:
                result['fitted'][model_name] = res_fcsts['fitted']
            logger.info(f'Computed forecasts for {model_name}.')
        if mode == 'cv':
            result['fcsts'] = {'y': test_values.flatten(), **result['fcsts']}
        return result

    def _data_parallel(self, h, test_size, step_size, input_size, xreg, fitted, level, mode='forecast'):
        result = {'fcsts': {}, 'fitted': {}}
        logger.info('Computing forecasts')
        gas = self.ga.split(self.n_jobs)
        if xreg is not None:
            xregs = xreg.split(self.n_jobs)
        else:
            from itertools import repeat

            xregs = repeat(None)

        if self.ray_address is not None:
            try:
                from ray.util.multiprocessing import Pool
            except ModuleNotFoundError as e:
                msg = (
                    f'{e}. To use a ray cluster you have to install '
                    'ray. Please run `pip install ray`. '
                )
                raise ModuleNotFoundError(msg) from e
            kwargs = dict(ray_address=self.ray_address)
        else:
            from multiprocessing import Pool
            kwargs = dict()

        with Pool(self.n_jobs, **kwargs) as executor:
            for model_args in self.models:
                model, *args = _as_tuple(model_args)
                model_name = _build_forecast_name(model, *args)
                futures = []
                for ga, xr in zip(gas, xregs):
                    if mode == 'forecast':
                        future = executor.apply_async(ga.compute_forecasts, (h, model, xr, fitted, level, *args,))
                    elif mode == 'cv':
                        future = executor.apply_async(ga.compute_cv, (h, test_size, model, step_size, input_size, fitted, *args))
                    futures.append(future)
                if mode == 'forecast':
                    res_fcsts = [f.get() for f in futures]
                    values = [d['forecasts'] for d in res_fcsts]
                    keys = [d['keys'] for d in res_fcsts]
                    keys = keys[0]
                elif mode == 'cv':
                    res_fcsts = [f.get() for f in futures]
                    values = [d['forecasts'] for d in res_fcsts]
                    test_values = [d['y'] for d in res_fcsts]
                    keys = None
                if keys is not None:
                    values = np.vstack(values)
                    for j, key in enumerate(keys):
                        result['fcsts'][f'{model_name}_{key}'] = values[:, j]
                else:
                    values = np.hstack([val.flatten() for val in values])
                    result['fcsts'][model_name] = values.flatten()
                if fitted:
                    res = {}
                    for k in res_fcsts[0]['fitted'].keys():
                        res[k] = np.concatenate([d['fitted'][k] for d in res_fcsts])
                    result['fitted'][model_name] = res
                logger.info(f'Computed forecasts for {model_name}.')
        if mode == 'cv':
            test_values = np.vstack(test_values)
            result['fcsts'] = {'y': test_values.flatten(), **result['fcsts']}
        return result