"""Draw a labeled quadrilateral sequence diagram for the plane depth sequencing paper."""
import matplotlib.pyplot as plt
import matplotlib.patches as patches
from matplotlib.patches import FancyArrowPatch
import numpy as np

sequence_data = [
    {"vertices": [-2, -1, 3, 7], "ordered": [-2, 7, -1, 3], "depths": [0, 1, 0, 1], "move": "Q_1", "type": "deep_diamond"},
    {"vertices": [-1, 3, 4, 8], "ordered": [3, -1, 4, 8], "depths": [1, 0, 1, 2], "move": "Q_2^{AD}", "type": "s_quad"},
    {"vertices": [2, 4, 5, 8], "ordered": [4, 2, 5, 8], "depths": [1, 0, 1, 2], "move": "Q_3^{AD}", "type": "s_quad"},
    {"vertices": [-1, 2, 4, 6], "ordered": [4, 2, 6, -1], "depths": [1, 0, 1, 0], "move": "Q_4^{LA}", "type": "shallow_diamond"},
]

COLORS = {
    "deep_diamond": "#B2DFDB",
    "shallow_diamond": "#FFE0B2",
    "s_quad": "#F8BBD0",
}

def draw_sequence_diagram(sequence_data: list[dict], output_path: str = "quad_sequence_diagram.png"):
    """Draw the full quadrilateral sequence diagram."""
    fig, ax = plt.subplots(figsize=(16, 5.5))

    step_width = 4.0
    y_center = 1.0
    x_positions = [i * step_width for i in range(len(sequence_data))]

    # Store vertex positions for connecting matching vertices
    quad_vertex_positions = []

    for i, quad_data in enumerate(sequence_data):
        x = x_positions[i]
        quad_type = quad_data["type"]
        ordered_verts = quad_data["ordered"]
        depths = quad_data["depths"]

        # Standard positions: top, right, bottom, left (or tweezers variant)
        if quad_type == "s_quad":
            # Tweezers shape: top, left-pinched-upper, right, left-pinched-lower
            r = 0.45
            pinch = 0.25
            std_positions = [
                (x, y_center + r),                    # e1: top
                (x - r + pinch, y_center + r * 0.35),  # a: left-upper
                (x + r, y_center),                    # e2: right
                (x - r + pinch, y_center - r * 0.35),  # b: left-lower
            ]
        else:
            # Diamond shape: top, right, bottom, left
            r = 0.45
            std_positions = [
                (x, y_center + r),        # e1: top
                (x + r, y_center),        # a: right
                (x, y_center - r),        # e2: bottom
                (x - r, y_center),        # b: left
            ]

        # Draw polygon
        polygon = patches.Polygon(
            std_positions,
            closed=True,
            facecolor=COLORS.get(quad_type, "#CCCCCC"),
            edgecolor="black",
            linewidth=1.5,
            alpha=0.75,
        )
        ax.add_patch(polygon)

        # Draw vertices and labels
        vertex_positions = {}
        for vert_idx, (vertex_id, (vx, vy), depth) in enumerate(zip(ordered_verts, std_positions, depths)):
            # Draw vertex point
            ax.plot(vx, vy, "ko", markersize=7, zorder=5)
            vertex_positions[vertex_id] = (vx, vy)

            # Draw vertex label below
            ax.text(vx, vy - 0.2, str(vertex_id), ha="center", va="top", fontsize=8, fontweight="bold",
                   bbox=dict(boxstyle="round,pad=0.15", facecolor="lightyellow", alpha=0.9, edgecolor="none"))

        # Draw quad label in center
        center_x = np.mean([p[0] for p in std_positions])
        center_y = np.mean([p[1] for p in std_positions])
        ax.text(center_x, center_y, f"${quad_data['move']}$", ha="center", va="center",
               fontsize=11, fontweight="bold",
               bbox=dict(boxstyle="round,pad=0.3", facecolor="white", edgecolor="black", alpha=0.85))

        # Draw depth labels on the left (only for first quad)
        if i == 0:
            unique_depths = sorted(set(depths))
            for d in unique_depths:
                label_y = y_center + 0.5 - d * 0.5
                ax.text(x - 1.0, label_y, f"d={d}", ha="right", va="center", fontsize=9, style="italic")

        quad_vertex_positions.append(vertex_positions)

        # Draw arrows connecting matching vertices to next quad
        if i < len(sequence_data) - 1:
            next_quad_verts = set(sequence_data[i + 1]["vertices"])
            for vertex_id, (vx, vy) in vertex_positions.items():
                if vertex_id in next_quad_verts:
                    next_quad_data = sequence_data[i + 1]
                    next_x = x_positions[i + 1]
                    next_idx = next_quad_data["ordered"].index(vertex_id)

                    if next_quad_data["type"] == "s_quad":
                        r = 0.45
                        pinch = 0.25
                        next_std_positions = [
                            (next_x, y_center + r),
                            (next_x - r + pinch, y_center + r * 0.35),
                            (next_x + r, y_center),
                            (next_x - r + pinch, y_center - r * 0.35),
                        ]
                    else:
                        r = 0.45
                        next_std_positions = [
                            (next_x, y_center + r),
                            (next_x + r, y_center),
                            (next_x, y_center - r),
                            (next_x - r, y_center),
                        ]

                    next_vx, next_vy = next_std_positions[next_idx]

                    # Draw connecting arrow
                    arrow = FancyArrowPatch(
                        (vx, vy), (next_vx, next_vy),
                        arrowstyle="->", color="gray", linewidth=0.8, alpha=0.5, zorder=1,
                        connectionstyle="arc3,rad=0.15"
                    )
                    ax.add_patch(arrow)

    # Set axis properties
    ax.set_xlim(-1.8, x_positions[-1] + 1.2)
    ax.set_ylim(-1.0, 2.2)
    ax.set_aspect("equal")
    ax.axis("off")

    plt.tight_layout()
    plt.savefig(output_path, dpi=150, bbox_inches="tight")
    print(f"Saved diagram to {output_path}")
    plt.show()


if __name__ == "__main__":
    draw_sequence_diagram(sequence_data, "quad_sequence_diagram.png")