Source code for mtnlion.formula

"""
The `Formula` class is the basis for defining models. The formulas define the name, applicable domains, and arguments
required to be able to fully assemble the formula for FFL.
"""
import abc
import collections
from typing import Callable, Dict, Iterable, List, NamedTuple, Optional, Tuple, Type, Union

from ufl.core.expr import Expr  # type: ignore

from mtnlion.cell import MeasureMap
from mtnlion.domain import Domain
from mtnlion.variable import Variable


class Arguments:
    """
    This class is designed to serve as a data class to pass variables, parameters, formulas, lambda functions, and time
    discretization schemes to any formula. This class also handles the assignment of named tuples from generic tuples
    and performs basic error checking.
    """

    def __init__(
        self,
        variables: Iterable[Variable] = (),
        parameters: Iterable = (),
        lambdas: Iterable[Callable] = (),
        time_discretization: Iterable = (),  # TODO: Need a base class for Rothes, should be optional with a Variable
    ):
        """
        Create an argument object with various required data for operation.

        :param variables: `Variable`'s required for a formula to evaluate
        :param parameters: Parameters for the formula to evaluate
        :param lambdas: Lambda functions that can be used to alter the behavior of the formulas
        :param time_discretization: Time discretization scheme, should be callable
        """

        # pylint: disable=too-many-arguments

        self.variables = variables
        self.parameters = parameters
        self.time_discretization = time_discretization
        self.lambdas = lambdas
        self._classes: Dict[str, Type[NamedTuple]] = {}

    def assign_argument_classes(
        self,
        variables: Type[NamedTuple],
        parameters: Type[NamedTuple],
        lambdas: Type[NamedTuple],
        time_discretization: Type[NamedTuple],
    ) -> None:
        """
        Assign the classes that the Iterable's should be converted to.

        :param variables: Named tuple for variables
        :param parameters: Named tuple for parameters
        :param lambdas: Named tuple for lambdas
        :param time_discretization: Named tuple for time_discretizations
        :return: None
        """

        # pylint: disable=too-many-arguments

        self._classes = {
            "variables": variables,
            "parameters": parameters,
            "lambdas": lambdas,
            "time_discretization": time_discretization,
        }

    def convert(self) -> None:
        """
        Convert the internal properties to the classes defined in _classes.
        :return: None
        """
        for arg_name, arg_type in self._classes.items():
            setattr(self, arg_name, self._prepare(arg_type, getattr(self, arg_name)))

    @staticmethod
    def _prepare(class_type: Type[NamedTuple], data: Iterable) -> NamedTuple:
        try:
            return class_type(*data)
        except KeyError as excpt:
            raise ArgumentError(f"Cannot find {excpt} in {class_type.__qualname__}")
        except Exception as excpt:
            raise type(excpt)(f"Could not fit data into the given datatype... type: {class_type.__qualname__}")

    @staticmethod
    def _check_exists(attr, field_names):
        for field in field_names:
            if field not in attr._fields:
                raise KeyError(f"Field {field} does not exist in {attr}")

    def _popper(self, attr_name: str, typename: str, field_names):
        attr = getattr(self, attr_name)
        self._check_exists(attr, field_names)
        if isinstance(attr, self._classes[attr_name]):
            fields = attr._fields

            keep_indices = [index for index in range(len(attr)) if fields[index] not in field_names]
            pop_indices = list(set(range(len(attr))) - set(keep_indices))
            pop_indices.sort()

            pop_type = collections.namedtuple(  # type: ignore
                typename, ", ".join(fields[index] for index in pop_indices)
            )
            new_type = collections.namedtuple(  # type: ignore
                typename, ", ".join(fields[index] for index in keep_indices)
            )

            pop_data = pop_type(*(attr[index] for index in pop_indices))
            setattr(self, attr_name, new_type(*(attr[index] for index in keep_indices)))

            return pop_data

        raise NotImplementedError

    def pop_variables(self, variables: Iterable[str]) -> NamedTuple:
        """
        Return a tuple containing the values of the requested arguments. The requested arguments are then removed from
        the object.

        :param variables: List of arguments to pop
        :return: Tuple
        """
        return self._popper("variables", "Variable", variables)

    def pop_parameters(self, parameters: Iterable[str]) -> NamedTuple:
        """
        Return a tuple containing the values of the requested arguments. The requested arguments are then removed from
        the object.

        :param parameters: List of arguments to pop
        :return: Tuple
        """
        return self._popper("parameters", "Parameter", parameters)

    def pop_time_derivatives(self, time_derivatives: Iterable[str]) -> NamedTuple:
        """
        Return a tuple containing the values of the requested arguments. The requested arguments are then removed from
        the object.

        :param time_derivatives: List of arguments to pop
        :return: Tuple
        """
        return self._popper("time_discretization", "TimeDiscretization", time_derivatives)


