import einx._src.namedtensor.stage1 as stage1
import einx._src.namedtensor.stage2 as stage2
import einx._src.namedtensor.stage3 as stage3
from einx._src.namedtensor import solve
from einx._src.namedtensor import NamedTensor
from einx._src.frontend.types import Tensor
import numpy as np
import functools
import types
import einx._src.tracer as tracer
import einx._src.adapter as adapter
import uuid
from functools import partial
from einx._src.frontend.errors import SemanticError
from collections import defaultdict
from einx._src.util.functools import use_name_of


def _tensor_to_string(tensor):
    if tensor.shape is None:
        return "Tensor factory"
    else:
        return f"Tensor with shape {tensor.shape}"


def _kwarg_to_string(v):
    s = str(v)
    if len(s) < 20:
        return s
    else:
        return "..."


class Invocation:
    def __init__(self, expression, name, tensors, kwargs):
        self.expression = expression
        self.name = name
        self.tensors = tensors
        self.kwargs = kwargs

        from einx._src.namedtensor import ExpressionIndicator

        self.indicator = ExpressionIndicator(expression)

    def to_call_signature_string(self):
        message = "The operation was called with the following arguments:\n"
        for i, tensor in enumerate(self.tensors):
            message += f"  - Positional argument #{i + 1}: {_tensor_to_string(tensor)}\n"
        for k, v in self.kwargs.items():
            message += f"  - Keyword argument '{k}': {_kwarg_to_string(v)}\n"
        return message


def _to_el_expr(expr):
    if isinstance(expr, list):
        newlist = [_to_el_expr(e) for e in expr]
        newlist = [expr for expr in newlist if expr.ndim != 0]
        return newlist
    elif isinstance(expr, stage1.List):
        return stage1.List.create(_to_el_expr(expr.children), expr.begin_pos, expr.end_pos)
    elif isinstance(expr, stage1.Axis):
        return stage1.List([])
    elif isinstance(expr, stage1.FlattenedAxis):
        inner = _to_el_expr(expr.inner)
        if inner.ndim == 0:
            return stage1.List([])
        else:
            return stage1.FlattenedAxis.create(inner, expr.begin_pos, expr.end_pos)
    elif isinstance(expr, stage1.ConcatenatedAxis):
        # ConcatenatedAxis cannot be used with brackets
        assert not any(stage1.is_in_brackets(c) for c in expr.nodes())
        return stage1.List([])
    elif isinstance(expr, stage1.Brackets):
        return expr.inner.__deepcopy__()
    elif isinstance(expr, stage1.Ellipsis):
        return stage1.Ellipsis.create(_to_el_expr(expr.inner), expr.begin_pos, expr.end_pos, expr.ellipsis_id)
    elif isinstance(expr, stage1.Op):
        return stage1.Op([_to_el_expr(c) for c in expr.children], expr.begin_pos, expr.end_pos)
    elif isinstance(expr, stage1.Args):
        return stage1.Args([_to_el_expr(c) for c in expr.children], expr.begin_pos, expr.end_pos)
    else:
        raise TypeError(f"Invalid expression type {type(expr)}")


