import typing
from abc import abstractmethod, ABC
from collections import OrderedDict
from collections.abc import Iterable
from enum import Enum
from typing import Optional, Union, Callable

import relationalai.early_access.builder as qb
from relationalai.early_access.builder import define, where, not_
from relationalai.early_access.builder.builder import Not, ConceptConstruct, Ref, RelationshipRef
from relationalai.early_access.dsl.bindings.common import Binding, BindableAttribute, IdentifierConceptBinding, \
    SubtypeConceptBinding, BindableTable, _PrimitiveType, FilterBy, ReferentConceptBinding, _get_transform_types, \
    BindableColumn
from relationalai.early_access.dsl.codegen.helpers import reference_entity, construct_entity
from relationalai.early_access.dsl.orm.constraints import Unique
from relationalai.early_access.dsl.orm.relationships import Role
from relationalai.early_access.metamodel.util import OrderedSet


def _get_map_madlib_and_name(name_prefix: str, value_concept: qb.Concept) -> tuple[str, str]:
    """
    Generates a name for the map based on the value concept.
    """
    name = f'{name_prefix}_row_to_{value_concept}'
    madlib = f'{name} {{row_id:RowId}} {{val:{value_concept}}}'
    return madlib, name


class GeneratedRelation(qb.Relationship):
    """
    A class representing a relation generated from a model.
    """
    def __init__(self, madlib, model, name):
        super().__init__(madlib, model=model.qb_model(), short_name=name)

    def __repr__(self):
        return f'GeneratedRelation({self._name})'


class InternallyGeneratedRelation(GeneratedRelation):
    """
    A class representing a relation generated by analyzing the model.
    """
    def __init__(self, madlib, dsl_model, name):
        super().__init__(madlib, dsl_model, name)


def _filtering_view(row, filter_by: Optional[FilterBy], column_ref=None):
    """
    Generates filtering atoms for the map if the binding has any filters, which must be simple
    filters on the attribute view of the same table.
    """
    if not filter_by:
        return where()
    if isinstance(filter_by, _PrimitiveType):
        assert column_ref is not None, 'Attribute must be provided if filter_by is a PrimitiveFilterBy'
        atoms = [_primitive_filtering_atom(column_ref, filter_by)]
    elif isinstance(filter_by, qb.Expression):
        atoms = [_expr_filtering_atom(row, filter_by)]
    elif isinstance(filter_by, Iterable):
        atoms = [_expr_filtering_atom(row, filter_expr) for filter_expr in filter_by]
    elif isinstance(filter_by, Not):
        # trivial negation is supported, either on a single column or a trivial expression
        atoms = [_negation_atom(row, filter_by)]
    else:
        raise TypeError(f'Expected a PrimitiveFilterBy, Expression, or Iterable of Expressions, got {type(filter_by)}')
    return where(*atoms)

def _primitive_filtering_atom(column_ref, condition: _PrimitiveType):
    """
    Generates a filtering atom for the map based on a primitive filter.
    """
    if isinstance(condition, Enum):
        condition = condition.name
    return where(
        column_ref == condition
    )

def _expr_from_filter(row, condition: Union[BindableColumn, qb.Expression]):
    """
    Generates a filtering atom for the map based on a condition.
    """
    if isinstance(condition, qb.Expression):
        return _expr_filtering_atom(row, condition)
    elif isinstance(condition, BindableColumn):
        return _column_filtering_atom(row, condition)
    else:
        raise TypeError(f'Expected an Expression or a BindableColumn, got {type(condition)}')

def _expr_filtering_atom(row, condition: qb.Expression):
    params = condition._params
    if len(params) != 2:
        raise ValueError(f'Expected a condition with two parameters, got {len(params)}: {condition}')
    column, value = params
    if isinstance(value, Enum):
        value = value.name
    orig = column.type().ref()
    return where(
        column(row, orig),
        qb.Expression(condition._op, orig, value)
    )

def _column_filtering_atom(row, column: BindableColumn):
    """
    Defines an expression which could be used to express `exists(val: column(row, val))` for a given column.
    """
    orig = column.type().ref()
    return where(
        column(row, orig)
    )

def _negation_atom(row, condition: Not):
    """
    Generates a negation atom based on the filter: either a column or a trivial expression.
    """
    conditions = condition._args
    return where(
        not_(
            *[_expr_from_filter(row, condition) for condition in conditions]
        )
    )


