from . import stage1, solver
import re
import tlib
import numpy as np
from collections import defaultdict


class Expression:
    def __init__(self, ellipsis_indices, begin_pos=None, end_pos=None):
        self.ellipsis_indices = ellipsis_indices
        self.parent = None
        self.begin_pos = begin_pos
        self.end_pos = end_pos

    @property
    def depth(self):
        return len(self.ellipsis_indices)

    @property
    def shape(self):
        return tuple(i[1] for i in self.ellipsis_indices) + (len(self),)


class Composition(Expression):
    def __init__(self, inner, ellipsis_indices, begin_pos=None, end_pos=None):
        Expression.__init__(self, ellipsis_indices, begin_pos=begin_pos, end_pos=end_pos)
        self.inner = inner
        inner.parent = self

    def __str__(self):
        return f"({self.inner})"

    def __len__(self):
        return 1

    def __iter__(self):
        yield self

    def __deepcopy__(self):
        return Composition(
            self.inner.__deepcopy__(),
            ellipsis_indices=self.ellipsis_indices,
            begin_pos=self.begin_pos,
            end_pos=self.end_pos,
        )

    def all(self):
        yield self
        yield from self.inner.all()


class List(Expression):
    @staticmethod
    def maybe(l, ellipsis_indices, begin_pos=None, end_pos=None):
        if len(l) == 1:
            return l[0]
        else:
            return List(l, ellipsis_indices, begin_pos=begin_pos, end_pos=end_pos)

    def __init__(self, children, ellipsis_indices, begin_pos=None, end_pos=None):
        Expression.__init__(self, ellipsis_indices, begin_pos=begin_pos, end_pos=end_pos)
        self.children = children
        for c in children:
            c.parent = self

    def __str__(self):
        return " ".join([str(c) for c in self.children])

    def __len__(self):
        return sum(len(c) for c in self.children)

    def __iter__(self):
        for c in self.children:
            yield from c

    def __deepcopy__(self):
        return List(
            [c.__deepcopy__() for c in self.children],
            ellipsis_indices=self.ellipsis_indices,
            begin_pos=self.begin_pos,
            end_pos=self.end_pos,
        )

    def all(self):
        yield self
        for c in self.children:
            yield from c.all()


class NamedAxis(Expression):
    def __init__(self, name, ellipsis_indices, begin_pos=None, end_pos=None):
        Expression.__init__(self, ellipsis_indices, begin_pos=begin_pos, end_pos=end_pos)
        self.name = name

        postfix = ""
        for idx, _num in self.ellipsis_indices:
            postfix = postfix + "." + str(idx)
        if not self.name.endswith(postfix):
            self.name = self.name + postfix

    def __str__(self):
        return self.name

    def __len__(self):
        return 1

    def __iter__(self):
        yield self

    def __deepcopy__(self):
        return NamedAxis(
            self.name,
            ellipsis_indices=self.ellipsis_indices,
            begin_pos=self.begin_pos,
            end_pos=self.end_pos,
        )

    def all(self):
        yield self


class UnnamedAxis(Expression):
    def __init__(self, value, ellipsis_indices, begin_pos=None, end_pos=None):
        Expression.__init__(self, ellipsis_indices, begin_pos=begin_pos, end_pos=end_pos)
        self.value = value

    def __str__(self):
        return str(self.value)

    def __len__(self):
        return 1

    def __iter__(self):
        yield self

    def __deepcopy__(self):
        return UnnamedAxis(
            self.value,
            ellipsis_indices=self.ellipsis_indices,
            begin_pos=self.begin_pos,
            end_pos=self.end_pos,
        )

    def all(self):
        yield self


class Concatenation(Expression):
    def __init__(self, children, ellipsis_indices, begin_pos=None, end_pos=None):
        Expression.__init__(self, ellipsis_indices, begin_pos=begin_pos, end_pos=end_pos)
        for c in children:
            if len(c) != 1:
                raise ValueError(
                    "Concatenation can only be used on expressions of length 1, " f"but got expression '{c}'"
                )
        self.children = children
        for c in children:
            c.parent = self

    def __str__(self):
        return "+".join([str(c) for c in self.children])

    def __len__(self):
        return 1

    def __iter__(self):
        yield self

    def __deepcopy__(self):
        return Concatenation(
            [c.__deepcopy__() for c in self.children],
            ellipsis_indices=self.ellipsis_indices,
            begin_pos=self.begin_pos,
            end_pos=self.end_pos,
        )

    def all(self):
        yield self
        for c in self.children:
            yield from c.all()


class Marker(Expression):
    @staticmethod
    def maybe(inner, ellipsis_indices, begin_pos=None, end_pos=None):
        if len(inner) == 0:
            return inner
        else:
            return Marker(inner, ellipsis_indices, begin_pos=None, end_pos=None)

    def __init__(self, inner, ellipsis_indices, begin_pos=None, end_pos=None):
        Expression.__init__(self, ellipsis_indices, begin_pos=begin_pos, end_pos=end_pos)
        self.inner = inner
        inner.parent = self
        assert len(inner) > 0

    def __str__(self):
        return f"[{self.inner}]"

    def __len__(self):
        return len(self.inner)

    def __iter__(self):
        yield from self.inner

    def __deepcopy__(self):
        return Marker(
            self.inner.__deepcopy__(),
            ellipsis_indices=self.ellipsis_indices,
            begin_pos=self.begin_pos,
            end_pos=self.end_pos,
        )

    def all(self):
        yield self
        yield from self.inner.all()


