"""Draw a graph with quadrilateral sequence diagram for the plane depth sequencing paper.

Usage:
  sage draw_quad_sequence.py --seed 42 --n 7 --output quad_sequence.png
"""
import os, sys
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", "..")))  # repo root for `lib`
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))  # sibling experiment modules
from PIL import Image, ImageDraw
import argparse
from pathlib import Path
from sage.all import graphs, Graph  # type: ignore
from sage.misc.randstate import set_random_seed  # type: ignore
from plane_depth_sequencing import (
    quadrilateral_sequencing,
    _quad_vertices,
    _level_edge_of_face,
    _quad_type,
    get_plane_depth_labelling,
)
from lib.tutte_embedding import tutte_embedding


def generate_sequence(seed: int, n: int) -> tuple[list[dict], dict, Graph, list, Graph]:
    """Generate a quadrilateral sequence and return (sequence_data, depth_labelling, original_graph, outer_cycle, deep_embedding_graph)."""
    set_random_seed(seed)
    g = graphs.RandomTriangulation(n)
    g.is_planar(set_embedding=True)
    embedding = g.get_embedding()
    faces = g.faces(embedding)
    outer_cycle = [u for u, _ in faces[0]]

    result = quadrilateral_sequencing(g, outer_cycle)
    sequence = result['sequence']
    move_codes = result['move_codes']
    depth_labelling = result['depth_labelling']
    g_prime = result['deep_embedding']

    move_names = {0: "AD", 1: "LA", 2: "J", 3: "RC"}
    quad_type_colors = {
        "deep_diamond": (178, 223, 219),
        "shallow_diamond": (255, 224, 178),
        "s_quad": (248, 187, 208),
    }

    # Pass 1: collect each quad's level-edge endpoints, apexes, and full vertex set.
    raw = []
    quads = sequence[:6]  # Limit to first 6 for readability
    for i, quad in enumerate(quads):
        quad_type = _quad_type(quad, depth_labelling)
        f1, f2 = list(quad)
        level_edge = _level_edge_of_face(f1, depth_labelling)
        p, q = list(level_edge)
        a = next(v for v in f1 if v not in level_edge)
        b = next(v for v in f2 if v not in level_edge)
        # Apex with the smaller depth goes on top, larger depth on bottom.
        top, bottom = (a, b) if depth_labelling[a] <= depth_labelling[b] else (b, a)
        move = move_names[move_codes[i - 1]] if i > 0 else ""
        move_label = f"Q_{i+1}" if i == 0 else f"Q_{i+1}^{{{move}}}"
        raw.append({
            "level": (p, q),
            "verts": {p, q, a, b},
            "top": top,
            "bottom": bottom,
            "move": move_label,
            "type": quad_type,
            "color": quad_type_colors[quad_type],
        })

    # Pass 2: chain the diamonds. The left/right corners of a diamond are its level
    # edge endpoints. A vertex shared with the PREVIOUS quad's level edge was that
    # quad's right corner, so it goes on the LEFT here; a vertex shared with the NEXT
    # quad's level edge goes on the RIGHT. Apex vertices are centered (top/bottom) and
    # impose no left/right constraint, so only level-edge sharing matters.
    sequence_data = []
    for i, r in enumerate(raw):
        p, q = r["level"]
        prev_level = set(raw[i - 1]["level"]) if i > 0 else set()
        next_level = set(raw[i + 1]["level"]) if i + 1 < len(raw) else set()

        p_prev, q_prev = p in prev_level, q in prev_level
        p_next, q_next = p in next_level, q in next_level

        if p_prev != q_prev:            # exactly one shared with prev -> left
            left = p if p_prev else q
            right = q if p_prev else p
        elif p_next != q_next:          # else exactly one shared with next -> right
            right = p if p_next else q
            left = q if p_next else p
        else:                           # no chain constraint: deterministic fallback
            right = max(p, q, key=str)
            left = q if right == p else p

        sequence_data.append({
            "left": left,
            "right": right,
            "top": r["top"],
            "bottom": r["bottom"],
            "depth_left": depth_labelling[left],
            "depth_right": depth_labelling[right],
            "depth_top": depth_labelling[r["top"]],
            "depth_bottom": depth_labelling[r["bottom"]],
            "move": r["move"],
            "type": r["type"],
            "color": r["color"],
        })

    return sequence_data, depth_labelling, g, outer_cycle, g_prime


