from manim import *
from matplotlib import animation
from manim_ml.neural_network.layers.parent_layers import VGroupNeuralNetworkLayer


class Convolutional2DLayer(VGroupNeuralNetworkLayer):
    def __init__(
        self,
        feature_map_height,
        feature_map_width,
        filter_width,
        filter_height,
        stride=1,
        cell_width=0.5,
        feature_map_color=BLUE,
        filter_color=ORANGE,
        **kwargs
    ):
        super(VGroupNeuralNetworkLayer, self).__init__(**kwargs)
        self.feature_map_height = feature_map_height
        self.feature_map_width = feature_map_width
        self.filter_width = filter_width
        self.filter_height = filter_height
        self.feature_map_color = feature_map_color
        self.filter_color = filter_color
        self.stride = stride
        self.cell_width = cell_width
        # Construct the input
        self.construct_feature_map()

    def construct_feature_map(self):
        """Makes feature map"""
        # Make feature map rectangle
        self.feature_map = Rectangle(
            width=self.feature_map_width * self.cell_width,
            height=self.feature_map_height * self.cell_width,
            color=self.feature_map_color,
            grid_xstep=self.cell_width,
            grid_ystep=self.cell_width,
        )

        self.add(self.feature_map)

    @override_animation(Create)
    def _create_override(self, **kwargs):
        return FadeIn(self.feature_map)

    def make_forward_pass_animation(self, **kwargs):
        """Make feed forward animation"""
        return AnimationGroup()