#!/usr/bin/env python
# -*- coding: utf-8 -*-

from typing import List, Optional

import torch
from torch import Tensor


@torch.jit.script
def fill_normal(
    inputs: Tensor,
    fill_mask: Optional[Tensor] = None,
    sample_mask: Optional[Tensor] = None,
    preserve_var: bool = True,
    unbiased: bool = True,
) -> Tensor:
    r"""Fills a tensor with samples from a normal distribution under optional masking constraints,
    preserving mean and optionally variance. This method is focused towards a use case where the
    parameters of the normal distribution are derived from the input tensor.

    .. note::
        This method may be significantly faster when ``sample_mask`` is ``None``.

    Args:
        inputs (:class:`torch.Tensor`):
            Input tensor to fill

        fill_mask (:class:`torch.Tensor`):
            Boolean mask of locations that should be filled with new values. By default, the entire
            tensor will be filled with new values.

        sample_mask (:class:`torch.Tensor`):
            Boolean mask of locations to include when calculating per channel mean/variance.
            By default, all locations are included in the mean/variance calculation.

        preserve_var (bool):
            If true, fill values for each channel will be generated by sampling values from a normal
            distribution parameterized by the calculated mean/variance of the channel. Otherwise,
            sample from a normal distribution centered at the per channel mean, but with zero variance.

        unbiased (bool):
            Whether to use an unbiased estimator for variance. See :func:`torch.var_mean`.

    Dropout Example::

        >>> t1 = torch.rand(4, 3, 10, 10)
        >>> # drop channels 0 and 1 while preserving channel mean/variance
        >>> t1[:, 0:2, :, :] = fill_normal(t1[:, 0:2, :, :])

    Shape
        * ``inputs`` - :math:`(N, C, *)`
        * ``fill_mask`` - Same as ``inputs``
        * ``sample_mask`` - Same as ``inputs``
    """
    assert fill_mask is None or fill_mask.sum()
    assert sample_mask is None or sample_mask.sum()

    # get mean / variance by reducing `dim`
    batch_dim, channel_dim = 0, 1
    batch, channel = inputs.shape[0], inputs.shape[1]
    if sample_mask is not None:
        var, mean = [], []
        sample_source = inputs.flatten(batch_dim, channel_dim)
        sample_mask = sample_mask.flatten(batch_dim, channel_dim)
        for src, mask in zip(sample_source, sample_mask):
            v, m = torch.var_mean(src[mask], dim=0, keepdim=True, unbiased=unbiased)
            var.append(v)
            mean.append(m)
        var = torch.stack(var, dim=0).view(batch, channel, 1)
        mean = torch.stack(mean, dim=0).view(batch, channel, 1)
    else:
        sample_source = inputs
        dims: List[int] = torch.arange(2, inputs.ndim).tolist()
        var, mean = torch.var_mean(sample_source, dim=dims, keepdim=True, unbiased=unbiased)

    # match dimensions with input
    while mean.ndim < inputs.ndim:
        mean.unsqueeze_(-1)
        var.unsqueeze_(-1)

    # get fill values, either with mean alone or sampled from variance preserving distribution
    if preserve_var:
        mean = mean.expand_as(inputs)
        std = var.sqrt().expand_as(inputs)
        fill_values = torch.normal(mean, std)
    else:
        fill_values = mean

    # apply fill values
    if fill_mask is not None:
        return torch.where(fill_mask, fill_values, inputs)
    else:
        return fill_values