def _parse_op(description, el_op, invocation, allow_concat=False, implicit_output=None, mark_reduced_axes=False, allow_duplicate_el_axes=True, keepdims=False):
    if not isinstance(description, str):
        raise ValueError("The operation description must be a string.")
    if keepdims is None:
        keepdims = False
    op = stage1.parse_op(description)

    if not allow_concat:
        # Disallow concatenation
        if any(isinstance(expr, stage1.ConcatenatedAxis) for expr in op.nodes()):
            raise SemanticError(
                invocation=invocation,
                pos=invocation.indicator.get_pos_for_concat(op),
                message="The concatenation operator (+) is not allowed in this function.\n%EXPR%",
            )

    exprs_in = op.children[0].children

    # Get signature that is expected by the elementary operation
    el_subop = _to_el_expr(op)
    if el_op is None:
        # No signature is explicitly provided -> extract the signature from the expression itself
        if len(el_subop.children) != 2:
            raise SemanticError(invocation=invocation, message="The operation expects an output expression, but no '->' was found.\n%EXPR%")
        el_op = el_subop
    elif isinstance(el_op, str):
        el_op = stage1.parse_op(el_op)
    elif callable(el_op):
        el_op = el_op(el_subop)
        el_op = stage1.parse_op(el_op)
    else:
        raise ValueError(f"Invalid type for el_op: {type(el_op)}.")
    assert len(el_op.children) == 2

    # Check number of input and output expressions
    if len(el_op.children[0].children) != len(el_subop.children[0].children):
        raise SemanticError(
            invocation=invocation,
            message=f"The operation expects {len(el_op.children[0].children)} input expression(s), but found {len(el_subop.children[0].children)}.\n%EXPR%",
        )
    if len(el_subop.children) > 1 and len(el_op.children[1].children) != len(el_subop.children[1].children):
        raise SemanticError(
            invocation=invocation,
            message=f"The operation expects {len(el_op.children[1].children)} output expression(s), but found {len(el_subop.children[1].children)}.\n%EXPR%",
        )

    # Determin output expression if it is not given
    if len(op.children) == 1:
        if implicit_output == "bijective":
            if len(el_op.children[0].children) == 1 and len(el_op.children[1].children) == 1 and el_op.children[0] == el_op.children[1]:
                # single input == single output
                # -> Use input expression as output expression
                exprs_out = [expr_in.__deepcopy__() for expr_in in exprs_in]
            elif len(el_op.children[0].children) == 1 and len(el_op.children[1].children) == 1 and el_op.children[1].children[0].ndim == 0:
                # Single input and single scalar output
                # -> Remove brackets from input expression for output expression
                if keepdims:

                    def _replace(expr):
                        if isinstance(expr, stage1.Brackets):
                            return stage1.FlattenedAxis.create(stage1.List.create([]))

                    exprs_out = [stage1.map(expr_in, _replace, include_children=True) for expr_in in exprs_in]
                else:
                    exprs_out = [stage1.remove(expr_in, stage1.Brackets, keep_children=False) for expr_in in exprs_in]
            elif all(c.ndim == 0 for c in list(el_op.children[0].children) + list(el_op.children[1].children)):
                # Scalar inputs and outputs
                # -> Find superset expression
                if len(exprs_in) == 1:
                    # Only one input -> use it as output
                    exprs_out = [exprs_in[0].__deepcopy__()]
                else:
                    # Use one of the input expression if it contains the axis names of
                    # all others and if this choice is unique
                    in_axis_names = [{expr.name for expr in root.nodes() if isinstance(expr, stage1.Axis) and expr.value != 1} for root in exprs_in]

                    valid_parents = set()
                    for i, parent in enumerate(in_axis_names):
                        for j, child in enumerate(in_axis_names):
                            if i != j and not child.issubset(parent):
                                break
                        else:
                            # Found valid parent
                            valid_parents.add(exprs_in[i])

                    if len(valid_parents) != 1:
                        raise SemanticError(
                            invocation=invocation,
                            message="The output expression is missing in this operation. einx allows implicitly determining the output expression"
                            " in this operation according to the rule: If one of the input expressions contains the axis names of all other input"
                            " expressions (excluding 1s) and if this choice is unique, then this expression is used as output expression. However, no unique"
                            " input expression was found.\n%EXPR%",
                        )
                    exprs_out = [valid_parents.pop().__deepcopy__()]
            elif len(el_op.children[0].children) == 1 and len(el_op.children[1].children) == 1:
                # Single input and single output
                # -> Replace single input bracket with single output bracket
                def _to_output(expr):
                    bracket_num = len([e for e in expr.nodes() if isinstance(e, stage1.Brackets)])
                    assert bracket_num > 0
                    if bracket_num == 1:

                        def _replace(expr):
                            if isinstance(expr, stage1.Brackets):
                                return stage1.Brackets.create(stage1.Axis.create("output.axis"))

                        return stage1.map(expr, _replace, include_children=False)
                    else:
                        raise SemanticError(
                            invocation=invocation,
                            pos=invocation.indicator.get_pos_for_brackets(expr),
                            message=(
                                "The output expression is missing in this operation. The output expression can only be determined implicitly, "
                                "if the input expression contains exactly one usage of brackets ([]).\n%EXPR%"
                            ),
                        )

                exprs_out = [_to_output(expr_in) for expr_in in exprs_in]
            else:
                raise SemanticError(invocation=invocation, message="The operation expects an output expression, but no '->' was found.\n%EXPR%")
        elif isinstance(implicit_output, int):
            exprs_out = [exprs_in[implicit_output].__deepcopy__()]
        elif isinstance(implicit_output, tuple | list) and all(isinstance(i, int) for i in implicit_output):
            exprs_out = [exprs_in[i].__deepcopy__() for i in implicit_output]
        elif implicit_output is None:
            raise SemanticError(invocation=invocation, message="The operation expects an output expression, but no '->' was found.\n%EXPR%")
        else:
            raise ValueError(f"Invalid value for implicit_output: {implicit_output}")
    else:
        exprs_out = op.children[1].children
    op = stage1.Op([stage1.Args(exprs_in), stage1.Args(exprs_out)])
    el_subop = _to_el_expr(op)
    assert len(el_op.children[0].children) == len(el_subop.children[0].children)
    assert len(el_op.children[1].children) == len(el_subop.children[1].children)

    # Check bracket usage
    def _to_ordinal_str(i):
        if i == 0:
            return "1st"
        elif i == 1:
            return "2nd"
        elif i == 2:
            return "3rd"
        else:
            return f"{i + 1}th"

    def check(i, el_arg, el_subarg, arg, inoutput):
        if el_arg.ndim == 0 and el_subarg.ndim != 0:
            raise SemanticError(
                invocation=invocation,
                pos=invocation.indicator.get_pos_for_brackets(arg),
                message=f"Brackets ([]) are not allowed in the {_to_ordinal_str(i)} {inoutput} expression of this operation.\n%EXPR%",
            )
        elif el_arg.ndim != 0 and el_subarg.ndim == 0:
            raise SemanticError(
                invocation=invocation,
                pos=invocation.indicator.get_pos_for_exprs(arg),
                message=f"The {_to_ordinal_str(i)} {inoutput} expression of this operation requires brackets, but no brackets were found.\n%EXPR%",
            )

    for i, (el_arg, el_subarg, arg) in enumerate(zip(el_op.children[0].children, el_subop.children[0].children, op.children[0].children, strict=False)):
        check(i, el_arg, el_subarg, arg, "input")
    for i, (el_arg, el_subarg, arg) in enumerate(zip(el_op.children[1].children, el_subop.children[1].children, op.children[1].children, strict=False)):
        check(i, el_arg, el_subarg, arg, "output")

    marking_reduced_axes = mark_reduced_axes and not any(isinstance(expr, stage1.Brackets) for expr_in in exprs_in for expr in expr_in.nodes())
    if marking_reduced_axes:
        # Check that no axis name is used twice per input
        for expr_in in exprs_in:
            counts = defaultdict(lambda: 0)
            for axis in expr_in.nodes():
                if isinstance(axis, stage1.Axis):
                    counts[axis.name] += 1
            duplicates = [name for name, count in counts.items() if count > 1]
            if len(duplicates) > 0:
                raise SemanticError(
                    invocation=invocation,
                    pos=invocation.indicator.get_pos_for_axisnames([expr_in], duplicates),
                    message=(
                        "If no axes are marked with brackets in this operation, brackets are placed automatically around all axes that do not appear "
                        "in the output expression. Since this mode is used here, no axis name may appear more than once in the same input tensor (to avoid "
                        f"confusion with diagonal extraction). However, the following axis names are used more than once: {', '.join(duplicates)}.\n%EXPR%"
                    ),
                )

        assert len(exprs_out) == 1
        expr_out = exprs_out[0]

        # If no brackets appear in exprs_in, mark all axes that don't appear in expr_out.
        axes_names_out = {axis.name for axis in expr_out.nodes() if isinstance(axis, stage1.Axis)}

        def _mark(expr):
            if isinstance(expr, stage1.Axis) and expr.name not in axes_names_out:
                return True
            else:
                return False

        exprs_in = [stage1.map(expr_in, lambda expr: stage1.Brackets(expr) if _mark(expr) else None, include_children=False) for expr_in in exprs_in]

    # Check that no two vectorized axes in any (unconcatenated) output have the same name
    for expr_out in exprs_out:
        for expr_out in stage1.split_concatenated_axes(expr_out):
            axis_names = [expr.name for expr in expr_out.nodes() if isinstance(expr, stage1.Axis) and not stage1.is_in_brackets(expr)]
            if len(axis_names) != len(set(axis_names)):
                duplicates = {name for name in axis_names if axis_names.count(name) > 1}
                raise SemanticError(
                    invocation=invocation,
                    pos=invocation.indicator.get_pos_for_axisnames([expr_out], duplicates),
                    message="The output expression must not contain multiple vectorized axes with the same name (after splitting concatenated axes).\n%EXPR%",
                )

    if not allow_duplicate_el_axes:
        # Check that no two marked axes in any tensor have the same name
        for expr in exprs_in + exprs_out:
            axis_names = [expr.name for expr in expr.nodes() if isinstance(expr, stage1.Axis) and stage1.is_in_brackets(expr)]
            if len(axis_names) != len(set(axis_names)):
                duplicates = {name for name in axis_names if axis_names.count(name) > 1}

                message = "The expression must not contain multiple axes with the same name in brackets ([]).\n%EXPR%"
                raise SemanticError(invocation=invocation, pos=invocation.indicator.get_pos_for_axisnames([expr], duplicates), message=message)

    return exprs_in, exprs_out


