generate_data.py

#!/usr/bin/env python3
"""
Knowledge distillation: run editorai:mini as teacher to generate gold training data for jtafull.
Outputs data/train.jsonl — one chat example per line.
"""
import json
import pathlib
import time
import ollama

TEACHER_MODEL = "entity12208/editorai:mini"

MERGED_SYSTEM_PROMPT = """You are editorai:jtafull — a Geometry Dash level designer and coding assistant AI.

Return ONLY valid JSON for level design requests — no markdown, no explanations, no code fences.
For coding requests, return clean Python or JavaScript code only.

Available objects: block_black_gradient_square, block_grid_patterned_top_square, block_grid_patterned_outer_corner_square, block_grid_patterned_inner_corner_square, block_grid_patterned_inner_square, block_grid_patterned_top_pillar_square, block_grid_patterned_pillar_square, spike_black_gradient_spike, spike_non_colorable_spike_black_pit_hazard, portal_normal_gravity_portal, portal_flipped_gravity_portal, portal_cube_portal, portal_ship_portal, decor_tall_rod, decor_medium_rod, decor_short_rod, spike_large_decorative_spikes, spike_medium_decorative_spikes, spike_small_decorative_spikes, spike_very_small_decorative_spikes, effect_no_enter_effect, effect_fade_bottom_enter_effect, effect_fade_top_enter_effect, effect_fade_left_enter_effect, effect_fade_right_enter_effect, effect_small_to_big_enter_effect, effect_big_to_small_enter_effect, effect_start_position, effect_enable_ghost_trail, effect_disable_ghost_trail, effect_10_level_end_trigger, jump_pad_yellow_jump_pad, jump_orb_yellow_jump_orb, spike_half_black_gradient_spike, block_black_gradient_single_slab, decor_tall_chain, portal_orange_mirror_portal, portal_blue_mirror_portal, portal_ball_portal, large_fading_cloud, small_fading_cloud, effect_pulsing_filled_circle, effect_pulsing_circle, effect_pulsing_heart, effect_pulsing_diamond, effect_pulsing_star, effect_chaotic_enter_effect, effect_halve_left_enter_effect, effect_halve_right_enter_effect, effect_halve_enter_effect, effect_inverse_halve_enter_effect, effect_pulsing_music_note, hazard_non_colorable_wavy_black_pit_hazard, block_wavy_black_slab, block_wavy_black_slab_outer_corner, block_wavy_black_slab_inner_corner, block_wavy_black_slab_right_edge, block_wavy_black_slab_left_edge, obj_blue_gravity_pad, block_wavy_black_single_slab, block_cross_black_square, block_beveled_top_square, block_beveled_outer_corner_square, block_beveled_inner_corner_square, block_beveled_inner_square, block_beveled_top_pillar_square, block_beveled_pillar_square, block_mechanical_square, block_mechanical_top_square, block_mechanical_outer_corner_square, block_mechanical_inner_square, block_mechanical_top_pillar_square, block_mechanical_pillar_square, block_grid_patterned_square, obj_blue_gravity_orb, decor_large_decorative_gear, decor_medium_decorative_gear, decor_small_decorative_gear, spike_large_black_sawblade, spike_medium_black_sawblade, block_black_square, block_black_top_square, block_black_outer_corner_square, block_black_inner_corner_square, block_black_inner_square, block_black_top_pillar_square, block_black_pillar_square, decor_very_small_decorative_gear, spike_small_black_sawblade, portal_green_size_portal, speed_portal_half, speed_portal_normal, speed_portal_double, speed_portal_triple, alpha_trigger, toggle_trigger, pulse_trigger, speed_portal_quadruple, rotate_trigger, hide_player_trigger, show_player_trigger, stop_trigger, spawn_trigger

JSON Format:
{
  "analysis": "Reasoning about layout, rhythm, visual style, and design choices",
  "objects": [
    {"type": "block_black_gradient_square", "x": 0, "y": 10},
    {"type": "spike_black_gradient_spike", "x": 50, "y": 0}
  ]
}

COORDINATE SYSTEM:
  X = horizontal (10 = 1 grid cell, increases right). Y = vertical (0 = ground, increases up, always >= 0).

SPACING BY DIFFICULTY:
  EASY=50-70  MEDIUM=30-50  HARD=20-30  EXTREME=10-20

GAME MODES:
  CUBE: gravity-based jumps, 1-3 blocks tall obstacles, jump pads/orbs for rhythm
  SHIP: continuous flight, narrow corridors, ceiling/floor hazards, smooth wave-like paths
  BALL: gravity-flip mechanic, use portal_normal/flipped_gravity_portal frequently
  WAVE: tight diagonal movement, very small gaps, spike walls

DESIGN PRINCIPLES:
  - Build in rhythmic patterns — obstacles should feel musical, not random
  - Use decorative objects (gears, rods, chains, clouds) to add visual depth between hazards
  - Layer blocks: use corner/pillar variants to create proper geometry, not just flat walls
  - Place speed portals to vary pacing — slow sections for build-up, fast for climax
  - Use effect objects (pulsing shapes, enter effects) at section transitions
  - Mirror or echo obstacle patterns for cohesion
  - Longer sections (20+ objects) are better than short ones

CODING CAPABILITY:
  - Can write Python or JavaScript to generate, parse, or transform GD level JSON
  - Can write scripts to mirror levels, scale coordinates, filter by object type, or batch-generate sections
  - For code requests: return only the code, no JSON wrapper

TERRAIN RULES:
  - NEVER place all objects at the same Y value — levels must have height variation
  - Ground floor objects at Y=0, obstacles at Y=10-40, ceiling hazards at Y=50-80+
  - Build elevated platforms using stacked blocks at different Y heights
  - Use pillar/top_pillar/inner_corner/outer_corner block variants to create proper 3D-looking terrain
  - Platforms should rise and fall — no perfectly flat stretches longer than 3 objects

RULES: Only use listed type names. Y >= 0. Vary object types richly. Never repeat the same pattern twice in a row."""

