from dataclasses import dataclass
from functools import reduce
import itertools
from operator import matmul, mul
from pprint import pformat
import re
from typing import Iterator, Literal, TypeAlias

import sympy

t_param = sympy.symbols("T")


@dataclass
class Polyhedron:
    vertex_count: int
    edge_count: int
    face_count: int

    def as_matrix(self) -> sympy.Matrix:
        return sympy.Matrix([[self.vertex_count, self.edge_count, self.face_count]]).T


def loeschian_norm(a: int, b: int) -> int:
    """Euclidean norm with respect to the sixth root of unity"""
    return a * a + a * b + b * b


def loeschian_numbers() -> Iterator[int]:
    """
    Numbers of the form a^2 + ab + b^2.
    No guarantees are made that this produces a sorted sequence without repeats.
    See https://oeis.org/A003136
    """
    for a in itertools.count(1):
        for b in range(a + 1):
            yield loeschian_norm(a, b)


GoldbergOperator: TypeAlias = Literal["c", "dk", "tk", "w", "g_T"]

goldberg_parameters_to_operators = {
    # Class I
    (1, 0): "",
    (2, 0): "dud = c",
    (3, 0): "du3d = tk",
    (4, 0): "du4d = cc",
    (5, 0): "du5d",
    (6, 0): "duu3d = ctk",
    # Class II
    (1, 1): "dk",
    (2, 2): "dkdud = dkc",
    (3, 3): "dkdu3d = dktk",
    (4, 4): "dkduud = dkcc",
    # Class III
    (2, 1): "w",
    (3, 1): "*",  # loeschian(3, 1) = 13 is prime
    (3, 2): "*",  # loeschian(3, 2) = 19 is prime
    (4, 1): "wdk",
    (4, 2): "wc",
}

# Reverse-lookup for the above table
goldberg_operators_to_parameters = {
    operation: parameter
    for parameter in goldberg_parameters_to_operators
    for operation in goldberg_parameters_to_operators[parameter].split(" = ")
    if operation != "*"
}

# Duals of operators appearing in the above table
_goldberg_operator_duals = {
    "k": "t",
    "c": "u2",
    "tk": "u3",
}
# Commutating pairs of operators
_goldberg_operator_commutations = {
    ("dk", "c"),
    ("dk", "w"),
}
# Multiply adjacent subdivisions
_gather_subdivisions = lambda formula: re.sub(
    "(u\\d*){2,}",
    lambda match: "u"
    + str(reduce(mul, [int(i) for i in match.group(0).split("u") if i])),
    formula,
)


def generalize_recipe(recipe: str | list[str], depth: int = 3) -> set[str]:
    """
    Rewrite a recipe by dualizing, commutating, and gathering operators together
    """
    # get rid of all of the regularization operators and double-dualizations
    if isinstance(recipe, str):
        recipe = recipe.split(" = ")

    ret = set(
        [
            # Need to run this twice to account for adjacent u's
            re.sub(
                "u(\\D)",
                "u2\\1",
                re.sub(
                    "u(\\D)",
                    "u2\\1",
                    re.sub("([KC]\\d*)", "", rec),
                ),
            ).replace("dd", "")
            for rec in recipe
        ]
    )
    prev = ret.copy()
    for _ in range(depth):
        new_recipes: set[str] = set(
            [
                dualize_polyhedron
                for op, dual in _goldberg_operator_duals.items()
                for left_comm, right_comm in _goldberg_operator_commutations
                for recipe in prev
                for dualize_operator in [
                    recipe.replace(op, f"d{dual}d").replace("dd", ""),
                    recipe.replace(dual, f"d{op}d").replace("dd", ""),
                ]
                for commute_operators in set(
                    [
                        dualize_operator,
                        dualize_operator.replace(
                            left_comm + right_comm, right_comm + left_comm
                        ),
                    ]
                )
                for subdivide_synonyms in set(
                    [
                        commute_operators,
                        _gather_subdivisions(commute_operators),
                    ]
                )
                # only implementing tetrahedron, dodecahedron, and icosahedron
                for dualize_polyhedron in set(
                    [
                        subdivide_synonyms,
                        subdivide_synonyms.replace("dT", "T"),
                        subdivide_synonyms.replace("dI", "D"),
                        subdivide_synonyms.replace("dD", "I"),
                    ]
                )
            ]
        )

        ret.update(new_recipes)
        prev = new_recipes

    return ret