def _semantic_checks_dot(exprs_in, exprs_out, invocation):
    expr_out = exprs_out[0]

    # Ensure at least two inputs
    if len(exprs_in) < 2:
        raise SemanticError(
            invocation=invocation, message=f"The dot operation requires at least 2 input expressions, but only {len(exprs_in)} is given.\n%EXPR%"
        )

    # Ensure that all marked axes appear in exactly two input expressions

    def is_marked_axis(expr):
        return isinstance(expr, stage3.Axis) and stage3.is_in_brackets(expr)

    marked_axis_names = {expr.name for expr_in in exprs_in for expr in expr_in.nodes() if is_marked_axis(expr)}
    invalid_axis_names = []
    for axis_name in marked_axis_names:
        count = 0
        for expr_in in exprs_in:
            if axis_name in {axis.name for axis in expr_in.nodes() if isinstance(axis, stage3.Axis)}:
                count += 1
        if count != 2:
            invalid_axis_names.append(axis_name)
    if len(invalid_axis_names) > 0:
        raise SemanticError(
            invocation=invocation,
            pos=invocation.indicator.get_pos_for_axisnames(exprs_in + [expr_out], invalid_axis_names),
            message="All contracted axes must appear in exactly two input expressions.\n%EXPR%",
        )


def _semantic_checks_get_at(exprs_in, exprs_out, invocation):
    expr_out = exprs_out[0]
    if len(exprs_in) < 2:
        raise SemanticError(invocation=invocation, message=f"The operation expects at least 2 input expressions, but found {len(exprs_in)}.\n%EXPR%")
    all_exprs = list(exprs_in) + [expr_out]
    tensor_expr = exprs_in[0]
    coords_exprs = exprs_in[1:]

    # Ensure that at most one axis is marked in each coordinate expression
    for coord_expr in coords_exprs:
        marked_axisnames = [expr.name for expr in coord_expr.nodes() if isinstance(expr, stage3.Axis) and stage3.is_in_brackets(expr)]
        if len(marked_axisnames) > 1:
            raise SemanticError(
                invocation=invocation,
                pos=invocation.indicator.get_pos_for_axisnames(all_exprs, marked_axisnames),
                message="Each coordinate expression must contain at most one axis in brackets.\n%EXPR%",
            )

    # Ensure that summed number of coordinates is equal to the number of marked axes in first input expression
    n = 0
    for coord_expr in coords_exprs:
        marked_axis = [expr for expr in coord_expr.nodes() if isinstance(expr, stage3.Axis) and stage3.is_in_brackets(expr)]
        assert len(marked_axis) <= 1
        if len(marked_axis) == 1:
            n += marked_axis[0].value
        else:
            n += 1
    marked_axes = [expr for expr in tensor_expr.nodes() if isinstance(expr, stage3.Axis) and stage3.is_in_brackets(expr)]
    if len(marked_axes) != n:
        raise SemanticError(
            invocation=invocation,
            pos=invocation.indicator.get_pos_for_axisnames(all_exprs, [expr.name for expr in marked_axes]),
            message="The number of coordinates must match the number of axes in brackets in the first input expression.\n%EXPR%",
        )


