from typing import Dict, Union, Optional
from copy import deepcopy

import numpy as np
import pandas as pd
from neuralforecast import NeuralForecast
from neuralforecast.models.nhits import NHITS
from neuralforecast.models.nbeats import NBEATS
from neuralforecast.losses.pytorch import MQLoss

from lightwood.helpers.log import log
from lightwood.mixer.base import BaseMixer
from lightwood.api.types import PredictionArguments
from lightwood.data.encoded_ds import EncodedDs, ConcatedEncodedDs


class NHitsMixer(BaseMixer):
    horizon: int
    target: str
    supports_proba: bool
    model_path: str
    hyperparam_search: bool
    default_config: dict
    SUPPORTED_MODELS = ('nhits', 'nbeats')

    def __init__(
            self,
            stop_after: float,
            target: str,
            horizon: int,
            window: int,
            dtype_dict: Dict,
            ts_analysis: Dict,
            pretrained: bool = False,
            train_args: Optional[Dict] = None,
    ):
        """
        Wrapper around an N-HITS deep learning model.
        
        :param stop_after: time budget in seconds.
        :param target: column to forecast.
        :param horizon: length of forecasted horizon.
        :param window: length of input data.
        :param ts_analysis: dictionary with miscellaneous time series info, as generated by 'lightwood.data.timeseries_analyzer'.
        :param train_args: arguments to steer the training process. 
            - `trainer_args`: all arguments for the PyTorchLightning trainer.
            - `conf_level`: level passed into MQLoss. Directly impacts prediction bounds. 
        """  # noqa
        super().__init__(stop_after)
        self.stable = False
        self.prepared = False
        self.supports_proba = False
        self.target = target
        self.window = window
        self.horizon = horizon
        self.dtype_dict = dtype_dict
        self.ts_analysis = ts_analysis
        self.grouped_by = ['__default'] if not ts_analysis['tss'].group_by else ts_analysis['tss'].group_by
        self.group_boundaries = {}  # stores last observed timestamp per series
        self.train_args = train_args.get('trainer_args', {}) if train_args else {}

        # we set a fairly aggressive training schedule by default
        self.train_args['early_stop_patience_steps'] = self.train_args.get('early_stop_patience_steps', 1)
        self.train_args['val_check_steps'] = self.train_args.get('val_check_steps', 10)
        self.train_args['learning_rate'] = self.train_args.get('learning_rate', 3e-3)
        self.train_args['mlp_units'] = self.train_args.get('mlp_units', [[128, 128], [128, 128]])
        self.train_args['random_seed'] = self.train_args.get('random_seed', 1)

        self.conf_level = self.train_args.pop('conf_level', [90])
        for level in self.conf_level:
            assert 0 <= level <= 100, f'A provided level is not in the [0, 100] range (found: {level})'
            assert isinstance(level, int), f'A provided level is not an integer (found: {level})'

        self.pretrained = pretrained
        self.base_url = 'https://nixtla-public.s3.amazonaws.com/transfer/pretrained_models/'
        self.freq_to_model = {
            'Y': 'yearly',
            'Q': 'monthly',
            'M': 'monthly',
            'W': 'daily',
            'D': 'daily',
            'H': 'hourly',
            'T': 'hourly',  # NOTE: use another pre-trained model once available
            'S': 'hourly'  # NOTE: use another pre-trained model once available
        }
        self.model = None
        self.model_class_str = self.train_args.get('model_class', 'nhits').lower()
        assert self.model_class_str in NHitsMixer.SUPPORTED_MODELS, f'Provided model class ({self.model_class_str}) is not supported. Supported models are: {NHitsMixer.SUPPORTED_MODELS}'  # noqa
        self.model_class = NBEATS if self.model_class_str == 'nbeats' else NHITS
        self.model_name = None
        self.model_names = {
            'nhits': {
                'hourly': 'nhits_m4_hourly.ckpt',  # hourly (non-tiny)
                'daily': 'nhits_m4_daily.ckpt',   # daily
                'monthly': 'nhits_m4_monthly.ckpt',  # monthly
                'yearly': 'nhits_m4_yearly.ckpt',  # yearly
            },
            'nbeats': {}  # TODO: complete
        }

    def fit(self, train_data: EncodedDs, dev_data: EncodedDs) -> None:
        """
        Fits the NeuralForecast model.
        """  # noqa
        log.info('Started fitting N-HITS forecasting model')

        # prepare data
        cat_ds = ConcatedEncodedDs([train_data, dev_data])
        oby_col = self.ts_analysis["tss"].order_by
        gby = self.ts_analysis["tss"].group_by if self.ts_analysis["tss"].group_by else []
        df = deepcopy(cat_ds.data_frame)
        Y_df, _ = self._make_initial_df(df, mode='train')
        self.group_boundaries = self._set_boundary(Y_df, gby)
        if gby:
            n_time = df[gby].value_counts().min()
        else:
            n_time = len(df[f'__mdb_original_{oby_col}'].unique())
        n_ts_val = max(int(.1 * n_time), self.horizon)  # at least self.horizon to validate on

        # train the model
        n_time_out = self.horizon
        if self.pretrained:
            # TODO: let user specify finetuning
            self.model_name = self.model_names.get(self.freq_to_model[self.ts_analysis['sample_freqs']['__default']],
                                                   None)
            self.model_name = self.model_names['hourly'] if self.model_name is None else self.model_name
            ckpt_url = self.base_url + self.model_name
            self.model = self.model_class.load_from_checkpoint(ckpt_url)

            if not self.window < self.model.hparams.n_time_in:
                log.info(f'NOTE: Provided window ({self.window}) is smaller than specified model input length ({self.model.hparams.n_time_in}). Will train a new model from scratch.')  # noqa
                self.pretrained = False
            if self.horizon > self.model.hparams.n_time_out:
                log.info(f'NOTE: Horizon ({self.horizon}) is bigger than that of the pretrained model ({self.model.hparams.n_time_out}). Will train a new model from scratch.')  # noqa
                self.pretrained = False
            if self.pretrained:
                log.info(f'Successfully loaded pretrained N-HITS forecasting model ({self.model_name})')

        if not self.pretrained:
            if self.window + self.horizon > n_time:
                new_window = max(1, n_time - self.horizon - 1)
                self.window = new_window
                log.info(f'Window {self.window} is too long for data provided (group: {df[gby].value_counts()[::-1].index[0]}), reducing window to {new_window}.')  # noqa
            model = self.model_class(h=n_time_out, input_size=self.window, **self.train_args, loss=MQLoss(level=self.conf_level))  # noqa
            self.model = NeuralForecast(models=[model], freq=self.ts_analysis['sample_freqs']['__default'],)
            self.model.fit(df=Y_df, val_size=n_ts_val)
            log.info('Successfully trained N-HITS forecasting model.')

    def partial_fit(self, train_data: EncodedDs, dev_data: EncodedDs, args: Optional[dict] = None) -> None:
        self.hyperparam_search = False
        self.train_args = args.get('trainer_args', {}) if args else {}  # NOTE: this replaces the original config
        self.fit(train_data, dev_data)
        self.prepared = True

    def __call__(self, ds: Union[EncodedDs, ConcatedEncodedDs],
                 args: PredictionArguments = PredictionArguments()) -> pd.DataFrame:
        """
        Calls the mixer to emit forecasts.
        
        NOTE: in the future we may support predicting every single row efficiently. For now, this mixer
        replicates the neuralforecast library behavior and returns a forecast strictly for the next `tss.horizon`
        timesteps after the end of the input dataframe.
        """  # noqa
        if args.predict_proba:
            log.warning('This mixer does not output probability estimates')

        # provided quantile must match one of the training levels, else we default to the largest one of these
        if args.fixed_confidence is not None and int(args.fixed_confidence * 100) in self.conf_level:
            level = int(args.fixed_confidence * 100)
        else:
            level = max(self.conf_level)

        target_cols = ['prediction', 'lower', 'upper']
        pred_cols = [
            f'{self.model_class_str.upper()}-median',
            f'{self.model_class_str.upper()}-lo-{level}',
            f'{self.model_class_str.upper()}-hi-{level}'
        ]

        input_df, idxs = self._make_initial_df(deepcopy(ds.data_frame))
        length = sum(ds.encoded_ds_lengths) if isinstance(ds, ConcatedEncodedDs) else len(ds)
        ydf = pd.DataFrame(0, index=range(length), columns=target_cols, dtype=object)

        # fill with zeroed arrays
        zero_array = [0 for _ in range(self.horizon)]
        for target_col in target_cols:
            ydf[target_col] = [zero_array] * len(ydf)

        grouper = input_df.groupby('unique_id')
        group_ends = grouper.last()['index'].values
        fcst = self.model.predict(input_df).reset_index()
        fcst['ds'] = fcst.groupby('unique_id').cumcount()
        horizons = pd.pivot_table(fcst, values=pred_cols, index='unique_id', columns='ds')

        temp_df = pd.DataFrame(0,  # zero-filled
                               index=range(len(horizons)),
                               columns=target_cols,
                               dtype=object)

        for pcol, tcol in zip(pred_cols, target_cols):
            temp_df[tcol] = horizons[pcol].values.tolist()

        for tcol in target_cols:
            ydf[tcol].iloc[group_ends] = temp_df[tcol]

        ydf['confidence'] = level / 100
        return ydf

    def _make_initial_df(self, df, mode='inference'):
        """
        Prepares a dataframe for the model according to what neuralforecast expects.

        If a per-group boundary exists, this method additionally drops out all observations prior to the cutoff.
        """  # noqa

        oby_col = self.ts_analysis["tss"].order_by
        # df = df.sort_values(by=f'__mdb_original_{oby_col}')  # TODO rm
        df[f'__mdb_parsed_{oby_col}'] = df.index
        df = df.reset_index(drop=True)

        Y_df = pd.DataFrame()
        Y_df['_index'] = np.arange(len(df))
        Y_df['y'] = df[self.target]
        Y_df['ds'] = df[f'__mdb_parsed_{oby_col}']

        if self.grouped_by != ['__default']:
            Y_df['unique_id'] = df[self.grouped_by].apply(lambda x: ','.join([elt for elt in x]), axis=1)
        else:
            Y_df['unique_id'] = '__default'

        Y_df = Y_df.reset_index()

        # filter if boundary exists
        if mode == 'train' and self.group_boundaries:
            filtered = []
            grouped = Y_df.groupby(by='unique_id')
            for group, sdf in grouped:
                if group in self.group_boundaries:
                    sdf = sdf[sdf['ds'].gt(self.group_boundaries[group])]
                    if sdf.shape[0] > 0:
                        filtered.append(sdf)
            if filtered:
                Y_df = pd.concat(filtered)

        filtered_idxs = Y_df.pop('_index').values
        return Y_df, filtered_idxs

    @staticmethod
    def _set_boundary(df: pd.DataFrame, gby: list) -> Dict[str, object]:
        """
        Finds last observation for every series in a pre-sorted `df` given a `gby` list of columns to group by.
        """
        if not gby:
            group_boundaries = {'__default': df.iloc[-1]['ds']}
        else:
            # could use groupby().transform('max'), but we leverage pre-sorting instead
            grouped_df = df.groupby(by='unique_id', as_index=False).last()
            group_boundaries = grouped_df[['unique_id', 'ds']].set_index('unique_id').to_dict()['ds']

        return group_boundaries
