


__all__ = ['masked_median', 'masked_mean', 'minmax_statistics', 'minmax1_statistics', 'std_statistics', 'robust_statistics',
           'invariant_statistics', 'identity_statistics', 'TemporalNorm']


import torch
import torch.nn as nn


def masked_median(x, mask, dim=-1, keepdim=True):
    """Masked Median

    Compute the median of tensor `x` along dim, ignoring values where
    `mask` is False. `x` and `mask` need to be broadcastable.

    Args:
        x (torch.Tensor): Tensor to compute median of along `dim` dimension.
        mask (torch.Tensor): Tensor bool with same shape as `x`, where `x` is valid and False
            where `x` should be masked. Mask should not be all False in any column of
            dimension dim to avoid NaNs from zero division.
        dim (int, optional): Dimension to take median of. Defaults to -1.
        keepdim (bool, optional): Keep dimension of `x` or not. Defaults to True.

    Returns:
        torch.Tensor: Normalized values.
    """
    x_nan = x.masked_fill(mask < 1, float("nan"))
    x_median, _ = x_nan.nanmedian(dim=dim, keepdim=keepdim)
    x_median = torch.nan_to_num(x_median, nan=0.0)
    return x_median


def masked_mean(x, mask, dim=-1, keepdim=True):
    """Masked  Mean

    Compute the mean of tensor `x` along dimension, ignoring values where
    `mask` is False. `x` and `mask` need to be broadcastable.

    Args:
        x (torch.Tensor): Tensor to compute mean of along `dim` dimension.
        mask (torch.Tensor): Tensor bool with same shape as `x`, where `x` is valid and False
            where `x` should be masked. Mask should not be all False in any column of
            dimension dim to avoid NaNs from zero division.
        dim (int, optional): Dimension to take mean of. Defaults to -1.
        keepdim (bool, optional): Keep dimension of `x` or not. Defaults to True.

    Returns:
        torch.Tensor: Normalized values.
    """
    x_nan = x.masked_fill(mask < 1, float("nan"))
    x_mean = x_nan.nanmean(dim=dim, keepdim=keepdim)
    x_mean = torch.nan_to_num(x_mean, nan=0.0)
    return x_mean


def minmax_statistics(x, mask, eps=1e-6, dim=-1):
    r"""MinMax Scaler

    Standardizes temporal features by ensuring its range dweels between
    [0,1] range. This transformation is often used as an alternative
    to the standard scaler. The scaled features are obtained as:

    ```math
    \mathbf{z} = (\mathbf{x}_{[B,T,C]}-\mathrm{min}({\mathbf{x}})_{[B,1,C]})/
        (\mathrm{max}({\mathbf{x}})_{[B,1,C]}- \mathrm{min}({\mathbf{x}})_{[B,1,C]})
    ```

    Args:
        x (torch.Tensor): Input tensor.
        mask (torch.Tensor): Tensor bool, same dimension as `x`, indicates where `x` is valid and False
            where `x` should be masked. Mask should not be all False in any column of
            dimension dim to avoid NaNs from zero division.
        eps (float, optional): Small value to avoid division by zero. Defaults to 1e-6.
        dim (int, optional): Dimension over to compute min and max. Defaults to -1.

    Returns:
        torch.Tensor: Same shape as `x`, except scaled.
    """
    mask = mask.clone()
    mask[mask == 0] = torch.inf
    mask[mask == 1] = 0
    x_max = torch.max(
        torch.nan_to_num(x - mask, nan=-torch.inf), dim=dim, keepdim=True
    )[0]
    x_min = torch.min(torch.nan_to_num(x + mask, nan=torch.inf), dim=dim, keepdim=True)[
        0
    ]
    x_max = x_max.type(x.dtype)
    x_min = x_min.type(x.dtype)

    # x_range and prevent division by zero
    x_range = x_max - x_min
    x_range[x_range == 0] = 1.0
    x_range = x_range + eps
    return x_min, x_range


def minmax_scaler(x, x_min, x_range):
    return (x - x_min) / x_range


def inv_minmax_scaler(z, x_min, x_range):
    return z * x_range + x_min


