class RMSELoss(LossBasics):
    def __init__(self, operator: torch.nn.Module = None) -> None:
        """Vanilla mean-squared error loss function

        Args:
            operator (torch.nn.Module): the operator used for evaluating
                the loss function (usually a neural network)
        """
        super().__init__()

        self.operator = operator
        self.loss_states = {"loss": list()}

    def _data_loss(
        self,
        output_tilde: torch.Tensor = None,
        norm_value: torch.Tensor = None,
        target_data_tensor: torch.Tensor = None,
    ) -> torch.Tensor:
        """It executes the evaluation of the data-driven mean-squared error

        Args:
            output_tilde (torch.Tensor): the output generated by
                self.operator
            norm_value (torch.Tensor): the value used for normalizing
                the loss evaluation
            target_data_tensor (torch.Tensor): the target tensor to be
                compared with output_tilde

        Returns:
            torch.Tensor: the loss function value for a given state
        """

        if norm_value is not None:
            data_loss = torch.mean(
                torch.square((output_tilde - target_data_tensor) / norm_value)
            )
        else:
            data_loss = torch.mean(torch.square((output_tilde - target_data_tensor)))

        return data_loss

    def __call__(
        self,
        input_data: Union[dict, torch.Tensor] = None,
        target_data: torch.Tensor = None,
        call_back: str = "",
        norm_value: list = None,
        lambda_1: float = 0.0,
        device: str = "cpu",
        lambda_2: float = 0.0,
    ) -> Callable:
        """Main function for generating complete loss function workflow

        Args:
            input_data (Union[dict, torch.Tensor]): the data used as
                input for self.operator
            target_data (torch.Tensor): the target data used for
                training self.oeprator
            call_back (str): a string used for composing the logging of
                the optimization process
            norm_value (list): a list of values used for normalizing the
                loss temrms
            lambda_1 (float): the penalty for the L^1  regularization
                term
            lambda_2 (float): the penalty for the L^2  regularization
                term
            device (str): the device in which the loss evaluation will
                be executed, 'cpu' or 'gpu'

        Returns:
            Callable: the closure function used for evaluating the loss
            value
        """

        l1_reg_multiplication = self._exec_multiplication_in_regularization(
            lambda_type=type(lambda_1), term_type=type(self.operator.weights_l1)
        )

        l2_reg_multiplication = self._exec_multiplication_in_regularization(
            lambda_type=type(lambda_2), term_type=type(self.operator.weights_l2)
        )

        def closure():
            output_tilde = self.operator.forward(**input_data)

            data_loss = self._data_loss(
                output_tilde=output_tilde,
                norm_value=norm_value,
                target_data_tensor=target_data,
            )

            # L² and L¹ regularization term
            weights_l2 = self.operator.weights_l2
            weights_l1 = self.operator.weights_l1

            # beta *||W||_2 + alpha * ||W||_1
            l2_reg = l2_reg_multiplication(lambda_2, weights_l2)
            l1_reg = l1_reg_multiplication(lambda_1, weights_l1)

            # Loss = ||Ũ_t - U_t||_2  +
            #         lambda_1 *||W||_2 + lambda_2 * ||W||_1

            loss = data_loss + l2_reg + l1_reg

            # Back-propagation
            loss.backward()

            self.loss_states["loss"].append(float(loss.detach().data))

            sys.stdout.write(("\rloss: {} {}").format(loss, call_back))
            sys.stdout.flush()

            return loss

        return closure