class AbstractMap(GeneratedRelation):
    """
    A class representing an abstract map.
    """
    def __init__(self, madlib, dsl_model, name):
        super().__init__(madlib, dsl_model, name)

    def row_ref(self):
        assert len(self._field_refs) == 2, 'AbstractMap must have exactly two field references'
        row_ref = self._field_refs[0]
        assert row_ref._thing._name == 'RowId', 'First field must be of type RowId'
        return row_ref

    def value_player(self) -> qb.Concept:
        """
        Returns the value player of the map.
        This is the concept that the map maps to.
        """
        raise NotImplementedError()

    def formula(self, row_ref):
        rez_val_ref = self.value_player().ref()
        return rez_val_ref, [self(row_ref, rez_val_ref)]


class RoleMap(AbstractMap):
    """
    A class representing a base role map relation.
    """

    def __init__(self, madlib, model, name, role):
        super().__init__(madlib, model, name)
        self._role = role

    def role(self):
        """
        Returns the role of the role map.
        """
        return self._role

    def value_player(self):
        """
        Returns the value player of the role map.
        """
        return self._role.player()

    @abstractmethod
    def column(self) -> BindableAttribute:
        """
        Returns the bindable column associated with this role map.
        """
        raise NotImplementedError("Subclasses must implement this method")

    def table(self) -> BindableTable:
        """
        Returns the table associated with the role map.
        """
        return self.column().table

    def __repr__(self):
        return f'RoleMap({self._name})'


class InlineValueMap:
    """
    A non-QB relationship class that represents an inline value map.
    """
    def __init__(self, model, binding: Binding, role):
        self._binding = binding
        self._role = role
        self._value_player = role.player()
        self._initialize()

    def _initialize(self):
        role = self._role
        if not role.player()._is_primitive():
            raise TypeError(f'Cannot construct a value map for {role}, concept is not a value type')

        table_name = self._binding.column.table.physical_name()
        concept = role.player()
        _, self._name = _get_map_madlib_and_name(table_name, concept)
        self._has_transforms = bool(self._binding.transform_with)

    def __call__(self, *args, **kwargs):
        """
        Returns a QB expression that can be used in a query.
        """
        # assert that len(args) is 2, as we expect row and value variables
        if len(args) != 2:
            raise ValueError(f'Expected 2 arguments, got {len(args)}: {args}')

        row_ref, rez_val_ref = args
        _, body_atoms = self.construct_body_formula(row_ref, rez_val_ref)
        return where(*body_atoms)

    def construct_body_formula(self, row_ref, rez_val_ref):
        """
        Constructs the body formula for the value map.
        This is used to generate the head and the body of the relation.
        """
        col_val_ref = self.column().type().ref() if self._has_transforms else rez_val_ref
        filter_atoms = _filtering_view(row_ref, self._binding.filter_by, col_val_ref)
        return self._construct_body_atoms(row_ref, col_val_ref, rez_val_ref, filter_atoms)

    def _construct_body_atoms(self, row_ref, col_val_ref, rez_val_ref, filter_atoms):
        result_type = self.value_player()
        if self._has_transforms:
            transform_atoms = self._generate_transform_body(col_val_ref, rez_val_ref)
            return [row_ref, result_type(rez_val_ref)], [
                where(
                    self._binding.column(row_ref, col_val_ref),
                    transform_atoms,
                    filter_atoms
                )]
        else:
            # col_val_ref and rez_val_ref are the same handles here
            return [row_ref, result_type(col_val_ref)], [
                self._binding.column(row_ref, col_val_ref),
                filter_atoms,
            ]

    def _generate_transform_body(self, col_val_ref, rez_val_ref):
        transformations = self._binding.transform_with
        if transformations:
            transformations = (transformations,) if not isinstance(transformations, tuple) else transformations
            return self._chain_transformations(col_val_ref, rez_val_ref, transformations)
        raise ValueError("Incorrect state, cannot apply binding transformations with empty value converter(s)")

    def _chain_transformations(self, input_val_ref, output_val_ref, transformations):
        if len(transformations) == 1:
            return self._apply_transformations(input_val_ref, output_val_ref, transformations[0])
        intermediate_vars = [self._get_transform_output_type(tr) for tr in transformations[:-1]]
        all_vars = [input_val_ref, *intermediate_vars, output_val_ref]
        return where(
            *[self._apply_transformations(all_vars[idx], all_vars[idx+1], tr)
              for idx, tr in enumerate(transformations)]
        )

    @staticmethod
    def _get_transform_output_type(transformer):
        _, output_type = _get_transform_types(transformer)
        return output_type

    @staticmethod
    def _apply_transformations(input_var, output_var, transformer):
        if isinstance(transformer, qb.Relationship):
            return transformer(input_var, output_var)
        elif isinstance(transformer, Callable):
            return output_var == transformer(input_var)
        else:
            raise TypeError(f'Expected a Relationship or Callable, got {type(transformer)}')

    def value_player(self):
        """
        Returns the value player of the inline value map.
        """
        return self._value_player

    def column(self) -> BindableAttribute:
        """
        Returns the bindable column associated with this inline value map.
        """
        return self._binding.column

    def table(self) -> BindableTable:
        """
        Returns the table associated with the inline value map.
        """
        return self.column().table

    def row_ref(self):
        """
        Returns the row reference for the inline value map.
        """
        return self.column().row_ref()

    def binding(self) -> Binding:
        """
        Returns the binding of the value map.
        """
        return self._binding

    def formula(self, row_ref):
        rez_val_ref = self.value_player().ref()
        return rez_val_ref, [self(row_ref, rez_val_ref)]

    def __repr__(self):
        return f'@inline ValueMap({self._name})'


