from __future__ import annotations

from collections import defaultdict
from dataclasses import dataclass
from typing import Optional, Union, Iterable, TypeVar, Tuple

from ...util import tracing
from ...util.error import err, exc, Source, Part
from ...util.naming import sanitize
from ...util.structures import OrderedSet

from . import metamodel as mm, builtins as bt
from .rewriter import Walker, Rewriter, NO_WALK
from .builtins import builtins as b

#--------------------------------------------------
# Typer
#--------------------------------------------------

T = TypeVar("T", bound=mm.Model | mm.Task)

class Typer:
    def __init__(self, enforce=False):
        self.enforce = enforce
        self.model_net: Optional[PropagationNetwork] = None
        # TODO: remove this once we are using spans for the debugger
        self.last_net: Optional[PropagationNetwork] = None

    def infer_model(self, model: mm.Model) -> mm.Model:
        """
        Infer types for the given model, returning a new model with updated types and
        storing the results within this typer so that subsequent calls can reuse them.
        """
        # create a brand new network with data from this model
        self.model_net = PropagationNetwork(model)
        return self._infer(model, self.model_net)

    def infer_query(self, query: mm.Task) -> mm.Task:
        """
        Infer types for the given query, returning a new query with updated types. If the
        typer was used to infer a model previously, the information from that model analysis
        will be reused.
        """
        if self.model_net is not None:
            net = PropagationNetwork(self.model_net.model)
            net.load_types(self.model_net.resolved_types)
        else:
            net = PropagationNetwork(mm.Model(root=query))
        self.last_net = net
        return self._infer(query, net)

    #--------------------------------------------------
    # Internal implementation
    #--------------------------------------------------

    def _infer(self, node: T, net: PropagationNetwork) -> T:
        # build the propagation network by analyzing the model
        with tracing.span("typer.analyze"):
            Analyzer(net).analyze(node)

        # propagate the types through the network
        with tracing.span("typer.propagate") as span:
            net.propagate()
            # span["type_graph"] = net.to_mermaid()

        # replace the fields in the model with the new types
        with tracing.span("typer.replace"):
            replacer = Replacer(net)
            final = replacer.rewrite(node)

        # report any errors found during typing
        if net.errors:
            for error in net.errors:
                error.report()
            if self.enforce:
                exc("TyperError", "Type errors detected during type inference.")
        return final


#--------------------------------------------------
# Propagation Network
#--------------------------------------------------

# The core idea of the typer is to build a propagation network that represents data
# dependencies between nodes. Nodes can be variables, literals, fields and tasks that refer
# to relations (lookups, aggregates, updates). An edge from node A to node B means that
# in order to resolve references and types for node B we need to know the type of node A.
#
# After building the network, we start with roots that have known types (literals, fields
# loaded from previous analysis and vars with concrete types) and propagate their types via
# the edges to other nodes. It is possible that the network contains cycles, so we iterate
# on the work list until fixpoint is reached, i.e. when we cannot propagate any new type.
#
# During the propagation, we are only gathering information about the types and references
# of nodes. Once we reach a fixpoint, we do a final pass to rewrite the model with the new
# type information, which may include adding casts to convert between types where needed.
#
Node = Union[mm.Var, mm.Field, mm.Literal, mm.Lookup, mm.Aggregate, mm.Update]

