
from escnn.gspaces import *
from escnn.nn import FieldType
from escnn.nn import GeometricTensor

from ..equivariant_module import EquivariantModule

import torch
import torch.nn.functional as F

from typing import List, Tuple, Any

import numpy as np

__all__ = ["PointwiseNonLinearity"]


class PointwiseNonLinearity(EquivariantModule):
    
    def __init__(self, in_type: FieldType, function: str = 'p_relu'):
        r"""
        
        Pointwise non-linearities.
        The same scalar function is applied to every channel independently.
        The input representation is preserved by this operation and, therefore, it equals the output representation.
        
        Only representations supporting pointwise non-linearities are accepted as input field type.
        
        Args:
            in_type (FieldType):  the input field type
            function (str): the identifier of the non-linearity.
                    It is used to specify which function to apply.
                    By default (``'p_relu'``), ReLU is used.
            
        """

        assert isinstance(in_type.gspace, GSpace)
        
        super(PointwiseNonLinearity, self).__init__()

        for r in in_type.representations:
            assert 'pointwise' in r.supported_nonlinearities, \
                'Error! Representation "{}" does not support "pointwise" non-linearity'.format(r.name)
        
        self.space = in_type.gspace
        self.in_type = in_type
        
        # the representation in input is preserved
        self.out_type = in_type
        
        # retrieve the activation function to apply
        if function == 'p_relu':
            self._function = F.relu
        elif function == 'p_elu':
            self._function = F.elu
        elif function == 'p_sigmoid':
            self._function = F.sigmoid
        elif function == 'p_tanh':
            self._function = F.tanh
        else:
            raise ValueError('Function "{}" not recognized!'.format(function))

    def forward(self, input: GeometricTensor) -> GeometricTensor:
        r"""

        Applies the pointwise activation function on the input fields

        Args:
            input (GeometricTensor): the input feature map

        Returns:
            the resulting feature map after the non-linearities have been applied

        """

        assert input.type == self.in_type
        
        # TODO - remove the 'contiguous()' call as soon as PyTorch's error is fixed
        # return GeometricTensor(self._function(input.tensor.contiguous()), self.out_type)
        return GeometricTensor(self._function(input.tensor), self.out_type, input.coords)

    def evaluate_output_shape(self, input_shape: Tuple[int, ...]) -> Tuple[int, ...]:

        assert len(input_shape) >= 2
        assert input_shape[1] == self.in_type.size

        b, c = input_shape[:2]
        spatial_shape = input_shape[2:]

        return (b, self.out_type.size, *spatial_shape)

    def check_equivariance(self, atol: float = 1e-6, rtol: float = 1e-5) -> List[Tuple[Any, float]]:
    
        c = self.in_type.size
    
        x = torch.randn(3, c, 10, 10)
    
        x = GeometricTensor(x, self.in_type)
    
        errors = []
    
        for el in self.space.testing_elements:
            out1 = self(x).transform_fibers(el)
            out2 = self(x.transform_fibers(el))
        
            errs = (out1.tensor - out2.tensor).detach().numpy()
            errs = np.abs(errs).reshape(-1)
            print(el, errs.max(), errs.mean(), errs.var())
        
            assert torch.allclose(out1.tensor, out2.tensor, atol=atol, rtol=rtol), \
                'The error found during equivariance check with element "{}" is too high: max = {}, mean = {} var ={}' \
                    .format(el, errs.max(), errs.mean(), errs.var())
        
            errors.append((el, errs.mean()))
    
        return errors