def solve(exprs1, exprs2, expansions1, expansions2, depths1, depths2, descs1, descs2, signature=None):
    exprs1 = list(exprs1)
    exprs2 = list(exprs2)
    expansions1 = list(expansions1)
    expansions2 = list(expansions2)
    depths1 = list(depths1)
    depths2 = list(depths2)
    descs1 = list(descs1)
    descs2 = list(descs2)
    if any(expr is not None and not isinstance(expr, stage1.Expression) for expr in exprs1 + exprs2):
        raise ValueError("Can only expand stage1.Expression")
    if (
        len(
            {
                len(exprs1),
                len(exprs2),
                len(expansions1),
                len(expansions2),
                len(depths1),
                len(depths2),
                len(descs1),
                len(descs2),
            }
        )
        != 1
    ):
        raise ValueError("Number of expressions, expansions and depths must be equal")

    expr_id_to_expr = {}
    ellipsis_id_to_exprs = defaultdict(list)
    axis_name_to_exprs = defaultdict(list)
    for root in exprs1 + exprs2:
        if root is not None:
            for expr in root.all():
                expr_id_to_expr[id(expr)] = expr
                if isinstance(expr, stage1.Ellipsis):
                    ellipsis_id_to_exprs[expr.ellipsis_id].append(expr)
                elif isinstance(expr, stage1.NamedAxis):
                    axis_name_to_exprs[expr.name].append(expr)

    # ##### 1. Find expression depths #####
    equations = []

    symbolic_expr_depths = {}
    for root in exprs1 + exprs2:
        if root is not None:
            for expr in root.all():
                symbolic_expr_depths[id(expr)] = solver.Variable(f"symbolic_expr_depths[{id(expr)}]", str(expr))

    # Add equations: Depth relations between subexpressions
    for root in exprs1 + exprs2:
        if root is not None:
            for expr in root.all():
                if isinstance(expr, stage1.Ellipsis):
                    # Ellipsis increases depth by one
                    equations.append(
                        (
                            symbolic_expr_depths[id(expr)] + 1,
                            symbolic_expr_depths[id(expr.inner)],
                        )
                    )
                else:
                    # All other expressions have the same depth as their children
                    for child in expr.direct_children:
                        equations.append(
                            (
                                symbolic_expr_depths[id(expr)],
                                symbolic_expr_depths[id(child)],
                            )
                        )

    # Add equations: Depth arguments
    for root, depth in zip(exprs1 + exprs2, depths1 + depths2):
        if root is not None and depth is not None:
            equations.append((symbolic_expr_depths[id(root)], depth))

    # Add equations: Root depths
    for root1, root2, expansion1, expansion2 in zip(exprs1, exprs2, expansions1, expansions2):
        if root1 is not None and root2 is not None and expansion1 is not None and expansion2 is not None:
            equations.append(
                (
                    symbolic_expr_depths[id(root1)] + len(expansion1),
                    symbolic_expr_depths[id(root2)] + len(expansion2),
                )
            )

    # Add equations: Multiple occurrences of the same named axis must have the same depth
    symbolic_axis_depths = {}
    for root in exprs1 + exprs2:
        if root is not None:
            for axis in root.all():
                if isinstance(axis, stage1.NamedAxis):
                    if axis.name not in symbolic_axis_depths:
                        symbolic_axis_depths[axis.name] = solver.Variable(
                            f"symbolic_axis_depths[{axis.name}]", axis.name
                        )
                    equations.append(
                        (
                            symbolic_expr_depths[id(axis)],
                            symbolic_axis_depths[axis.name],
                        )
                    )

    # Add equations: Ellipses with the same id must have the same depth
    symbolic_ellipsis_depths = {}
    for root in exprs1 + exprs2:
        if root is not None:
            for ellipsis in root.all():
                if isinstance(ellipsis, stage1.Ellipsis):
                    if ellipsis.ellipsis_id not in symbolic_ellipsis_depths:
                        symbolic_ellipsis_depths[ellipsis.ellipsis_id] = solver.Variable(
                            f"symbolic_ellipsis_depths[{ellipsis.ellipsis_id}]", str(ellipsis)
                        )
                    equations.append(
                        (
                            symbolic_expr_depths[id(ellipsis)],
                            symbolic_ellipsis_depths[ellipsis.ellipsis_id],
                        )
                    )

    # Solve
    def solve(equations):
        solutions = solver.solve(equations)

        expr_depths = {}
        for k, v in solutions.items():
            if k.startswith("symbolic_expr_depths["):
                expr_depths[int(k[len("symbolic_expr_depths[") : -1])] = int(v)

        # Raise exception on negative depths
        failed_exprs = set()
        for root in exprs1 + exprs2:
            if root is not None:
                for expr in root.all():
                    if id(expr) in expr_depths and expr_depths[id(expr)] < 0:
                        failed_exprs.add(str(expr))
        if len(failed_exprs) > 0:
            raise solver.SolveExceptionNoSolution()

        # Raise exception on missing depths
        failed_exprs = set()
        for root in exprs1 + exprs2:
            if root is not None:
                for expr in root.all():
                    if id(expr) not in expr_depths:
                        failed_exprs.add(str(expr))
        if len(failed_exprs) > 0:
            raise solver.SolveExceptionTooManySolutions()

        return expr_depths

    illegal_depth_message = (
        "Please check the following:\n"
        " - Each axis name may be used either with or without an ellipsis, "
        "but not both.\n - The rank of a constraint "
        "must be equal to or less than the number of ellipses around the corresponding axis."
    )
    try:
        expr_depths = solve(equations)
    except solver.SolveExceptionTooManySolutions as e:
        raise DimensionError(
            f"Found an invalid usage of ellipses and/or constraints.\n{illegal_depth_message}",
            text=signature.text,
            constraints=[
                (expr1, expr2) for expr1, expr2 in zip(exprs1, exprs2) if expr1 is not None and expr2 is not None
            ],
        ) from e
    except solver.SolveExceptionNoSolution as e:
        # Check which axes are leading to contradiction
        def axisnames_in_equation(eq):
            axisnames = set()
            for term in eq:
                if isinstance(term, solver.Expression):
                    for var in term:
                        if isinstance(var, solver.Variable):
                            if var.id.startswith("symbolic_axis_depths["):
                                axisnames.add(var.id[len("symbolic_axis_depths[") : -1])
                            elif var.id.startswith("symbolic_expr_depths["):
                                expr_id = var.id[len("symbolic_expr_depths[") : -1]
                                expr = expr_id_to_expr.get(int(expr_id), None)
                                for expr in expr.all():
                                    if isinstance(expr, stage1.NamedAxis):
                                        axisnames.add(expr.name)
            return list(axisnames)

        def equation_contains_axisconsistency(eq, axis_name):
            for term in eq:
                if isinstance(term, solver.Expression):
                    for var in term:
                        if isinstance(var, solver.Variable):
                            if var.id == f"symbolic_axis_depths[{axis_name}]":
                                return True
            return False

        contradicting_axis_names = []
        for axis_name in axis_name_to_exprs.keys():
            try_equations = [
                eq
                for eq in equations
                if [axis_name] != axisnames_in_equation(eq) and not equation_contains_axisconsistency(eq, axis_name)
            ]
            try:
                solve(try_equations)
                still_contradicts = False
            except solver.SolveExceptionNoSolution:
                still_contradicts = True
            except solver.SolveExceptionTooManySolutions:
                still_contradicts = False
            if not still_contradicts:
                contradicting_axis_names.append(axis_name)

        if len(contradicting_axis_names) > 0:
            contradicting_axis_names = sorted(contradicting_axis_names)
            if len(contradicting_axis_names) == 1:
                axis = f"axis {contradicting_axis_names[0]}"
            else:
                axis = f"axes {', '.join(contradicting_axis_names)}"
            raise tlib.DimensionError(
                f"Found an invalid usage of ellipses and/or constraints for the {axis}:",
                postfix=illegal_depth_message,
                text=signature.text,
                pos=signature.get_pos_for_axisnames(exprs1 + exprs2, contradicting_axis_names),
                constraints=[
                    (expr1, expr2) for expr1, expr2 in zip(exprs1, exprs2) if expr1 is not None and expr2 is not None
                ],
            ) from e
        else:
            raise tlib.DimensionError(
                "Found an invalid usage of ellipses and/or constraints.",
                postfix=illegal_depth_message,
                text=signature.text,
                constraints=[
                    (expr1, expr2) for expr1, expr2 in zip(exprs1, exprs2) if expr1 is not None and expr2 is not None
                ],
            ) from e

    # Expand ellipses and add missing dimensions to expansions
    for exprs, expansions, _depths in zip([exprs1, exprs2], [expansions1, expansions2], [depths1, depths2]):
        for i in range(len(exprs)):
            if exprs[i] is not None:
                missing_depth = expr_depths[id(exprs[i])]
                assert missing_depth >= 0

                # Add missing dimensions to expansions
                if expansions[i] is not None:
                    assert len(expansions[i]) >= 1
                    if missing_depth > 0:
                        expansions[i] = [None] * missing_depth + list(expansions[i])

                # Add missing ellipses around root expressions
                if missing_depth > 0:
                    for _ in range(missing_depth):
                        exprs[i] = stage1.Ellipsis(exprs[i], -1, -1)
                        expr_depths[id(exprs[i])] = expr_depths[id(exprs[i].inner)] - 1

    # ##### 2. Find ellipsis expansions #####
    equations = []

    symbolic_expr_expansions = {}
    for root in exprs1 + exprs2:
        if root is not None:
            for expr in root.all():
                for depth in range(expr_depths[id(expr)] + 1):
                    key = (id(expr), depth)
                    symbolic_expr_expansions[key] = solver.Variable(
                        f"symbolic_expr_expansions[{id(expr)},{depth}]", f"{expr} at depth {depth}"
                    )

    # Add equations: Expansion of an expression at depth d (less than own depth)
    # is equal to the expansion of each child at depth d
    for root in exprs1 + exprs2:
        if root is not None:
            for expr in root.all():
                for depth in range(expr_depths[id(expr)]):
                    for child in expr.direct_children:
                        equations.append(
                            (
                                symbolic_expr_expansions[(id(expr), depth)],
                                symbolic_expr_expansions[(id(child), depth)],
                            )
                        )

    # Add equations: Relations between expressions and their children
    for root in exprs1 + exprs2:
        if root is not None:
            for expr in root.all():
                depth = expr_depths[id(expr)]
                if isinstance(expr, stage1.List):
                    v = sum(symbolic_expr_expansions[(id(child), depth)] for child in expr.children)
                elif isinstance(expr, stage1.Concatenation):
                    v = 1
                elif isinstance(expr, stage1.NamedAxis):
                    v = 1
                elif isinstance(expr, stage1.UnnamedAxis):
                    v = 1
                elif isinstance(expr, stage1.Composition):
                    v = 1
                elif isinstance(expr, stage1.Marker):
                    v = symbolic_expr_expansions[(id(expr.inner), depth)]
                elif isinstance(expr, stage1.Ellipsis):
                    v = symbolic_expr_expansions[(id(expr.inner), depth)]
                else:
                    raise AssertionError(f"{expr}")
                equations.append((symbolic_expr_expansions[(id(expr), depth)], v))

    # Add equations: Expansions stored in "expansions"
    for expansion1, expansion2, expr1, expr2, desc1, desc2 in zip(
        expansions1, expansions2, exprs1, exprs2, descs1, descs2
    ):
        if expansion1 is not None and expansion2 is not None:
            if len(expansion1) != len(expansion2) or any(
                e1 is not None and e2 is not None and e1 != e2 for e1, e2 in zip(expansion1, expansion2)
            ):
                raise tlib.DimensionError(
                    f"The number of dimensions of the {desc1} does not match the number of "
                    f"dimensions of the {desc2}",
                    text=signature.text,
                    constraints=[
                        (expr1, expr2)
                        for expr1, expr2 in zip(exprs1, exprs2)
                        if expr1 is not None and expr2 is not None
                    ],
                )

        if expansion1 is not None and expansion2 is not None:
            expansion = [e1 if e1 is not None else e2 for e1, e2 in zip(expansion1, expansion2)]
        elif expansion1 is not None:
            expansion = expansion1
        elif expansion2 is not None:
            expansion = expansion2
        else:
            expansion = None

        if expansion is not None:
            for depth, e in enumerate(expansion):
                if e is not None:
                    if expr1 is not None and depth <= expr_depths[id(expr1)]:
                        equations.append((symbolic_expr_expansions[(id(expr1), depth)], int(e)))
                    if expr2 is not None and depth <= expr_depths[id(expr2)]:
                        equations.append((symbolic_expr_expansions[(id(expr2), depth)], int(e)))

    # Add equations: Multiple occurrences of the same named axis must have the same expansions
    symbolic_axis_expansions = {}
    for root in exprs1 + exprs2:
        if root is not None:
            for axis in root.all():
                if isinstance(axis, stage1.NamedAxis):
                    for depth in range(expr_depths[id(axis)] + 1):
                        if axis.name not in symbolic_axis_expansions:
                            symbolic_axis_expansions[(axis.name, depth)] = solver.Variable(
                                f"symbolic_axis_expansions[{axis.name},{depth}]",
                                f"{axis.name} at depth {depth}",
                            )
                        equations.append(
                            (
                                symbolic_expr_expansions[(id(axis), depth)],
                                symbolic_axis_expansions[(axis.name, depth)],
                            )
                        )

    # Add equations: Ellipses with the same id must have the same expansions
    symbolic_ellipsis_expansions = {}
    for root in exprs1 + exprs2:
        if root is not None:
            for ellipsis in root.all():
                if isinstance(ellipsis, stage1.Ellipsis):
                    for depth in range(expr_depths[id(ellipsis)] + 1):
                        if ellipsis.ellipsis_id not in symbolic_ellipsis_expansions:
                            symbolic_ellipsis_expansions[(ellipsis.ellipsis_id, depth)] = solver.Variable(
                                f"symbolic_ellipsis_expansions[{ellipsis.ellipsis_id},{depth}]",
                                f"{ellipsis} at depth {depth}",
                            )
                        equations.append(
                            (
                                symbolic_expr_expansions[(id(ellipsis), depth)],
                                symbolic_ellipsis_expansions[(ellipsis.ellipsis_id, depth)],
                            )
                        )

    # Add equations: Same root expansions
    for root1, root2 in zip(exprs1, exprs2):
        if root1 is not None and root2 is not None:
            assert expr_depths[id(root1)] == expr_depths[id(root2)]
            for depth in range(expr_depths[id(root1)] + 1):
                equations.append(
                    (
                        symbolic_expr_expansions[(id(root1), depth)],
                        symbolic_expr_expansions[(id(root2), depth)],
                    )
                )

    # Solve
    def solve(equations):
        solutions = solver.solve(equations)

        expansion_values = {}
        for k, v in solutions.items():
            if k.startswith("symbolic_expr_expansions["):
                k = k[len("symbolic_expr_expansions[") : -1]
                id_expr, depth = str(k).split(",")
                try:
                    id_expr = int(id_expr)
                except ValueError:
                    continue
                depth = int(depth)
                expansion_values[(id_expr, depth)] = int(v)

        # Raise exception on negative expansions
        failed_exprs = set()
        for root in exprs1 + exprs2:
            if root is not None:
                for expr in root.all():
                    depth = expr_depths[id(expr)]
                    key = (id(expr), depth)
                    if key in expansion_values and expansion_values[key] < 0:
                        failed_exprs.add(expr)
        if len(failed_exprs) > 0:
            raise solver.SolveExceptionNoSolution()

        # Raise exception on missing expansions for expressions not wrapped in composition
        def is_at_root(expr):
            parent = expr.parent
            while parent is not None:
                if isinstance(parent, stage1.Composition):
                    return False
                parent = parent.parent
            return True

        failed_exprs = set()
        for root in exprs1 + exprs2:
            if root is not None:
                for expr in root.all():
                    if not is_at_root(expr):
                        continue
                    depth = expr_depths[id(expr)]
                    key = (id(expr), depth)
                    if key not in expansion_values:
                        failed_exprs.add(expr)
        if len(failed_exprs) > 0:
            raise solver.SolveExceptionTooManySolutions()

        return expansion_values

    try:
        expansion_values = solve(equations)
    except solver.SolveExceptionTooManySolutions as e:
        raise tlib.DimensionError(
            "Failed to uniquely determine the expansion of ellipses in the expression. Please "
            "provide more constraints.",
            text=signature.text,
            pos=signature.get_pos_for_ellipses(exprs1 + exprs2),
            constraints=[
                (expr1, expr2) for expr1, expr2 in zip(exprs1, exprs2) if expr1 is not None and expr2 is not None
            ],
        ) from e
    except solver.SolveExceptionNoSolution as e:
        raise tlib.DimensionError(
            "Failed to find an expansion of ellipses in the expression such that the number of "
            "dimensions matches all given constraints.",
            text=signature.text,
            pos=signature.get_pos_for_ellipses(exprs1 + exprs2),
            constraints=[
                (expr1, expr2) for expr1, expr2 in zip(exprs1, exprs2) if expr1 is not None and expr2 is not None
            ],
        ) from e

    # Expand ellipses and map stage1 expressions to stage2 expressions
    def is_unnamed(expr):
        for expr in expr.all():
            if isinstance(expr, stage1.NamedAxis):
                return False
        return True

    def get_unnamed_value(expr):
        if isinstance(expr, stage1.List):
            return np.prod([get_unnamed_value(child) for child in expr.children]).astype("int")
        elif isinstance(expr, stage1.Concatenation):
            return np.sum([get_unnamed_value(child) for child in expr.children])
        elif isinstance(expr, stage1.NamedAxis):
            raise AssertionError()
        elif isinstance(expr, stage1.UnnamedAxis):
            return expr.value
        elif isinstance(expr, stage1.Composition):
            return get_unnamed_value(expr.inner)
        elif isinstance(expr, stage1.Marker):
            return get_unnamed_value(expr.inner)
        elif isinstance(expr, stage1.Ellipsis):
            value = get_unnamed_value(expr.inner)
            if value != 1:  # TODO: implement this
                raise NotImplementedError(
                    f"Found unnamed and unexpanded ellipsis '{expr}'. We currently disallow this "
                    "case, since it could can take on multiple values ('2...' could have values "
                    "2, 4, ...) that should be resolved in the solver and then checked to be "
                    "consistent with these constraints."
                )
            return 1
        else:
            raise AssertionError(f"{expr}")

    def map(expr, ellipsis_indices):
        if isinstance(expr, list):
            return [c for expr in expr for c in map(expr, ellipsis_indices=ellipsis_indices)]
        elif isinstance(expr, stage1.NamedAxis):
            return [
                NamedAxis(
                    expr.name,
                    ellipsis_indices=ellipsis_indices,
                    begin_pos=expr.begin_pos,
                    end_pos=expr.end_pos,
                )
            ]
        elif isinstance(expr, stage1.UnnamedAxis):
            return [
                UnnamedAxis(
                    expr.value,
                    ellipsis_indices=ellipsis_indices,
                    begin_pos=expr.begin_pos,
                    end_pos=expr.end_pos,
                )
            ]
        elif isinstance(expr, stage1.List):
            return map(expr.children, ellipsis_indices=ellipsis_indices)
        elif isinstance(expr, stage1.Concatenation):
            return [
                Concatenation(
                    [
                        List.maybe(
                            map(c, ellipsis_indices=ellipsis_indices),
                            ellipsis_indices=ellipsis_indices,
                        )
                        for c in expr.children
                    ],
                    ellipsis_indices=ellipsis_indices,
                    begin_pos=expr.begin_pos,
                    end_pos=expr.end_pos,
                )
            ]
        elif isinstance(expr, stage1.Composition):
            return [
                Composition(
                    List.maybe(
                        map(expr.inner, ellipsis_indices=ellipsis_indices),
                        ellipsis_indices=ellipsis_indices,
                    ),
                    ellipsis_indices=ellipsis_indices,
                    begin_pos=expr.begin_pos,
                    end_pos=expr.end_pos,
                )
            ]
        elif isinstance(expr, stage1.Marker):
            return [
                Marker.maybe(
                    List.maybe(
                        map(expr.inner, ellipsis_indices=ellipsis_indices),
                        ellipsis_indices=ellipsis_indices,
                    ),
                    ellipsis_indices=ellipsis_indices,
                    begin_pos=expr.begin_pos,
                    end_pos=expr.end_pos,
                )
            ]
        elif isinstance(expr, stage1.Ellipsis):
            key = (id(expr), expr_depths[id(expr)])
            if key in expansion_values:
                # Ellipsis is expanded
                expansion = expansion_values[key]
                assert expansion >= 0
                return [
                    c
                    for i in range(expansion)
                    for c in map(expr.inner, ellipsis_indices=ellipsis_indices + [(i, expansion)])
                ]
            else:
                # Ellipsis is not expanded
                if is_unnamed(expr):
                    # Contains no named axes -> convert to unnamed axis
                    return [
                        UnnamedAxis(
                            get_unnamed_value(expr),
                            ellipsis_indices=ellipsis_indices,
                            begin_pos=expr.begin_pos,
                            end_pos=expr.end_pos,
                        )
                    ]
                else:
                    # Contains named axes -> convert to named axis
                    return [
                        NamedAxis(
                            str(expr),
                            ellipsis_indices=ellipsis_indices,
                            begin_pos=expr.begin_pos,
                            end_pos=expr.end_pos,
                        )
                    ]
        else:
            raise AssertionError(f"{expr}")

    exprs1 = [
        List.maybe(map(root, ellipsis_indices=[]), ellipsis_indices=[]) if root is not None else None for root in exprs1
    ]
    exprs2 = [
        List.maybe(map(root, ellipsis_indices=[]), ellipsis_indices=[]) if root is not None else None for root in exprs2
    ]

    return exprs1, exprs2