class PropagationNetwork():
    def __init__(self, model: mm.Model):
        # the model that is either being analyzed or that is the context for a query
        self.model = model

        # the resolved types of nodes in the network
        self.resolved_types: dict[Node, mm.Type] = {}

        # map from unresolved placeholder relations to their potential target replacements
        self.potential_targets: dict[mm.Relation, list[mm.Relation]] = {}

        # nodes loaded from previous analysis, which we use as roots to start propagation
        self.loaded_roots = set()

        # edges in the propagation network, from one node to potentially many
        self.edges:dict[Node, OrderedSet[Node]] = defaultdict(lambda: OrderedSet())
        self.back_edges:dict[Node, OrderedSet[Node]] = defaultdict(lambda: OrderedSet())

        # all errors collected during inference
        self.errors:list[TyperError] = []

        # overloads resolved for a lookup/update/aggregate, by node id. This is only for
        # relations that declare overloads
        self.resolved_overload:dict[int, mm.Overload] = {}

        # placeholders resolved for a lookup, by node id. This is only for relations that
        # are placeholders (i.e. only Any fields) and will be replaced by references to
        # these concrete relations. E.g. a query for "name(Any, Any)" may be replaced by
        # the union of "name(Dog, String)" and name(Cat, String)".
        self.resolved_placeholder:dict[int, list[mm.Relation]] = {}

        # for a given lookup/update/aggregate that involves numbers, the specific number
        # type resolved for it.
        self.resolved_number:dict[int, mm.NumberType] = {}

    #--------------------------------------------------
    # Error reporting
    #--------------------------------------------------

    def type_mismatch(self, node: Node|mm.Update, expected: mm.Type, actual: mm.Type):
        self.errors.append(TypeMismatch(node, expected, actual))

    def invalid_type(self, node: Node, type: mm.Type):
        self.errors.append(InvalidType(node, type))

    def unresolved_overload(self, node: mm.Lookup|mm.Aggregate):
        # TODO - consider renaming this to UnresolvedReference
        self.errors.append(UnresolvedOverload(node, [self.resolve(a) for a in node.args]))

    def unresolved_type(self, node: Node):
        # TODO - this is not being used yet, we need a pass at the end to check for any node
        # that we could not resolve
        self.errors.append(UnresolvedType(node))

    def has_errors(self, node: Node) -> bool:
        for mismatch in self.errors:
            if mismatch.node == node:
                return True
        return False

    #--------------------------------------------------
    # Types and Edges
    #--------------------------------------------------

    def add_edge(self, source: Node, target: Node):
        self.edges[source].add(target)
        self.back_edges[target].add(source)

    def add_resolved_type(self, node: Node, type: mm.Type):
        """ Register that this node was resolved to have this type. """
        if node in self.resolved_types:
            self.resolved_types[node] = merge_types(self.resolved_types[node], type)
        else:
            self.resolved_types[node] = type

    #--------------------------------------------------
    # Load types from a previous analysis
    #--------------------------------------------------

    def load_types(self, type_dict: dict[Node, mm.Type]):
        for node, type in type_dict.items():
            if isinstance(node, (mm.Field)):
                self.add_resolved_type(node, type)
                self.loaded_roots.add(node)

    #--------------------------------------------------
    # Resolve Values
    #--------------------------------------------------

    def resolve(self, value: mm.Value) -> mm.Type:
        if isinstance(value, (mm.Var, mm.Literal)):
            return self.resolved_types.get(value) or to_type(value)
        if isinstance(value, mm.Field):
            return self.resolved_types.get(value) or value.type
        assert not isinstance(value, (mm.Task)), "Should never try to resolve a task"
        return to_type(value)

    #--------------------------------------------------
    # Resolve References
    #--------------------------------------------------

    def all_dependencies_resolved(self, op:mm.Lookup|mm.Aggregate):
        """ True iff all dependencies required to resolve this reference are met. """
        rel = get_relation(op)
        # if this is a placeholder, eq or cast, we assume all possible args were resolved
        if bt.is_placeholder(rel) or rel == b.core.eq or rel == b.core.cast:
            return True

        # else, find whether all back-edges were resolved
        for node in self.back_edges[op]:
            # cannot resolve if a required var, literal or input field is still abstract
            if isinstance(node, (mm.Var, mm.Literal)) or (isinstance(node, mm.Field) and node.input):
                node_type = self.resolve(node)
                if bt.is_abstract(node_type):
                    return False
        return True

    def resolve_reference(self, op: mm.Lookup|mm.Aggregate) -> Optional[mm.Overload|list[mm.Relation]]:
        # check if all dependencies required to resolve this reference are met
        if not self.all_dependencies_resolved(op):
            return None

        relation = get_relation(op)
        if bt.is_placeholder(relation):
            # when replacing a placeholder, we may have multiple matches
            matches = []
            resolved_args = [self.resolve(arg) for arg in op.args]
            for target in self.potential_targets.get(relation, []):
                fields = get_relation_fields(target, relation.name)
                if all(type_matches(arg, self.resolve(field))
                        for arg, field in zip(resolved_args, fields)):
                    matches.append(target)
            return matches

        elif relation.overloads:
            # when resolving an overload for a concrete relation, we can only have one
            resolved_args = [self.resolve(arg) for arg in op.args]
            is_function = bt.is_function(relation)
            for overload in relation.overloads:
                # note that for functions we only consider input fields when matching
                if all(type_matches(arg, field_type) or conversion_allowed(arg, field_type)
                       for arg, field, field_type in zip(resolved_args, relation.fields, overload.types)
                       if field.input or not is_function):
                    self.resolved_overload[op.id] = overload
                    return overload
            return []  # no matches found

        else:
            # this is a relation with type vars or numbers that needs to be specialized
            return [relation]

    #--------------------------------------------------
    # Propagation
    #--------------------------------------------------

    def propagate(self):
        edges = self.edges
        work_list = []

        # start with the loaded roots + all literals + sources of edges without back edges
        work_list.extend(self.loaded_roots)
        for source in self.edges.keys():
            if isinstance(source, (mm.Literal)) or not source in self.back_edges:
                work_list.append(source)

        # limit the number of iterations to avoid infinite loops
        i = 0
        max_iterations = 100 * len(self.edges)

        # propagate types until we reach a fixed point
        while work_list:
            i += 1
            if i > max_iterations:
                err("Infinite Loop", "Infinite loop detected in the typer. Please, report this as a bug.")
                break

            node = work_list.pop(0)
            next = None
            if isinstance(node, mm.Field):
                # this is a field loaded from a previous analysis, so all we need to do is
                # propagate its type to its output edges
                next = edges.get(node, [])

            elif isinstance(node, (mm.Lookup, mm.Aggregate)):
                # this is a lookup/aggregate that may be overloaded or a placeholder; try
                # to resolve its reference, i.e. determine which specific relation or
                # overload it refers to
                found = self.resolve_reference(node)
                if found is not None:
                    # if found is None it means that we need more info to resolve the
                    # reference; otherwise it was possible to resolve it (even if to no matches)

                    # keep the resolved args before propagation to see if they will change
                    resolved_args = [self.resolve(arg) for arg in node.args]
                    # propagate the reference resolution
                    self.propagate_reference(node, resolved_args, found)
                    if found:
                        # the next nodes to process are all the outgoing edges plus any arg that
                        # changed during propagation
                        next = OrderedSet()
                        next.update(edges.get(node, []))
                        next.update([arg for idx, arg in enumerate(node.args) if isinstance(arg, mm.Var) and not (self.resolve(arg) == resolved_args[idx])])
                    else:
                        # the reference is unresolved, so remove from the worklist the args
                        # because they depend on this being resolved
                        for arg in node.args:
                            if arg in work_list:
                                work_list.remove(arg)
            else:
                assert isinstance(node, (mm.Var, mm.Literal))
                resolved = self.resolve(node)
                # if we ended up with an invalid type, report it on a next edge (which is
                # where the var or literal is going to be used)
                if invalid_type(resolved):
                    self.invalid_type(edges.get(node, [node])[0], resolved)

                # if the var is still abstract, add it to the work list to try again later
                if isinstance(node, mm.Var) and bt.is_abstract(resolved):
                    # but if everything that is left are vars, we cannot make progress
                    if not all(isinstance(x, mm.Var) for x in work_list):
                        next = [node]
                else:
                    # otherwise, check all outgoing edges
                    next = []
                    for out in edges.get(node, []):
                        if isinstance(out, mm.Update):
                            # this is an update to a field. We have to check that the type is
                            # valid for the field and if the field is abstract we can refine the
                            # type. But if the field is concrete we have to check that the type
                            # matche or can be converted.
                            is_population_lookup = isinstance(out.relation, mm.TypeNode)
                            for field in get_update_fields(node, out):
                                if not type_matches(resolved, field.type, accept_expected_super_types=is_population_lookup) and not conversion_allowed(resolved, field.type):
                                    self.type_mismatch(out, field.type, resolved)
                                elif bt.is_abstract(field.type):
                                    # if the field type changed, propagate further
                                    if resolved != self.resolve(field):
                                        next.append(field)
                                    self.add_resolved_type(field, resolved)
                        else:
                            # this is an arg flowing into a task, so resolve the task next
                            next.append(out)

            # add the new nodes to the work list
            if next is not None:
                for n in next:
                    if n not in work_list:
                        work_list.append(n)


    def propagate_reference(self, task:mm.Lookup|mm.Aggregate, resolved_args:list[mm.Type], references:mm.Overload|list[mm.Relation]):
        # TODO: distinguish between overloads and placeholders better when raising errors
        if not references:
            return self.unresolved_overload(task)

        relation = get_relation(task)

        if relation == b.core.cast:
            # cast is special cased: we just propagate the type to the target
            cast_type, target = task.args[0], task.args[2]
            assert isinstance(cast_type, mm.Type)
            assert isinstance(target, mm.Var)
            self.add_resolved_type(target, cast_type)
            return

        if bt.is_placeholder(relation):
            assert(references and isinstance(references, list))
            # we've resolved the placeholder, so store that
            self.resolved_placeholder[task.id] = references

            # we need to determine the final types of our args by taking all the references
            # and adding the type of their fields back to the args.
            for reference in references:
                resolved_fields = [self.resolve(f) for f in get_relation_fields(reference, relation.name)]
                for field_type, arg_type, arg in zip(resolved_fields, resolved_args, task.args):
                    if bt.is_abstract(arg_type) and isinstance(arg, mm.Var):
                        self.add_resolved_type(arg, field_type)
        else:
            if isinstance(references, mm.Overload):
                # we resolved an overload, use its types
                types = list(references.types)
            else:
                # we resolved to a single concrete relation that contains type vars or needs
                # number specialization, so use that relation's field types
                types = list([self.resolve(f) for f in relation.fields])

            # if our overload preserves types, we check to see if there's a preserved output
            # type given the inputs and if so, shadow the field type with the preserved type.
            # this is only attempted if all input types match the field types, i.e. no
            # conversions are needed
            resolved_fields = types
            if bt.is_function(relation) and len(set(resolved_fields)) == 1 and not relation in self.NON_TYPE_PRESERVERS and\
                all(type_matches(arg_type, field_type) for arg_type, field_type, field in zip(resolved_args, types, relation.fields) if field.input):

                input_types = set([arg_type for field, arg_type
                                    in zip(relation.fields, resolved_args) if field.input])
                if out_type := self.try_preserve_type(input_types):
                    resolved_fields = [field_type if field.input else out_type
                                        for field, field_type in zip(relation.fields, types)]


            # eq is special cased because we don't want to specialize a number for it, as it
            # can be just comparing numbers of different types.
            if relation != b.core.eq and (b.core.Number in types or (b.core.TypeVar in types and any(bt.is_number(t) for t in resolved_args))):
                # this relation contains generic numbers or typevars bound to numbers, so
                # find which specific type of number to use given the arguments being passed
                number, resolved_fields = self.specialize_number(relation, resolved_fields, resolved_args)
                self.resolved_number[task.id] = number
                # push the specialized number to the outputs
                for field, field_type, arg in zip(relation.fields, resolved_fields, task.args):
                    if not field.input and isinstance(arg, mm.Var):
                        self.add_resolved_type(arg, field_type)

            elif b.core.TypeVar in types:
                # this relation contains type vars, so we have to make sure that all args
                # bound to the same type var are consistent

                # find which arg is bound to the type var and check that they are all consistent
                typevar_type = None
                for arg_type, field_type in zip(resolved_args, types):
                    if field_type == b.core.TypeVar:
                        if typevar_type is None:
                            typevar_type = arg_type
                        else:
                            typevar_type = merge_types(typevar_type, arg_type)
                assert typevar_type is not None

                # compute the final arg types by replacing type vars with the typevar_type
                computed_arg_types = []
                for arg_type, field_type in zip(resolved_args, types):
                    # check that the arg type matches the type var type
                    if not type_matches(typevar_type, arg_type) and not conversion_allowed(arg_type, typevar_type):
                        self.type_mismatch(task, typevar_type, arg_type)
                    if field_type == b.core.TypeVar:
                        computed_arg_types.append(typevar_type)
                    else:
                        computed_arg_types.append(arg_type)

                # if no mismatches were found, propagate the computed arg types back to the args
                if len(computed_arg_types) == len(task.args):
                    for computed_type, arg, arg_type in zip(computed_arg_types, task.args, resolved_args):
                        # TODO: we could allow for non-numeric/string literals because this means
                        # the literal is being used as a value type, but the backend emitters
                        # would have to deal with that.
                        if isinstance(arg, mm.Var) or (isinstance(arg, mm.Literal) and (bt.is_numeric(computed_type) or computed_type == b.core.String)):
                            self.add_resolved_type(arg, computed_type)

            else:
                # no typevar or number specialization shenanigans, just propagate field types to args
                for field, field_type, arg, arg_type in zip(relation.fields, resolved_fields, task.args, resolved_args):
                    # add a resolved type if we learned more about the arg's type
                    if isinstance(arg, mm.Var) and (bt.is_abstract(arg_type) or bt.extends(field_type, arg_type)):
                        self.add_resolved_type(arg, field_type)
                    elif not type_matches(arg_type, field_type) and not conversion_allowed(arg_type, field_type):
                        self.type_mismatch(task, field_type, arg_type)


    # we try to preserve types for relations that are functions (i.e. potentially multiple
    # input but a single output) and where all types are the same. However, there are some
    # exceptions to this rule, e.g. range() always returns int regardless of whether the
    # input type is an int or a value type that extends int.
    NON_TYPE_PRESERVERS = [
        b.common.range
    ]

    def try_preserve_type(self, types:set[mm.Type]) -> Optional[mm.Type]:
        # we keep the input type as the output type if either all inputs are the exact same
        # type or there's one nominal and its base primitive type, e.g. USD + Decimal
        if len(types) == 1:
            return next(iter(types))
        if len(types) == 2:
            t1, t2 = types
            t1_base = bt.get_primitive_supertype(t1)
            t2_base = bt.get_primitive_supertype(t2)
            if t1_base is None or t2_base is None:
                base_equivalent = type_matches(t1, t2, accept_expected_super_types=True)
            else:
                base_equivalent = type_matches(t1_base, t2_base)
            if base_equivalent:
                # as long as one of the types is a base primitive, we can use the
                # other type as final preserved type
                if bt.is_primitive(t1):
                    return t2
                elif bt.is_primitive(t2):
                    return t1
        return None

    def specialize_number(self, op, field_types:list[mm.Type], arg_types:list[mm.Type]) -> Tuple[mm.NumberType, list[mm.Type]]:
        """
        Find the number type to use for an overload that has Number in its field_types, and
        which is being referred to with these arg_types.

        Return a tuple where the first element is the specialized number type, and the second
        element is a new list that contains the same types as field_types but with
        Number replaced by this specialized number.
        """
        # special case a few operators according to Snowflake's rules in
        # https://docs.snowflake.com/en/sql-reference/operators-arithmetic#scale-and-precision-in-arithmetic-operations
        if op == b.core.div:
            # see https://docs.snowflake.com/en/sql-reference/operators-arithmetic#division
            numerator, denominator = get_number_type(arg_types[0]), get_number_type(arg_types[1])
            s = max(numerator.scale, min(numerator.scale + 6, 12))
            number = mm.NumberType(name=f"Number(38,{s})", precision=38, scale=s, source=op.source, super_types=(b.core.Numeric,))
            return number, [numerator, denominator, number]
        elif op == b.core.mul:
            # see https://docs.snowflake.com/en/sql-reference/operators-arithmetic#multiplication
            t1, t2 = get_number_type(arg_types[0]), get_number_type(arg_types[1])
            S1 = t1.scale
            S2 = t2.scale
            s = min(S1 + S2, max(S1, S2, 12))
            number = mm.NumberType(name=f"Number(38,{s})", precision=38, scale=s, source=op.source, super_types=(b.core.Numeric,))
            return number, [t1, t2, number]
        elif op == b.aggregates.avg:
            # TODO!! - implement proper avg specialization
            pass

        # fall back to the the current specialization policy, which is to select the number
        # with largest scale and, if there multiple with the largest scale, the one with the
        # largest precision. This is safe because when converting a number to the
        # specialized number, we never truncate fractional digits (because we selected the
        # largest scale) and, if the non-fractional digits are too large to fit the
        # specialized number, we will have a runtime overflow, which should alert the user
        # of the problem.
        number = None
        for arg_type in arg_types:
            x = bt.get_number_supertype(arg_type)
            if isinstance(x, mm.NumberType):
                if number is None or x.scale > number.scale or (x.scale == number.scale and x.precision > number.precision):
                    number = x
        assert(number is not None)
        return number, [number if bt.is_number(t) else t for t in field_types]


    #--------------------------------------------------
    # Display
    #--------------------------------------------------

    # draw the network as a mermaid graph for the debugger
    def to_mermaid(self, max_edges=500) -> str:
        # add links for edges while collecting nodes
        nodes = OrderedSet()
        link_strs = []
        # edges
        for src, dsts in self.edges.items():
            nodes.add(src)
            for dst in dsts:
                if len(link_strs) > max_edges:
                    break
                nodes.add(dst)
                link_strs.append(f"n{src.id} --> n{dst.id}")
            if len(link_strs) > max_edges:
                break

        def type_span(t:mm.Type) -> str:
            type_str = t.name if isinstance(t, mm.ScalarType) else str(t)
            return f"<span style='color:cyan;'>{type_str.strip()}</span>"

        def reference_span(rel:mm.Relation, arg_types:list[mm.Type]) -> str:
            args = []
            for field, arg_type in zip(rel.fields, arg_types):
                field_type = self.resolve(field)
                if not type_matches(arg_type, field_type) and not conversion_allowed(arg_type, field_type) and not bt.is_abstract(field_type):
                    args.append(f"<span style='color:yellow;'>{str(arg_type).strip()} -> {str(field_type).strip()}</span>")
                elif isinstance(arg_type, mm.UnionType):
                    args.append(type_span(field_type))
                else:
                    args.append(type_span(arg_type))
            return f'{rel.name}({", ".join(args)})'

        resolved = self.resolved_types
        node_strs = []
        for node in nodes:
            klass = ""
            if isinstance(node, mm.Var):
                ir_type = resolved.get(node) or self.resolve(node)
                type_str = type_span(ir_type)
                label = f'(["{node.name}:{type_str}"])'
            elif isinstance(node, mm.Literal):
                ir_type = resolved.get(node) or self.resolve(node)
                type_str = type_span(ir_type)
                klass = ":::literal"
                label = f'(["{node.value}: {type_str}"])'
            elif isinstance(node, mm.Field):
                ir_type = resolved.get(node) or self.resolve(node)
                type_str = type_span(ir_type)
                klass = ":::field"
                rel = node._relation
                if rel is not None:
                    rel = str(node._relation)
                    label = f'[/"{node.name}:{type_str}\nfrom {rel}"\\]'
                else:
                    label = f'[/"{node.name}:\n{type_str}"\\]'
            elif isinstance(node, (mm.Lookup, mm.Aggregate, mm.Update)):
                arg_types = [self.resolve(arg) for arg in node.args]
                if node.id in self.resolved_placeholder:
                    overloads = self.resolved_placeholder[node.id]
                    content = "<br/>".join([reference_span(o, arg_types) for o in overloads])
                else:
                    content = reference_span(get_relation(node), arg_types)
                if isinstance(node, mm.Update):
                    klass = ":::update"
                    label = f'{{{{"{content}"}}}}'
                else:
                    klass = ":::reference"
                    label = f'[/"{content}"/]'
            else:
                raise NotImplementedError(f"Unknown node type: {type(node)}")
            if self.has_errors(node):
                klass = ":::error"
            node_strs.append(f'n{node.id}{label}{klass}')

        node_str = "\n                ".join(node_strs)
        link_str = "\n                ".join(link_strs)
        template = f"""
            %%{{init: {{'theme':'dark', 'flowchart':{{'useMaxWidth':false, 'htmlLabels': true}}}}}}%%
            flowchart TD
                linkStyle default stroke:#666
                classDef field fill:#245,stroke:#478
                classDef update fill:#245,stroke:#478
                classDef literal fill:#452,stroke:#784
                classDef error fill:#624,stroke:#945,color:#f9a
                classDef default stroke:#444,stroke-width:2px, font-size:12px

                %% nodes
                {node_str}

                %% edges
                {link_str}
        """
        return template