def _semantic_checks_sort(exprs_in, exprs_out, invocation):
    assert len(exprs_in) == 1 and len(exprs_out) == 1
    for expr in exprs_in + exprs_out:
        marked_axisnames = [expr.name for expr in expr.nodes() if isinstance(expr, stage3.Axis) and stage3.is_in_brackets(expr)]
        if len(marked_axisnames) != 1:
            raise SemanticError(
                invocation=invocation,
                pos=invocation.indicator.get_pos_for_axisnames(exprs_in + exprs_out, marked_axisnames),
                message="The expression for this operation must contain exactly one axis in brackets.\n%EXPR%",
            )


def _semantic_checks_update_at(exprs_in, exprs_out, invocation):
    expr_out = exprs_out[0]
    if len(exprs_in) < 2:
        raise SemanticError(invocation=invocation, message=f"The operation expects at least 3 input expressions, but found {len(exprs_in)}.\n%EXPR%")

    all_exprs = list(exprs_in) + [expr_out]
    tensor_expr = exprs_in[0]
    coords_exprs = exprs_in[1:-1]
    update_expr = exprs_in[-1]

    # Ensure same set of axes is marked in first input and output expression
    input_axisnames = {expr.name for expr in tensor_expr.nodes() if isinstance(expr, stage3.Axis) and stage3.is_in_brackets(expr)}
    output_axisnames = {expr.name for expr in expr_out.nodes() if isinstance(expr, stage3.Axis) and stage3.is_in_brackets(expr)}
    if input_axisnames != output_axisnames:
        raise SemanticError(
            invocation=invocation,
            pos=invocation.indicator.get_pos_for_axisnames(all_exprs, input_axisnames.symmetric_difference(output_axisnames)),
            message="The first input and output expressions must have the same set of axes in brackets.\n%EXPR%",
        )

    # Ensure that at most one axis is marked in each coordinate expression (aside from reduced axes that are also marked in updates)
    for coord_expr in coords_exprs:
        marked_axisnames = {expr.name for expr in coord_expr.nodes() if isinstance(expr, stage3.Axis) and stage3.is_in_brackets(expr)}
        if len(marked_axisnames) > 1:
            raise SemanticError(
                invocation=invocation,
                pos=invocation.indicator.get_pos_for_axisnames(all_exprs, marked_axisnames),
                message="Each coordinate expression must contain at most one coordinate axis in brackets.\n%EXPR%",
            )

    # Ensure marked axes in (1) target tensor, and (2) coordinate and update expressions are non-overlapping
    target_axisnames = {expr.name for expr in tensor_expr.nodes() if isinstance(expr, stage3.Axis) and stage3.is_in_brackets(expr)}
    coordupdate_axisnames = {
        expr.name for expr in coords_exprs + [update_expr] for expr in expr.nodes() if isinstance(expr, stage3.Axis) and stage3.is_in_brackets(expr)
    }
    intersection = target_axisnames.intersection(coordupdate_axisnames)
    if len(intersection) > 0:
        raise SemanticError(
            invocation=invocation,
            pos=invocation.indicator.get_pos_for_axisnames(all_exprs, intersection),
            message="Axes may not appear in brackets both in the first input tensor and in a coordinate or update expression.\n%EXPR%",
        )

    # Ensure that summed number of coordinates is equal to the number of marked axes in first input expression
    n = 0
    for coord_expr in coords_exprs:
        marked_axis = {expr for expr in coord_expr.nodes() if isinstance(expr, stage3.Axis) and stage3.is_in_brackets(expr)}
        assert len(marked_axis) <= 1
        if len(marked_axis) == 1:
            n += marked_axis.pop().value
        else:
            n += 1
    marked_axes = [expr for expr in tensor_expr.nodes() if isinstance(expr, stage3.Axis) and stage3.is_in_brackets(expr)]
    if len(marked_axes) != n:
        raise SemanticError(
            invocation=invocation,
            pos=invocation.indicator.get_pos_for_axisnames(all_exprs, [expr.name for expr in marked_axes]),
            message="The number of coordinates must match the number of axes in brackets in the first input expression.\n%EXPR%",
        )