def cse(expressions, cse_concat=True, cse_in_markers=False, verbose=False):
    expressions = list(expressions)
    if any(expr is not None and not isinstance(expr, Expression) for expr in expressions):
        raise TypeError("Expected expressions to be of type Expression")

    # Find possible expressions, identified by their string representation
    str_to_common_expr = defaultdict(list)
    for root in expressions:
        if root is not None:
            for expr in root.all():
                if expr.parent is not None:
                    str_expr = str(expr)
                    str_to_common_expr[str_expr].append([expr])

                    if isinstance(expr, List):
                        for start_index in range(len(expr.children)):
                            for end_index in range(start_index, len(expr.children)):
                                children = expr.children[start_index : end_index + 1]
                                str_expr = " ".join([str(c) for c in children])
                                str_to_common_expr[str_expr].append(children)

    if verbose:
        print("CSE: All subexpressions")
        for k in str_to_common_expr.keys():
            print(f"    {k}")

    # Keep only expressions
    # 1. with at least one named axis
    # 2. where named axes are not also used outside the expression
    common_exprs = set()
    for str_expr in str_to_common_expr.keys():
        used_axis_ids = set()
        used_axis_names = set()
        for exprlist in str_to_common_expr[str_expr]:
            for expr in exprlist:
                for v in expr.all():
                    if isinstance(v, NamedAxis):
                        used_axis_ids.add(id(v))
                        used_axis_names.add(v.name)

        if len(used_axis_ids) == 0:
            continue

        axes_used_only_in_this_subexpression = True
        for root in expressions:
            if root is not None:
                for global_axis in root.all():
                    if isinstance(global_axis, NamedAxis) and global_axis.name in used_axis_names:
                        axes_used_only_in_this_subexpression = (
                            axes_used_only_in_this_subexpression and id(global_axis) in used_axis_ids
                        )

        if axes_used_only_in_this_subexpression:
            common_exprs.add(str_expr)

    common_exprs = [str_to_common_expr[k] for k in common_exprs]  # list of common_expr(=list of exprlist)

    if verbose:
        print("CSE: Removed expressions with axes that are also used outside the expression")
        for v in common_exprs:
            print(f"    {[' '.join([str(y) for y in x]) for x in v]}")

    def remove_duplicates(common_expr):
        new_common_expr = []
        for exprlist1 in common_expr:
            is_duplicate = False
            for exprlist2 in new_common_expr:
                is_duplicate = is_duplicate or (
                    len(exprlist1) == len(exprlist2)
                    and all(id(expr1) == id(expr2) for expr1, expr2 in zip(exprlist1, exprlist2))
                )
            if not is_duplicate:
                new_common_expr.append(exprlist1)
        return new_common_expr

    common_exprs = [remove_duplicates(exprlists) for exprlists in common_exprs]

    if verbose:
        print("CSE: Removed duplicates")
        for v in common_exprs:
            print(f"    {[' '.join([str(y) for y in x]) for x in v]} " f"{[[id(y) for y in x] for x in v]}")

    # Remove singletons
    def is_singleton(expr):
        if isinstance(expr, list):
            return len(expr) == 1 and is_singleton(expr[0])
        elif isinstance(expr, List):
            return is_singleton(expr.children)
        elif isinstance(expr, NamedAxis):
            return True
        elif isinstance(expr, UnnamedAxis):
            return True
        elif isinstance(expr, Marker):
            return is_singleton(expr.inner)
        else:
            return False

    common_exprs = [common_expr for common_expr in common_exprs if not is_singleton(common_expr[0])]

    if verbose:
        print("CSE: Removed singletons")
        for v in common_exprs:
            print(f"    {[' '.join([str(y) for y in x]) for x in v]}")

    # Remove expressions with/ in markers
    if cse_in_markers:
        common_exprs = [
            common_expr
            for common_expr in common_exprs
            if not any(isinstance(expr, Marker) for exprlist in common_expr for expr in exprlist for expr in expr.all())
        ]
    else:
        common_exprs = [
            common_expr
            for common_expr in common_exprs
            if not any(
                tlib.expr.stage2.is_marked(expr) for exprlist in common_expr for expr in exprlist for expr in expr.all()
            )
        ]

    # Remove expressions that contain concatenations
    if not cse_concat:
        common_exprs = [
            common_expr
            for common_expr in common_exprs
            if not any(
                isinstance(expr, Concatenation) for exprlist in common_expr for expr in exprlist for expr in expr.all()
            )
        ]

    if verbose:
        print("CSE: Removed expressions with markers")
        for v in common_exprs:
            print(f"    {[' '.join([str(y) for y in x]) for x in v]}")

    # Remove expressions at root level with len > 1
    common_exprs = [
        common_expr
        for common_expr in common_exprs
        if not (is_at_root(common_expr[0][0]) and (len(common_expr[0]) > 1 or len(common_expr[0][0]) > 1))
    ]

    if verbose:
        print("CSE: Removed subexpressions of root with len > 1")
        for v in common_exprs:
            print(f"    {[' '.join([str(y) for y in x]) for x in v]}")

    # Remove subexpressions of subexpressions
    def any_is_parent_of(parent, child):
        if isinstance(parent, list):
            return any(any_is_parent_of(p, child) for p in parent)
        elif isinstance(child, list):
            return any(any_is_parent_of(parent, c) for c in child)
        else:
            return child.parent is not None and (
                id(child.parent) == id(parent) or any_is_parent_of(parent, child.parent)
            )

    common_exprs = [
        common_expr
        for common_expr in common_exprs
        if not any(
            id(common_expr) != id(common_expr2) and any_is_parent_of(common_expr2, common_expr)
            for common_expr2 in common_exprs
        )
    ]

    if verbose:
        print("CSE: Removed subexpressions of subexpressions")
        for v in common_exprs:
            print(f"    {[' '.join([str(y) for y in x]) for x in v]}")

    # All subexpressions have been found. Now replace them with new Axis objects.
    def replace(expr):
        if isinstance(expr, list) and len(expr) == 1:
            return replace(expr[0])
        if not isinstance(expr, list):
            for idx, common_expr in enumerate(common_exprs):
                for exprlist in common_expr:
                    if len(exprlist) == 1 and id(expr) == id(exprlist[0]):
                        return [
                            NamedAxis(
                                f"cse.{idx}",
                                expr.ellipsis_indices,
                                begin_pos=expr.begin_pos,
                                end_pos=expr.end_pos,
                            )
                        ]

        if isinstance(expr, list):
            result = []
            i = 0
            while i < len(expr):
                # Check if a subexpression starts at position i
                exprlist_found = None
                for idx, common_expr in enumerate(common_exprs):  # noqa: B007
                    for exprlist in common_expr:
                        for j in range(len(exprlist)):
                            if i + j >= len(expr) or id(exprlist[j]) != id(expr[i + j]):
                                break
                        else:
                            exprlist_found = exprlist
                    if exprlist_found is not None:
                        break
                exprlist = exprlist_found

                if exprlist is not None:
                    assert len(exprlist) > 0
                    result.append(
                        NamedAxis(
                            f"cse.{idx}",
                            exprlist[0].ellipsis_indices,
                            begin_pos=exprlist[0].begin_pos,
                            end_pos=exprlist[-1].end_pos,
                        )
                    )
                    i += len(exprlist)
                else:
                    result.extend(replace(expr[i]))
                    i += 1

            return result
        elif isinstance(expr, NamedAxis):
            return [expr.__deepcopy__()]
        elif isinstance(expr, UnnamedAxis):
            return [expr.__deepcopy__()]
        elif isinstance(expr, List):
            return replace(expr.children)
        elif isinstance(expr, Concatenation):
            return [
                Concatenation(
                    [c2 for c1 in expr.children for c2 in replace(c1)],
                    expr.ellipsis_indices,
                    begin_pos=expr.begin_pos,
                    end_pos=expr.end_pos,
                )
            ]
        elif isinstance(expr, Marker):
            return [
                Marker.maybe(
                    List.maybe(replace(expr.inner), expr.ellipsis_indices),
                    expr.ellipsis_indices,
                    begin_pos=expr.begin_pos,
                    end_pos=expr.end_pos,
                )
            ]
        elif isinstance(expr, Composition):
            return [
                Composition(
                    List.maybe(replace(expr.inner), expr.ellipsis_indices),
                    expr.ellipsis_indices,
                    begin_pos=expr.begin_pos,
                    end_pos=expr.end_pos,
                )
            ]
        else:
            raise AssertionError()

    return [List.maybe(replace(root), ellipsis_indices=[]) if root is not None else None for root in expressions]