class ValueMap(RoleMap):
    """
    A class representing a value map relation.
    """

    def __init__(self, model, binding: Binding, role):
        # =
        # Skip QB relationship initialization if *inline* is True, as __call__ will return a QB expression
        # =
        madlib, name = self._handle_params(binding, role)
        super().__init__(madlib, model, name, role)
        self._binding = binding
        self._inline_value_map = InlineValueMap(model, binding, role)
        self._generate_body()

    @staticmethod
    def _handle_params(binding, role):
        if not role.player()._is_primitive():
            raise TypeError(f'Cannot construct a value map for {role}, concept is not a value type')

        table_name = binding.column.table.physical_name()
        concept = role.player()
        return _get_map_madlib_and_name(table_name, concept)

    def _generate_body(self):
        row_ref = self.row_ref()
        has_transforms = bool(self._binding.transform_with)
        rez_val_ref = self.value_player().ref() if has_transforms else self.column().type().ref()

        head_vars, body_atoms = self._inline_value_map.construct_body_formula(row_ref, rez_val_ref)

        define(self(*head_vars)).where(*body_atoms)

    def value_player(self):
        return self._inline_value_map.value_player()

    def column(self) -> BindableAttribute:
        return self._inline_value_map.column()

    def binding(self) -> Binding:
        return self._inline_value_map.binding()

    def __repr__(self):
        return f'ValueMap({self._name})'


class MaterializedEntityMap(AbstractMap):
    """
    A class representing a materialized entity map relation.
    This is used to materialize an entity map from a set of role maps.
    """

    def __init__(self, model, inline_entity_map: 'AbstractInlineEntityMap'):
        super().__init__(inline_entity_map.madlib(), model, inline_entity_map.short_name())
        self._inline_entity_map = inline_entity_map
        self._generate_body()

    def _generate_body(self):
        # materialize the population of the entity map as necessary
        self._inline_entity_map.materialize_population()

        row_ref = self.row_ref()
        rez_val_ref, body_formula = self._inline_entity_map.formula(row_ref)
        where(*body_formula).define(self(row_ref, rez_val_ref))

    def value_player(self):
        """
        Returns the result value player of the materialized entity map.
        This is the same as the value player of the inline entity map.
        """
        return self._inline_entity_map.value_player()

    def column(self) -> BindableAttribute:
        return self._inline_entity_map.column()

    def table(self) -> BindableTable:
        """
        Returns the table associated with the materialized entity map.
        """
        return self._inline_entity_map.table()

    def __repr__(self):
        # replace @inline with @materialized
        return self._inline_entity_map.__repr__().replace('@inline', '@materialized')


