origami_server/tasks.py

"""Task definitions for origami RL training.

Each task defines a target shape as a reference FOLD crease pattern.
The LLM must discover a crease pattern that folds into the same shape.
"""

TASKS: dict[str, dict] = {
    "triangle": {
        "name": "triangle",
        "description": "Fold the paper in half diagonally to make a triangle",
        "difficulty": 1,
        "paper": {"width": 1.0, "height": 1.0},
        "target_fold": {
            "vertices_coords": [[0, 0], [1, 0], [1, 1], [0, 1]],
            "edges_vertices": [[0, 1], [1, 2], [2, 3], [3, 0], [0, 2]],
            "edges_assignment": ["B", "B", "B", "B", "V"],
            "edges_foldAngle": [0, 0, 0, 0, 180],
            "faces_vertices": [[0, 1, 2], [0, 2, 3]],
        },
    },
    "half_fold": {
        "name": "half_fold",
        "description": "Fold the paper in half horizontally",
        "difficulty": 1,
        "paper": {"width": 1.0, "height": 1.0},
        "target_fold": {
            "vertices_coords": [
                [0, 0], [1, 0], [1, 1], [0, 1], [0, 0.5], [1, 0.5],
            ],
            "edges_vertices": [
                [0, 1], [1, 5], [5, 2], [2, 3], [3, 4], [4, 0],
                [4, 5],
            ],
            "edges_assignment": ["B", "B", "B", "B", "B", "B", "V"],
            "edges_foldAngle": [0, 0, 0, 0, 0, 0, 180],
            "faces_vertices": [[0, 1, 5, 4], [4, 5, 2, 3]],
        },
    },
    "quarter_fold": {
        "name": "quarter_fold",
        "description": "Fold the paper into quarters (two perpendicular folds)",
        "difficulty": 2,
        "paper": {"width": 1.0, "height": 1.0},
        "target_fold": {
            "vertices_coords": [
                [0, 0], [0.5, 0], [1, 0],
                [0, 0.5], [0.5, 0.5], [1, 0.5],
                [0, 1], [0.5, 1], [1, 1],
            ],
            "edges_vertices": [
                [0, 1], [1, 2], [2, 5], [5, 8], [8, 7], [7, 6], [6, 3], [3, 0],
                [1, 4], [4, 7],
                [3, 4], [4, 5],
            ],
            "edges_assignment": [
                "B", "B", "B", "B", "B", "B", "B", "B",
                "V", "V", "V", "V",
            ],
            "edges_foldAngle": [
                0, 0, 0, 0, 0, 0, 0, 0,
                180, 180, 180, 180,
            ],
            "faces_vertices": [
                [0, 1, 4, 3],
                [1, 2, 5, 4],
                [3, 4, 7, 6],
                [4, 5, 8, 7],
            ],
        },
    },
    "letter_fold": {
        "name": "letter_fold",
        "description": "Tri-fold the paper like a letter (two parallel folds)",
        "difficulty": 2,
        "paper": {"width": 1.0, "height": 1.0},
        "target_fold": {
            "vertices_coords": [
                [0, 0], [1, 0],
                [0, 1/3], [1, 1/3],
                [0, 2/3], [1, 2/3],
                [0, 1], [1, 1],
            ],
            "edges_vertices": [
                [0, 1], [1, 3], [3, 5], [5, 7], [7, 6], [6, 4], [4, 2], [2, 0],
                [2, 3],
                [4, 5],
            ],
            "edges_assignment": [
                "B", "B", "B", "B", "B", "B", "B", "B",
                "V", "M",
            ],
            "edges_foldAngle": [
                0, 0, 0, 0, 0, 0, 0, 0,
                180, -180,
            ],
            "faces_vertices": [
                [0, 1, 3, 2],
                [2, 3, 5, 4],
                [4, 5, 7, 6],
            ],
        },
    },
    # --- Optigami patterns (from lib/patterns.ts) ---
    "waterbomb": {
        "name": "waterbomb",
        "description": "Create a waterbomb base with valley folds on diagonals and mountain folds on midlines of a square sheet",
        "difficulty": 3,
        "paper": {"width": 2.0, "height": 2.0},
        "target_fold": {
            "vertices_coords": [
                [-1, 1], [0, 1], [1, 1],
                [-1, 0], [0, 0], [1, 0],
                [-1, -1], [0, -1], [1, -1],
            ],
            "edges_vertices": [
                [0, 1], [1, 2], [2, 5], [5, 8], [8, 7], [7, 6], [6, 3], [3, 0],
                [0, 4], [2, 4], [6, 4], [8, 4],
                [1, 4], [3, 4], [5, 4], [7, 4],
            ],
            "edges_assignment": [
                "B", "B", "B", "B", "B", "B", "B", "B",
                "V", "V", "V", "V",
                "M", "M", "M", "M",
            ],
            "edges_foldAngle": [
                0, 0, 0, 0, 0, 0, 0, 0,
                180, 180, 180, 180,
                -180, -180, -180, -180,
            ],
            "faces_vertices": [
                [0, 1, 4], [1, 2, 4],
                [2, 5, 4], [5, 8, 4],
                [8, 7, 4], [7, 6, 4],
                [6, 3, 4], [3, 0, 4],
            ],
        },
    },
    "accordion": {
        "name": "accordion",
        "description": "Make an accordion (zig-zag) fold with alternating mountain and valley creases like a paper fan",
        "difficulty": 2,
        "paper": {"width": 2.0, "height": 2.0},
        "target_fold": {
            "vertices_coords": [
                [-1, 1], [-0.5, 1], [0, 1], [0.5, 1], [1, 1],
                [-1, -1], [-0.5, -1], [0, -1], [0.5, -1], [1, -1],
            ],
            "edges_vertices": [
                [0, 1], [1, 2], [2, 3], [3, 4],
                [5, 6], [6, 7], [7, 8], [8, 9],
                [0, 5], [4, 9],
                [1, 6], [2, 7], [3, 8],
            ],
            "edges_assignment": [
                "B", "B", "B", "B",
                "B", "B", "B", "B",
                "B", "B",
                "V", "M", "V",
            ],
            "edges_foldAngle": [
                0, 0, 0, 0,
                0, 0, 0, 0,
                0, 0,
                180, -180, 180,
            ],
            "faces_vertices": [
                [0, 1, 6, 5], [1, 2, 7, 6],
                [2, 3, 8, 7], [3, 4, 9, 8],
            ],
        },
    },
    "miura_ori": {
        "name": "miura_ori",
        "description": "Create a Miura-ori tessellation fold pattern on a 2x2 grid with offset zigzag vertices",
        "difficulty": 3,
        "paper": {"width": 2.0, "height": 2.0},
        "target_fold": {
            "vertices_coords": [
                [-1, 1], [0, 1.2], [1, 1],
                [-1, 0], [0, 0.2], [1, 0],
                [-1, -1], [0, -0.8], [1, -1],
            ],
            "edges_vertices": [
                [0, 1], [1, 2], [2, 5], [5, 8], [8, 7], [7, 6], [6, 3], [3, 0],
                [1, 4], [4, 7],
                [3, 4], [4, 5],
            ],
            "edges_assignment": [
                "B", "B", "B", "B", "B", "B", "B", "B",
                "M", "M",
                "V", "M",
            ],
            "edges_foldAngle": [
                0, 0, 0, 0, 0, 0, 0, 0,
                -180, -180,
                180, -180,
            ],
            "faces_vertices": [
                [0, 1, 4, 3], [1, 2, 5, 4],
                [3, 4, 7, 6], [4, 5, 8, 7],
            ],
        },
    },
}


def get_task(name: str | None = None) -> dict:
    """Get a task by name. Defaults to 'triangle'."""
    if name is None:
        name = "triangle"
    if name not in TASKS:
        raise ValueError(f"Unknown task '{name}'. Available: {list(TASKS.keys())}")
    return TASKS[name]


def list_tasks() -> list[str]:
    """List all available task names."""
    return list(TASKS.keys())