import itertools
from dataclasses import dataclass
from typing import Iterable

import numpy as np
from numpy.typing import NDArray
import matplotlib.pyplot as plt

root3 = 0.5 - (3**0.5) / 2j
embed_triangular = lambda i: (i[0] + i[1] * root3.real, i[1] * root3.imag)


def plot_triangular(plottype="plot"):
    """
    Plot an iterable of coordinates in a triangular grid using the member `plottype` of `pyplot`.
    Coordinates are in the form `(x, y)`.
    Extra arguments are passed to the plotting method.
    """

    def ret(coords: Iterable[tuple[int, int] | NDArray], *args, **kwargs):
        getattr(plt, plottype)(
            *list(zip(*map(embed_triangular, coords))), *args, **kwargs
        )

    return ret


def hex_positions(n: int, offset: int = 0) -> list[tuple[int, int]]:
    """
    List of coordinates which, in a triangular lattice, form the vertices of hexagons.

    The points (1, 1), (3, 0), and (0, 3) are all at the center of hexagons,
        so linear combinations of these are NOT part of the grid.
    """
    coordinates = [
        (i + offset, j + offset)
        for i in range(n)
        for j in range(n)
        if i + j < n and (i - j) % 3 != 0
    ]
    return coordinates


def draw_hex_lines(n: int, offset: int = 0):
    # vertices of hexagons in rightmost third of the the upper half plane
    hex_coordinates = hex_positions(n, offset)
    # find connections between a vertex and its (up to 3) neighbors
    lines = [
        list(
            itertools.chain.from_iterable(
                zip(
                    # the point being considered
                    itertools.repeat((i, j)),
                    # its neighbors which exist in the grid
                    filter(
                        lambda x: x in hex_coordinates,
                        [
                            (i + 1, j),
                            (i, j + 1),
                            (i + 1, j - 1),
                        ],
                    ),
                )
            )
        )
        for i, j in hex_coordinates
    ]
    for line in lines:
        # draw lines between connected points
        plot_triangular("plot")(line, "ko-")


def draw_triangular_grid(n: int, offset: int = 0, show_hexes: bool = False):
    plot_triangular("triplot")(
        [(i + offset, j + offset) for i in range(n) for j in range(n) if i + j < n],
        "-",
        color="0.7",
        lw=0.75,
    )
    if show_hexes:
        draw_hex_lines(n, offset=offset)
    plt.gca().axis("off")


class Triangle:
    def __init__(
        self,
        left: tuple[int, int] | NDArray,
        right: tuple[int, int] | NDArray,
        top: tuple[int, int] | NDArray,
    ):
        self.left = np.array(left)
        self.right = np.array(right)
        self.top = np.array(top)

    def __iter__(self):
        yield self.left
        yield self.right
        yield self.top
        yield self.left


def draw_class_i(a: int, fill_all: bool = False):
    """
    Plot a Class I Goldberg-Coxeter sector with parameter (a, 0)
    """
    origin = Triangle(
        left=(0, 0),
        right=(1, 0),
        top=(0, 1),
    )
    plot_triangular("fill")(origin, "r")

    terminal = Triangle(
        right=(a, a),
        left=(a - 1, a),
        top=(a, a - 1),  # actually the bottom
    )
    plot_triangular("fill")(terminal, "b")

    bottom_sector = Triangle(
        left=origin.left,
        right=(a, 0),
        top=(0, a),
    )
    top_sector = Triangle(
        left=(0, a),
        right=terminal.right,
        top=(a, 0),
    )

    # outline sector
    plot_triangular("plot")(
        [
            origin.top,
            origin.right,
            bottom_sector.right,
            terminal.top,
            terminal.left,
            bottom_sector.top,
            origin.top,
        ],
        color="0.25",
        lw=2,
    )

    if fill_all:
        plot_triangular("fill")(bottom_sector, "r")
        plot_triangular("fill")(top_sector, "b")


def draw_trapezoids(a: int):
    """Draw trapezoidal regions in a Class II Goldberg-Coxeter sector"""
    plot_triangular("fill")(
        [(0, 0), (0, a), (a, a), (2 * a, 0)],
        "r",
    )
    plot_triangular("fill")(
        [(0, a), (a, a), (a, 2 * a), (0, 3 * a)],
        "g",
    )
    plot_triangular("fill")(
        [(a, a), (a, 2 * a), (3 * a, 0), (2 * a, 0)],
        "b",
    )


