# AUTOGENERATED FROM: babilim/model/layers/pooling.ipynb

# Cell: 0
"""doc
# babilim.model.layers.pooling

> Pooling operations.
"""

# Cell: 1
from babilim.core.annotations import RunOnlyOnce
from babilim.core.module_native import ModuleNative


# Cell: 2
class MaxPooling1D(ModuleNative):
    def __init__(self, pool_size=2, stride=None):
        """
        A N max pooling layer.
    
        Computes the max of a N region with stride S.
        This divides the feature map size by S.

        :param pool_size: Size of the region over which is pooled.
        :param stride: The stride defines how the top left corner of the pooling moves across the image. If None then it is same to pool_size resulting in zero overlap between pooled regions.
        """
        super().__init__()
        self.pool_size = pool_size
        self.stride = stride
        if self.stride is None:
            self.stride = self.pool_size
        
    @RunOnlyOnce
    def _build_pytorch(self, features):
        pass
        
    def _call_pytorch(self, features):
        from torch.nn.functional import max_pool1d as _MaxPooling1D
        return _MaxPooling1D(features, self.pool_size, stride=self.stride)
    
    @RunOnlyOnce
    def _build_tf(self, features):
        from tensorflow.keras.layers import MaxPooling1D as MaxPooling1D
        self.pool = _MaxPooling1D(pool_size=self.pool_size, strides=self.stride)
        
    def _call_tf(self, features):
        return self.pool(features)


# Cell: 3
class MaxPooling2D(ModuleNative):
    def __init__(self, pool_size=(2, 2), stride=None):
        """
        A NxN max pooling layer.
    
        Computes the max of a NxN region with stride S.
        This divides the feature map size by S.

        :param pool_size: Size of the region over which is pooled.
        :param stride: The stride defines how the top left corner of the pooling moves across the image. If None then it is same to pool_size resulting in zero overlap between pooled regions.
        """
        super().__init__()
        self.pool_size = pool_size
        self.stride = stride
        if self.stride is None:
            self.stride = self.pool_size
        
    @RunOnlyOnce
    def _build_pytorch(self, features):
        pass
        
    def _call_pytorch(self, features):
        from torch.nn.functional import max_pool2d as _MaxPooling2D
        return _MaxPooling2D(features, self.pool_size, stride=self.stride)
    
    @RunOnlyOnce
    def _build_tf(self, features):
        from tensorflow.keras.layers import MaxPooling2D as MaxPooling2D
        self.pool = _MaxPooling2D(pool_size=self.pool_size, strides=self.stride)
        
    def _call_tf(self, features):
        return self.pool(features)


# Cell: 4
class GlobalAveragePooling1D(ModuleNative):
    def __init__(self):
        """
        A global average pooling layer.
    
        This computes the global average in N dimension (B, N, C), so that the result is of shape (B, C).
        """
        super().__init__()
        
    @RunOnlyOnce
    def _build_pytorch(self, features):
        from babilim.model.layers.flatten import Flatten
        self.flatten = Flatten()
        
    def _call_pytorch(self, features):
        from torch.nn.functional import avg_pool1d as _AveragePooling1D
        return self.flatten(_AveragePooling1D(features, features.size()[2:]))
    
    @RunOnlyOnce
    def _build_tf(self, features):
        from tensorflow.keras.layers import GlobalAveragePooling1D as _GlobalAveragePooling1D
        self.pool = _GlobalAveragePooling1D()
        
    def _call_tf(self, features):
        return self.pool(features)


# Cell: 5
class GlobalAveragePooling2D(ModuleNative):
    def __init__(self):
        """
        A global average pooling layer.
    
        This computes the global average in W, H dimension, so that the result is of shape (B, C).
        """
        super().__init__()
        
    @RunOnlyOnce
    def _build_pytorch(self, features):
        from babilim.model.layers.flatten import Flatten
        self.flatten = Flatten()
        
    def _call_pytorch(self, features):
        from torch.nn.functional import avg_pool2d as _AveragePooling2D
        return self.flatten(_AveragePooling2D(features, features.size()[2:]))
    
    @RunOnlyOnce
    def _build_tf(self, features):
        from tensorflow.keras.layers import GlobalAveragePooling2D as _GlobalAveragePooling2D
        self.pool = _GlobalAveragePooling2D()
        
    def _call_tf(self, features):
        return self.pool(features)