def draw_graph(g_prime: Graph, outer_cycle: list, depth_labelling: dict, width: int, height: int) -> Image.Image:
    """Draw the deep embedding graph with Tutte embedding and depth labels in a square."""
    img = Image.new('RGB', (width, height), 'white')
    draw = ImageDraw.Draw(img)

    # Get Tutte embedding for the deep embedding graph
    pos = tutte_embedding(g_prime, outer_cycle)

    # Get data bounds
    xs = [p[0] for p in pos.values()]
    ys = [p[1] for p in pos.values()]
    x_min, x_max = min(xs), max(xs)
    y_min, y_max = min(ys), max(ys)
    x_range = x_max - x_min if x_max > x_min else 1
    y_range = y_max - y_min if y_max > y_min else 1

    # Create a square region in the center of the image
    margin = 30
    max_size = min(width, height) - 2 * margin
    graph_size = max_size

    # Center the square graph horizontally
    graph_x_min = (width - graph_size) // 2
    graph_x_max = graph_x_min + graph_size
    graph_y_min = margin
    graph_y_max = margin + graph_size

    def data_to_pixel(x: float, y: float) -> tuple[int, int]:
        px = graph_x_min + (x - x_min) / x_range * (graph_x_max - graph_x_min)
        py = graph_y_min + (y - y_min) / y_range * (graph_y_max - graph_y_min)
        return int(px), int(py)

    # Draw edges
    for u, v in g_prime.edges(labels=False):
        px1, py1 = data_to_pixel(pos[u][0], pos[u][1])
        px2, py2 = data_to_pixel(pos[v][0], pos[v][1])
        # Color level edges differently
        if depth_labelling[u] == depth_labelling[v]:
            draw.line([(px1, py1), (px2, py2)], fill=(150, 150, 150), width=1)
        else:
            draw.line([(px1, py1), (px2, py2)], fill=(50, 50, 50), width=1)

    # Draw vertices
    for v, (x, y) in pos.items():
        px, py = data_to_pixel(x, y)
        depth = depth_labelling[v]
        if v in outer_cycle:
            color = (25, 118, 210)  # blue
        else:
            color = (0, 0, 0)  # black
        draw.ellipse([px-5, py-5, px+5, py+5], fill=color)
        # Draw label with vertex ID and depth
        label = f"{v}^{depth}"
        draw.text((px-12, py-20), label, fill='black')

    return img