def draw_classes_ii_iii(parameter: tuple[int, int], fill_all: int):
    """
    Plot a Class II or III Goldberg-Coxeter sector with parameter (a, b)

    Parameters where the second term is larger are NOT plotted correctly by fill_all.
    """
    # (a, b) specifies coordinates to the center of a hexagon
    a, b = parameter
    # (a, a) is the hexagon we turn at,
    # (2b, -b) is the length of the path after the 60 degree turn, so
    # (a+2b, a-b) is the center of the hexagon
    w = a + 2 * b - 1
    l = a - b

    # origin
    origin = Triangle(
        left=(0, 0),
        right=(1, 0),
        top=(0, 1),
    )
    plot_triangular("fill")(origin, "r")

    # terminal triangle #1
    terminal1 = Triangle(right=(w + 1, l), left=(w, l), top=(w, l + 1))
    plot_triangular("fill")(terminal1, "b")

    # terminal triangle #2
    terminal2 = Triangle(top=(-l, w + l + 1), left=(-l, w + l), right=(-l + 1, w + l))
    plot_triangular("fill")(terminal2, "g")

    # lines connecting triangles
    plot_triangular("plot")(
        [
            origin.top,
            origin.right,
            terminal1.left,
            terminal1.top,
            terminal2.right,
            terminal2.left,
            origin.top,
        ],
        color="0.25",
        lw=2,
    )

    if fill_all:
        if a == b:
            draw_trapezoids(a)
        else:
            plot_triangular("fill")(
                [
                    origin.top,
                    origin.right,
                    terminal1.left,
                    terminal1.top,
                    terminal2.right,
                    terminal2.left,
                    origin.top,
                ],
                color="0.4",
            )

            # lower parallelogram
            plot_triangular("fill")(
                [
                    origin.right,
                    (1, l),
                    terminal1.left,
                    terminal1.left + (0, -l),
                    origin.right,
                ],
                color="xkcd:light red",
            )

            # upper right parallelogram
            plot_triangular("fill")(
                [
                    terminal1.top,
                    terminal1.top + (-l, 0),
                    terminal2.right,
                    terminal2.right + (l, 0),
                    terminal1.top,
                ],
                color="xkcd:light blue",
            )

            # upper left parallelogram
            plot_triangular("fill")(
                [
                    terminal2.left,
                    terminal2.left + (l, -l),
                    origin.top,
                    origin.top + (-l, l),
                    terminal2.left,
                ],
                color="xkcd:light green",
            )

            # center triangle
            plot_triangular("fill")(
                [
                    (0, a - b),
                    (3 * b, a - b),
                    (0, a + 2 * b),
                    (0, a - b),
                ],
                color="0.8",
            )


def draw_path(a: int, b: int):
    plot_triangular("plot")(
        [(i, i) for i in range(a + 1)],
        marker="o",
        color="xkcd:light red",
        mfc="k",
        lw=3,
    )
    plt.text(
        *embed_triangular((a / 2 - 1, a / 2 + 1.5)),
        f"a = {a}",
        bbox={
            "facecolor": "xkcd:light red",
            "pad": 6,
        },
    )

    plot_triangular("plot")(
        [(a + 2 * i, a - i) for i in range(b + 1)],
        marker="o",
        color="xkcd:light blue",
        mfc="k",
        lw=3,
    )
    plt.text(
        *embed_triangular((a + b, a - b / 2 + 0.75)),
        f"b = {b}",
        bbox={
            "facecolor": "xkcd:light blue",
            "pad": 6,
        },
    )


def recenter_axes(a: int, b: int):
    minx, maxx = plt.xlim()
    miny, maxy = plt.ylim()
    xspan = maxx - minx
    yspan = maxy - miny

    extent_x = a + 2 * b + (a - b) * root3.real  # rightward extent due to blue triangle
    if a > b and min(a, b) != 0:
        extent_x += (
            -(a - b) + (1 + a - b) * root3.real
        )  # leftward extent due to green parallelogram

    extent_y = (2 * max(a, b) + min(a, b)) * root3.imag

    if min(a, b) == 0 or (extent_x * yspan / xspan) > extent_y:
        span_diff = extent_x * yspan / xspan - extent_y
        plt.xlim(-0.1, extent_x + 0.1)
        plt.ylim(span_diff / 2, (extent_x + 0.2) * yspan / xspan + span_diff / 2)
    else:
        span_diff = extent_y * xspan / yspan - extent_x
        # the figure extends below the y axis if a < b
        if a >= b:
            plt.ylim(-0.1, extent_y + 0.1)
        else:
            plt.ylim(
                -extent_y / (2 * a + 1) - 0.1, extent_y * (1 - (1 / (2 * a + 1))) + 0.1
            )

        plt.xlim(-span_diff / 2, (extent_y + 0.2) * xspan / yspan - span_diff / 2)


def draw_sector(parameter: tuple[int, int], fill_all: bool = False, path: bool = False):
    """
    Draw a sector from the Goldberg-Coxeter construction with a certain `parameter`.
    `offset` will translate the figure to the point `(offset, offset)` in a triangular plane.
    `fill_all` controls whether the sector should be filled.
    `path` controls whether to label the path described by `parameter`
    """
    draw_triangular_grid(50, offset=-10, show_hexes=True)
    # (a, b) specifies coordinates to the center of a hexagon
    a, b = parameter

    if min(parameter) == 0:
        draw_class_i(max(parameter), fill_all)
    else:
        draw_classes_ii_iii(parameter, fill_all)

    if path:
        draw_path(*parameter)

    recenter_axes(*parameter)


# if __name__ == "__main__":
#     draw_triangular_grid(24, True)
#     # draw_sector((4,2), point=False, path=True)
#     draw_sector((4, 2), offset=5, point=False, path=True)
#     plt.show()