def minmax1_statistics(x, mask, eps=1e-6, dim=-1):
    r"""MinMax1 Scaler

    Standardizes temporal features by ensuring its range dweels between
    [-1,1] range. This transformation is often used as an alternative
    to the standard scaler or classic Min Max Scaler.
    The scaled features are obtained as:

    ```math
    \mathbf{z} = 2 (\mathbf{x}_{[B,T,C]}-\mathrm{min}({\mathbf{x}})_{[B,1,C]})/ (\mathrm{max}({\mathbf{x}})_{[B,1,C]}- \mathrm{min}({\mathbf{x}})_{[B,1,C]})-1
    ```


    Args:
        x (torch.Tensor): Input tensor.
        mask (torch.Tensor): Tensor bool, same dimension as `x`, indicates where `x` is valid and False
            where `x` should be masked. Mask should not be all False in any column of
            dimension dim to avoid NaNs from zero division.
        eps (float, optional): Small value to avoid division by zero. Defaults to 1e-6.
        dim (int, optional): Dimension over to compute min and max. Defaults to -1.

    Returns:
        torch.Tensor: Same shape as `x`, except scaled.
    """
    # Mask values (set masked to -inf or +inf)
    mask = mask.clone()
    mask[mask == 0] = torch.inf
    mask[mask == 1] = 0
    x_max = torch.max(
        torch.nan_to_num(x - mask, nan=-torch.inf), dim=dim, keepdim=True
    )[0]
    x_min = torch.min(torch.nan_to_num(x + mask, nan=torch.inf), dim=dim, keepdim=True)[
        0
    ]
    x_max = x_max.type(x.dtype)
    x_min = x_min.type(x.dtype)

    # x_range and prevent division by zero
    x_range = x_max - x_min
    x_range[x_range == 0] = 1.0
    x_range = x_range + eps
    return x_min, x_range


def minmax1_scaler(x, x_min, x_range):
    x = (x - x_min) / x_range
    z = x * (2) - 1
    return z


def inv_minmax1_scaler(z, x_min, x_range):
    z = (z + 1) / 2
    return z * x_range + x_min


def std_statistics(x, mask, dim=-1, eps=1e-6):
    r"""Standard Scaler

    Standardizes features by removing the mean and scaling
    to unit variance along the `dim` dimension.

    For example, for `base_windows` models, the scaled features are obtained as (with dim=1):

    ```math
    \mathbf{z} = (\mathbf{x}_{[B,T,C]}-\bar{\mathbf{x}}_{[B,1,C]})/\hat{\sigma}_{[B,1,C]}
    ```

    Args:
        x (torch.Tensor): Input tensor.
        mask (torch.Tensor): Tensor bool, same dimension as `x`, indicates where `x` is valid and False
            where `x` should be masked. Mask should not be all False in any column of
            dimension dim to avoid NaNs from zero division.
        eps (float, optional): Small value to avoid division by zero. Defaults to 1e-6.
        dim (int, optional): Dimension over to compute mean and std. Defaults to -1.

    Returns:
        torch.Tensor: Same shape as `x`, except scaled.
    """
    x_means = masked_mean(x=x, mask=mask, dim=dim)
    x_stds = torch.sqrt(masked_mean(x=(x - x_means) ** 2, mask=mask, dim=dim))

    # Protect against division by zero
    x_stds[x_stds == 0] = 1.0
    x_stds = x_stds + eps
    return x_means, x_stds


def std_scaler(x, x_means, x_stds):
    return (x - x_means) / x_stds


def inv_std_scaler(z, x_mean, x_std):
    return (z * x_std) + x_mean


def robust_statistics(x, mask, dim=-1, eps=1e-6):
    r"""Robust Median Scaler

    Standardizes features by removing the median and scaling
    with the mean absolute deviation (mad) a robust estimator of variance.
    This scaler is particularly useful with noisy data where outliers can
    heavily influence the sample mean / variance in a negative way.
    In these scenarios the median and amd give better results.

    For example, for `base_windows` models, the scaled features are obtained as (with dim=1):

    ```math
    \mathbf{z} = (\mathbf{x}_{[B,T,C]}-\textrm{median}(\mathbf{x})_{[B,1,C]})/\textrm{mad}(\mathbf{x})_{[B,1,C]}
    ```

    ```math
    \textrm{mad}(\mathbf{x}) = \frac{1}{N} \sum_{}|\mathbf{x} - \mathrm{median}(x)|
    ```

    Args:
        x (torch.Tensor): Input tensor.
        mask (torch.Tensor): Tensor bool, same dimension as `x`, indicates where `x` is valid and False
            where `x` should be masked. Mask should not be all False in any column of
            dimension dim to avoid NaNs from zero division.
        eps (float, optional): Small value to avoid division by zero. Defaults to 1e-6.
        dim (int, optional): Dimension over to compute median and mad. Defaults to -1.

    Returns:
        torch.Tensor: Same shape as `x`, except scaled.
    """
    x_median = masked_median(x=x, mask=mask, dim=dim)
    x_mad = masked_median(x=torch.abs(x - x_median), mask=mask, dim=dim)

    # Protect x_mad=0 values
    # Assuming normality and relationship between mad and std
    x_means = masked_mean(x=x, mask=mask, dim=dim)
    x_stds = torch.sqrt(masked_mean(x=(x - x_means) ** 2, mask=mask, dim=dim))
    x_mad_aux = x_stds * 0.6744897501960817
    x_mad = x_mad * (x_mad > 0) + x_mad_aux * (x_mad == 0)

    # Protect against division by zero
    x_mad[x_mad == 0] = 1.0
    x_mad = x_mad + eps
    return x_median, x_mad


