class AutoencoderMLP(NetworkTemplate):
    r"""This is an implementation of a Fully-connected AutoEncoder as
    Reduced Order Model;

    A MLP autoencoder architecture consists of two stages:

    - Fully-connected encoder
    - Fully connected decoder

    Graphical scheme:

                |         |
                |  |   |  |
        Z ->    |  | | |  |  -> Z_til
                |  |   |  |
                |         |

    ENCODER       DECODER

    """

    def __init__(
        self,
        encoder: DenseNetwork = None,
        decoder: DenseNetwork = None,
        input_dim: Optional[int] = None,
        output_dim: Optional[int] = None,
        latent_dim: Optional[int] = None,
        activation: Optional[Union[list, str]] = None,
        shallow: Optional[bool] = False,
        devices: Union[str, list] = "cpu",
        name: str = None,
    ) -> None:
        """Initialize the AutoencoderMLP network

        Args:
            encoder (DenseNetwork, optional): The encoder network architecture. (Default value = None)
            decoder (DenseNetwork, optional): The decoder network architecture. (Default value = None)
            input_dim (Optional[int], optional): The input dimensions of the data, by default None.
            output_dim (Optional[int], optional): The output dimensions of the data, by default None.
            latent_dim (Optional[int], optional): The dimensions of the latent space, by default None.
            activation (Optional[Union[list, str]], optional): The activation functions used by the network, by default None.
            shallow (Optional[bool], optional): Whether the network should be shallow or not, by default False.
            devices (Union[str, list], optional): The device(s) to be used for allocating subnetworks, by default "cpu".
            name (str, optional): The name of the network, by default None.

        """

        super(AutoencoderMLP, self).__init__(name=name)

        self.weights = list()

        # This option is used when no network is provided
        # and it uses default choices for the architectures
        if encoder == None and decoder == None:
            encoder, decoder = mlp_autoencoder_auto(
                input_dim=input_dim,
                latent_dim=latent_dim,
                output_dim=output_dim,
                activation=activation,
                shallow=shallow,
            )

        # Determining the kind of device to be used for allocating the
        # subnetworks used in the DeepONet model
        self.device = self._set_device(devices=devices)

        self.encoder = self.to_wrap(entity=encoder, device=self.device)
        self.decoder = self.to_wrap(entity=decoder, device=self.device)

        self.add_module("encoder", self.encoder)
        self.add_module("decoder", self.decoder)

        self.weights += self.encoder.weights
        self.weights += self.decoder.weights

        self.last_encoder_channels = None

        self.shapes_dict = dict()

    def summary(self) -> None:
        """Prints the summary of the network architecture"""
        self.encoder.summary()
        self.decoder.summary()

    def projection(
        self, input_data: Union[np.ndarray, torch.Tensor] = None
    ) -> torch.Tensor:
        """Project the input dataset into the latent space.

        Args:
            input_data (Union[np.ndarray, torch.Tensor], optional): The dataset to be projected, by default None.

        Returns:
            torch.Tensor: The dataset projected over the latent space.

        """
        latent = self.encoder.forward(input_data=input_data)

        return latent

    def reconstruction(
        self, input_data: Union[torch.Tensor, np.ndarray] = None
    ) -> torch.Tensor:
        """Reconstruct the latent dataset to the original one.

        Args:
            input_data (Union[torch.Tensor, np.ndarray], optional): The dataset to be reconstructed, by default None.

        Returns:
            torch.Tensor: The dataset reconstructed.

        """
        reconstructed = self.decoder.forward(input_data=input_data)

        return reconstructed

    def forward(
        self, input_data: Union[np.ndarray, torch.Tensor] = None
    ) -> torch.Tensor:
        """Execute the complete projection/reconstruction pipeline.

        Args:
            input_data (Union[np.ndarray, torch.Tensor], optional): The input dataset, by default None.

        Returns:
            torch.Tensor: The dataset reconstructed.

        """
        latent = self.projection(input_data=input_data)
        reconstructed = self.reconstruction(input_data=latent)

        return reconstructed

    def eval_projection(
        self, input_data: Union[np.ndarray, torch.Tensor] = None
    ) -> np.ndarray:
        """Evaluate the projection of the input dataset into the latent space.

        Args:
            input_data (Union[np.ndarray, torch.Tensor], optional): The dataset to be projected, by default None.

        Returns:
            np.ndarray: The dataset projected over the latent space.

        """
        return self.projection(input_data=input_data).detach().numpy()