PROMPTS = [
    # --- CUBE sections ---
    "Design an easy cube section with 5 obstacles",
    "Design a medium difficulty cube section with 10 obstacles",
    "Design a hard cube section with tight spike patterns",
    "Design an extreme demon cube section with 15 obstacles",
    "Create a cube section with jump pad rhythm pattern",
    "Create a cube section using jump orbs for a multi-jump sequence",
    "Design a cube section with a tall 3-block wall jump",
    "Design a cube section with alternating ground and ceiling spikes",
    "Create an easy auto cube section using only jump pads",
    "Design a cube section with a fake-out spike then safe path",
    "Create a cube section transitioning from easy to hard difficulty",
    "Design a 20-object cube section with mechanical block theme",
    "Create a cube section with beveled block architecture and decorative gears",
    "Design a cube section featuring wavy black slab terrain",
    "Create a cube section with a speed portal transition mid-section",
    "Design a cube section with grid patterned blocks forming a staircase",
    "Create an easy cube section for a beginner-friendly level",
    "Design a cube section with cross black square pillars",
    "Create a cube section with ghost trail effects",
    "Design a cube section with pulsing heart decorations",

    # --- SHIP sections ---
    "Design a medium ship section through a narrow cave",
    "Create a hard ship section with ceiling and floor spikes",
    "Design an easy ship section with wide open space",
    "Create a ship section with mirror portals for a flip effect",
    "Design a ship section with gravity portals for inversion",
    "Create an extreme ship section with tiny corridor and sawblades",
    "Design a ship section featuring wavy terrain on both sides",
    "Create a ship section with speed changes using speed portals",
    "Design a ship section with decorative chains on the walls",
    "Create a ship section transitioning to ball mode",

    # --- BALL sections ---
    "Design a medium ball section with gravity flips",
    "Create a hard ball section with precise gravity portal timing",
    "Design an easy ball section using blue gravity orbs",
    "Create a ball section with alternating floor and ceiling blocks",
    "Design a ball section with tight spike corridors",
    "Create a ball section transitioning to cube mode mid-section",
    "Design a ball section with blue gravity pads on the floor",

    # --- Multi-mode transitions ---
    "Design a section transitioning from cube to ship to ball",
    "Create a level section that goes ship → cube → ship",
    "Design a section with a gravity portal mid-cube section",
    "Create a section with a size portal making the cube small",

    # --- Themed / visual sections ---
    "Design a mechanical-themed section with gears and metal blocks",
    "Create a dark cave section with wavy black terrain and spikes",
    "Design a sky section with fading clouds and beveled blocks",
    "Create a section with pulsing star decorations throughout",
    "Design a section with a musical rhythm feel using pulsing effects",
    "Create a section with tall chain decorations on the walls",
    "Design a section with decorative rod pillars between obstacles",
    "Create a section with a chaotic enter effect at the start",
    "Design a section that uses halve enter effects for dramatic entry",
    "Create a dense decorative section with gears of all sizes",

    # --- Specific object focus ---
    "Design a section using only sawblades as hazards",
    "Create a section where all hazards are half spikes",
    "Design a section using pit hazards instead of spikes",
    "Create a section using only jump pads for navigation",
    "Design a section with only decorative objects, no hazards",
    "Create a section using triggers: alpha, toggle, pulse, rotate",
    "Design a section that uses spawn_trigger for timed objects",
    "Create a section with hide_player and show_player triggers",
    "Design a section using stop_trigger and rotate_trigger together",

    # --- Difficulty progressions ---
    "Design a 30-object section that starts easy and ends extreme",
    "Create a section with 3 difficulty waves: easy, hard, easy",
    "Design a section that builds tension with increasing speed portals",
    "Create a demon-level section with maximum obstacle density",
    "Design a 1-star auto level section requiring no player input",

    # --- Longer/complex sections ---
    "Design a complete 40-object level section with intro, buildup, and climax",
    "Create a 25-object mechanical factory section with gears and metal blocks",
    "Design a 30-object cave section with varying ceiling heights",
    "Create a 20-object section that perfectly mirrors itself",
    "Design a 35-object ship section through a winding tunnel",

    # --- Specific GD style references ---
    "Design a section in the style of a classic GD stereo madness level",
    "Create a section that feels like a geometry dash theory of everything level",
    "Design a section inspired by clubstep: dark mechanical theme, hard",
    "Create a section inspired by deadlocked: sci-fi, extreme difficulty",
    "Design a section inspired by fingerdash: colorful, fast, medium difficulty",

    # --- Coding prompts ---
    "Write a Python function that takes a GD level JSON and returns all spike positions",
    "Write a Python script to mirror a GD level JSON horizontally",
    "Write a Python function to scale all x coordinates in a level by a factor",
    "Write a Python script that generates a random easy cube section as JSON",
    "Write a JavaScript function to filter level objects by type prefix",
    "Write a Python function to find the rightmost object x position in a level",
    "Write a Python script to validate that all object types in a level are in the allowed list",
    "Write a Python function to shift all objects in a level by a given x offset",
    "Write a Python script to merge two GD level JSON objects arrays",
    "Write a Python function to count objects by type in a level JSON",
    "Write a JavaScript function to sort level objects by x position",
    "Write a Python script to find and remove duplicate objects at the same position",
    "Write a Python function to generate a staircase pattern of blocks",
    "Write a Python script to convert a level JSON to a simple ASCII art preview",
    "Write a Python function to generate a spike gauntlet at a given difficulty",

    # --- Mixed prompts ---
    "Design a hard cube section then write a Python function to mirror it",
    "Create a medium ship section and write a script to double its length",
    "Design a spike pattern and write Python code to repeat it 5 times with offset",
    "Create an extreme section then write code to validate its object types",
    "Design a 10-object section then write JavaScript to render it as HTML divs",
]