def _semantic_checks_argfind(exprs_in, exprs_out, invocation):
    if len(exprs_in) != 1:
        raise SemanticError(invocation=invocation, message=f"The operation expects exactly one input expression, but found {len(exprs_in)}.\n%EXPR%")
    if len(exprs_out) != 1:
        raise SemanticError(invocation=invocation, message=f"The operation expects exactly one output expression, but found {len(exprs_in)}.\n%EXPR%")
    expr_in = exprs_in[0]
    expr_out = exprs_out[0]
    all_exprs = [expr_in, expr_out]

    # Ensure that at most one axis is marked in output expression
    marked_output_axes = [expr for expr in expr_out.nodes() if isinstance(expr, stage3.Axis) and stage3.is_in_brackets(expr)]
    if len(marked_output_axes) > 1:
        raise SemanticError(
            invocation=invocation,
            pos=invocation.indicator.get_pos_for_axisnames(all_exprs, [a.name for a in marked_output_axes]),
            message="The output expression must contain at most one axis in brackets.\n%EXPR%",
        )
    marked_output_axis = marked_output_axes[0] if len(marked_output_axes) == 1 else None

    # Ensure that number of marked axes in input expression is equal to value of marked axis in output expression
    marked_input_axes = [expr for expr in expr_in.nodes() if isinstance(expr, stage3.Axis) and stage3.is_in_brackets(expr)]
    if marked_output_axis is None:
        if len(marked_input_axes) != 1:
            raise SemanticError(
                invocation=invocation,
                pos=invocation.indicator.get_pos_for_axisnames(all_exprs, [expr.name for expr in marked_input_axes]),
                message=(
                    f"If no axis is marked in the output expression, exactly one axis must be marked in the input expression, "
                    f"but found {len(marked_input_axes)} marked axes.\n%EXPR%"
                ),
            )
    else:
        if len(marked_input_axes) != marked_output_axis.value:
            raise SemanticError(
                invocation=invocation,
                pos=invocation.indicator.get_pos_for_axisnames(all_exprs, [a.name for a in marked_input_axes] + [marked_output_axis.name]),
                message=(
                    f"The number of axes in brackets in the input expression ({len(marked_input_axes)}) must match the value "
                    "of the marked axis in the output expression ({marked_output_axis.value}).\n%EXPR%"
                ),
            )