#--------------------------------------------------
# Analyzer
#--------------------------------------------------

class Analyzer(Walker):
    """ Walks the metamodel and builds the propagation network. """

    def __init__(self, net:PropagationNetwork):
        super().__init__()
        self.net = net

    def analyze(self, node: mm.Node):
        self(node)

    # TODO - ignoring requires for now because the typing of constraints seems incorrect
    def enter_require(self, require: mm.Require):
        return NO_WALK

    def compute_potential_targets(self, relation: mm.Relation):
        # register potential targets for placeholders
        if bt.is_placeholder(relation):
            self.net.potential_targets[relation] = get_potential_targets(self.net.model, relation)

    #--------------------------------------------------
    # Walk Update
    #--------------------------------------------------

    def update(self, node: mm.Update):
        rel = node.relation
        self.compute_potential_targets(rel)

        # arg is flowing into a field
        for arg, field in zip(node.args, rel.fields):
            if isinstance(arg, (mm.Var, mm.Literal)):
                self.net.add_edge(arg, node)

    #--------------------------------------------------
    # Walk Lookups + Aggregates
    #--------------------------------------------------

    def lookup(self, task: mm.Lookup):
        self.compute_potential_targets(task.relation)
        self.visit_rel_op(task)

    def aggregate(self, task: mm.Aggregate):
        self.visit_rel_op(task)

    def visit_rel_op(self, task: mm.Lookup|mm.Aggregate):
        relation = get_relation(task)

        is_placeholder = bt.is_placeholder(relation)
        for field, arg in zip(relation.fields, task.args):
            if isinstance(arg, (mm.Var, mm.Literal)):
                if field.input:
                    # we need to resolve all inputs before resolving the relation
                    self.net.add_edge(arg, task)
                else:
                    # placeholders also need the output to be resolved
                    if is_placeholder:
                        self.net.add_edge(arg, task)
                    else:
                        # args bound to outputs can be resolved after
                        self.net.add_edge(task, arg)
                        # if the field is abstract, it needs to be resolved before we can
                        # resolve the task.
                        if bt.is_abstract(self.net.resolve(field)):
                            self.net.add_edge(field, task)


