from dataclasses import dataclass
import sympy
from sympy.abc import x


def display_integral_root(
    root: int, multiplicity: int, pad_to: int | None = None
) -> str:
    return (
        "("
        # denote sign
        + ("\\pm" if root != 0 else "")
        + str(root)
        + ")^{"
        + str(multiplicity)
        # pad multiplicity
        + (
            ("\\phantom{" + ("0" * (pad_to - len(str(multiplicity))) + "}"))
            if pad_to is not None and len(str(multiplicity)) < pad_to
            else ""
        )
        + "}"
    )


@dataclass
class SymmetricSpectrum:
    """
    Dataclass for symmetric spectra -- those for which, if a spectrum contains a root,
    it also contains its negative.

    If a root polynomial is not symmetric, then its symmetric counterpart will *not* be specified.
    """

    # Either the integral root or a polynomial with the root, and the multiplicity
    roots: list[tuple[sympy.Expr | int, int]]
    not_shown_count: int = 0

    def to_markdown(self) -> str:
        """
        Display a spectrum as a markdown string
        """
        integer_roots = [
            (root, multiplicity)
            for root, multiplicity in self.roots
            if isinstance(root, int)
        ]
        # symmetrify polynomials
        polynomial_roots = [
            (root, multiplicity)
            for base_root, multiplicity in self.roots
            if isinstance(base_root, sympy.Expr)
            for root in set(
                [base_root, (-1) ** (sympy.degree(base_root)) * base_root.subs(x, -x)]
            )
        ]

        shown_roots = len(polynomial_roots) > 0 or len(integer_roots) > 0
        not_shown_string = (
            (
                f"*Not shown: {self.not_shown_count} other roots*"
                if shown_roots
                else f"*Not shown: all {self.not_shown_count} roots*"
            )
            if self.not_shown_count > 0
            else ""
        )

        # strictly integral spectra
        if len(polynomial_roots) == 0 and len(integer_roots) != 0:
            longest_integer = max(
                len(str(multiplicity)) for _, multiplicity in integer_roots
            )
            return (
                (
                    "$$\\begin{gather*}"
                    + " \\\\ ".join(
                        display_integral_root(*rm, pad_to=longest_integer)
                        for rm in integer_roots
                    )
                    + "\\end{gather*}$$"
                )
                if shown_roots
                else ""
            ) + not_shown_string

        # conditionally split integral roots onto multiple lines
        integral_roots_lines = (
            " ".join(display_integral_root(*rm) for rm in integer_roots)
            if len(integer_roots) < 5
            else (
                " ".join(
                    display_integral_root(*rm)
                    for rm in integer_roots[: len(integer_roots) // 2]
                )
                + "\\\\"
                + " ".join(
                    display_integral_root(*rm)
                    for rm in integer_roots[len(integer_roots) // 2 :]
                )
            )
        )

        return (
            (
                "$$\\begin{gather*}"
                + integral_roots_lines
                + " \\\\ "
                + " \\\\ ".join(
                    sympy.latex(root**multiplicity)
                    for root, multiplicity in polynomial_roots
                )
                + "\\end{gather*}$$"
            )
            if shown_roots
            else ""
        ) + not_shown_string


def count_roots_spectrum(spectrum: SymmetricSpectrum) -> int:
    return (
        sum(
            [
                (
                    (multiplicity * 2 if root != 0 else multiplicity)
                    if isinstance(root, int)
                    else (int(sympy.degree(root, x)) * multiplicity)
                )
                for root, multiplicity in spectrum.roots
            ]
        )
        + spectrum.not_shown_count
    )


@dataclass
class GraphData:
    spectrum: SymmetricSpectrum
    vertex_count: int | None = None
    edge_count: int | None = None
    distance_classes: list[int] | None = None