def _cast_shape(tensor, expr):
    if tensor.shape is None:
        assert isinstance(tensor, tracer.signature.classical.ConvertibleTensor)
        return tracer.cast(tensor, lambda origin: tracer.signature.classical.ConvertibleTensor(origin, tensor.concrete, expr.shape))
    else:
        assert tuple(tensor.shape) == tuple(expr.shape)
        return tensor


def op(
    op,
    el_op=None,
    allow_concat=False,
    implicit_output=None,
    mark_reduced_axes=False,
    add_keepdims_param=False,
    check=None,
    cse_in_brackets=True,
    iskwarg=lambda name: False,
    equations_stage3=None,
    allow_nontrivial_unmarked_reduced_axes=False,
    allow_duplicate_el_axes=True,
    no_el_axis_permute=False,
):
    if isinstance(iskwarg, list | tuple | set):
        iskwarg = lambda name, iskwarg=iskwarg: name in iskwarg

    def inner(description, *tensors: Tensor, **kwargs):
        invocation = Invocation(
            description, name=op.__name__ if hasattr(op, "__name__") else "operation", tensors=tensors, kwargs=kwargs
        )  # Used for error reporting

        if add_keepdims_param:
            keepdims = {"keepdims": kwargs["keepdims"] if "keepdims" in kwargs else False}
            kwargs.pop("keepdims", None)
        else:
            keepdims = {}

        exprs_in, exprs_out = _parse_op(
            description,
            el_op,
            invocation=invocation,
            allow_concat=allow_concat,
            implicit_output=implicit_output,
            mark_reduced_axes=mark_reduced_axes,
            allow_duplicate_el_axes=allow_duplicate_el_axes,
            **keepdims,
        )

        if len(exprs_in) != len(tensors):

            def choose_s(i):
                return "s" if i != 1 else ""

            raise ValueError(
                f"The operation is defined with {len(exprs_in)} input expression{choose_s(len(exprs_in))}, but {len(tensors)} "
                f"input tensor{choose_s(len(tensors))} {'are' if len(tensors) != 1 else 'is'} given as argument{choose_s(len(tensors))}."
            )

        used_axis_names = {expr.name for expr in exprs_in + exprs_out for expr in expr.nodes() if isinstance(expr, stage1.Axis)}
        if any(iskwarg(name) for name in used_axis_names):
            invalid_kwargnames = [name for name in used_axis_names if iskwarg(name)]
            raise SemanticError(
                invocation=invocation,
                pos=invocation.indicator.get_pos_for_axisnames(exprs_in + exprs_out, invalid_kwargnames),
                message=(
                    f"The following axis names may not be used in the expression, since they are keyword arguments of the elementary "
                    f"operation: {', '.join(invalid_kwargnames)}.\n%EXPR%"
                ),
            )

        # Split kwargs from axis parameters
        parameters = {}
        new_kwargs = {}
        for key, value in kwargs.items():
            if iskwarg(key):
                new_kwargs[key] = value
            else:
                parameters[key] = value
        kwargs = new_kwargs
        del new_kwargs

        if no_el_axis_permute:
            for expr_in, expr_out in zip(exprs_in, exprs_out, strict=False):
                axisnames_in = []
                for axis in expr_in.nodes():
                    if isinstance(axis, stage1.Axis) and stage1.is_in_brackets(axis):
                        axisnames_in.append(axis.name)
                axisnames_out = []
                for axis in expr_out.nodes():
                    if isinstance(axis, stage1.Axis) and stage1.is_in_brackets(axis):
                        axisnames_out.append(axis.name)
                if axisnames_in != axisnames_out:
                    raise SemanticError(
                        invocation=invocation,
                        pos=invocation.indicator.get_pos_for_axisnames(exprs_in + exprs_out, set(axisnames_in).union(set(axisnames_out))),
                        message="All axes marked with brackets must appear in the same order in the inputs and outputs of this operation.\n%EXPR%",
                    )

        exprs_in, exprs_out = solve(
            exprs_in,
            exprs_out,
            [tensor.shape for tensor in tensors],
            invocation,
            parameters,
            cse_in_brackets=cse_in_brackets,
            equations_stage3=partial(equations_stage3, invocation=invocation) if equations_stage3 is not None else None,
        )

        if not allow_nontrivial_unmarked_reduced_axes:
            axis_names_in = {
                axis.name
                for expr_in in exprs_in
                for axis in expr_in.nodes()
                if isinstance(axis, stage3.Axis) and not stage3.is_in_brackets(axis) and axis.value != 1
            }
            axis_names_out = {
                axis.name
                for expr_out in exprs_out
                for axis in expr_out.nodes()
                if isinstance(axis, stage3.Axis) and not stage3.is_in_brackets(axis) and axis.value != 1
            }
            axis_names_reduced = axis_names_in - axis_names_out
            if len(axis_names_reduced) > 0:
                raise SemanticError(
                    invocation=invocation,
                    pos=invocation.indicator.get_pos_for_axisnames(exprs_in + exprs_out, axis_names_reduced),
                    message=f"The input axes {axis_names_reduced} must appear in the output expression.\n%EXPR%",
                )

        if check is not None:
            check(exprs_in, exprs_out, invocation)

        tensors = [_cast_shape(tensor, expr_in) for tensor, expr_in in zip(tensors, exprs_in, strict=False)]

        tensors = [NamedTensor(tensor, expr_in) for tensor, expr_in in zip(tensors, exprs_in, strict=False)]
        try:
            tensors = op(*tensors, out=exprs_out[0] if len(exprs_out) == 1 else exprs_out, **kwargs)
        except SemanticError as e:
            raise SemanticError(invocation=invocation, pos=e.pos, message=str(e)) from e

        if len(exprs_out) > 1:
            return tuple(t.value for t in tensors)
        else:
            return tensors.value

    inner.__name__ = op.__name__
    inner.__qualname__ = op.__qualname__
    inner.__doc__ = op.__doc__
    inner.__module__ = op.__module__
    return inner