class WRMSELoss(LossBasics):
    def __init__(self, operator=None):
        """Weighted mean-squared error loss function

        Args:
            operator (torch.nn.Module): the operator used for evaluating
                the loss function (usually a neural network)
        """

        super().__init__()

        self.operator = operator
        self.split_dim = 1
        self.tol = 1e-25

        self.loss_evaluator = None
        self.norm_evaluator = None

        self.axis_loss_evaluator = lambda res: torch.mean(torch.square((res)), dim=1)

        self.loss_states = {"loss": list()}

    def _data_loss(
        self,
        output_tilde: torch.Tensor = None,
        weights: list = None,
        target_data_tensor: torch.Tensor = None,
        axis: int = -1,
    ) -> List:
        """It executes the evaluation of the data-driven mean-squared error

        Args:
            output_tilde (torch.Tensor): the output generated by
                self.operator
            norm_value (torch.Tensor): the value used for normalizing
                the loss evaluation
            target_data_tensor (torch.Tensor): the target tensor to be
                compared with output_tilde

        Returns:
            torch.Tensor: the loss function value for a given state
        """

        target_split = torch.split(target_data_tensor, self.split_dim, dim=axis)
        output_split = torch.split(output_tilde, self.split_dim, dim=axis)[:len(target_split)]

        data_losses = [
            weights[i]
            * self.loss_evaluator(out_split - tgt_split)
            / self.norm_evaluator(tgt_split)
            for i, (out_split, tgt_split) in enumerate(zip(output_split, target_split))
        ]

        return data_losses

    def _no_data_loss_wrapper(
        self,
        output_tilde: torch.Tensor = None,
        weights: list = None,
        target_data_tensor: torch.Tensor = None,
        axis: int = -1,
    ) -> torch.Tensor:
        """It executes the evaluation of the data-driven mean-squared error without considering causality preserving

        Args:
            output_tilde (torch.Tensor): the output generated by
                self.operator
            weights (list): weights for rescaling each variable
                outputted by self.operator
            target_data_tensor (torch.Tensor): the target tensor to be
                compared with output_tilde
            axis (int): the axis in which the variables are split

        Returns:
            torch.Tensor: the loss function value for a given state
        """

        return self.data_loss(
            output_tilde=output_tilde,
            weights=weights,
            target_data_tensor=target_data_tensor,
            axis=axis,
        )

    def __call__(
        self,
        input_data: Union[dict, torch.Tensor] = None,
        target_data: torch.Tensor = None,
        call_back: str = "",
        lambda_1: float = 0.0,
        lambda_2: float = 0.0,
        axis: int = -1,
        relative: bool = False,
        device: str = "cpu",
        weights: list = None,
        use_mean: bool = True,
    ) -> Callable:
        """Main function for generating complete loss function workflow

        Args:
            input_data (Union[dict, torch.Tensor]): the data used as
                input for self.operator
            target_data (torch.Tensor): the target data used for
                training self.oeprator
            call_back (str): a string used for composing the logging of
                the optimization process
            norm_value (list): a list of values used for normalizing the
                loss temrms
            lambda_1 (float): the penalty for the L^1  regularization
                term
            lambda_2 (float): the penalty for the L^2  regularization
                term
            device (str): the device in which the loss evaluation will
                be executed, 'cpu' or 'gpu'
            weights (list): a list of weights for rescaling each
                variable outputted by self.operator
            use_mean (bool): use mean for evaluating the losses or not
                (the alternative is sum)

        Returns:
            Callable: the closure function used for evaluating the loss
            value
        """

        self.data_loss = self._data_loss

        l1_reg_multiplication = self._exec_multiplication_in_regularization(
            lambda_type=type(lambda_1), term_type=type(self.operator.weights_l1)
        )

        l2_reg_multiplication = self._exec_multiplication_in_regularization(
            lambda_type=type(lambda_2), term_type=type(self.operator.weights_l2)
        )

        # Using mean evaluation or not
        if use_mean == True:
            self.loss_evaluator = lambda res: torch.mean(torch.square((res)))
        else:
            self.loss_evaluator = lambda res: torch.sum(torch.square((res)))

        # Relative norm or not
        if relative == True:
            if use_mean == True:
                self.norm_evaluator = lambda ref: torch.mean(torch.square((ref)))
            else:
                self.norm_evaluator = lambda ref: torch.sum(torch.square((ref)))
        else:
            self.norm_evaluator = lambda ref: 1

        self.data_loss_wrapper = self._no_data_loss_wrapper

        def closure():
            output_tilde = self.operator.forward(**input_data)

            data_losses = self.data_loss_wrapper(
                output_tilde=output_tilde,
                weights=weights,
                target_data_tensor=target_data,
                axis=axis,
            )

            # L² and L¹ regularization term
            weights_l2 = self.operator.weights_l2
            weights_l1 = self.operator.weights_l1

            # beta *||W||_2 + alpha * ||W||_1
            l2_reg = l2_reg_multiplication(lambda_2, weights_l2)
            l1_reg = l1_reg_multiplication(lambda_1, weights_l1)

            # Loss = ||Ũ_t - U_t||_2  +
            #         lambda_1 *||W||_2 + lambda_2 * ||W||_1
            loss = sum(data_losses) + l2_reg + l1_reg

            # Back-propagation
            loss.backward()

            self.loss_states["loss"].append(float(loss.detach().data))

            sys.stdout.write(("\rloss: {} {}").format(loss, call_back))
            sys.stdout.flush()

        return closure