def expr_map(f):
    def outer(expr, *args, **kwargs):
        # Wrap the user function to return a list of expressions
        def f2(expr):
            t = f(expr, *args, **kwargs)
            if t is None:
                return None, expr_map.CONTINUE
            expr, signal = t

            if isinstance(expr, list) or expr is None:
                return expr, signal
            if isinstance(expr, List):
                return expr.children, signal
            elif isinstance(expr, Expression):
                return [expr], signal
            else:
                raise TypeError(f"Invalid return type {type(expr)}")

        return List.maybe(_expr_map(expr, f2))

    return outer


expr_map.CONTINUE = 1
expr_map.COPY_AND_STOP = 2
expr_map.REPLACE_AND_STOP = 3
expr_map.REPLACE_AND_CONTINUE = 4


def _expr_map(expr, f):
    exprs, signal = f(expr)
    if signal == expr_map.REPLACE_AND_STOP:
        assert isinstance(exprs, list)
        return exprs
    elif signal == expr_map.COPY_AND_STOP:
        return [expr.__deepcopy__()]
    elif signal == expr_map.REPLACE_AND_CONTINUE:
        return [c for expr in exprs for c in _expr_map(expr, f)]

    if isinstance(expr, NamedAxis):
        return [expr.__deepcopy__()]
    elif isinstance(expr, UnnamedAxis):
        return [expr.__deepcopy__()]
    elif isinstance(expr, Composition):
        return [Composition(List.maybe(_expr_map(expr.inner, f)))]
    elif isinstance(expr, List):
        return [c2 for c1 in expr.children for c2 in _expr_map(c1, f)]
    elif isinstance(expr, Concatenation):
        return [Concatenation([List.maybe(_expr_map(c, f)) for c in expr.children])]
    elif isinstance(expr, Marker):
        x = _expr_map(expr.inner, f)
        if len(x) == 0:
            # Drop empty marker
            return []
        else:
            return [Marker.maybe(List.maybe(x))]
    else:
        raise TypeError(f"Invalid expression type {type(expr)}")