def verify_recipes(
    left: dict[tuple[int, int], str], right: dict[tuple[int, int], str], message: str
) -> None:
    """Ensure all entries in `left` also exist in `right` and agree on at least one generalizationed recipe"""
    for i, left_recipe in left.items():
        assert i in right, f"Could not find tuple {i} on the RHS"
        if left_recipe.find("*") == 0 or right[i].find("*") == 0:
            continue

        left_recipes = generalize_recipe(left_recipe)
        right_recipes = generalize_recipe(right[i])
        assert (
            left_recipes & right_recipes
        ), f"{message}\nLeft recipes: {pformat(left_recipes)}, Right recipes: {pformat(right_recipes)}"


def from_same_eigenspace(
    uneigen: sympy.Matrix, values: list[int | sympy.Symbol]
) -> sympy.Matrix:
    """Diagonalize `uneigen`, then apply its eigenvectors to `values` as a diagonal matrix."""
    eigenvectors = uneigen.diagonalize()[0]
    return eigenvectors @ sympy.diag(*values) @ eigenvectors.inv()


# Matrices applied over column vectors of the form [v, e, f]
goldberg_matrix_operators: dict[GoldbergOperator, sympy.Matrix] = {
    "dk": sympy.Matrix(
        [
            [0, 2, 0],
            [0, 3, 0],
            [1, 0, 1],
        ]
    ),
    "c": sympy.Matrix(
        [
            [1, 2, 0],
            [0, 4, 0],
            [0, 1, 1],
        ]
    ),
    "w": sympy.Matrix(
        [
            [1, 4, 0],
            [0, 7, 0],
            [0, 2, 1],
        ]
    ),
}
goldberg_matrix_operators["tk"] = (
    goldberg_matrix_operators["dk"] @ goldberg_matrix_operators["dk"]
)
goldberg_matrix_operators["g_T"] = from_same_eigenspace(
    goldberg_matrix_operators["dk"], [t_param % 3, 1, t_param]
)


def apply_goldberg_operator(
    operator: GoldbergOperator | tuple[int, int] | str,
    poly: Polyhedron,
) -> Polyhedron:
    """Apply a (list of) operator(s) to a polyhedron. Operator can also be given as a parameter."""
    # Naive combinatorial method
    if isinstance(operator, tuple):
        return Polyhedron(
            *(goldberg_matrix_operators["g_T"] @ poly.as_matrix()).subs(
                t_param, loeschian_norm(*operator)
            )
        )

    # Try lookup parameter
    if (parameter := goldberg_operators_to_parameters.get(operator)) is not None:
        return Polyhedron(
            *(goldberg_matrix_operators["g_T"] @ poly.as_matrix()).subs(
                t_param, loeschian_norm(*parameter)
            )
        )
    elif (
        matrix := goldberg_matrix_operators.get(
            operator  # pyright: ignore[reportArgumentType]
        )
    ) is not None:
        return Polyhedron(*(matrix @ poly.as_matrix()))

    # Partition the recipe into operators for which we have the matrix
    partitioned = []
    for recipe in generalize_recipe(operator):
        partitioned.clear()
        bad_recipe = False
        while recipe != "":
            for op in goldberg_matrix_operators.keys():
                if recipe.find(op) == 0:
                    partitioned.append(op)
                    recipe = recipe[len(op) :]
                    break
            else:
                bad_recipe = True
                partitioned.clear()
                break
        if not bad_recipe:
            break

    if partitioned:
        # Technically you could just pointwise multiply the eigenvalues
        matrix = reduce(matmul, (goldberg_matrix_operators[i] for i in partitioned))
        return Polyhedron(*(matrix @ poly.as_matrix()))

    raise ValueError("Could not partition operators!")