def generate(prompt: str) -> str | None:
    try:
        resp = ollama.chat(
            model=TEACHER_MODEL,
            messages=[{"role": "user", "content": prompt}],
            options={"temperature": 0.8, "num_predict": 4096},
        )
        return resp["message"]["content"].strip()
    except Exception as e:
        print(f"  ERROR: {e}")
        return None


def is_valid_response(text: str, prompt: str) -> bool:
    if not text:
        return False
    # coding prompts — just check it returned something non-trivial
    if any(kw in prompt.lower() for kw in ["write a python", "write a javascript", "write code"]):
        return len(text) > 50
    # level design — must be parseable JSON with objects key
    try:
        data = json.loads(text)
        return "objects" in data and len(data["objects"]) > 0
    except json.JSONDecodeError:
        return False


def main():
    out_path = pathlib.Path("data/train.jsonl")
    out_path.parent.mkdir(exist_ok=True)

    existing = 0
    if out_path.exists():
        existing = sum(1 for _ in out_path.open())
        print(f"Resuming — {existing} examples already in {out_path}")

    written = existing
    with out_path.open("a") as f:
        for i, prompt in enumerate(PROMPTS):
            if i < existing:
                continue
            print(f"[{i+1}/{len(PROMPTS)}] {prompt[:60]}...")
            response = generate(prompt)
            if is_valid_response(response, prompt):
                example = {
                    "messages": [
                        {"role": "system", "content": MERGED_SYSTEM_PROMPT},
                        {"role": "user", "content": prompt},
                        {"role": "assistant", "content": response},
                    ]
                }
                f.write(json.dumps(example) + "\n")
                f.flush()
                written += 1
                print(f"  OK ({written} total)")
            else:
                print(f"  SKIP — invalid response")
            time.sleep(0.2)

    print(f"\nDone. {written} examples written to {out_path}")


if __name__ == "__main__":
    main()