@expr_map
def demark(expr):
    if isinstance(expr, Marker):
        return expr.inner, expr_map.REPLACE_AND_CONTINUE


def any_parent_is(expr, pred, include_self=True):
    if not include_self:
        if expr.parent is None:
            return False
        expr = expr.parent
    while expr is not None:
        if pred(expr):
            return True
        expr = expr.parent
    return False


def is_at_root(expr):
    return not any_parent_is(expr, lambda expr: isinstance(expr, Composition), include_self=False)


def is_marked(expr):
    return any_parent_is(expr, lambda expr: isinstance(expr, Marker))


def _get_marked(expr):
    if isinstance(expr, NamedAxis):
        return []
    elif isinstance(expr, UnnamedAxis):
        return []
    elif isinstance(expr, Marker):
        return [expr.inner.__deepcopy__()]
    elif isinstance(expr, Concatenation):
        return [Concatenation.maybe([x for c in expr.children for x in _get_marked(c)])]
    elif isinstance(expr, Composition):
        return [Composition(List.maybe(_get_marked(expr.inner)))]
    elif isinstance(expr, List):
        return [List.maybe([x for c in expr.children for x in _get_marked(c)])]
    else:
        raise TypeError(f"Invalid expression type {type(expr)}")


def get_marked(expr):
    return List.maybe(_get_marked(expr))


def get_unmarked(expr):
    return remove(expr, lambda expr: not is_marked(expr))


@expr_map
def replace(expr, f):
    expr = f(expr)
    if expr is not None:
        return expr, expr_map.REPLACE_AND_STOP


@expr_map
def remove(expr, pred):
    if pred(expr):
        return [], expr_map.REPLACE_AND_STOP