def id(op, **kwargs):
    if "allow_concat" not in kwargs:
        kwargs["allow_concat"] = True

    def el_op(op):
        n_in = len(op.children[0].children)
        n_out = n_in if len(op.children) == 1 else len(op.children[1].children)
        args_in = ", ".join("" for i in range(n_in))
        args_out = ", ".join("" for i in range(n_out))
        return f"{args_in} -> {args_out}"

    return globals()["op"](op, el_op=el_op, implicit_output="bijective", **kwargs)


def elementwise(op, **kwargs):
    def el_op(op):
        n_in = len(op.children[0].children)
        args_in = ", ".join("" for i in range(n_in))
        return f"{args_in} ->"

    return globals()["op"](op, el_op=el_op, implicit_output="bijective", **kwargs)


def dot(op, **kwargs):
    def el_op(op):
        return f"{op.children[0]} ->"

    return globals()["op"](
        op, el_op=el_op, implicit_output="bijective", mark_reduced_axes=True, allow_duplicate_el_axes=False, check=_semantic_checks_dot, **kwargs
    )


def _equations_stage3_index_at(exprs_in, exprs_out, invocation, is_update):
    if len(exprs_in) <= 1:
        return []

    tensor_expr = exprs_in[0]
    coord_exprs = exprs_in[1:-1] if is_update else exprs_in[1:]
    coords_axes = []
    for coord_expr in coord_exprs:
        marked_axes = [expr for expr in coord_expr.nodes() if isinstance(expr, stage2.Axis) and stage2.is_in_brackets(expr)]
        if len(marked_axes) == 0:
            coords_axes.append(stage2.Axis(f"unnamed.{uuid.uuid4().int}", 1, ellipsis_indices=[]))
        elif len(marked_axes) == 1:
            coords_axes.append(marked_axes[0])
        else:
            raise SemanticError(
                invocation=invocation,
                pos=invocation.indicator.get_pos_for_axisnames(exprs_in + exprs_out, [expr.name for expr in marked_axes]),
                message="Each coordinate expression must contain at most one axis in brackets.\n%EXPR%",
            )
    marked_coord_axis = stage2.ConcatenatedAxis.create(coords_axes, ellipsis_indices=[])
    marked_axes_in = [expr for expr in tensor_expr.nodes() if isinstance(expr, stage2.Axis) and stage2.is_in_brackets(expr)]

    if marked_coord_axis.value is not None and marked_coord_axis.value != len(marked_axes_in):
        raise SemanticError(
            invocation=invocation,
            pos=invocation.indicator.get_pos_for_axisnames(exprs_in + exprs_out, [marked_coord_axis.name] + [expr.name for expr in marked_axes_in]),
            message=(
                f"The sum of the lengths of marked coordinate axes ({marked_coord_axis.value}) must match the number of marked axes "
                f"in the first input expression ({len(marked_axes_in)}).\n%EXPR%"
            ),
        )

    return [stage3.Equation(marked_coord_axis, stage2.Axis(f"unnamed.{uuid.uuid4().int}", len(marked_axes_in), ellipsis_indices=[]))]


def get_at(op, **kwargs):
    def el_op(op):
        return ", ".join(str(c) for c in op.children[0].children) + " ->"

    return globals()["op"](
        op,
        el_op=el_op,
        implicit_output="bijective",
        check=_semantic_checks_get_at,
        equations_stage3=partial(_equations_stage3_index_at, is_update=False),
        **kwargs,
    )