class Formula(abc.ABC):
    """
    `Formula` defines the name, domains of evaluation, and relevant dependencies in order to define the FFL formulation.

    The method `Formula.form` must be overridden to return the fully specified form given the arguments. The arguments
    are defined by overriding one or more of `Formula.Variables`, `Formula.Formulas`, `Formula.Parameters`,
    `Formula.Lambdas`, and `Formula.TimeDiscretization`.

    `Formula.Variables`: Must contain the names of variables defined in a given model
    `Formula.Formulas`: Must contain the names of `Formulas` defined in a given model
    `Formula.Parameters`: Contains the names of parameters that will be provided later
    `Formula.Lambdas`: Contains the names of lambda functions that will be provided later
    `Formula.TimeDiscretization`: Contains the names of time derivatives that will be defined later. This can be either
    a time stepping scheme for Rothes, or a member of the solution vector for MOL.
    """

    def __init__(self, name: Optional[str] = None, domains: Iterable[str] = (),) -> None:
        """
        Create a Formula object with an optional name and a list of domains.

        If no name is provided, the name from the inheriting class will be used.

        :param name: name of the `Formula`
        :param domains: List of domains in which the `Formula` is defined
        """
        self.domains = domains
        self.name = name if name is not None else type(self).__name__
        self.Variables = self.typedef("Variables", "")  # pylint: disable=invalid-name
        self.Parameters = self.typedef("Parameters", "")  # pylint: disable=invalid-name
        self.Lambdas = self.typedef("Lambdas", "")  # pylint: disable=invalid-name
        self.TimeDiscretization = self.typedef("TimeDiscretization", "")  # pylint: disable=invalid-name

    @staticmethod
    def typedef(name: str, parameter_string: Union[str, Iterable]) -> Type[NamedTuple]:
        """
        Creates a `namedtuple` class with a given name and parameter string.

        :param name: Name of the new class
        :param parameter_string: string of parameter names for each tuple element. Comma or space delimited.
        :return: `namedtuple`
        """
        return collections.namedtuple(name, parameter_string)  # type: ignore

    @abc.abstractmethod
    def form(self, arguments: Arguments, domain: str) -> Expr:
        """
        This method must be overloaded to define the form of the `Formula`.

        :param arguments: All arguments defined by overriding one or more of `Formula.Variables`, `Formula.Formulas`,
            `Formula.Parameters`, `Formula.Lambdas`, and `Formula.TimeDiscretization`
        :param domain: The current domain in which the function is being evaluated
        :return: FFL form
        """

    @property
    def variables(self) -> Tuple[str]:
        """
        Return a tuple of argument names.

        :return: Names of the required variables
        """
        return self.Variables._fields  # type: ignore

    @property
    def parameters(self) -> Tuple[str]:
        """
        Return a tuple of argument names.

        :return: Names of the required parameters
        """
        return self.Parameters._fields  # type: ignore

    @property
    def time_discretizations(self) -> Tuple[str]:
        """
        Return a tuple of argument names.

        :return: Names of the required time_discretizations
        """
        return self.TimeDiscretization._fields  # type: ignore

    def append_arguments(self, type_: Type[NamedTuple], arguments: Iterable) -> Type[NamedTuple]:
        """
        Append argument names to the given type definition.

        :param type_: Existing type
        :param arguments: New arguments to add
        :return: NamedTuple
        """
        return self.typedef(type_.__name__, type_._fields + tuple(arguments))

    def formulate(self, arguments: Arguments, domain: str) -> Expr:
        """
        Validate arguments and perform type conversions necessary for evaluating `Formula.form`.

        :param arguments: Arguments required to evaluate
        :param domain: Domain in which the form is evaluated
        :return: FFL form
        """
        arguments.assign_argument_classes(self.Variables, self.Parameters, self.Lambdas, self.TimeDiscretization)
        arguments.lambdas = tuple(func(arguments) for func in arguments.lambdas)
        arguments.convert()

        return self.form(arguments, domain)

    def __repr__(self):
        names = ", ".join(self.Variables._fields + self.Parameters._fields + self.TimeDiscretization._fields)

        return f"{self.name}({names})"


class FormMap:
    """
    A mapping between formula specifications (unformulated), the formulated FFL representation, the domain/boundary
    measure map, and evaluation domain.
    """

    def __init__(
        self,
        formula: Formula,
        measure_map: MeasureMap,
        eval_domain: str,
        formulation: Optional[Domain[str, List]] = None,
        name: Optional[str] = None,
    ) -> None:
        """
        Create a mapping between formula definitions and implementation details.

        :param formula: `Formula` definition
        :param measure_map: Mapping of domain measure with primary domains and a multiple
        :param eval_domain: Domain of evaluation for the formulation
        :param formulation: FFL formulation
        :param name: Name of the mapping
        """

        # pylint: disable=too-many-arguments

        if not isinstance(formula, Formula):
            raise AttributeError("Only Formula types are permitted for formulas")

        self.formula = formula
        self.measure_map = measure_map
        self.eval_domain = eval_domain
        self.formulation = formulation

        if name is not None:
            self.name = name

    @property
    def name(self) -> str:
        """
        Name of the formula
        """
        return self.formula.name

    @name.setter
    def name(self, name: str) -> None:
        """
        Name of the formula
        """
        self.formula.name = name

    @property
    def domains(self) -> Iterable[str]:
        """
        Domains that the formula is defined in
        """
        return self.formula.domains

    @property
    def variables(self) -> Tuple[str]:
        """
        List of test functions in the formula
        """
        return self.formula.variables

    @property
    def primary_domain(self) -> Optional[str]:
        """
        Return the value of the measure's primary (parent) domain. This is usually set if the current measure is a
        boundary.

        :return: Primary domain
        """
        return self.measure_map.primary_domain

    def __repr__(self) -> str:
        """
        Nice printout for the formula, indicates if it has been formulated
        """
        primary_domain_str = f"{self.measure_map.primary_domain} -> " if self.measure_map.primary_domain else ""
        formulated_str = "formulated" if self.formulation else "unformulated"
        return f"{self.formula.name}[{primary_domain_str}{self.eval_domain}]: {formulated_str}"


class ArgumentError(Exception):
    """Raised when an argument doesn't match the class specified"""