def draw_diagram(sequence_data: list[dict], depth_labelling: dict, output_path: str):
    """Draw and save the sequence diagram."""
    width = 300 + len(sequence_data) * 400
    height = 380
    img = Image.new('RGB', (width, height), 'white')
    draw = ImageDraw.Draw(img)

    # Data-to-pixel coordinate mapping
    data_x_min = -2.0
    data_x_max = 0.5 + len(sequence_data) * 3.5
    data_y_min = -1.1
    data_y_max = 2.3
    pixel_x_min = 80
    pixel_x_max = width - 40
    pixel_y_min = 40
    pixel_y_max = height - 40

    def data_to_pixel(x: float, y: float) -> tuple[int, int]:
        px = pixel_x_min + (x - data_x_min) / (data_x_max - data_x_min) * (pixel_x_max - pixel_x_min)
        py = pixel_y_max - (y - data_y_min) / (data_y_max - data_y_min) * (pixel_y_max - pixel_y_min)
        return int(px), int(py)

    step_width = 3.5
    r = 0.45

    # Depth -> data-y: smaller depth higher on screen (larger data-y).
    all_depths = [
        d for quad in sequence_data
        for d in (quad["depth_left"], quad["depth_right"], quad["depth_top"], quad["depth_bottom"])
    ]
    min_d, max_d = min(all_depths), max(all_depths)
    y_of = lambda d: 1.0 - (d - min_d) * 0.5

    for i, quad_data in enumerate(sequence_data):
        x = i * step_width
        color = quad_data["color"]

        # Place the four corners. The level edge {left,right} is horizontal (both at
        # the same depth); the apexes sit above (smaller depth) and below (larger).
        p_left = data_to_pixel(x - r, y_of(quad_data["depth_left"]))
        p_right = data_to_pixel(x + r, y_of(quad_data["depth_right"]))
        p_top = data_to_pixel(x, y_of(quad_data["depth_top"]))
        p_bottom = data_to_pixel(x, y_of(quad_data["depth_bottom"]))

        # Polygon in cyclic order so the diamond is non-self-intersecting.
        corners = [p_left, p_top, p_right, p_bottom]
        draw.polygon(corners, fill=color, outline='black', width=2)
        for px, py in corners:
            draw.ellipse([px-5, py-5, px+5, py+5], fill='black')

        # Vertex labels, offset away from the diamond body.
        labels = [
            (p_left, quad_data["left"], (-18, -6)),
            (p_right, quad_data["right"], (10, -6)),
            (p_top, quad_data["top"], (-4, -18)),
            (p_bottom, quad_data["bottom"], (-4, 8)),
        ]
        for (px, py), vid, (dx, dy) in labels:
            draw.text((px + dx, py + dy), str(vid), fill='black')

        # Quad move label at the centroid.
        cx = sum(p[0] for p in corners) // 4
        cy = sum(p[1] for p in corners) // 4
        draw.text((cx-15, cy-6), quad_data["move"], fill='black')

        # Depth gridline labels on the left of the first quad.
        if i == 0:
            for d in range(min_d, max_d + 1):
                label_px, label_py = data_to_pixel(x - 1.2, y_of(d))
                draw.text((label_px-28, label_py-6), f"d={d}", fill=(80, 80, 80))

    if output_path:
        img.save(output_path)
        print(f"Saved diagram to {output_path}")
    print(f"Sequence diagram size: {img.size}")
    return img


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--seed", type=int, default=42, help="Random seed for graph generation")
    parser.add_argument("--n", type=int, default=7, help="Number of vertices in the random triangulation")
    parser.add_argument("--output", type=str, default="quad_sequence_diagram.png", help="Output file path")
    args = parser.parse_args()

    print(f"Generating sequence for n={args.n} with seed={args.seed}")
    sequence_data, depth_labelling, g, outer_cycle, g_prime = generate_sequence(args.seed, args.n)
    print(f"Generated {len(sequence_data)} quadrilaterals")
    print(f"Deep embedding has {g_prime.order()} vertices (original had {g.order()})")

    # Draw deep embedding graph with depths
    graph_width = 300 + len(sequence_data) * 400
    graph_height = 500
    graph_img = draw_graph(g_prime, outer_cycle, depth_labelling, graph_width, graph_height)

    # Draw sequence
    seq_width = graph_width
    seq_height = 380
    seq_img = draw_diagram(sequence_data, depth_labelling, None)  # Return image instead of saving

    # Combine images vertically
    combined_width = graph_width
    combined_height = graph_height + seq_height + 20  # 20px gap
    combined_img = Image.new('RGB', (combined_width, combined_height), 'white')
    combined_img.paste(graph_img, (0, 0))
    combined_img.paste(seq_img, (0, graph_height + 20))

    combined_img.save(args.output)
    print(f"Saved diagram to {args.output}")
    print(f"Combined image size: {combined_img.size}")