# -*- coding: utf-8 -*-
from typing import Tuple

import torch
from torch import Tensor
from torch.nn import init
from torch.nn.parameter import Parameter
from .. import functional


class BatchNorm1dDuo(torch.nn.Module):
    __constant__ = ['in_features']
    in_features: int

    def __init__(self, in_features: int, bias_cov: bool = False, special_init: bool = True) -> None:
        super(BatchNorm1dDuo, self).__init__()
        self.in_features = in_features
        self.bn_mean = torch.nn.BatchNorm1d(in_features)
        if special_init:
            self.bn_mean.weight.data.fill_(2.5)
            self.bn_mean.bias.data.fill_(2.5)
        self.scale = in_features ** -0.5
        if bias_cov:
            self.bias_cov = Parameter(torch.Tensor(in_features, in_features))
        else:
            self.register_parameter("bias_cov", None)
        self.reset_parameters()

    def reset_parameters(self) -> None:
        if self.bias_cov is not None:
            init.uniform_(self.bias_cov, 1, 3)

    def forward(self, u: Tensor, cov: Tensor) -> Tuple[Tensor, Tensor]:
        return functional.mnn_bn1d_forward(self.bn_mean, u, cov, self.bn_mean.weight, self.bias_cov, self.scale)


class BatchNorm1dNoRho(torch.nn.Module):
    __constant__ = ['in_features']
    in_features: int

    def __init__(self, in_features: int, bias_std: bool = False, special_init: bool = True) -> None:
        super(BatchNorm1dNoRho, self).__init__()
        self.in_features = in_features
        self.special_init = special_init
        self.bn_mean = torch.nn.BatchNorm1d(in_features)
        if special_init:
            self.bn_mean.weight.data.fill_(2.5)
            self.bn_mean.bias.data.fill_(2.5)

        if bias_std:
            self.bias_std = Parameter(torch.Tensor(in_features))
        else:
            self.register_parameter("bias_std", None)
        self.reset_parameters()

    def reset_parameters(self) -> None:
        if self.bias_std is not None:
            if self.special_init:
                init.uniform_(self.bias_std, 2, 10)
            else:
                init.zeros_(self.bias_std)

    def forward(self, u: Tensor, s: Tensor) -> Tuple[Tensor, Tensor]:
        uhat = self.bn_mean(u)
        if self.bn_mean.training:
            var = torch.pow(s, 2) * torch.pow(self.bn_mean.weight, 2) / \
                  (torch.var(u, dim=0, keepdim=True) + self.bn_mean.eps)

        else:
            if self.bn_mean.track_running_stats:
                var = torch.pow(s, 2) * torch.pow(self.bn_mean.weight, 2) / \
                      (self.bn_mean.running_var + self.bn_mean.eps)
            else:
                var = torch.pow(s, 2) * torch.pow(self.bn_mean.weight, 2) / \
                      (torch.var(u, dim=0, keepdim=True) + self.bn_mean.eps)

        if self.bias_std is not None:
            var += torch.pow(self.bias_std, 2)

        return uhat, torch.sqrt(var)



# -*- coding: utf-8 -*-
from typing import Tuple

import torch
from torch import Tensor
from torch.nn import init
from torch.nn.parameter import Parameter
from .. import functional


class BatchNorm1dDuo(torch.nn.Module):
    __constant__ = ['in_features']
    in_features: int

    def __init__(self, in_features: int, bias_cov: bool = False, special_init: bool = True) -> None:
        super(BatchNorm1dDuo, self).__init__()
        self.in_features = in_features
        self.bn_mean = torch.nn.BatchNorm1d(in_features)
        if special_init:
            self.bn_mean.weight.data.fill_(2.5)
            self.bn_mean.bias.data.fill_(2.5)
        self.scale = in_features ** -0.5
        if bias_cov:
            self.bias_cov = Parameter(torch.Tensor(in_features, in_features))
        else:
            self.register_parameter("bias_cov", None)
        self.reset_parameters()

    def reset_parameters(self) -> None:
        if self.bias_cov is not None:
            init.uniform_(self.bias_cov, 1, 3)

    def forward(self, u: Tensor, cov: Tensor) -> Tuple[Tensor, Tensor]:
        return functional.mnn_bn1d_forward(self.bn_mean, u, cov, self.bn_mean.weight, self.bias_cov, self.scale)


class BatchNorm1dNoRho(torch.nn.Module):
    __constant__ = ['in_features']
    in_features: int

    def __init__(self, in_features: int, bias_std: bool = False, special_init: bool = True) -> None:
        super(BatchNorm1dNoRho, self).__init__()
        self.in_features = in_features
        self.special_init = special_init
        self.bn_mean = torch.nn.BatchNorm1d(in_features)
        if special_init:
            self.bn_mean.weight.data.fill_(2.5)
            self.bn_mean.bias.data.fill_(2.5)

        if bias_std:
            self.bias_std = Parameter(torch.Tensor(in_features))
        else:
            self.register_parameter("bias_std", None)
        self.reset_parameters()

    def reset_parameters(self) -> None:
        if self.bias_std is not None:
            if self.special_init:
                init.uniform_(self.bias_std, 2, 10)
            else:
                init.zeros_(self.bias_std)

    def forward(self, u: Tensor, s: Tensor) -> Tuple[Tensor, Tensor]:
        uhat = self.bn_mean(u)
        if self.bn_mean.training:
            var = torch.pow(s, 2) * torch.pow(self.bn_mean.weight, 2) / \
                  (torch.var(u, dim=0, keepdim=True) + self.bn_mean.eps)

        else:
            if self.bn_mean.track_running_stats:
                var = torch.pow(s, 2) * torch.pow(self.bn_mean.weight, 2) / \
                      (self.bn_mean.running_var + self.bn_mean.eps)
            else:
                var = torch.pow(s, 2) * torch.pow(self.bn_mean.weight, 2) / \
                      (torch.var(u, dim=0, keepdim=True) + self.bn_mean.eps)

        if self.bias_std is not None:
            var += torch.pow(self.bias_std, 2)

        return uhat, torch.sqrt(var)



# -*- coding: utf-8 -*-