class AbstractInlineEntityMap(ABC):
    """
    A non-QB relationship class that represents an inline entity map.
    This is used to construct entities from a set of role maps.
    """

    def __init__(self, model, madlib, name):
        self._model = model
        self._madlib = madlib
        self._name = name

    def __call__(self, *args, **kwargs):
        if len(args) != 1:
            raise ValueError(f'Expected a single argument to a Concept Map, got {len(args)}: {args}')

        row_ref = args[0]
        return self._construct_body_formula(row_ref)

    def madlib(self):
        """
        Returns the madlib of the inline entity map.
        """
        return self._madlib

    def short_name(self):
        """
        Returns the name of the inline entity map.
        """
        return self._name

    @abstractmethod
    def value_player(self):
        """
        Returns the value player of the role map.
        """
        raise NotImplementedError("Subclasses must implement this method")

    @abstractmethod
    def column(self) -> BindableAttribute:
        """
        Returns the bindable column associated with this role map.
        """
        raise NotImplementedError("Subclasses must implement this method")

    def table(self) -> BindableTable:
        """
        Returns the table associated with the role map.
        """
        return self.column().table

    @abstractmethod
    def row_ref(self):
        """
        Returns the row reference for the inline entity map, should be typically a row reference of an arbitrary
        inline entity or value map that is used to construct the entity, all the way to a column.
        """
        raise NotImplementedError("Subclasses must implement this method")

    def formula(self, row_ref):
        return self._construct_body_formula(row_ref)

    @abstractmethod
    def _construct_body_formula(self, row_ref) -> tuple[Union[ConceptConstruct, Ref, RelationshipRef], list[typing.Any]]:
        """
        Constructs the body formula for the inline entity map.
        This is used to generate the body of the relation.
        """
        pass

    def materialize(self):
        return MaterializedEntityMap(self._model, self)

    @abstractmethod
    def materialize_population(self):
        """
        Materializes the population of the entity map, using the body of the entity map.
        This has to check the conditions and only update the population if those are met.
        """
        pass


class SimpleConstructorEntityMap(AbstractInlineEntityMap):
    """
    A class representing an entity map relation.
    """

    def __init__(self, model, binding: Binding, role: Role, identifier_to_role_map: OrderedDict[Unique, 'RoleMap']):
        madlib, name = self._handle_params(binding, role)
        super().__init__(model, madlib, name)
        self._binding = binding
        self._role = role
        self._identifier_to_role_map = identifier_to_role_map
        self._reference_role_map = list(identifier_to_role_map.items())[-1][1]  # last role map is the reference one

    @staticmethod
    def _handle_params(binding, role):
        if role.player()._is_primitive():
            raise TypeError(f'Cannot construct an entity map for {role}, concept is not an entity type')

        table_name = binding.column.table.physical_name()
        concept = role.player()
        return _get_map_madlib_and_name(table_name, concept)

    def _construct_body_formula(self, row_ref):
        concept = self._role.player()
        # and populate the entity map, role_map is always the last one in the identifier_to_role_map
        value = self._reference_role_map.value_player().ref()
        body_atoms = [
            self._reference_role_map(row_ref, value),
            entity_ref := reference_entity(concept, value),
            self._concept_population_atom(entity_ref)
        ]
        return entity_ref, body_atoms

    def _concept_population_atom(self, entity):
        """
        Generates the *optional* population atom for the entity map.
        Only used for bindings that have a FK.
        """
        if self._should_reference():
            return where(
                self._role.player()(entity)
            )
        else:
            return where()

    def _should_reference(self) -> bool:
        return self._reference_role_map.column().references_column is not None

    def value_player(self):
        """
        Returns the value player of the entity map.
        """
        return self._role.player()

    def column(self) -> BindableAttribute:
        """
        Returns the bindable column associated with this entity map.
        """
        return self._binding.column

    def row_ref(self):
        return self._reference_role_map.row_ref()

    def materialize_population(self):
        row_ref = self.row_ref()
        concept = self._role.player()
        if not self._should_reference():
            # then create entities
            role_maps = self._identifier_to_role_map.values()
            map_value_refs = [role_map.value_player().ref() for role_map in role_maps]
            where(
                *[role_map(row_ref, value) for role_map, value in zip(role_maps, map_value_refs)]
            ).define(construct_entity(concept, *map_value_refs))

    def __repr__(self):
        return f'@inline CtorEntityMap({self._name})'