def robust_scaler(x, x_median, x_mad):
    return (x - x_median) / x_mad


def inv_robust_scaler(z, x_median, x_mad):
    return z * x_mad + x_median


def invariant_statistics(x, mask, dim=-1, eps=1e-6):
    r"""Invariant Median Scaler

    Standardizes features by removing the median and scaling
    with the mean absolute deviation (mad) a robust estimator of variance.
    Aditionally it complements the transformation with the arcsinh transformation.

    For example, for `base_windows` models, the scaled features are obtained as (with dim=1):

    ```math
    \mathbf{z} = (\mathbf{x}_{[B,T,C]}-\textrm{median}(\mathbf{x})_{[B,1,C]})/\textrm{mad}(\mathbf{x})_{[B,1,C]}
    ```

    ```math
    \mathbf{z} = \textrm{arcsinh}(\mathbf{z})
    ```



    Args:
        x (torch.Tensor): Input tensor.
        mask (torch.Tensor): Tensor bool, same dimension as `x`, indicates where `x` is valid and False
            where `x` should be masked. Mask should not be all False in any column of
            dimension dim to avoid NaNs from zero division.
        eps (float, optional): Small value to avoid division by zero. Defaults to 1e-6.
        dim (int, optional): Dimension over to compute median and mad. Defaults to -1.

    Returns:
        torch.Tensor: Same shape as `x`, except scaled.
    """
    x_median = masked_median(x=x, mask=mask, dim=dim)
    x_mad = masked_median(x=torch.abs(x - x_median), mask=mask, dim=dim)

    # Protect x_mad=0 values
    # Assuming normality and relationship between mad and std
    x_means = masked_mean(x=x, mask=mask, dim=dim)
    x_stds = torch.sqrt(masked_mean(x=(x - x_means) ** 2, mask=mask, dim=dim))
    x_mad_aux = x_stds * 0.6744897501960817
    x_mad = x_mad * (x_mad > 0) + x_mad_aux * (x_mad == 0)

    # Protect against division by zero
    x_mad[x_mad == 0] = 1.0
    x_mad = x_mad + eps
    return x_median, x_mad


def invariant_scaler(x, x_median, x_mad):
    return torch.arcsinh((x - x_median) / x_mad)


def inv_invariant_scaler(z, x_median, x_mad):
    return torch.sinh(z) * x_mad + x_median


def identity_statistics(x, mask, dim=-1, eps=1e-6):
    """Identity Scaler

    A placeholder identity scaler, that is argument insensitive.

    Args:
        x (torch.Tensor): Input tensor.
        mask (torch.Tensor): Tensor bool, same dimension as `x`, indicates where `x` is valid and False
            where `x` should be masked. Mask should not be all False in any column of
            dimension dim to avoid NaNs from zero division.
        eps (float, optional): Small value to avoid division by zero. Defaults to 1e-6.
        dim (int, optional): Dimension over to compute median and mad. Defaults to -1.

    Returns:
        torch.Tensor: Original `x`.
    """
    # Collapse dim dimension
    shape = list(x.shape)
    shape[dim] = 1

    x_shift = torch.zeros(shape, device=x.device)
    x_scale = torch.ones(shape, device=x.device)

    return x_shift, x_scale


def identity_scaler(x, x_shift, x_scale):
    return x


def inv_identity_scaler(z, x_shift, x_scale):
    return z