#--------------------------------------------------
# Replacer
#--------------------------------------------------

# Once we've pushed all the types through the network, we need to replace the types of
# fields and vars that we may have discovered. We also need to replace placeholder lookups
# with the chosen relations and do any conversions that are needed.
class Replacer(Rewriter):
    def __init__(self, net:PropagationNetwork):
        super().__init__()
        self.net = net

    def rewrite(self, model: T) -> T:
        return self(model) # type: ignore

    def logical(self, logical: mm.Logical):
        if len(logical.body) == 0:
            return logical
        # inline logicals that are just there to group other nodes during rewrite
        body = []
        for child in logical.body:
            if isinstance(child, mm.Logical) and not child.optional and not child.scope:
                body.extend(child.body)
            else:
                body.append(child)
        return logical.mut(body = tuple(body))

    #--------------------------------------------------
    # Rewriter handlers
    #--------------------------------------------------

    def field(self, node: mm.Field):
        # TODO - this is only modifying the relation in the model, but then we have a new
        # relation there, which is different than the object referenced by tasks.
        if node in self.net.resolved_types:
            return mm.Field(node.name, self.net.resolved_types[node], node.input, _relation = node._relation)
        return node

    def var(self, node: mm.Var):
        if node in self.net.resolved_types:
            return mm.Var(self.net.resolved_types[node], node.name)
        return node

    def literal(self, node: mm.Literal):
        if node in self.net.resolved_types:
            return mm.Literal(self.net.resolved_types[node], node.value)
        return node

    def update(self, node: mm.Update):
        return self.convert_arguments(node, node.relation)

    def lookup(self, node: mm.Lookup):
        # We need to handle eq specially because its arguments can be converted symmetrically
        if node.relation == b.core.eq:
            return self.visit_eq_lookup(node)

        args = types = None
        if node.id in self.net.resolved_placeholder:
            # placeholder resolved to multiple relations
            resolved_relations = self.net.resolved_placeholder[node.id]
            args = get_lookup_args(node, resolved_relations[0])
            types = [f.type for f in resolved_relations[0].fields]
        elif node.id in self.net.resolved_overload:
            # overload resolved to a specific relation
            resolved_relations = [node.relation]
            types = self.net.resolved_overload[node.id].types
        else:
            # single relation
            resolved_relations = [node.relation]

        if len(resolved_relations) == 1:
            # single relation, just convert arguments
            x = self.convert_arguments(node, resolved_relations[0], args, types)
            if isinstance(x, mm.Logical) and len(x.body) == 1:
                return x.body[0]
            else:
                return x

        # multiple relations, create a union
        branches:list = []
        for target in resolved_relations:
            args = get_lookup_args(node, target)
            types = [f.type for f in get_relation_fields(resolved_relations[0], node.relation.name)]
            # adding this logical to avoid issues in the old backend
            branches.append(mm.Logical((self.convert_arguments(node, target, args, types=types, force_copy=True),)))
        return mm.Union(tuple(branches))

    def convert_arguments(self, node: mm.Lookup|mm.Update, relation: mm.Relation, args: Iterable[mm.Value]|None=None, types: Iterable[mm.Type]|None=None, force_copy=False) -> mm.Logical|mm.Lookup|mm.Update:
        """ This node was resolved to target this relation using these args, which should
        have these types. Convert any arguments as needed and return a new node with the
        proper relation and converted args. If multiple conversions are needed, return a
        logical that contains all the conversion tasks plus the final node. """
        args = args or node.args
        types = types or [self.net.resolve(f) for f in relation.fields]
        number_type = self.net.resolved_number.get(node.id)
        is_function = bt.is_function(relation)
        tasks = []
        final_args = []
        for arg, field, field_type in zip(args, relation.fields, types):
            if isinstance(arg, (mm.Var, mm.Literal)) and (not is_function or field.input):
                arg_type = to_type(arg)
                if number_type and bt.is_number(arg_type) and not arg_type == b.core.ScaledNumber:
                    field_type = number_type
                # the typer previously made sure that this should be valid so a type mismatch
                # means we need to convert
                if not type_matches(arg_type, field_type):
                    final_args.append(convert(arg, field_type, tasks))
                else:
                    final_args.append(arg)
            else:
                final_args.append(arg)
        # add the original node with the proper target relation and converted args;
        # here we want to use mut because we keep information about nodes based on their ids
        # (e.g. we store resolved types based on ids), so we want to keep the same id. But
        # when a lookup is being converted as part of a union (i.e. multiple targets for a
        # placeholder), we need to create nodes with new ids to avoid conflicts.
        if force_copy:
            tasks.append(node.replace(relation = relation, args = tuple(final_args)))
        else:
            tasks.append(node.mut(relation = relation, args = tuple(final_args)))
        # if we need conversion tasks, wrap in a logical
        if len(tasks) == 1:
            return tasks[0]
        return mm.Logical(tuple(tasks))

    def visit_eq_lookup(self, node: mm.Lookup):
        (left, right) = node.args
        left_type = to_type(left)
        right_type = to_type(right)

        if type_matches(left_type, right_type):
            return node

        assert isinstance(left, (mm.Var, mm.Literal)) and isinstance(right, (mm.Var, mm.Literal))
        final_args = []
        tasks = []
        if conversion_allowed(left_type, right_type):
            final_args = [convert(left, right_type, tasks), right]
        elif conversion_allowed(right_type, left_type):
            final_args = [left, convert(right, left_type, tasks)]
        else:
            # this type mismatch was reported during propagation, so just return the node
            return node

        tasks.append(mm.Lookup(b.core.eq, tuple(final_args)))
        return mm.Logical(tuple(tasks))

