Source code for mtnlion.model

"""
Provides a class interface for defining FEM models.
"""
import warnings
from typing import Iterable, List, Tuple, Type

from mtnlion.formula import Formula
from mtnlion.variable import Variable


class Model:
    """
    Abstract class for defining FEM models.
    """

    def __init__(self, minimum_rothes_order: int = 1):
        self.formulas: List[Formula] = []
        self.variables: List[Variable] = []
        self.minimum_rothes_order = minimum_rothes_order

    def exclude_formulas(self, formula_names: Iterable[str] = (), formula_types: Iterable[Type] = ()):
        """
        Exclude a formula definition from assembly.

        This function will remove any object instances from the assembly list that match the given name or class
        definition.

        :param formula_names: Name of the formula(s) to remove
        :param formula_types: Class of the formula(s) to remove
        :return: None
        """
        formula_names = tuple(formula_names)

        # Check existance
        for name in formula_names:
            self.formula_by_name(name)

        self.formulas = [formula for formula in self.formulas if formula.name not in formula_names]
        self.formulas = [
            formula
            for formula in self.formulas
            if not any(filter(lambda x, form=formula: isinstance(form, x), formula_types))  # type: ignore
        ]

    def replace_formulas(self, *formulas: Formula) -> "Model":
        """
        Replace an existing formula with a new formula. The formulas are replaced by name.

        :param formulas: New formula
        :return: None
        """
        self.exclude_formulas(form.name for form in formulas)
        self.formulas += formulas

        return self

    def add_formula(self, *forms: Formula) -> "Model":
        """
        Add a formulae to the model
        :param forms: formulas to add
        """
        self.formulas += forms

        return self

    def formula_by_name(self, name: str) -> Formula:
        """
        Retrieve a `Formula` object by name from the list of formulas.

        :param name: name of the formula
        :return: formula object
        """

        try:
            return next(filter(lambda x: x.name == name, self.formulas))
        except StopIteration:
            raise KeyError(f"Cannot find formula: {name}")

    def variable_by_name(self, name: str) -> Variable:
        """
        Retrieve a `Variable` object by name from the list of variables.

        :param name: name of the variable
        :return: variable object
        """
        try:
            return next(filter(lambda x: x.name == name, self.variables))
        except StopIteration:
            raise KeyError(f"Cannot find variable: {name}")

    def rename_variable(self, old: str, new: str) -> None:
        """
        Rename a given variable defined in the model to a new name.

        :param old: Parameter to rename
        :param new: New name
        :return: None
        """
        self.variable_by_name(old).name = new

        for formula in self.formulas:
            if old in formula.variables:
                variables = (new if variable == old else variable for variable in formula.variables)
                formula.Variables = formula.typedef(formula.Variables.__name__, ", ".join(variables))

    def rename_parameter(self, old: str, new: str) -> None:
        """
        Rename a given parameter defined in the model to a new name.

        :param old: Parameter to rename
        :param new: New name
        :return: None
        """
        for formula in self.formulas:
            if old in formula.parameters:
                parameters = (new if parameter == old else parameter for parameter in formula.parameters)
                formula.Parameters = formula.typedef(formula.Parameters.__name__, ", ".join(parameters))

    def _check_variable_names(self) -> None:
        """
        Check variable names in formulas against the model variables to ensure that variables that are used in formulas
        exist in the model variable definitions and vise versa.

        :return: None
        """
        form_vars = set(field for obj in self.formulas for field in obj.variables)
        def_vars = set(var.name for var in self.variables)

        if form_vars != def_vars:
            if form_vars - def_vars:
                raise LookupError(
                    f"Variables {form_vars - def_vars} are used in formulas but are undefined in the model."
                )
            if def_vars - form_vars:
                warnings.warn(f"Variables {def_vars - form_vars} are defined but are unused.")
            else:
                raise ValueError(
                    f"Unknown error occured. Formula variables and defined variables differ: {form_vars} -- {def_vars}."
                )

    def _check_name_uniqueness(self) -> None:
        """
        Check for naming collisions in variables and formulas.

        :return: None
        """
        form_names = [form.name for form in self.formulas]
        var_names = [var.name for var in self.variables]

        if len(form_names) != len(set(form_names)):
            duplicates = set(x for x in form_names if form_names.count(x) > 1)
            raise ReferenceError(f"Duplicate formula names exist: {duplicates}")

        if len(var_names) != len(set(var_names)):
            duplicates = set(x for x in var_names if var_names.count(x) > 1)
            raise ReferenceError(f"Duplicate variable names exist: {duplicates}")

    @property
    def domains(self) -> Tuple[str, ...]:
        """
        Retrieve a list of domains used in this model.

        :return: Tuple of sorted domain strings
        """
        form_domains = set(domain for form in self.formulas for domain in form.domains)
        var_domains = set(domain for var in self.variables for domain in var.domains)

        return tuple(sorted(form_domains.union(var_domains)))

    def check(self) -> None:
        """
        Perform high-level sanity checks to ensure the model is correctly defined.

        :return: None
        """
        self._check_variable_names()
        self._check_name_uniqueness()