class ReferentEntityMap(AbstractInlineEntityMap):
    """
    A class representing a referent entity map relation.
    """

    def __init__(self, model, binding: Binding, role: Role, constructing_role_map: 'Map'):
        madlib, name = self._handle_params(binding, role)
        super().__init__(model, madlib, name)
        self._binding = binding
        self._role = role
        self._constructing_role_map = constructing_role_map

    @staticmethod
    def _handle_params(binding, role):
        if role.player()._is_primitive():
            raise TypeError(f'Cannot construct an entity map for {role}, concept is not an entity type')

        table_name = binding.column.table.physical_name()
        concept = role.player()
        return _get_map_madlib_and_name(f'{table_name}_ref', concept)

    def _construct_body_formula(self, row_ref):
        concept = self._role.player()
        orig_entity_ref, subformula_atoms = self._constructing_role_map.formula(row_ref)

        entity_ref = reference_entity(concept, orig_entity_ref)
        return entity_ref, [
            *subformula_atoms,
            entity_ref,
            concept(entity_ref)
        ]

    def value_player(self):
        return self._role.player()

    def binding(self) -> Binding:
        """
        Returns the binding of the entity map.
        """
        return self._binding

    def column(self) -> BindableAttribute:
        """
        Returns the bindable column associated with this entity map.
        """
        return self._binding.column

    def row_ref(self):
        return self._constructing_role_map.row_ref()

    def materialize_population(self):
        """
        Referent entity maps do not need to materialize population, as they always reference existing entities.
        """
        pass

    def __repr__(self):
        return f'@inline ReferentEntityMap({self._name})'


class CompositeEntityMap(AbstractInlineEntityMap):
    """
    A class representing a composite entity map relation.

    Takes a sequence of EntityMaps and constructs a composite entity type.
    """

    def __init__(self, model, entity_concept: qb.Concept, *role_maps: 'Map'):
        madlib, name = self._handle_params(entity_concept, *role_maps)
        super().__init__(model, madlib, name)
        self._entity_concept = entity_concept
        self._role_maps = list(role_maps)

    def _handle_params(self, entity_concept: qb.Concept, *role_maps: 'Map'):
        if entity_concept._is_primitive():
            raise TypeError(f'Cannot construct a composite entity map for {entity_concept},'
                            f' concept is not an entity type')
        if len(role_maps) < 2:
            raise ValueError('CompositeEntityMap requires at least two EntityMaps')

        role_map = role_maps[0]
        table = role_map.table()
        self._table = table
        return _get_map_madlib_and_name(table.physical_name(), entity_concept)

    def _construct_body_formula(self, row_ref):
        var_refs, body_atoms = self._body_formula(row_ref)

        # the rule to populate the entity map
        entity_ref = reference_entity(self._entity_concept, *var_refs)
        body_atoms.append(entity_ref)  # TODO should this be added?
        return entity_ref, body_atoms

    @staticmethod
    def _should_reference() -> bool:
        """
        Determines if the composite entity map should reference a column.
        Currently not supported fully.
        """
        return False

    def _body_formula(self, row_ref):
        var_refs, subformulas = zip(*(role_map.formula(row_ref) for role_map in self._role_maps))
        return var_refs, [
            *(atom for subformula in subformulas for atom in subformula)
        ]

    def value_player(self):
        """
        Returns the value player of the composite entity map.
        """
        return self._entity_concept

    def table(self) -> BindableTable:
        """
        Returns the table associated with the composite entity map.
        """
        return self._table

    def column(self) -> BindableAttribute:
        raise ValueError('CompositeEntityMap does not index a single column')

    def row_ref(self):
        # take arbitrary row reference from the first role map
        role_map = self._role_maps[0]
        return role_map.row_ref()

    def materialize_population(self):
        row_ref = self.row_ref()
        var_refs, body_atoms = self._body_formula(row_ref)
        if not self._should_reference():
            # construct entities
            where(
                *body_atoms,
                entity_ref := construct_entity(self._entity_concept, *var_refs),
            ).define(entity_ref)

    def __repr__(self):
        return f'@inline CompositeEntityMap({self._name})'


AbstractEntityMap = Union['SimpleConstructorEntityMap', 'ReferentEntityMap', 'CompositeEntityMap', 'UnionEntityMap']
Map = Union['ValueMap', 'InlineValueMap', 'UnionEntityMap', 'AbstractInlineEntityMap', 'MaterializedEntityMap']