#--------------------------------------------------
# Helpers
#--------------------------------------------------

def get_relation(node: mm.Lookup|mm.Update|mm.Aggregate) -> mm.Relation:
    if isinstance(node, mm.Aggregate):
        return node.aggregation
    return node.relation

def get_name(type: mm.Type) -> str:
    if isinstance(type, mm.ScalarType):
        return type.name
    elif isinstance(type, mm.UnionType):
        return '|'.join([get_name(t) for t in type.types])
    elif isinstance(type, mm.ListType):
        return f'List[{get_name(type.element_type)}]'
    elif isinstance(type, mm.TupleType):
        return f'Tuple[{", ".join([get_name(t) for t in type.element_types])}]'
    else:
        raise TypeError(f"Unknown type: {type}")

#--------------------------------------------------
# Type and Relation helpers
#--------------------------------------------------

def get_update_fields(arg, update: mm.Update) -> Iterable[mm.Field]:
    """ Get the fields of the relation being updated by this arg. Note that an arg can be
    bound to multiple fields at the same time. """
    if arg in update.args:
        for x, field in zip(update.args, update.relation.fields):
            if arg == x:
                yield field
    return []

def get_relation_fields(relation: mm.Relation, name: str) -> Iterable[mm.Field]:
    """ Get the fields of this relation, potentially reordered to match the reading with the given name."""
    if name == relation.name:
        return relation.fields
    for reading in relation.readings:
        if reading.name == name:
            # reorder the fields to match the correct ordering
            fields = []
            for idx in reading.field_order:
                fields.append(relation.fields[idx])
            return fields
    return []