class AutoencoderCNN(NetworkTemplate):
    r"""This is an implementation of a convolutional autoencoder as Reduced Order Model.
    An autoencoder architecture consists of three stages:

    - The convolutional encoder
    - The bottleneck stage, subdivided in:
        - Fully-connected encoder
        - Fully connected decoder
    - The convolutional decoder

    Graphical scheme

        Z -> [Conv] -> [Conv] -> ... [Conv] -> |  | | |  | -> [Conv.T] -> [Conv.T] -> ... [Conv.T] -> Z_til


                        ENCODER               DENSE BOTTLENECK           DECODER

    """

    def __init__(
        self,
        encoder: ConvolutionalNetwork = None,
        bottleneck_encoder: Linear = None,
        bottleneck_decoder: Linear = None,
        decoder: ConvolutionalNetwork = None,
        encoder_activation: str = "relu",
        input_dim: Optional[Tuple[int, ...]] = None,
        output_dim: Optional[Tuple[int, ...]] = None,
        latent_dim: Optional[int] = None,
        kernel_size: Optional[int] = None,
        activation: Optional[Union[list, str]] = None,
        channels: Optional[int] = None,
        case: Optional[str] = None,
        shallow: Optional[bool] = False,
        devices: Union[str, list] = "cpu",
        name: str = None,
        **kwargs,
    ) -> None:
        """Initialize the AutoencoderCNN network.

        Args:
            encoder (ConvolutionalNetwork, optional): The encoder network architecture, by default None.
            bottleneck_encoder (Linear, optional): The bottleneck encoder network architecture, by default None.
            bottleneck_decoder (Linear, optional): The bottleneck decoder network architecture, by default None.
            decoder (ConvolutionalNetwork, optional): The decoder network architecture, by default None.
            encoder_activation (str, optional): The activation function used by the encoder network, by default 'relu'.
            input_dim (Optional[Tuple[int, ...]], optional): The input dimensions of the data, by default None.
            output_dim (Optional[Tuple[int, ...]], optional): The output dimensions of the data, by default None.
            latent_dim (Optional[int], optional): The dimensions of the latent space, by default None.
            kernel_size (Optional[int], optional):  (Default value = None)
            activation (Optional[Union[list, str]], optional): The activation functions used by the network, by default None.
            channels (Optional[int], optional): The number of channels of the convolutional layers, by default None.
            case (Optional[str], optional): The type of convolutional encoder and decoder to be used, by default None.
            shallow (Optional[bool], optional): Whether the network should be shallow or not, by default False.
            devices (Union[str, list], optional): The device(s) to be used for allocating subnetworks, by default 'cpu'.
            name (str, optional): The name of the network, by default None.
            **kwargs

        """

        super(AutoencoderCNN, self).__init__(name=name)

        self.weights = list()

        # Determining the kind of device to be used for allocating the
        # subnetworks
        self.device = self._set_device(devices=devices)

        self.input_dim = None

        # If not network is provided, the automatic generation
        # pipeline is activated.
        if all(
            [
                isn == None
                for isn in [encoder, decoder, bottleneck_encoder, bottleneck_decoder]
            ]
        ):
            self.input_dim = input_dim

            (
                encoder,
                decoder,
                bottleneck_encoder,
                bottleneck_decoder,
            ) = cnn_autoencoder_auto(
                input_dim=input_dim,
                latent_dim=latent_dim,
                output_dim=output_dim,
                activation=activation,
                kernel_size=kernel_size,
                channels=channels,
                case=case,
                shallow=shallow,
            )

        self.encoder = self.to_wrap(entity=encoder, device=self.device)
        self.bottleneck_encoder = self.to_wrap(
            entity=bottleneck_encoder, device=self.device
        )
        self.bottleneck_decoder = self.to_wrap(
            entity=bottleneck_decoder, device=self.device
        )
        self.decoder = self.to_wrap(entity=decoder, device=self.device)

        self.add_module("encoder", self.encoder)
        self.add_module("bottleneck_encoder", self.bottleneck_encoder)
        self.add_module("bottleneck_decoder", self.bottleneck_decoder)
        self.add_module("decoder", self.decoder)

        self.weights += self.encoder.weights
        self.weights += self.bottleneck_encoder.weights
        self.weights += self.bottleneck_decoder.weights
        self.weights += self.decoder.weights

        self.last_encoder_channels = None
        self.before_flatten_dimension = None

        self.encoder_activation = self._get_operation(operation=encoder_activation)

        self.shapes_dict = dict()

    def summary(
        self,
        input_data: Union[np.ndarray, torch.Tensor] = None,
        input_shape: list = None,
        verbose: bool = True,
    ) -> torch.Tensor:
        """Prints the summary of the network architecture.

        Args:
            input_data (Union[np.ndarray, torch.Tensor], optional): The input dataset. (Default value = None)
            input_shape (list, optional): The shape of the input data. (Default value = None)
            verbose (bool, optional):  (Default value = True)

        Returns:
            torch.Tensor: The dataset projected over the latent space.

        """

        if verbose == True:
            if self.input_dim != None:
                input_shape = self.input_dim
            else:
                pass

            self.encoder.summary(
                input_data=input_data, input_shape=input_shape, device=self.device
            )

            if isinstance(input_data, np.ndarray):
                btnk_input = self.encoder.forward(input_data=input_data)
            else:
                assert (
                    input_shape
                ), "It is necessary to have input_shape when input_data is None."
                input_shape = self.encoder.input_size
                input_shape[0] = 1

                input_data = self.to_wrap(
                    entity=torch.ones(input_shape), device=self.device
                )

                btnk_input = self.encoder.forward(input_data=input_data)

            before_flatten_dimension = tuple(btnk_input.shape[1:])
            btnk_input = btnk_input.reshape((-1, np.prod(btnk_input.shape[1:])))

            latent = self.bottleneck_encoder.forward(input_data=btnk_input)

            self.bottleneck_encoder.summary()
            self.bottleneck_decoder.summary()

            bottleneck_output = self.encoder_activation(
                self.bottleneck_decoder.forward(input_data=latent)
            )

            bottleneck_output = bottleneck_output.reshape(
                (-1, *before_flatten_dimension)
            )

            self.decoder.summary(input_data=bottleneck_output, device=self.device)

            # Saving the content of the subnetworks to the overall architecture dictionary
            self.shapes_dict.update({"encoder": self.encoder.shapes_dict})
            self.shapes_dict.update(
                {"bottleneck_encoder": self.bottleneck_encoder.shapes_dict}
            )
            self.shapes_dict.update(
                {"bottleneck_decoder": self.bottleneck_decoder.shapes_dict}
            )
            self.shapes_dict.update({"decoder": self.decoder.shapes_dict})

        else:
            print(self)

    @as_tensor
    def projection(self, input_data: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
        """Project input dataset into the latent space.

        Args:
            input_data (Union[np.ndarray, torch.Tensor]): The dataset to be projected.

        Returns:
            torch.Tensor: The dataset projected over the latent space.

        """

        btnk_input = self.encoder.forward(input_data=input_data)

        self.before_flatten_dimension = tuple(btnk_input.shape[1:])

        btnk_input = btnk_input.reshape((-1, np.prod(self.before_flatten_dimension)))

        latent = self.bottleneck_encoder.forward(input_data=btnk_input)

        return latent

    @as_tensor
    def reconstruction(
        self, input_data: Union[torch.Tensor, np.ndarray]
    ) -> torch.Tensor:
        """Reconstruct the latent dataset to the original one.

        Args:
            input_data (Union[torch.Tensor, np.ndarray]): The dataset to be reconstructed.

        Returns:
            torch.Tensor: The reconstructed dataset.

        """

        bottleneck_output = self.encoder_activation(
            self.bottleneck_decoder.forward(input_data=input_data)
        )

        bottleneck_output = bottleneck_output.reshape(
            (-1,) + self.before_flatten_dimension
        )

        reconstructed = self.decoder.forward(input_data=bottleneck_output)

        return reconstructed

    def forward(self, input_data: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
        """Execute the complete projection/reconstruction pipeline.

        Args:
            input_data (Union[np.ndarray, torch.Tensor]): The input dataset.

        Returns:
            torch.Tensor: The reconstructed dataset.

        """

        latent = self.projection(input_data=input_data)
        reconstructed = self.reconstruction(input_data=latent)

        return reconstructed

    def eval(self, input_data: Union[np.ndarray, torch.Tensor] = None) -> np.ndarray:
        """Evaluate the autoencoder on the given dataset.

        Args:
            input_data (Union[np.ndarray, torch.Tensor], optional): The dataset to be evaluated, by default None.

        Returns:
            np.ndarray: The dataset projected over the latent space.

        """

        if isinstance(input_data, np.ndarray):
            input_data = torch.from_numpy(input_data.astype(ARRAY_DTYPE))

        input_data = input_data.to(self.device)

        return super().eval(input_data=input_data)

    def project(self, input_data: Union[np.ndarray, torch.Tensor] = None) -> np.ndarray:
        """Project the input dataset into the latent space.

        Args:
            input_data (Union[np.ndarray, torch.Tensor], optional): The dataset to be projected, by default None.

        Returns:
            np.ndarray: The dataset projected over the latent space.

        """

        projected_data = self.projection(input_data=input_data)

        return projected_data.cpu().detach().numpy()

    def reconstruct(
        self, input_data: Union[np.ndarray, torch.Tensor] = None
    ) -> np.ndarray:
        """Reconstructs the latent dataset to the original one.

        Args:
            input_data (Union[np.ndarray, torch.Tensor], optional): The dataset to be reconstructed. If not provided, uses the original input data, by default None.

        Returns:
            np.ndarray: The reconstructed dataset.

        """
        reconstructed_data = self.reconstruction(input_data=input_data)
        return reconstructed_data.cpu().detach().numpy()

class AutoencoderKoopman(NetworkTemplate):
    r"""This is an implementation of a Koopman autoencoder as a Reduced Order Model.

    A Koopman autoencoder architecture consists of five stages:

    - The convolutional encoder [Optional]
    - Fully-connected encoder
    - Koopman operator
    - Fully connected decoder
    - The convolutional decoder [Optional]

    Graphical scheme

                                                (Koopman OPERATOR)
                                                         ^
                                                  |      |      |
                                                  |  |   |   |  |
           Z -> [Conv] -> [Conv] -> ... [Conv] -> |  | | - | |  | -> [Conv.T] -> [Conv.T] -> ... [Conv.T] -> Z_til
                                                  |  |       |  |
                                                  |             |

                                ENCODER          DENSE BOTTLENECK        DECODER

    """

    def __init__(
        self,
        encoder: Union[ConvolutionalNetwork, DenseNetwork] = None,
        bottleneck_encoder: Optional[Union[Linear, DenseNetwork]] = None,
        bottleneck_decoder: Optional[Union[Linear, DenseNetwork]] = None,
        decoder: Union[ConvolutionalNetwork, DenseNetwork] = None,
        input_dim: Optional[Tuple[int, ...]] = None,
        output_dim: Optional[Tuple[int, ...]] = None,
        latent_dim: Optional[int] = None,
        activation: Optional[Union[list, str]] = None,
        channels: Optional[int] = None,
        case: Optional[str] = None,
        architecture: Optional[str] = None,
        shallow: Optional[bool] = False,
        use_batch_norm: Optional[bool] = False,
        encoder_activation: str = "relu",
        devices: Union[str, list] = "cpu",
        name: str = None,
    ) -> None:
        """Constructs a new instance of the Autoencoder

        Args:
            encoder (Union[ConvolutionalNetwork, DenseNetwork], optional): The encoder network. Defaults to None.
            bottleneck_encoder (Optional[Union[Linear, DenseNetwork]], optional): The bottleneck encoder network. Defaults to None.
            bottleneck_decoder (Optional[Union[Linear, DenseNetwork]], optional): The bottleneck decoder network. Defaults to None.
            decoder (Union[ConvolutionalNetwork, DenseNetwork], optional): The decoder network. Defaults to None.
            input_dim (Optional[Tuple[int, ...]], optional): The input dimensions. Used for automatic network generation. Defaults to None.
            output_dim (Optional[Tuple[int, ...]], optional): The output dimensions. Used for automatic network generation. Defaults to None.
            latent_dim (Optional[int], optional): The latent dimensions. Used for automatic network generation. Defaults to None.
            activation (Optional[Union[list, str]], optional): The activation functions for each layer. Used for automatic network generation. Defaults to None.
            channels (Optional[int], optional): The number of channels. Used for automatic network generation. Defaults to None.
            case (Optional[str], optional): The type of problem. Used for automatic network generation. Defaults to None.
            architecture (Optional[str], optional): The network architecture. Used for automatic network generation. Defaults to None.
            shallow (Optional[bool], optional): Whether to use shallow or deep network. Used for automatic network generation. Defaults to False.
            use_batch_norm (Optional[bool], optional):  (Default value = False)
            encoder_activation (str, optional): The activation function for the encoder. Defaults to "relu".
            devices (Union[str, list], optional): The devices to use. Defaults to "cpu".
            name (str, optional): The name of the autoencoder. Defaults to None.

        """
        super(AutoencoderKoopman, self).__init__(name=name)

        self.weights = list()

        # Determining the kind of device to be used for allocating the
        # subnetworks
        self.device = self._set_device(devices=devices)

        self.input_dim = None

        # If not network is provided, the automatic generation
        # pipeline is activated.
        if all(
            [
                isn == None
                for isn in [encoder, decoder, bottleneck_encoder, bottleneck_decoder]
            ]
        ):
            self.input_dim = input_dim

            encoder, decoder, bottleneck_encoder, bottleneck_decoder = autoencoder_auto(
                input_dim=input_dim,
                latent_dim=latent_dim,
                output_dim=output_dim,
                activation=activation,
                channels=channels,
                architecture=architecture,
                case=case,
                shallow=shallow,
                use_batch_norm=use_batch_norm,
            )

        self.encoder = encoder.to(self.device)
        self.decoder = decoder.to(self.device)

        self.add_module("encoder", self.encoder)
        self.add_module("decoder", self.decoder)

        self.weights += self.encoder.weights
        self.weights += self.decoder.weights

        # These subnetworks are optional
        if bottleneck_encoder is not None and bottleneck_decoder is not None:
            self.bottleneck_encoder = self.to_wrap(
                entity=bottleneck_encoder, device=self.device
            )
            self.bottleneck_decoder = self.to_wrap(
                entity=bottleneck_decoder, device=self.device
            )

            self.add_module("bottleneck_encoder", self.bottleneck_encoder)
            self.add_module("bottleneck_decoder", self.bottleneck_decoder)

            self.weights += self.bottleneck_encoder.weights
            self.weights += self.bottleneck_decoder.weights

        # These subnetworks are optional
        if bottleneck_encoder is not None and bottleneck_decoder is not None:
            self.bottleneck_encoder = self.to_wrap(
                entity=bottleneck_encoder, device=self.device
            )
            self.bottleneck_decoder = self.to_wrap(
                entity=bottleneck_decoder, device=self.device
            )

            self.add_module("bottleneck_encoder", self.bottleneck_encoder)
            self.add_module("bottleneck_decoder", self.bottleneck_decoder)

            self.weights += self.bottleneck_encoder.weights
            self.weights += self.bottleneck_decoder.weights

        if bottleneck_encoder is not None and bottleneck_decoder is not None:
            self.projection = self._projection_with_bottleneck
            self.reconstruction = self._reconstruction_with_bottleneck
        else:
            self.projection = self._projection
            self.reconstruction = self._reconstruction

        self.last_encoder_channels = None
        self.before_flatten_dimension = None

        self.latent_dimension = None

        if bottleneck_encoder is not None:
            self.latent_dimension = bottleneck_encoder.output_size
        else:
            self.latent_dimension = self.encoder.output_size

        self.K_op = self.to_wrap(
            entity=torch.nn.Linear(
                self.latent_dimension, self.latent_dimension, bias=False
            ).weight,
            device=self.device,
        )

        self.encoder_activation = self._get_operation(operation=encoder_activation)

        self.shapes_dict = dict()

    def summary(
        self,
        input_data: Union[np.ndarray, torch.Tensor] = None,
        input_shape: list = None,
        verbose: bool = True,
    ) -> torch.Tensor:
        if verbose == True:
            if self.input_dim != None:
                input_shape = list(self.input_dim)
            else:
                pass

            self.encoder.summary(
                input_data=input_data, input_shape=input_shape, device=self.device
            )

            self.before_flatten_dimension = tuple(self.encoder.output_size[1:])

            if isinstance(input_data, np.ndarray):
                btnk_input = self.encoder.forward(input_data=input_data)
            else:
                assert (
                    input_shape
                ), "It is necessary to have input_shape when input_data is None."
                input_shape = self.encoder.input_size
                input_shape[0] = 1

                input_data = self.to_wrap(
                    entity=torch.ones(input_shape), device=self.device
                )

                btnk_input = self.encoder.forward(input_data=input_data)

            before_flatten_dimension = tuple(btnk_input.shape[1:])
            btnk_input = btnk_input.reshape((-1, np.prod(btnk_input.shape[1:])))

            latent = self.bottleneck_encoder.forward(input_data=btnk_input)

            self.bottleneck_encoder.summary()

            print(f"The Koopman Operator has shape: {self.K_op.shape} ")

            self.bottleneck_decoder.summary()

            bottleneck_output = self.encoder_activation(
                self.bottleneck_decoder.forward(input_data=latent)
            )

            bottleneck_output = bottleneck_output.reshape(
                (-1, *before_flatten_dimension)
            )

            self.decoder.summary(input_data=bottleneck_output, device=self.device)

            # Saving the content of the subnetworks to the overall architecture dictionary
            self.shapes_dict.update({"encoder": self.encoder.shapes_dict})
            self.shapes_dict.update(
                {"bottleneck_encoder": self.bottleneck_encoder.shapes_dict}
            )
            self.shapes_dict.update(
                {"bottleneck_decoder": self.bottleneck_decoder.shapes_dict}
            )
            self.shapes_dict.update({"decoder": self.decoder.shapes_dict})

        else:
            print(self)

    @as_tensor
    def _projection_with_bottleneck(
        self, input_data: Union[np.ndarray, torch.Tensor] = None
    ) -> torch.Tensor:
        """Computes the projection of the input data onto the bottleneck encoder.

        Args:
            input_data (Union[np.ndarray, torch.Tensor], optional): The input data. Defaults to None.

        Returns:
            torch.Tensor: The projected latent representation.

        """
        btnk_input = self.encoder.forward(input_data=input_data)

        self.before_flatten_dimension = tuple(btnk_input.shape[1:])

        btnk_input = btnk_input.reshape((-1, np.prod(self.before_flatten_dimension)))

        latent = self.bottleneck_encoder.forward(input_data=btnk_input)

        return latent

    @as_tensor
    def _projection(
        self, input_data: Union[np.ndarray, torch.Tensor] = None
    ) -> torch.Tensor:
        """Computes the projection of the input data onto the encoder.

        Args:
            input_data (Union[np.ndarray, torch.Tensor], optional): The input data. Defaults to None.

        Returns:
            torch.Tensor: The projected latent representation.

        """
        latent = self.encoder.forward(input_data=input_data)

        return latent

    @as_tensor
    def _reconstruction_with_bottleneck(
        self, input_data: Union[torch.Tensor, np.ndarray] = None
    ) -> torch.Tensor:
        """Reconstructs the input data using the bottleneck decoder.

        Args:
            input_data (Union[torch.Tensor, np.ndarray], optional): The input data. Defaults to None.

        Returns:
            torch.Tensor: The reconstructed data.

        """
        bottleneck_output = self.encoder_activation(
            self.bottleneck_decoder.forward(input_data=input_data)
        )

        bottleneck_output = bottleneck_output.reshape(
            (-1,) + self.before_flatten_dimension
        )

        reconstructed = self.decoder.forward(input_data=bottleneck_output)

        return reconstructed

    @as_tensor
    def _reconstruction(
        self, input_data: Union[torch.Tensor, np.ndarray] = None
    ) -> torch.Tensor:
        """Reconstructs the input data using the decoder.

        Args:
            input_data (Union[torch.Tensor, np.ndarray], optional): The input data. Defaults to None.

        Returns:
            torch.Tensor: The reconstructed data.

        """
        reconstructed = self.decoder.forward(input_data=input_data)

        return reconstructed

    def latent_forward_m(
        self, input_data: Union[np.ndarray, torch.Tensor] = None, m: int = 1
    ) -> torch.Tensor:
        """Evaluates the operation $u^{u+m} = K^m u^{i}$

        Args:
            input_data (Union[np.ndarray, torch.Tensor], optional): The input data. Defaults to None.
            m (int, optional): The number of Koopman iterations. Defaults to 1.

        Returns:
            torch.Tensor: The computed latent representation.

        """
        return torch.matmul(input_data, torch.pow(self.K_op.T, m))

    def latent_forward(
        self, input_data: Union[np.ndarray, torch.Tensor] = None
    ) -> torch.Tensor:
        """Evaluates the operation u^{u+1} = K u^{i}

        Args:
            input_data (Union[np.ndarray, torch.Tensor], optional): The input data. Defaults to None.

        Returns:
            torch.Tensor: The computed latent representation.

        """
        return torch.matmul(input_data, self.K_op.T)

    def reconstruction_forward(
        self, input_data: Union[np.ndarray, torch.Tensor] = None
    ) -> torch.Tensor:
        """Evaluates the operation Ũ = D(E(U))

        Args:
            input_data (Union[np.ndarray, torch.Tensor], optional): The input data. Defaults to None.

        Returns:
            torch.Tensor: The reconstructed data.

        """
        latent = self.projection(input_data=input_data)
        reconstructed = self.reconstruction(input_data=latent)

        return reconstructed

    def reconstruction_forward_m(
        self, input_data: Union[np.ndarray, torch.Tensor] = None, m: int = 1
    ) -> torch.Tensor:
        """Evaluates the operation Ũ_m = D(K^m E(U))

        Args:
            input_data (Union[np.ndarray, torch.Tensor], optional): The input data. Defaults to None.
            m (int, optional): The number of Koopman iterations. Defaults to 1.

        Returns:
            torch.Tensor: The reconstructed data.

        """
        latent = self.projection(input_data=input_data)
        latent_m = self.latent_forward_m(input_data=latent, m=m)
        reconstructed_m = self.reconstruction(input_data=latent_m)

        return reconstructed_m

    def predict(
        self, input_data: Union[np.ndarray, torch.Tensor] = None, n_steps: int = 1
    ) -> np.ndarray:
        """Predicts the reconstructed data for the input data after n_steps extrapolation in the latent space.

        Args:
            input_data (Union[np.ndarray, torch.Tensor], optional): The input data. Defaults to None.
            n_steps (int, optional): The number of extrapolations to perform. Defaults to 1.

        Returns:
            np.ndarray: The predicted reconstructed data.

        """
        if isinstance(input_data, np.ndarray):
            input_data = torch.from_numpy(input_data.astype(ARRAY_DTYPE))

        predictions = list()
        latent = self.projection(input_data=input_data)
        init_latent = latent

        # Extrapolating in the latent space over n_steps steps
        for s in range(n_steps):
            latent_s = self.latent_forward(input_data=init_latent)
            init_latent = latent_s
            predictions.append(latent_s)

        predictions = torch.vstack(predictions)

        reconstructed_predictions = self.reconstruction(input_data=predictions)

        return reconstructed_predictions.detach().numpy()

    def project(self, input_data: Union[np.ndarray, torch.Tensor] = None) -> np.ndarray:
        """Projects the input data into the latent space.

        Args:
            input_data (Union[np.ndarray, torch.Tensor], optional): The input data. Defaults to None.

        Returns:
            np.ndarray: The projected data.

        """
        projected_data = self.projection(input_data=input_data)

        return projected_data.cpu().detach().numpy()

    def reconstruct(
        self, input_data: Union[np.ndarray, torch.Tensor] = None
    ) -> np.ndarray:
        """Reconstructs the input data.

        Args:
            input_data (Union[np.ndarray, torch.Tensor], optional): The input data. Defaults to None.

        Returns:
            np.ndarray: The reconstructed data.

        """
        reconstructed_data = self.reconstruction(input_data=input_data)

        return reconstructed_data.cpu().detach().numpy()