class EntitySubtypeMap(AbstractInlineEntityMap):
    """
    A class representing an entity subtype map relation.
    """

    def __init__(self, model, binding: Union[IdentifierConceptBinding, SubtypeConceptBinding, ReferentConceptBinding],
                 ctor_entity_map: Union[AbstractInlineEntityMap, MaterializedEntityMap]):
        # type of the entity subtype map is the parent type coming from the ctor_entity_map
        madlib, name = self._handle_params(binding, ctor_entity_map.value_player())
        super().__init__(model, madlib, name)
        self._binding = binding
        self._ctor_entity_map = ctor_entity_map

    @staticmethod
    def _handle_params(binding: Union[IdentifierConceptBinding, SubtypeConceptBinding, ReferentConceptBinding],
                       result_concept: qb.Concept):
        table = binding.column.table
        entity_concept = binding.entity_type
        # check that the entity concept is a subtype of the result concept
        assert entity_concept._isa(result_concept), f"Entity concept {entity_concept} must be a subtype of result concept {result_concept}"
        name = f'{table.physical_name()}_row_to_{entity_concept}'
        madlib = f'{name} {{row_id:RowId}} {{val:{result_concept}}}'
        return madlib, name

    def _construct_body_formula(self, row_ref):
        parent_type_ref, body_atoms = self._generate_body_atoms(row_ref)

        # note: have to return the parent type here, as otherwise it gets a join on the subtype population
        #       but doesn't cast to the type of the subtype...
        return parent_type_ref, body_atoms

    def _generate_body_atoms(self, row_ref):
        filtering_atom = _filtering_view(row_ref, self._binding.filter_by)
        parent_type_ref, subformula_atoms = self._ctor_entity_map.formula(row_ref)
        return parent_type_ref, [
            *subformula_atoms,
            filtering_atom
        ]

    def binding(self):
        """
        Returns the binding of the entity subtype map.
        """
        return self._binding

    def value_player(self):
        """
        Returns the subtype of the entity subtype map.
        """
        return self._binding.entity_type

    def column(self) -> BindableAttribute:
        """
        Returns the bindable column associated with this entity subtype map.
        """
        return self._binding.column

    def table(self) -> BindableTable:
        """
        Returns the table associated with the entity subtype map.
        """
        return self._binding.column.table

    def row_ref(self):
        return self._ctor_entity_map.row_ref()

    def materialize_population(self):
        row_ref = self.row_ref()
        subtype = self._binding.entity_type
        parent_type_ref, body_atoms = self._generate_body_atoms(row_ref)
        # TODO instead of Person.new maybe try Employee.new to construct these?
        # derive subtype population
        where(*body_atoms).define(subtype(parent_type_ref))

    def __repr__(self):
        return f'@inline EntitySubtypeMap({self._name})'


class UnionEntityMap(AbstractMap):
    """
    A class representing a union entity map relation.
    """

    def __init__(self, model, entity_concept: qb.Concept, *entity_maps: AbstractEntityMap, generate_population: bool=False):
        madlib, name = self._handle_params(entity_concept, *entity_maps)
        super().__init__(madlib, model, name)
        self._entity_type = entity_concept
        self._entity_maps = OrderedSet().update(entity_maps)
        self._generate_population = generate_population
        self._generate_body()

    def _handle_params(self, entity_concept: qb.Concept, *entity_maps: AbstractEntityMap):
        if len(entity_maps) == 0:
            raise ValueError('UnionEntityMap requires at least one EntityMap')
        # pick an arbitrary entity map to get the table, as they must all have the same
        table = entity_maps[0].table()
        self._table = table
        return _get_map_madlib_and_name(table.physical_name(), entity_concept)

    def _generate_body(self):
        for entity_map in self._entity_maps:
            self._generate_body_rule(entity_map)

    def _generate_body_rule(self, entity_map: AbstractEntityMap):
        row_ref = entity_map.row_ref()
        rez_val_ref, subformula_atoms = entity_map.formula(row_ref)

        # derive union entity map
        where(
            *subformula_atoms
        ).define(self(row_ref, rez_val_ref))

    def value_player(self):
        """
        Returns the type of the entity map.
        """
        return self._entity_type

    def table(self):
        """
        Returns the table associated with the entity map.
        """
        return self._table

    def update(self, entity_map: AbstractEntityMap):
        """
        Updates the union entity map with a new entity map.
        """
        if entity_map in self._entity_maps:
            return
        self._entity_maps.add(entity_map)
        self._generate_body_rule(entity_map)

    def materialize_population(self):
        if self._generate_population:
            # derive type population
            for entity_map in self._entity_maps:
                row_ref = entity_map.row_ref()
                rez_val_ref, subformula_atoms = entity_map.formula(row_ref)
                where(
                    *subformula_atoms
                ).define(self._entity_type(rez_val_ref))

    def __repr__(self):
        return f'UnionEntityMap({self._name})'