def get_lookup_args(node: mm.Lookup, target: mm.Relation):
    """ Get the args of this lookup, potentially reordered to match the reading with the given name."""
    for reading in target.readings:
        if reading.name == node.relation.name:
            # reorder the args to match the correct ordering
            args = []
            for idx in reading.field_order:
                args.append(node.args[idx])
            return args
    return node.args

def get_number_type(t: mm.Type) -> mm.NumberType:
    # Get a number type from the given type, if it is Number return the default number
    x = bt.get_number_supertype(t)
    if isinstance(x, mm.NumberType):
        return x
    return b.core.DefaultNumber

def is_potential_target(placeholder: mm.Relation, target: mm.Relation) -> bool:
    """ Whether this target is matches the placeholder signature and, thus, can be a potential target. """
    if placeholder != target and len(placeholder.fields) == len(target.fields) and not bt.is_placeholder(target):
        return placeholder.name == target.name or any(placeholder.name == reading.name for reading in target.readings)
    return False

def get_potential_targets(model: mm.Model, placeholder: mm.Relation) -> list[mm.Relation]:
    """ Get all potential target relations in the model that match the placeholder signature. """
    return list(filter(lambda r: is_potential_target(placeholder, r), model.relations))

def to_type(value: mm.Value|mm.Field|mm.Literal) -> mm.Type:
    if isinstance(value, (mm.Var, mm.Literal)):
        return value.type

    if isinstance(value, mm.Type):
        return b.core.Type

    if isinstance(value, mm.Field):
        return b.core.Field

    if isinstance(value, tuple):
        return mm.TupleType(element_types=tuple(to_type(v) for v in value))

    raise TypeError(f"Cannot determine IR type for value: {value} of type {type(value).__name__}")

