# Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License").
# You may not use this file except in compliance with the License.
# A copy of the License is located at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# or in the "license" file accompanying this file. This file is distributed
# on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
# express or implied. See the License for the specific language governing
# permissions and limitations under the License.

from typing import Any, Dict, List, Optional, Iterable

import lightning.pytorch as pl
import torch
import numpy as np

from gluonts_customized.core.component import validated
from gluonts_customized.dataset.common import Dataset
from gluonts_customized.dataset.field_names import FieldName
from gluonts_customized.dataset.loader import as_stacked_batches
from gluonts_customized.itertools import Cyclic
from gluonts_customized.torch.model.estimator import PyTorchLightningEstimator
from gluonts_customized.torch.model.predictor import PyTorchPredictor
from gluonts_customized.time_feature import (
    get_seasonality,
    time_features_from_frequency_str,
    TimeFeature,
)
from gluonts_customized.transform import (
    Transformation,
    AddAgeFeature,
    AddObservedValuesIndicator,
    AddTimeFeatures,
    AsNumpyArray,
    Chain,
    ExpectedNumInstanceSampler,
    InstanceSampler,
    InstanceSplitter,
    SetField,
    RemoveFields,
    QuantizeMeanScaled,
    VstackFeatures,
    Identity,
    TestSplitSampler,
    ValidationSplitSampler,
)
from gluonts_customized.maybe import unwrap_or

from .module import WaveNet
from .lightning_module import WaveNetLightningModule

PREDICTION_INPUT_NAMES = [
    "feat_static_cat",
    "feat_static_real",
    "past_target",
    "past_observed_values",
    "past_time_feat",
    "future_time_feat",
    "scale",
]

TRAINING_INPUT_NAMES = PREDICTION_INPUT_NAMES + [
    "future_target",
    "future_observed_values",
]