def update_at(op, **kwargs):
    def el_op(op):
        return ", ".join(str(c) for c in op.children[0].children[:-1]) + f", -> {op.children[0].children[0]}"

    op = globals()["op"](
        op,
        el_op=el_op,
        implicit_output=0,
        check=_semantic_checks_update_at,
        equations_stage3=partial(_equations_stage3_index_at, is_update=True),
        allow_nontrivial_unmarked_reduced_axes=True,
        no_el_axis_permute=True,
        **kwargs,
    )

    @use_name_of(op)
    def op_with_zerosized_args(description, *tensors: Tensor, **kwargs):
        if len(tensors) > 0 and any(int(i) == 0 for tensor in [t for t in tensors[1:] if t.shape is not None] for i in tensor.shape):
            return tensors[0]  # TODO: remove this when dynamic dimensions are supported?
        else:
            return op(description, *tensors, **kwargs)

    return op_with_zerosized_args


def reduce(op, **kwargs):
    def el_op(op):
        return f"{op.children[0].children[0]} ->"

    return globals()["op"](op, el_op=el_op, implicit_output="bijective", mark_reduced_axes=True, add_keepdims_param=True, **kwargs)


def _equations_stage3_argfind(exprs_in, exprs_out, invocation):
    marked_axes_in = [expr for expr in exprs_in[0].nodes() if isinstance(expr, stage2.Axis) and stage2.is_in_brackets(expr)]
    marked_axes_out = [expr for expr in exprs_out[0].nodes() if isinstance(expr, stage2.Axis) and stage2.is_in_brackets(expr)]
    if len(marked_axes_out) == 0:
        if len(marked_axes_in) != 1:
            raise SemanticError(
                invocation=invocation,
                pos=invocation.indicator.get_pos_for_axisnames(exprs_in + exprs_out, [a.name for a in marked_axes_in]),
                message=(
                    f"If no axis is marked with brackets in the output expression, exactly one axis must be marked in the "
                    f"input expression, but found {len(marked_axes_in)} marked axes.\n%EXPR%"
                ),
            )
        return []
    elif len(marked_axes_out) == 1:
        marked_axis_out = marked_axes_out[0]
        if marked_axis_out.value is not None and marked_axis_out.value != len(marked_axes_in):
            raise SemanticError(
                invocation=invocation,
                pos=invocation.indicator.get_pos_for_axisnames(exprs_in + exprs_out, [marked_axis_out.name] + [a.name for a in marked_axes_in]),
                message=(
                    f"The value of the marked axis in the output expression ({marked_axis_out.value}) must match the number of marked "
                    f"axes in the input expression ({len(marked_axes_in)}).\n%EXPR%"
                ),
            )
        return [stage3.Equation(marked_axis_out, stage2.Axis(f"unnamed.{uuid.uuid4().int}", len(marked_axes_in), ellipsis_indices=[]))]
    else:
        raise SemanticError(
            invocation=invocation,
            pos=invocation.indicator.get_pos_for_axisnames(exprs_in + exprs_out, [a.name for a in marked_axes_out]),
            message="The output expression must contain at most one axis in brackets.\n%EXPR%",
        )


def argfind(op, **kwargs):
    def el_op(op):
        if len(op.children) == 1 or op.children[1].children[0].ndim != 0:
            return f"{op.children[0].children[0]} -> a{uuid.uuid4().int}"
        else:
            return f"{op.children[0].children[0]} ->"

    return globals()["op"](op, el_op=el_op, implicit_output="bijective", check=_semantic_checks_argfind, equations_stage3=_equations_stage3_argfind, **kwargs)


def preserve_shape(op, **kwargs):
    def el_op(op):
        return f"{op.children[0].children[0]} -> {op.children[0].children[0]}"

    return globals()["op"](op, el_op=el_op, implicit_output="bijective", no_el_axis_permute=True, **kwargs)


_name_to_op = (
    {name: elementwise for name in adapter.ops.elementwise}
    | {name: partial(reduce, cse_in_brackets=True) for name in adapter.ops.reduce}
    | {name: update_at for name in adapter.ops.update_at}
    | {name: argfind for name in adapter.ops.argfind}
    | {
        "get_at": get_at,
        "dot": dot,
        "id": id,
        "roll": partial(preserve_shape, iskwarg=["shift"]),
        "flip": preserve_shape,
        "sort": partial(preserve_shape, check=_semantic_checks_sort),
        "argsort": partial(preserve_shape, check=_semantic_checks_sort),
        "softmax": preserve_shape,
        "log_softmax": preserve_shape,
    }
)


def ops(namedtensor_ops, **kwargs):
    return {name: _name_to_op[name](namedtensor_ops[name], **kwargs) for name in namedtensor_ops.keys()}