def convert(value: mm.Var|mm.Literal, to_type: mm.Type, tasks: list[mm.Task]) -> mm.Value:
    # if the arg is a literal, we can just change its type
    # TODO - we may want to check that the value is actually convertible
    if isinstance(value, mm.Literal):
        return mm.Literal(to_type, value.value)

    # otherise we need to add a cast
    name = sanitize(value.name + "_" + get_name(to_type))
    to_type_base = bt.get_primitive_supertype(to_type) or to_type
    new_value = mm.Var(to_type_base, name)
    tasks.append(mm.Lookup(b.core.cast, (to_type_base, value, new_value)))
    return new_value

def conversion_allowed(from_type: mm.Type, to_type: mm.Type) -> bool:
    # value type conversion is allowed only if the value types are related by inheritance
    if bt.is_value_type(from_type) and bt.is_value_type(to_type) and not bt.extends(from_type, to_type):
        return False
    # value type conversion is allowed
    x = bt.get_primitive_supertype(from_type)
    y = bt.get_primitive_supertype(to_type)
    if x and y and (x != from_type or y != to_type) and conversion_allowed(x, y):
        return True

    # numbers can be converted to floats
    if bt.is_numeric(from_type) and to_type == b.core.Float:
        return True

    # a number can be converted to another number of larger scale
    if isinstance(from_type, mm.NumberType) and isinstance(to_type, mm.NumberType):
        if to_type.scale >= from_type.scale:
            return True

    if from_type == b.core.Number and isinstance(to_type, mm.NumberType):
        return True

    return False

def type_matches(actual: mm.Type, expected: mm.Type, accept_expected_super_types=False) -> bool:
    """
    True iff we can use a value of the actual type when expecting the expected type, without
    conversions.

    Any super-type of `actual` can match `expected`. For example if we expect a `Person`, we
    can use an `Employee` if `Employee < Person`.

    In general, the other way around is not true: if we expect an `Employee` we cannot use a
    `Person` instead.

    However, when the relation is a Type (previsously known as "population relations"), it
    is valid to provide sub-types of the expected type. For example, `Employee(Person)` is
    a valid way to check that `Person` is an `Employee` on a `Lookup`, or to assert that a
    particular `Person` is an `Employee` on an `Update`.
    """
    # exact match
    if actual == expected:
        return True

    # any matches anything
    if actual == b.core.Any or expected == b.core.Any:
        return True

    # type vars match anything
    if expected == b.core.TypeVar:
        return True

    # if an entity type var or any entity is expected, it matches any actual entity type
    if (expected == b.core.EntityTypeVar or bt.extends(expected, b.core.AnyEntity)) and not bt.is_primitive(actual):
        return True

    # the abstract Number type and the number type variable match any number type
    if (expected == b.core.Number) and bt.is_number(actual):
        return True

    if (expected == b.core.Numeric) and bt.is_numeric(actual):
        return True

    # different value types never match
    if bt.is_value_type(actual) and bt.is_value_type(expected) and not bt.extends(actual, expected):
        return False

    # if actual is scalar, any of its parents may match the expected type
    if isinstance(actual, mm.ScalarType) and any([type_matches(parent, expected) for parent in actual.super_types]):
        return True

    # if expected is a value type or this is a check for a type relation, any of the expected type's parents may match the actual type
    if (accept_expected_super_types or bt.is_value_type(expected)) and isinstance(expected, mm.ScalarType) and any([type_matches(actual, parent, accept_expected_super_types) for parent in expected.super_types]):
        return True

    # if we expect a union, the actual can match any of its types
    if isinstance(expected, mm.UnionType):
        for t in expected.types:
            if type_matches(t, actual, accept_expected_super_types):
                return True

    # if actual is a union, every one of its types must match the expected type
    # if isinstance(actual, mm.UnionType):
    #     for t in actual.types:
    #         if not type_matches(t, expected, accept_expected_super_types):
    #             return False
    #     return True
    # TODO - we have to distinguish between when we are checking that a specific arg matches
    # a relation vs when we are selecting relations for the placeholders; then we have to
    # decide between the above and this.
    if isinstance(actual, mm.UnionType):
        for t in actual.types:
            if type_matches(t, expected, accept_expected_super_types):
                return True

    # a list type matches if their element types match
    if isinstance(actual, (mm.ListType, mm.TupleType)) and isinstance(expected, mm.ListType):
        if isinstance(actual, mm.TupleType):
            for et in actual.element_types:
                if not type_matches(et, expected.element_type):
                    return False
            return True
        return type_matches(actual.element_type, expected.element_type)

    # a tuple types match if any of all their types match
    if isinstance(actual, mm.TupleType) and isinstance(expected, mm.TupleType):
        return all([type_matches(ae, ee) for ae, ee in zip(actual.element_types, expected.element_types)])

    # accept tuples with a single element type to match a list with that type
    if isinstance(actual, mm.TupleType) and isinstance(expected, mm.ListType):
        return type_matches(actual.element_types[0], expected.element_type)

    # otherwise no match
    return False