class PIRMSELoss(LossBasics):
    def __init__(self, operator: torch.nn.Module = None) -> None:
        """Physics-Informed mean-squared error loss function

        Args:
            operator (torch.nn.Module): the operator used for evaluating
                the loss function (usually a neural network)
        """

        super().__init__()

        self.split_dim = 1
        self.operator = operator
        self.loss_evaluator = None
        self.residual = None
        self.tol = 1e-15
        self.device = None

        self.axis_loss_evaluator = lambda res: torch.mean(torch.square((res)), dim=1)

        self.loss_states = {
            "pde": list(),
            "init": list(),
            "bound": list(),
            "extra_data": list(),
        }
        self.loss_tags = list(self.loss_states.keys())
        self.hybrid_data_pinn = False

        self.losses_terms_indices = {
            "pde": 0,
            "init": 1,
            "bound": 2,
            "extra_data": 3,
            "causality_weights": 4,
        }

    def _convert(
        self, input_data: Union[dict, np.ndarray] = None, device: str = None
    ) -> Union[dict, torch.Tensor]:
        """It converts a dataset to the proper format (torch.Tensor) and send it to
        the chosen execution device ('gpu' or 'cpu')

        Args:
            input_data (Union[dict, np.ndarray]): the data structure to
                be converted
            device: the device in which the converted dataset must be
                placed

        Returns:
            Union[dict, torch.Tensor]: the converted data structure
        """

        if type(input_data) == dict:
            return {
                key: torch.from_numpy(item.astype(ARRAY_DTYPE)).to(device)
                for key, item in input_data.items()
            }

        else:
            return torch.from_numpy(input_data.astype(ARRAY_DTYPE)).to(device)

    def _to_tensor(self, *args, device: str = "cpu") -> List[torch.Tensor]:
        """It converted a size indefined list of arrays to tensors

        Args:
            *args: list of arrays to be converted
            input_data (Union[dict, np.ndarray])
            device: the device in which the converted dataset must be
                placed
        :type np.array, np.array, ..., np.array

        Returns:
            List[torch.Tensor]: a list of tensors
        """
        return [self._convert(input_data=arg, device=device) for arg in args]

    def _data_loss(
        self,
        output_tilde: torch.Tensor = None,
        target_data_tensor: torch.Tensor = None,
        weights: List[float] = None,
    ) -> torch.Tensor:
        """It executes the evaluation of the data-driven mean-squared error

        Args:
            output_tilde (torch.Tensor): the output generated by
                self.operator
            target_data_tensor (torch.Tensor): the target tensor to be
                compared with output_tilde

        Returns:
            torch.Tensor: the loss function value
        """

        target_split = torch.split(target_data_tensor, self.split_dim, dim=-1)
        output_split = torch.split(output_tilde, self.split_dim, dim=-1)[:len(target_split)]

        data_losses = [
            self.loss_evaluator_data((out_split, tgt_split))
            / (self.norm_evaluator(tgt_split) or torch.tensor(1.0).to(self.device))
            for i, (out_split, tgt_split) in enumerate(zip(output_split, target_split))
        ]

        return self.weighted_loss_evaluator(data_losses, weights)

    def _data_loss_adaptive(
        self,
        output_tilde: torch.Tensor = None,
        target_data_tensor: torch.Tensor = None,
        **kwargs,
    ) -> torch.Tensor:
        """It executes the evaluation of the data-driven mean-squared error

        Args:
            output_tilde (torch.Tensor): the output generated by
                self.operator
            target_data_tensor (torch.Tensor): the target tensor to be
                compared with output_tilde

        Returns:
            torch.Tensor: the loss function value
        """

        output_split = torch.split(output_tilde, self.split_dim, dim=-1)
        target_split = torch.split(target_data_tensor, self.split_dim, dim=-1)

        data_discrepancy = [
            out_split - tgt_split
            for i, (out_split, tgt_split) in enumerate(zip(output_split, target_split))
        ]

        weights = self.data_weights_estimator(
            residual=data_discrepancy,
            loss_evaluator=self.loss_evaluator,
            loss_history=self.loss_states,
            operator=self.operator,
        )

        data_losses = [
            weights[i] * self.loss_evaluator_data((out_split, tgt_split))
            for i, (out_split, tgt_split) in enumerate(zip(output_split, target_split))
        ]

        return [sum(data_losses)]

    def _global_weights_bypass(
        self, initial_penalty: float = None, **kwargs
    ) -> List[float]:
        return [1.0, initial_penalty, 1.0, 1.0]

    def _global_weights_estimator(self, **kwargs) -> List[float]:
        weights = self.global_weights_estimator(**kwargs)

        return weights

    def _residual_loss(
        self, residual_approximation: List[torch.Tensor] = None, weights: list = None
    ) -> List[torch.Tensor]:
        """It evaluates the physics-driven residual loss

        Args:
            residual_approximation (List[torch.Tensor]): a list of
                tensors containing the evaluation for the physical
                residual for each sample in the dataset
            weights (list): a list of weights used for rescaling the
                residuals of each variable

        Returns:
            torch.Tensor: the list of residual losses
        """
        residual_losses = [self.loss_evaluator(res) for res in residual_approximation]

        return self.weighted_loss_evaluator(residual_losses, weights)

    def _residual_loss_adaptive(
        self, residual_approximation: List[torch.Tensor] = None, weights: list = None
    ) -> List[torch.Tensor]:
        """It evaluates the physics-driven residual loss

        Args:
            residual_approximation (List[torch.Tensor]): a list of
                tensors containing the evaluation for the physical
                residual for each sample in the dataset
            weights (list): a list of weights used for rescaling the
                residuals of each variable

        Returns:
            torch.Tensor: the list of residual losses
        """

        weights = self.residual_weights_estimator(
            residual=residual_approximation,
            loss_evaluator=self.loss_evaluator,
            loss_history=self.loss_states,
            operator=self.operator,
        )

        residual_loss = [
            weight * self.loss_evaluator(res)
            for weight, res in zip(weights, residual_approximation)
        ]

        return [sum(residual_loss)]

    def _extra_data(
        self, input_data: torch.Tensor = None, target_data: torch.Tensor = None, weights :list = None, 
    ) -> torch.Tensor:
        # Evaluating data for the initial condition
        output_tilde = self.operator(input_data=input_data)

        # Evaluating loss approximation for extra data
        data_loss = self._data_loss(
            output_tilde=output_tilde,
            target_data_tensor=target_data, 
            weights=weights,
        )

        return data_loss

    def _boundary_penalisation(
        self, boundary_input: dict = None, residual: SymbolicOperator = None
    ) -> List[torch.Tensor]:
        """It applies the boundary conditions

        Args:
            boundary_input (dict): a dictionary containing the
                coordinates of the boundaries
            residual (SymbolicOperator): a symbolic expression for the
                boundary condition

        Returns:
            list: the evaluation of each boundary condition
        """
        return [
            residual.eval_expression(k, boundary_input[k])
            for k in boundary_input.keys()
        ]

    def _no_boundary_penalisation(
        self, boundary_input: dict = None, residual: object = None
    ) -> List[torch.Tensor]:
        """It is used for cases in which no boundary condition is applied

        """

        return [torch.Tensor([0.0]).to(self.device) for k in boundary_input.keys()]

    def _no_boundary(
        self, boundary_input: dict = None, residual: object = None
    ) -> List[torch.Tensor]:
        """It is used for cases where there are not boundaries

        """

        return torch.Tensor([0.0]).to(self.device)

    def _no_extra_data(
        self, input_data: torch.Tensor = None, target_data: torch.Tensor = None, weights: list=None,
    ) -> torch.Tensor:
        return torch.Tensor([0.0]).to(self.device)

    def _no_residual_wrapper(self, input_data: torch.Tensor = None) -> torch.Tensor:
        return self.residual(input_data)

    def _causality_preserving_residual_wrapper(
        self, input_data: torch.Tensor = None
    ) -> List:
        return self.causality_preserving(self.residual(input_data))

    def _filter_necessary_loss_terms(self, residual: SymbolicOperator = None):
        tags = ["pde", "init"]
        indices = [0, 1]

        if residual.g_expressions:
            tags.append("bound")
            indices.append(2)
        else:
            pass

        if self.hybrid_data_pinn:
            tags.append("extra_data")
            indices.append(3)
        else:
            pass

        return tags, indices

    def _losses_states_str(self, tags: List[str] = None):
        losses_str = "\r"
        for item in tags:
            losses_str += f"{item}:{{}} "

        return losses_str

    def __call__(
        self,
        input_data: Union[dict, torch.Tensor] = None,
        target_data: Union[dict, torch.Tensor] = None,
        verbose: bool = False,
        call_back: str = "",
        residual: Callable = None,
        initial_input: Union[dict, torch.Tensor] = None,
        initial_state: Union[dict, torch.Tensor] = None,
        boundary_input: dict = None,
        boundary_penalties: list = [1],
        extra_input_data: Union[dict, torch.Tensor] = None,
        extra_target_data: Union[dict, torch.Tensor] = None,
        initial_penalty: float = 1,
        axis: int = -1,
        relative: bool = False,
        lambda_1: float = 0.0,
        lambda_2: float = 0.0,
        weights=None,
        weights_residual=None,
        weights_extra_data=None,
        device: str = "cpu",
        split_losses: bool = False,
        causality_preserving: Callable = None,
        global_weights_estimator: Callable = None,
        residual_weights_estimator: Callable = None,
        data_weights_estimator: Callable = None,
        use_mean: bool = True,
        use_data_log: bool = False,
    ) -> Callable:
        self.residual = residual

        self.device = device

        self.causality_preserving = causality_preserving

        # Handling expection when AnnealingWeights and split_losses
        # are used together.
        if isinstance(global_weights_estimator, AnnealingWeights):
            if split_losses:
                raise RuntimeError(
                    "Global weights estimator, AnnealingWeights, is not"
                    + "compatible with split loss terms."
                )
            else:
                pass

        self.global_weights_estimator = global_weights_estimator

        self.residual_weights_estimator = residual_weights_estimator

        self.data_weights_estimator = data_weights_estimator

        if split_losses:
            self.weighted_loss_evaluator = self._bypass_weighted_loss
        else:
            self.weighted_loss_evaluator = self._eval_weighted_loss

        if (
            isinstance(extra_input_data, np.ndarray)
            == isinstance(extra_target_data, np.ndarray)
            == True
        ):
            self.hybrid_data_pinn = True
        else:
            pass

        # When no weight is provided, they are
        # set to the default choice
        if weights is None:
            weights = len(residual.output_names) * [1]

        if weights_residual is None:
            weights_residual = len(residual.output_names) * [1]

        loss_tags, loss_indices = self._filter_necessary_loss_terms(residual=residual)
        loss_str = self._losses_states_str(tags=loss_tags)

        # Boundary conditions are optional, since they are not
        # defined in some cases, as ODE, for example.
        if residual.g_expressions:
            boundary = self._boundary_penalisation
        else:
            if boundary_input == None:
                boundary = self._no_boundary
            else:
                boundary = self._no_boundary_penalisation

        if self.causality_preserving:
            call_back = f", causality_weights: {self.causality_preserving.call_back}"
            self.residual_wrapper = self._causality_preserving_residual_wrapper

        else:
            self.residual_wrapper = self._no_residual_wrapper

        l1_reg_multiplication = self._exec_multiplication_in_regularization(
            lambda_type=type(lambda_1), term_type=type(self.operator.weights_l1)
        )

        l2_reg_multiplication = self._exec_multiplication_in_regularization(
            lambda_type=type(lambda_2), term_type=type(self.operator.weights_l2)
        )
        if type(input_data) is dict:
            try:
                input_data = input_data["input_data"]
            except Exception:
                pass

        initial_input, initial_state = self._to_tensor(
            initial_input, initial_state, device=device
        )

        # Preparing extra data, when necessary
        if self.hybrid_data_pinn:
            extra_input_data, extra_target_data = self._to_tensor(
                extra_input_data, extra_target_data, device=device
            )
            self.extra_data = self._extra_data
        else:
            self.extra_data = self._no_extra_data

        if use_data_log == True:
            self.inner_square = self._two_term_log_loss
        else:
            self.inner_square = self._two_term_loss

        if use_mean == True:
            self.loss_evaluator = lambda res: torch.mean(self._single_term_loss(res))
        else:
            self.loss_evaluator = lambda res: torch.sum(self._single_term_loss(res))

        if use_mean == True:
            self.loss_evaluator_data = lambda res: torch.mean(self.inner_square(*res))
        else:
            self.loss_evaluator_data = lambda res: torch.sum(self.inner_square(*res))

        # Relative norm or not
        if relative == True:
            if use_mean == True:
                self.norm_evaluator = lambda ref: torch.mean(torch.square((ref)))
            else:
                self.norm_evaluator = lambda ref: torch.sum(torch.square((ref)))
        else:
            self.norm_evaluator = lambda ref: 1

        # Determing the usage of special residual loss weighting
        if residual_weights_estimator:
            self.residual_loss = self._residual_loss_adaptive
        else:
            self.residual_loss = self._residual_loss

        # Determing the usage of special data loss weighting
        if data_weights_estimator:
            self.data_loss = self._data_loss_adaptive
        else:
            self.data_loss = self._data_loss

        # Determining the usage of special global loss weighting
        if global_weights_estimator:
            self.global_weights = self._global_weights_estimator
        else:
            self.global_weights = self._global_weights_bypass

        if verbose:
            self.pprint = self._pprint_verbose
        else:
            self.pprint = self._pprint_simple

        def closure():
            # Executing the symbolic residual evaluation
            residual_approximation = self.residual_wrapper(input_data)

            # Boundary, if appliable
            boundary_approximation = boundary(
                boundary_input=boundary_input, residual=residual
            )

            # Evaluating data for the initial condition
            initial_output_tilde = self.operator(input_data=initial_input)

            # Evaluating loss function for residual
            residual_loss = self.residual_loss(
                residual_approximation=residual_approximation, weights=weights_residual
            )

            # Evaluating loss for the boundary approaximation, if appliable
            boundary_loss = self._residual_loss(
                residual_approximation=boundary_approximation,
                weights=boundary_penalties,
            )

            # Evaluating loss approximation for initial condition
            initial_data_loss = self.data_loss(
                output_tilde=initial_output_tilde,
                target_data_tensor=initial_state,
                weights=weights,
            )

            # Evaluating extra data loss, when appliable
            extra_data = self.extra_data(
                input_data=extra_input_data,
                target_data=extra_target_data,
                weights=weights_extra_data,
            )

            # L² and L¹ regularization term
            weights_l2 = self.operator.weights_l2
            weights_l1 = self.operator.weights_l1

            # beta *||W||_2 + alpha * ||W||_1
            l2_reg = l2_reg_multiplication(lambda_2, weights_l2)
            l1_reg = l1_reg_multiplication(lambda_1, weights_l1)

            # The complete loss function
            pde = residual_loss
            init = initial_data_loss
            bound = boundary_loss

            loss_terms = self._aggregate_terms(*pde, *init, *bound, *extra_data)

            # Updating the loss weights if necessary
            loss_weights = self.global_weights(
                initial_penalty=initial_penalty,
                operator=self.operator,
                loss_evaluator=self.loss_evaluator,
                residual=loss_terms,
            )
            # Overall loss function
            loss = (
                sum(self._eval_weighted_loss(loss_terms, loss_weights))
                + l2_reg
                + l1_reg
            )

            # Back-propagation
            loss.backward()

            pde_detach = float(sum(pde).detach().data)
            init_detach = float(sum(init).detach().data)
            bound_detach = float(sum(bound).detach().data)
            extra_data_detach = float(sum(extra_data).detach().data)

            self.loss_states["pde"].append(pde_detach)
            self.loss_states["init"].append(init_detach)
            self.loss_states["bound"].append(bound_detach)
            self.loss_states["extra_data"].append(extra_data_detach)

            losses_list = np.array(
                [pde_detach, init_detach, bound_detach, extra_data_detach]
            )

            self.pprint(
                loss_str=loss_str,
                losses_list=losses_list,
                call_back=call_back,
                loss_indices=loss_indices,
                loss_terms=loss_terms,
                loss_weights=loss_weights,
            )

            _current_loss = loss

            return _current_loss

        return closure