class TemporalNorm(nn.Module):
    r"""Temporal Normalization

    Standardization of the features is a common requirement for many
    machine learning estimators, and it is commonly achieved by removing
    the level and scaling its variance. The `TemporalNorm` module applies
    temporal normalization over the batch of inputs as defined by the type of scaler.

    ```math
    \mathbf{z}_{[B,T,C]} = \textrm{Scaler}(\mathbf{x}_{[B,T,C]})
    ```


    If `scaler_type` is `revin` learnable normalization parameters are added on top of
    the usual normalization technique, the parameters are learned through scale decouple
    global skip connections. The technique is available for point and probabilistic outputs.

    ```math
    \mathbf{\hat{z}}_{[B,T,C]} = \boldsymbol{\hat{\gamma}}_{[1,1,C]} \mathbf{z}_{[B,T,C]} +\boldsymbol{\hat{\beta}}_{[1,1,C]}
    ```


    Args:
        scaler_type (str, optional): Defines the type of scaler used by TemporalNorm. Available [`identity`, `standard`, `robust`, `minmax`, `minmax1`, `invariant`, `revin`]. Defaults to "robust".
    `dim` (int, optional): Dimension over to compute scale and shift. Defaults to -1.
        eps (float, optional): Small value to avoid division by zero. Defaults to 1e-6.
        num_features (int, optional): For RevIN-like learnable affine parameters initialization. Defaults to None.

    References:
        - [Kin G. Olivares, David Luo, Cristian Challu, Stefania La Vattiata, Max Mergenthaler, Artur Dubrawski (2023). "HINT: Hierarchical Mixture Networks For Coherent Probabilistic Forecasting". Neural Information Processing Systems, submitted. Working Paper version available at arxiv.](https://arxiv.org/abs/2305.07089)
    """

    def __init__(self, scaler_type="robust", dim=-1, eps=1e-6, num_features=None):
        super().__init__()
        compute_statistics = {
            None: identity_statistics,
            "identity": identity_statistics,
            "standard": std_statistics,
            "revin": std_statistics,
            "robust": robust_statistics,
            "minmax": minmax_statistics,
            "minmax1": minmax1_statistics,
            "invariant": invariant_statistics,
        }
        scalers = {
            None: identity_scaler,
            "identity": identity_scaler,
            "standard": std_scaler,
            "revin": std_scaler,
            "robust": robust_scaler,
            "minmax": minmax_scaler,
            "minmax1": minmax1_scaler,
            "invariant": invariant_scaler,
        }
        inverse_scalers = {
            None: inv_identity_scaler,
            "identity": inv_identity_scaler,
            "standard": inv_std_scaler,
            "revin": inv_std_scaler,
            "robust": inv_robust_scaler,
            "minmax": inv_minmax_scaler,
            "minmax1": inv_minmax1_scaler,
            "invariant": inv_invariant_scaler,
        }
        assert scaler_type in scalers.keys(), f"{scaler_type} not defined"
        if (scaler_type == "revin") and (num_features is None):
            raise Exception("You must pass num_features for ReVIN scaler.")

        self.compute_statistics = compute_statistics[scaler_type]
        self.scaler = scalers[scaler_type]
        self.inverse_scaler = inverse_scalers[scaler_type]
        self.scaler_type = scaler_type
        self.dim = dim
        self.eps = eps

        if scaler_type == "revin":
            self._init_params(num_features=num_features)

    def _init_params(self, num_features):
        # Initialize RevIN scaler params to broadcast:
        if self.dim == 1:  # [B,T,C]  [1,1,C]
            self.revin_bias = nn.Parameter(torch.zeros(1, 1, num_features, 1))
            self.revin_weight = nn.Parameter(torch.ones(1, 1, num_features, 1))
        elif self.dim == -1:  # [B,C,T]  [1,C,1]
            self.revin_bias = nn.Parameter(torch.zeros(1, num_features, 1, 1))
            self.revin_weight = nn.Parameter(torch.ones(1, num_features, 1, 1))

    # @torch.no_grad()
    def transform(self, x, mask):
        """Center and scale the data.

        Args:
            x (torch.Tensor): Tensor shape [batch, time, channels].
            mask (torch.Tensor): Tensor bool, shape  [batch, time] where `x` is valid and False
                where `x` should be masked. Mask should not be all False in any column of
                dimension dim to avoid NaNs from zero division.

        Returns:
            torch.Tensor: Same shape as `x`, except scaled.
        """
        x_shift, x_scale = self.compute_statistics(
            x=x, mask=mask, dim=self.dim, eps=self.eps
        )
        self.x_shift = x_shift
        self.x_scale = x_scale

        # Original Revin performs this operation
        # z = self.revin_weight * z
        # z = z + self.revin_bias
        # However this is only valid for point forecast not for
        # distribution's scale decouple technique.
        if self.scaler_type == "revin":
            self.x_shift = self.x_shift + self.revin_bias
            self.x_scale = self.x_scale * (torch.relu(self.revin_weight) + self.eps)

        z = self.scaler(x, x_shift, x_scale)
        return z

    # @torch.no_grad()
    def inverse_transform(self, z, x_shift=None, x_scale=None):
        """Scale back the data to the original representation.

        Args:
            z (torch.Tensor): Tensor shape [batch, time, channels], scaled.
            x_shift (torch.Tensor, optional): Tensor shape [1, 1, channels], shift. Defaults to None.
            x_scale (torch.Tensor, optional): Tensor shape [1, 1, channels], scale. Defaults to None.

        Returns:
            torch.Tensor: Original data.
        """

        if x_shift is None:
            x_shift = self.x_shift
        if x_scale is None:
            x_scale = self.x_scale

        # Original Revin performs this operation
        # z = z - self.revin_bias
        # z = (z / (self.revin_weight + self.eps))
        # However this is only valid for point forecast not for
        # distribution's scale decouple technique.

        x = self.inverse_scaler(z, x_shift, x_scale)
        return x

    def forward(self, x):
        # The gradients are optained from BaseWindows/BaseRecurrent forwards.
        pass