def merge_numeric_types(type1: mm.Type, type2: mm.Type) -> Optional[mm.Type]:
    # if one of them is the abstract Number type, pick the other
    if type1 == b.core.Number and isinstance(type2, mm.NumberType):
        return type2
    if type2 == b.core.Number and isinstance(type1, mm.NumberType):
        return type1

    # if both are number types, pick the one with larger scale/precision
    if isinstance(type1, mm.NumberType) and isinstance(type2, mm.NumberType):
        if type1.scale > type2.scale or (type1.scale == type2.scale and type1.precision > type2.precision):
            return type1
        else:
            return type2

    # if we are overriding a number with a float, pick float
    if isinstance(type1, mm.NumberType) and type2 == b.core.Float:
        return type2
    if isinstance(type2, mm.NumberType) and type1 == b.core.Float:
        return type1

    return None

def merge_types(type1: mm.Type, type2: mm.Type) -> mm.Type:
    if type1 == type2:
        return type1
    if bt.is_type_var(type1):
        return type2
    if bt.is_type_var(type2):
        return type1

    types_to_process = [type1, type2]

    numeric_merge = merge_numeric_types(type1, type2)
    if numeric_merge is not None:
        return numeric_merge

    # if one extends the other, pick the most specific one
    if bt.extends(type1, type2):
        return type1
    if bt.extends(type2, type1):
        return type2

    # give precedence to nominal types (e.g. merging USD(decimal) with decimal gives USD(decimal))
    base_primitive_type1 = bt.get_primitive_supertype(type1)
    base_primitive_type2 = bt.get_primitive_supertype(type2)
    if base_primitive_type1 == base_primitive_type2:
        if bt.is_primitive(type1):
            return type2
        elif bt.is_primitive(type2):
            return type1

    if base_primitive_type1 and base_primitive_type2:
        numeric_merge = merge_numeric_types(base_primitive_type1, base_primitive_type2)
        if numeric_merge == base_primitive_type1:
            return type1
        elif numeric_merge == base_primitive_type2:
            return type2

    combined = OrderedSet()
    # Iterative flattening of union types
    while types_to_process:
        t = types_to_process.pop()
        if isinstance(t, mm.UnionType):
            types_to_process.extend(t.types)
        else:
            combined.add(t)

    # If we have multiple types and Any or AnyEntity is one of them, remove Any
    if len(combined) > 1:
        if b.core.Any in combined:
            combined.remove(b.core.Any)
        if b.core.AnyEntity in combined:
            combined.remove(b.core.AnyEntity)

    # If we still have multiple types, make sure supertypes are removed to keep only the
    # most specific types
    if len(combined) > 1:
        to_remove = set()
        for t1 in combined:
            for t2 in combined:
                if t1 != t2 and bt.extends(t1, t2):
                    to_remove.add(t2)
        for r in to_remove:
            combined.remove(r)

    # Return single type or create a union
    return next(iter(combined)) if len(combined) == 1 else mm.UnionType(types=tuple(combined))

def invalid_type(type:mm.Type) -> bool:
    if isinstance(type, mm.UnionType):
        # if there are multiple primitives, or a primitive and a non-primitive
        # then we have an invalid type
        if len(type.types) > 1:
            return any([bt.is_primitive(t) for t in type.types])
    return False


#--------------------------------------------------
# Type Errors
#--------------------------------------------------

@dataclass
class TyperError():
    node: mm.Node

    def report(self):
        err(self.name(), self.message(), self.parts())

    def name(self) -> str:
        return type(self).__name__

    def message(self) -> str:
        raise NotImplementedError()

    def parts(self) -> list[Part]:
        if self.node.source is None:
            return [str(self.node)]
        return [str(self.node), Source(self.node.source)]

@dataclass
class TypeMismatch(TyperError):
    expected: mm.Type
    actual: mm.Type

    def message(self) -> str:
        return f"Expected {get_name(self.expected)}, got {get_name(self.actual)}"

@dataclass
class InvalidType(TyperError):
    type: mm.Type

    def message(self) -> str:
        return f"Incompatible types infered: {get_name(self.type)}"

@dataclass
class UnresolvedOverload(TyperError):
    arg_types: list[mm.Type]

    def message(self) -> str:
        assert isinstance(self.node, (mm.Lookup, mm.Update, mm.Aggregate))
        rel = get_relation(self.node)
        types = ', '.join([get_name(t) for t in self.arg_types])
        return f"Unresolved overload: {rel.name}({types})"

@dataclass
class UnresolvedType(TyperError):

    def message(self) -> str:
        return "Unable to determine concrete type."