class WaveNetEstimator(PyTorchLightningEstimator):
    @validated()
    def __init__(
        self,
        freq: str,
        prediction_length: int,
        num_bins: int = 1024,
        num_residual_channels: int = 24,
        num_skip_channels: int = 32,
        dilation_depth: Optional[int] = None,
        num_stacks: int = 1,
        temperature: float = 1.0,
        num_feat_dynamic_real: int = 0,
        num_feat_static_cat: int = 0,
        num_feat_static_real: int = 0,
        cardinality: List[int] = [1],
        seasonality: Optional[int] = None,
        embedding_dimension: int = 5,
        use_log_scale_feature: bool = True,
        time_features: Optional[List[TimeFeature]] = None,
        lr: float = 1e-3,
        weight_decay: float = 1e-8,
        train_sampler: Optional[InstanceSampler] = None,
        validation_sampler: Optional[InstanceSampler] = None,
        batch_size: int = 32,
        num_batches_per_epoch: int = 50,
        num_parallel_samples: int = 100,
        negative_data: bool = False,
        trainer_kwargs: Optional[Dict[str, Any]] = None,
    ) -> None:
        """
        WaveNet estimator that uses the architecture proposed in.

        [Oord et al., 2016] with quantized targets. The model is trained
        using the cross-entropy loss.

        [Oord et al., 2016] "Wavenet: A generative model for raw audio."
        arXiv preprint arXiv:1609.03499 (2016).

        Parameters
        ----------
        freq
            Frequency of the time series
        prediction_length
            Length of the prediction horizon
        num_bins, optional
            Number ofum bins used for quantization of the time series,
            by default 1024
        num_residual_channels, optional
            Number of residual channels in wavenet architecture, by default 24
        num_skip_channels, optional
            Number of skip channels in wavenet architecture, by default 32
        dilation_depth, optional
            Number of dilation layers in wavenet architecture. If set to None
            (default), dialation_depth is set such that the receptive length
            is at least as long as typical seasonality for the frequency and
            at least 2 * prediction_length, by default None
        num_stacks, optional
            Number of dilation stacks in wavenet architecture, by default 1
        temperature, optional
            Temparature used for sampling from softmax distribution,
            by default 1.0
        num_feat_dynamic_real, optional
            The number of dynamic real features, by default 0
        num_feat_static_cat, optional
            The number of static categorical features, by default 0
        num_feat_static_real, optional
            The number of static real features, by default 0
        cardinality, optional
            The cardinalities of static categorical features, by default [1]
        seasonality, optional
            The seasonality of the time series. If set to None, seasonality
            is set to default value based on the `freq`, by default None
        embedding_dimension, optional
            The dimension of the embeddings for categorical features,
            by default 5
        use_log_scale_feature, optional
            If True, logarithm of the scale of the past data will be used as an
            additional static feature,
            by default True
        time_features, optional
            List of time features, from :py:mod:`gluonts_customized.time_feature`,
            by default None
        lr, optional
            Learning rate, by default 1e-3
        weight_decay, optional
            Weight decay, by default 1e-8
        train_sampler, optional
            Controls the sampling of windows during training,
            by default None
        validation_sampler, optional
            Controls the sampling of windows during validation,
            by default None
        batch_size, optional
            The size of the batches to be used for training, by default 32
        num_batches_per_epoch, optional
            Number of batches to be processed in each training epoch,
            by default 50
        num_parallel_samples, optional
            The number of parallel samples to generate during inference.
            This parameter is only used in inference mode, by default 100
        negative_data, optional
            Flag indicating whether the time series take negative values,
            by default False
        trainer_kwargs, optional
            Additional arguments to provide to ``pl.Trainer`` for construction,
            by default None
        """
        default_trainer_kwargs = {
            "max_epochs": 100,
            "gradient_clip_val": 10.0,
        }
        if trainer_kwargs is not None:
            default_trainer_kwargs.update(trainer_kwargs)
        super().__init__(trainer_kwargs=default_trainer_kwargs)

        self.freq = freq
        self.prediction_length = prediction_length
        self.num_feat_dynamic_real = num_feat_dynamic_real
        self.num_feat_static_cat = num_feat_static_cat
        self.num_feat_static_real = num_feat_static_real
        self.cardinality = cardinality
        self.embedding_dimension = embedding_dimension
        self.num_bins = num_bins
        self.num_residual_channels = num_residual_channels
        self.num_skip_channels = num_skip_channels
        self.num_stacks = num_stacks
        self.use_log_scale_feature = use_log_scale_feature
        self.time_features = unwrap_or(
            time_features, time_features_from_frequency_str(freq)
        )
        self.temperature = temperature
        self.lr = lr
        self.weight_decay = weight_decay
        self.train_sampler = unwrap_or(
            train_sampler,
            ExpectedNumInstanceSampler(
                num_instances=1.0, min_future=self.prediction_length
            ),
        )
        self.validation_sampler = unwrap_or(
            validation_sampler,
            ValidationSplitSampler(min_future=self.prediction_length),
        )
        self.batch_size = batch_size
        self.num_batches_per_epoch = num_batches_per_epoch
        self.num_parallel_samples = num_parallel_samples
        self.negative_data = negative_data
        low = -10.0 if self.negative_data else 0
        high = 10.0
        bin_centers = np.linspace(low, high, self.num_bins)
        bin_edges = np.concatenate(
            [[-1e20], (bin_centers[1:] + bin_centers[:-1]) / 2.0, [1e20]]
        )
        self.bin_centers = bin_centers.tolist()
        self.bin_edges = bin_edges.tolist()

        self.seasonality = unwrap_or(
            seasonality,
            get_seasonality(
                freq,
                {
                    "H": 7 * 24,
                    "D": 7,
                    "W": 52,
                    "M": 12,
                    "B": 7 * 5,
                    "T": 24 * 60,
                    "S": 60 * 60,
                },
            ),
        )

        goal_receptive_length = max(
            2 * self.seasonality, 2 * self.prediction_length
        )
        if dilation_depth is None:
            d = 1
            while (
                WaveNet.get_receptive_field(
                    dilation_depth=d, num_stacks=num_stacks
                )
                < goal_receptive_length
            ):
                d += 1
            self.dilation_depth = d
        else:
            self.dilation_depth = dilation_depth
        self.context_length = WaveNet.get_receptive_field(
            dilation_depth=self.dilation_depth, num_stacks=num_stacks
        )

    def create_transformation(self) -> Transformation:
        remove_field_names = []
        if self.num_feat_static_real == 0:
            remove_field_names.append(FieldName.FEAT_STATIC_REAL)
        if self.num_feat_dynamic_real == 0:
            remove_field_names.append(FieldName.FEAT_DYNAMIC_REAL)
        return Chain(
            [
                RemoveFields(field_names=remove_field_names),
                (
                    SetField(output_field=FieldName.FEAT_STATIC_CAT, value=[0])
                    if self.num_feat_static_cat == 0
                    else Identity()
                ),
                (
                    SetField(
                        output_field=FieldName.FEAT_STATIC_REAL, value=[0.0]
                    )
                    if self.num_feat_static_real == 0
                    else Identity()
                ),
                AsNumpyArray(
                    field=FieldName.FEAT_STATIC_CAT, expected_ndim=1, dtype=int
                ),
                AsNumpyArray(
                    field=FieldName.FEAT_STATIC_REAL,
                    expected_ndim=1,
                    dtype=np.float32,
                ),
                AsNumpyArray(field=FieldName.TARGET, expected_ndim=1),
                AddObservedValuesIndicator(
                    target_field=FieldName.TARGET,
                    output_field=FieldName.OBSERVED_VALUES,
                ),
                AddTimeFeatures(
                    start_field=FieldName.START,
                    target_field=FieldName.TARGET,
                    output_field=FieldName.FEAT_TIME,
                    time_features=self.time_features,
                    pred_length=self.prediction_length,
                ),
                AddAgeFeature(
                    target_field=FieldName.TARGET,
                    output_field=FieldName.FEAT_AGE,
                    pred_length=self.prediction_length,
                ),
                VstackFeatures(
                    output_field=FieldName.FEAT_TIME,
                    input_fields=[FieldName.FEAT_TIME, FieldName.FEAT_AGE]
                    + (
                        [FieldName.FEAT_DYNAMIC_REAL]
                        if self.num_feat_dynamic_real > 0
                        else []
                    ),
                ),
                AsNumpyArray(
                    FieldName.FEAT_TIME, expected_ndim=2, dtype=np.float32
                ),
            ]
        )

    def _create_instance_splitter(self, mode: str):
        assert mode in ["training", "validation", "test"]

        instance_sampler = {
            "training": self.train_sampler,
            "validation": self.validation_sampler,
            "test": TestSplitSampler(),
        }[mode]

        return InstanceSplitter(
            target_field=FieldName.TARGET,
            is_pad_field=FieldName.IS_PAD,
            start_field=FieldName.START,
            forecast_start_field=FieldName.FORECAST_START,
            instance_sampler=instance_sampler,
            past_length=self.context_length,
            future_length=self.prediction_length,
            output_NTC=False,
            time_series_fields=[
                FieldName.FEAT_TIME,
                FieldName.OBSERVED_VALUES,
            ],
        ) + QuantizeMeanScaled(bin_edges=self.bin_edges)

    def create_training_data_loader(
        self,
        data: Dataset,
        module: WaveNetLightningModule,
        shuffle_buffer_length: Optional[int] = None,
        **kwargs,
    ) -> Iterable:
        data = Cyclic(data).stream()
        instances = self._create_instance_splitter("training").apply(
            data, is_train=True
        )
        return as_stacked_batches(
            instances,
            batch_size=self.batch_size,
            shuffle_buffer_length=shuffle_buffer_length,
            field_names=TRAINING_INPUT_NAMES,
            output_type=torch.tensor,
            num_batches_per_epoch=self.num_batches_per_epoch,
        )

    def create_validation_data_loader(
        self, data: Dataset, module: WaveNetLightningModule, **kwargs
    ) -> Iterable:
        instances = self._create_instance_splitter("validation").apply(
            data, is_train=True
        )
        return as_stacked_batches(
            instances,
            batch_size=self.batch_size,
            field_names=TRAINING_INPUT_NAMES,
            output_type=torch.tensor,
        )

    def create_lightning_module(self) -> pl.LightningModule:
        return WaveNetLightningModule(
            lr=self.lr,
            weight_decay=self.weight_decay,
            model_kwargs=dict(
                bin_values=self.bin_centers,
                num_residual_channels=self.num_residual_channels,
                num_skip_channels=self.num_skip_channels,
                dilation_depth=self.dilation_depth,
                num_stacks=self.num_stacks,
                num_feat_dynamic_real=1
                + self.num_feat_dynamic_real
                + len(self.time_features),
                num_feat_static_real=max(1, self.num_feat_static_real),
                cardinality=self.cardinality,
                embedding_dimension=self.embedding_dimension,
                pred_length=self.prediction_length,
                num_parallel_samples=self.num_parallel_samples,
                temperature=self.temperature,
                use_log_scale_feature=self.use_log_scale_feature,
            ),
        )

    def create_predictor(
        self, transformation: Transformation, module: WaveNetLightningModule
    ) -> PyTorchPredictor:
        prediction_splitter = self._create_instance_splitter("test")

        return PyTorchPredictor(
            input_transform=transformation + prediction_splitter,
            input_names=PREDICTION_INPUT_NAMES,
            prediction_net=module,
            batch_size=self.batch_size,
            prediction_length=self.prediction_length,
            device="auto",
        )
