python_tools/cadquery_env.py

from __future__ import annotations

import argparse
import json
import math
import re
import shutil
import subprocess
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any

import numpy as np
import trimesh
from PIL import Image, ImageDraw
from scipy.spatial import cKDTree


APP_ROOT = Path(__file__).resolve().parents[1]
REPO_ROOT = APP_ROOT.parent
MODELS_ROOT = REPO_ROOT / "3d-models"
RUNS_ROOT = APP_ROOT / "runs" / "cadquery-env"
REFERENCE_ROOT = APP_ROOT / "runs" / "cadquery-reference"
RUNNER = APP_ROOT / "python_tools" / "cadquery_code_runner.py"
DEFAULT_GLB = MODELS_ROOT / "ikea_markus_office_chair.glb"
DEFAULT_IDEAL_CODE = MODELS_ROOT / "ikea_markus_idealish_code.md"
DEFAULT_TASKS = APP_ROOT / "data" / "cad_tasks.json"

VIEWS = {
    "front": {"axes": (0, 2), "flip_x": False},
    "back": {"axes": (0, 2), "flip_x": True},
    "left": {"axes": (1, 2), "flip_x": False},
    "right": {"axes": (1, 2), "flip_x": True},
    "top": {"axes": (0, 1), "flip_x": False},
}

SEMANTIC_HINTS = [
    "seat",
    "backrest",
    "headrest",
    "armrest",
    "gas_cylinder",
    "central_column",
    "star_base",
    "caster",
    "lumbar",
    "mechanism",
]


@dataclass(frozen=True)
class MeshBundle:
    mesh: trimesh.Trimesh
    normalized: trimesh.Trimesh
    bbox: dict[str, float]


def safe_slug(value: str) -> str:
    slug = re.sub(r"[^a-zA-Z0-9_.-]+", "-", value).strip("-")
    return slug[:80] or "run"


def extract_code(text: str) -> str:
    match = re.search(r"```(?:python|py)?\s*\n(.*?)```", text, re.IGNORECASE | re.DOTALL)
    return (match.group(1) if match else text).strip()


def read_code(path: Path) -> str:
    return extract_code(path.read_text())


def read_task_spec(path_or_id: str | None) -> dict[str, Any] | None:
    if not path_or_id:
        return None
    path = Path(path_or_id)
    if path.exists():
        return json.loads(path.read_text())
    if DEFAULT_TASKS.exists():
        for task in json.loads(DEFAULT_TASKS.read_text()):
            if task.get("id") == path_or_id:
                return task
    raise FileNotFoundError(f"Task spec not found: {path_or_id}")


def concise_error(stdout: str = "", stderr: str = "") -> str:
    text = (stderr or stdout or "").strip()
    if not text:
        return ""
    lines = [line.strip() for line in text.splitlines() if line.strip()]
    for line in reversed(lines):
        if re.search(r"(Error|Exception|Traceback|NameError|TypeError|ValueError|AttributeError|ImportError|RuntimeError)", line):
            return line[:500]
    return lines[-1][:500] if lines else ""


def mesh_from_file(path: Path) -> trimesh.Trimesh:
    loaded = trimesh.load(path, force="scene")
    if isinstance(loaded, trimesh.Scene):
        mesh = loaded.to_geometry() if hasattr(loaded, "to_geometry") else loaded.dump(concatenate=True)
    else:
        mesh = loaded
    if isinstance(mesh, list):
        mesh = trimesh.util.concatenate(mesh)
    if not isinstance(mesh, trimesh.Trimesh):
        raise ValueError(f"Could not load mesh from {path}")
    mesh = mesh.copy()
    mesh.remove_unreferenced_vertices()
    return mesh


def canonicalize_largest_axis_to_z(mesh: trimesh.Trimesh) -> trimesh.Trimesh:
    canonical = mesh.copy()
    bounds = np.asarray(canonical.bounds, dtype=float)
    dims = bounds[1] - bounds[0]
    up_axis = int(np.argmax(dims))
    if up_axis == 2:
        return canonical

    vertices = np.asarray(canonical.vertices, dtype=float).copy()
    if up_axis == 1:
        # Common GLB convention for this asset: Y is vertical.
        vertices = vertices[:, [0, 2, 1]]
    elif up_axis == 0:
        vertices = vertices[:, [1, 2, 0]]
    canonical.vertices = vertices
    canonical.remove_unreferenced_vertices()
    return canonical


def bbox_dict(mesh: trimesh.Trimesh) -> dict[str, float]:
    bounds = np.asarray(mesh.bounds, dtype=float)
    dims = bounds[1] - bounds[0]
    return {
        "xmin": float(bounds[0, 0]),
        "xmax": float(bounds[1, 0]),
        "ymin": float(bounds[0, 1]),
        "ymax": float(bounds[1, 1]),
        "zmin": float(bounds[0, 2]),
        "zmax": float(bounds[1, 2]),
        "xlen": float(dims[0]),
        "ylen": float(dims[1]),
        "zlen": float(dims[2]),
    }


def normalize_mesh(mesh: trimesh.Trimesh, target_height: float = 1000.0) -> trimesh.Trimesh:
    normalized = mesh.copy()
    bounds = np.asarray(normalized.bounds, dtype=float)
    dims = bounds[1] - bounds[0]
    height = float(max(dims[2], 1e-9))
    scale = target_height / height
    normalized.apply_scale(scale)
    bounds = np.asarray(normalized.bounds, dtype=float)
    center_xy = (bounds[0, :2] + bounds[1, :2]) / 2.0
    normalized.apply_translation([-center_xy[0], -center_xy[1], -bounds[0, 2]])
    normalized.remove_unreferenced_vertices()
    return normalized


def mesh_bundle(path: Path) -> MeshBundle:
    mesh = mesh_from_file(path)
    normalized = normalize_mesh(mesh)
    return MeshBundle(mesh=mesh, normalized=normalized, bbox=bbox_dict(mesh))


def run_cadquery(code: str, out_dir: Path, name: str, timeout: int = 180) -> dict[str, Any]:
    out_dir.mkdir(parents=True, exist_ok=True)
    started = time.time()
    proc = subprocess.run(
        [sys.executable, str(RUNNER), "--out-dir", str(out_dir), "--name", safe_slug(name)],
        cwd=APP_ROOT,
        input=json.dumps({"code": code}),
        text=True,
        capture_output=True,
        timeout=timeout,
        env={
            **dict(**__import__("os").environ),
            "PYTHONPATH": str(APP_ROOT / "python_tools"),
            "XDG_CACHE_HOME": str(APP_ROOT / ".cache"),
        },
    )
    elapsed_ms = int((time.time() - started) * 1000)
    output: dict[str, Any] | None = None
    if proc.returncode == 0:
        lines = [line for line in proc.stdout.splitlines() if line.strip()]
        if lines:
            output = json.loads(lines[-1])
    return {
        "ok": proc.returncode == 0 and output is not None,
        "returncode": proc.returncode,
        "stdout": proc.stdout,
        "stderr": proc.stderr,
        "elapsed_ms": elapsed_ms,
        "output": output,
    }


def edge_counts(mesh: trimesh.Trimesh) -> tuple[int, int]:
    if len(mesh.faces) == 0:
        return 0, 0
    inverse = mesh.edges_unique_inverse
    counts = np.bincount(inverse, minlength=len(mesh.edges_unique))
    boundary_edges = int(np.sum(counts == 1))
    non_manifold_edges = int(np.sum(counts > 2))
    return boundary_edges, non_manifold_edges


def topology_metrics(mesh: trimesh.Trimesh) -> dict[str, Any]:
    components = mesh.split(only_watertight=False)
    boundary_edges, non_manifold_edges = edge_counts(mesh)
    area_faces = np.asarray(mesh.area_faces) if len(mesh.faces) else np.array([])
    degenerate_faces = int(np.sum(area_faces < 1e-8))
    return {
        "vertices": int(len(mesh.vertices)),
        "faces": int(len(mesh.faces)),
        "components": int(len(components)),
        "floating_components": int(max(0, len(components) - 1)),
        "watertight": bool(mesh.is_watertight),
        "winding_consistent": bool(mesh.is_winding_consistent),
        "boundary_edges": boundary_edges,
        "non_manifold_edges": non_manifold_edges,
        "degenerate_faces": degenerate_faces,
    }


def topology_reward(metrics: dict[str, Any]) -> dict[str, Any]:
    # Markus is an assembly benchmark: a good chair can contain many valid solids.
    # For monolithic hooks/brackets we can add a stricter single-body reward later.
    component_count = metrics["components"]
    component_score = 0.0
    if 1 <= component_count <= 120:
        component_score = 0.25
    elif component_count <= 250:
        component_score = 0.12

    face_count = max(1, metrics["faces"])
    boundary_ratio = metrics["boundary_edges"] / max(1, metrics["faces"])
    non_manifold_ratio = metrics["non_manifold_edges"] / face_count
    degenerate_ratio = metrics["degenerate_faces"] / face_count

    score = 0.0
    score += component_score
    score += 0.15 if metrics["watertight"] else 0.05
    score += 0.20 * max(0.0, 1.0 - boundary_ratio * 250.0)
    score += 0.15 * max(0.0, 1.0 - non_manifold_ratio * 250.0)
    score += 0.15 * max(0.0, 1.0 - degenerate_ratio * 250.0)
    score += 0.10 if 50 <= metrics["faces"] <= 750000 else 0.03
    penalties = {
        "too_many_components_penalty": -0.20 if component_count > 250 else 0.0,
        "high_boundary_ratio_penalty": -0.10 if boundary_ratio > 0.02 else 0.0,
        "high_non_manifold_ratio_penalty": -0.10 if non_manifold_ratio > 0.02 else 0.0,
        "high_degenerate_ratio_penalty": -0.10 if degenerate_ratio > 0.02 else 0.0,
    }
    return {"score": float(max(0.0, min(1.0, score + sum(penalties.values())))), "penalties": penalties}


def bbox_gap(a: np.ndarray, b: np.ndarray) -> float:
    lower_gap = np.maximum(a[0] - b[1], 0.0)
    upper_gap = np.maximum(b[0] - a[1], 0.0)
    return float(np.linalg.norm(np.maximum(lower_gap, upper_gap)))


def contact_metrics(mesh: trimesh.Trimesh) -> dict[str, Any]:
    components = list(mesh.split(only_watertight=False))
    if len(components) <= 1:
        return {
            "score": 1.0,
            "components_checked": len(components),
            "mean_gap_ratio": 0.0,
            "max_gap_ratio": 0.0,
            "large_gap_components": 0,
        }

    bounds = []
    root_bbox = bbox_dict(mesh)
    height = max(root_bbox["zlen"], 1e-9)
    for component in components:
        bbox = np.asarray(component.bounds, dtype=float)
        dims = bbox[1] - bbox[0]
        diag = float(np.linalg.norm(dims))
        # Ignore tiny mesh shards from tessellation/import noise.
        if diag >= 0.015 * height:
            bounds.append((bbox, diag))
    if len(bounds) <= 1:
        return {
            "score": 1.0,
            "components_checked": len(bounds),
            "mean_gap_ratio": 0.0,
            "max_gap_ratio": 0.0,
            "large_gap_components": 0,
        }

    gaps = []
    for i, (bbox, _diag) in enumerate(bounds):
        nearest = min(bbox_gap(bbox, other_bbox) for j, (other_bbox, _other_diag) in enumerate(bounds) if i != j)
        gaps.append(nearest / height)
    gaps_array = np.asarray(gaps, dtype=float)
    mean_gap = float(np.mean(gaps_array))
    max_gap = float(np.max(gaps_array))
    large_gap_components = int(np.sum(gaps_array > 0.08))
    score = math.exp(-mean_gap * 28.0) * math.exp(-max_gap * 4.0)
    score -= min(0.40, large_gap_components * 0.04)
    return {
        "score": float(max(0.0, min(1.0, score))),
        "components_checked": len(bounds),
        "mean_gap_ratio": mean_gap,
        "max_gap_ratio": max_gap,
        "large_gap_components": large_gap_components,
    }


def project_vertices(vertices: np.ndarray, view: str) -> np.ndarray:
    if view == "isometric":
        rz = math.radians(42)
        rx = math.radians(58)
        rot_z = np.array([[math.cos(rz), -math.sin(rz), 0], [math.sin(rz), math.cos(rz), 0], [0, 0, 1]])
        rot_x = np.array([[1, 0, 0], [0, math.cos(rx), -math.sin(rx)], [0, math.sin(rx), math.cos(rx)]])
        rotated = vertices @ (rot_z @ rot_x).T
        return rotated[:, [0, 2]]
    spec = VIEWS[view]
    pts = vertices[:, list(spec["axes"])]
    if spec["flip_x"]:
        pts = pts.copy()
        pts[:, 0] *= -1
    return pts


def silhouette_mask(mesh: trimesh.Trimesh, view: str, size: int = 512, padding: int = 28) -> Image.Image:
    pts = project_vertices(np.asarray(mesh.vertices, dtype=float), view)
    mins = pts.min(axis=0)
    maxs = pts.max(axis=0)
    span = np.maximum(maxs - mins, 1e-9)
    scale = (size - 2 * padding) / float(max(span))
    xy = (pts - mins) * scale + padding
    xy[:, 1] = size - xy[:, 1]
    image = Image.new("L", (size, size), 0)
    draw = ImageDraw.Draw(image)
    for face in mesh.faces:
        poly = [(float(xy[i, 0]), float(xy[i, 1])) for i in face]
        draw.polygon(poly, fill=255)
    return image


def view_depth(vertices: np.ndarray, view: str) -> np.ndarray:
    if view == "front":
        return vertices[:, 1]
    if view == "back":
        return -vertices[:, 1]
    if view == "left":
        return vertices[:, 0]
    if view == "right":
        return -vertices[:, 0]
    if view == "top":
        return vertices[:, 2]
    return vertices @ np.array([0.42, -0.55, 0.72])


def projected_image_points(mesh: trimesh.Trimesh, view: str, size: int = 720, padding: int = 42) -> np.ndarray:
    pts = project_vertices(np.asarray(mesh.vertices, dtype=float), view)
    mins = pts.min(axis=0)
    maxs = pts.max(axis=0)
    span = np.maximum(maxs - mins, 1e-9)
    scale = (size - 2 * padding) / float(max(span))
    xy = (pts - mins) * scale + padding
    xy[:, 1] = size - xy[:, 1]
    return xy


def color_render_image(mesh: trimesh.Trimesh, view: str, size: int = 720, padding: int = 42) -> Image.Image:
    vertices = np.asarray(mesh.vertices, dtype=float)
    xy = projected_image_points(mesh, view, size=size, padding=padding)
    depths = view_depth(vertices, view)
    face_depths = depths[mesh.faces].mean(axis=1)
    order = np.argsort(face_depths)
    z_values = vertices[:, 2]
    zmin = float(z_values.min())
    zspan = float(max(z_values.max() - zmin, 1e-9))

    image = Image.new("RGB", (size, size), (244, 247, 250))
    draw = ImageDraw.Draw(image)
    grid = (222, 228, 235)
    for value in range(0, size, 72):
        draw.line([(value, 0), (value, size)], fill=grid)
        draw.line([(0, value), (size, value)], fill=grid)

    palette = [
        np.array([82, 111, 135]),
        np.array([94, 147, 113]),
        np.array([198, 95, 73]),
        np.array([207, 151, 63]),
        np.array([92, 124, 181]),
    ]
    for face_index in order:
        face = mesh.faces[face_index]
        center_z = float(vertices[face, 2].mean())
        height_mix = (center_z - zmin) / zspan
        base = palette[int(height_mix * (len(palette) - 1)) % len(palette)]
        shade = 0.68 + 0.32 * height_mix
        color = tuple(np.clip(base * shade + 20, 0, 255).astype(int).tolist())
        poly = [(float(xy[i, 0]), float(xy[i, 1])) for i in face]
        draw.polygon(poly, fill=color, outline=(63, 78, 92))

    draw.rectangle([0, 0, size - 1, size - 1], outline=(210, 218, 228), width=2)
    return image


def save_silhouettes(mesh: trimesh.Trimesh, out_dir: Path) -> dict[str, str]:
    out_dir.mkdir(parents=True, exist_ok=True)
    paths: dict[str, str] = {}
    for view in [*VIEWS.keys(), "isometric"]:
        image = silhouette_mask(mesh, view)
        path = out_dir / f"{view}.png"
        image.save(path)
        paths[view] = str(path)
    return paths


def save_color_renders(mesh: trimesh.Trimesh, out_dir: Path) -> dict[str, str]:
    out_dir.mkdir(parents=True, exist_ok=True)
    paths: dict[str, str] = {}
    for view in [*VIEWS.keys(), "isometric"]:
        image = color_render_image(mesh, view)
        path = out_dir / f"{view}.png"
        image.save(path)
        paths[view] = str(path)
    return paths


def mask_iou(a_path: Path, b_path: Path) -> float:
    a = np.asarray(Image.open(a_path).convert("L")) > 0
    b = np.asarray(Image.open(b_path).convert("L")) > 0
    union = np.logical_or(a, b).sum()
    if union == 0:
        return 0.0
    return float(np.logical_and(a, b).sum() / union)


def sample_points(mesh: trimesh.Trimesh, count: int = 3000) -> np.ndarray:
    state = np.random.get_state()
    np.random.seed(7)
    try:
        points = mesh.sample(count)
    finally:
        np.random.set_state(state)
    return np.asarray(points, dtype=np.float32)


def save_reference_mesh(name: str, mesh: trimesh.Trimesh, root: Path) -> dict[str, Any]:
    ref_dir = root / name
    ref_dir.mkdir(parents=True, exist_ok=True)
    normalized = normalize_mesh(mesh)
    stl_path = ref_dir / f"{name}_normalized.stl"
    normalized.export(stl_path)
    silhouettes = save_silhouettes(normalized, ref_dir / "silhouettes")
    points = sample_points(normalized)
    points_path = ref_dir / "points.npy"
    np.save(points_path, points)
    metrics = {
        "name": name,
        "stl_path": str(stl_path),
        "points_path": str(points_path),
        "silhouettes": silhouettes,
        "bbox": bbox_dict(mesh),
        "normalized_bbox": bbox_dict(normalized),
        "topology": topology_metrics(normalized),
    }
    (ref_dir / "metrics.json").write_text(json.dumps(metrics, indent=2))
    return metrics


def preprocess_reference(glb_path: Path, ideal_code_path: Path | None = DEFAULT_IDEAL_CODE, out_root: Path = REFERENCE_ROOT) -> dict[str, Any]:
    out_root.mkdir(parents=True, exist_ok=True)
    glb_mesh = canonicalize_largest_axis_to_z(mesh_from_file(glb_path))
    glb_metrics = save_reference_mesh("glb_reference", glb_mesh, out_root)

    ideal_metrics = None
    correlation = None
    if ideal_code_path is not None and ideal_code_path.exists():
        ideal_code = read_code(ideal_code_path)
        ideal_run_dir = out_root / "ideal_cadquery_run"
        ideal_run = run_cadquery(ideal_code, ideal_run_dir, "ideal_cadquery")
        if not ideal_run["ok"]:
            raise RuntimeError(f"Ideal CadQuery reference failed: {ideal_run['stderr'] or ideal_run['stdout']}")
        ideal_stl = Path(ideal_run["output"]["stl_path"])
        ideal_mesh = mesh_from_file(ideal_stl)
        ideal_metrics = save_reference_mesh("ideal_cadquery", ideal_mesh, out_root)
        correlation = compare_to_references(normalize_mesh(ideal_mesh), out_root, include_ideal=False, mode="full")

    summary = {
        "source_glb": str(glb_path),
        "source_ideal_code": str(ideal_code_path) if ideal_code_path else None,
        "glb_reference": glb_metrics,
        "ideal_cadquery": ideal_metrics,
        "ideal_vs_glb_correlation": correlation,
        "semantic_hints": SEMANTIC_HINTS,
    }
    (out_root / "reference_summary.json").write_text(json.dumps(summary, indent=2))
    return summary


def load_reference_summary(root: Path = REFERENCE_ROOT, auto_default: bool = True) -> dict[str, Any]:
    path = root / "reference_summary.json"
    if not path.exists():
        if not auto_default:
            return {}
        return preprocess_reference(DEFAULT_GLB, DEFAULT_IDEAL_CODE, root)
    return json.loads(path.read_text())


def bbox_similarity(candidate_bbox: dict[str, float], ref_bbox: dict[str, float]) -> float:
    cand = np.array([candidate_bbox["xlen"], candidate_bbox["ylen"], candidate_bbox["zlen"]], dtype=float)
    ref = np.array([ref_bbox["xlen"], ref_bbox["ylen"], ref_bbox["zlen"]], dtype=float)
    cand = cand / max(cand[2], 1e-9)
    ref = ref / max(ref[2], 1e-9)
    err = float(np.mean(np.abs(cand - ref) / np.maximum(ref, 1e-9)))
    return float(max(0.0, 1.0 - err))


def chamfer_score(candidate_points: np.ndarray, ref_points_path: Path) -> float:
    ref_points = np.load(ref_points_path)
    cand_tree = cKDTree(candidate_points)
    ref_tree = cKDTree(ref_points)
    d1 = cand_tree.query(ref_points, k=1, workers=-1)[0]
    d2 = ref_tree.query(candidate_points, k=1, workers=-1)[0]
    chamfer = float((np.mean(d1) + np.mean(d2)) / 2.0)
    normalized = chamfer / 1000.0
    return float(math.exp(-normalized / 0.10))


def silhouette_scores(candidate_silhouettes: dict[str, str], ref_silhouettes: dict[str, str]) -> dict[str, Any]:
    per_view = {}
    for view, path in candidate_silhouettes.items():
        if view in ref_silhouettes:
            per_view[view] = mask_iou(Path(path), Path(ref_silhouettes[view]))
    weights = {"front": 0.18, "back": 0.12, "left": 0.18, "right": 0.18, "top": 0.14, "isometric": 0.20}
    total = sum(per_view.get(view, 0.0) * weight for view, weight in weights.items())
    return {"score": float(total), "per_view": per_view}


def compare_to_reference_name(
    candidate: trimesh.Trimesh,
    candidate_silhouettes: dict[str, str] | None,
    root: Path,
    name: str,
    mode: str,
) -> dict[str, Any]:
    metrics = json.loads((root / name / "metrics.json").read_text())
    bbox = bbox_similarity(bbox_dict(candidate), metrics["normalized_bbox"])
    if mode == "fast":
        return {
            "bbox": bbox,
            "chamfer": bbox,
            "silhouette": {"score": bbox, "per_view": {}},
        }
    if candidate_silhouettes is None:
        raise ValueError("Full reference comparison requires candidate silhouettes.")
    cand_points = sample_points(candidate, count=1500)
    return {
        "bbox": bbox,
        "chamfer": chamfer_score(cand_points, Path(metrics["points_path"])),
        "silhouette": silhouette_scores(candidate_silhouettes, metrics["silhouettes"]),
    }


def compare_to_references(
    candidate: trimesh.Trimesh,
    root: Path,
    include_ideal: bool = True,
    mode: str = "full",
    candidate_silhouettes: dict[str, str] | None = None,
) -> dict[str, Any]:
    if mode not in {"fast", "full"}:
        raise ValueError(f"Unknown reward mode: {mode}")
    if mode == "full" and candidate_silhouettes is None:
        temp_dir = root / "_tmp_compare"
        if temp_dir.exists():
            shutil.rmtree(temp_dir)
        candidate_silhouettes = save_silhouettes(candidate, temp_dir / "silhouettes")
    glb = compare_to_reference_name(candidate, candidate_silhouettes, root, "glb_reference", mode)
    result = {"glb_reference": glb}
    if include_ideal and (root / "ideal_cadquery" / "metrics.json").exists():
        ideal = compare_to_reference_name(candidate, candidate_silhouettes, root, "ideal_cadquery", mode)
        result["ideal_cadquery"] = ideal
    return result


def semantic_reward(code: str, mesh: trimesh.Trimesh, task_spec: dict[str, Any] | None = None) -> dict[str, Any]:
    hints = list((task_spec or {}).get("semantic_hints") or SEMANTIC_HINTS)
    lowered = code.lower()
    code_hits = {hint: hint.lower() in lowered or hint.lower().replace("_", "") in lowered.replace("_", "") for hint in hints}
    code_score = sum(1 for ok in code_hits.values() if ok) / max(1, len(hints))

    bbox = bbox_dict(mesh)
    height = max(bbox["zlen"], 1e-9)
    if task_spec and task_spec.get("bbox_mm"):
        target = np.asarray(task_spec["bbox_mm"], dtype=float)
        target = target / max(target[2], 1e-9)
        actual = np.asarray([bbox["xlen"], bbox["ylen"], bbox["zlen"]], dtype=float)
        actual = actual / max(actual[2], 1e-9)
        ratio_score = float(max(0.0, 1.0 - np.mean(np.abs(actual - target) / np.maximum(target, 1e-9))))
        geometry_score = ratio_score
    else:
        width_ratio = bbox["xlen"] / height
        depth_ratio = bbox["ylen"] / height
        chair_ratio_score = float(max(0.0, 1.0 - abs(width_ratio - 0.55) - abs(depth_ratio - 0.60)))

        vertices = np.asarray(mesh.vertices)
        lower = vertices[vertices[:, 2] < bbox["zmin"] + 0.25 * height]
        upper = vertices[vertices[:, 2] > bbox["zmin"] + 0.55 * height]
        lower_radius = 0.0 if len(lower) == 0 else float(np.percentile(np.linalg.norm(lower[:, :2], axis=1), 90) / height)
        upper_height_presence = float(len(upper) > max(20, len(vertices) * 0.08))
        base_spread_score = float(min(1.0, lower_radius / 0.30))
        geometry_score = 0.45 * chair_ratio_score + 0.35 * base_spread_score + 0.20 * upper_height_presence

    functions = len(re.findall(r"^\s*def\s+\w+\s*\(", code, re.MULTILINE))
    assembly_score = min(1.0, functions / 6.0)
    score = 0.35 * code_score + 0.45 * geometry_score + 0.20 * assembly_score
    return {
        "score": float(max(0.0, min(1.0, score))),
        "code_hits": code_hits,
        "code_score": float(code_score),
        "geometry_score": float(geometry_score),
        "function_count": functions,
        "assembly_score": float(assembly_score),
    }


def editability_reward(code: str) -> dict[str, Any]:
    functions = len(re.findall(r"^\s*def\s+\w+\s*\(", code, re.MULTILINE))
    named_values = len(re.findall(r"^\s*[a-zA-Z_][a-zA-Z0-9_]*\s*=\s*[-+]?\d", code, re.MULTILINE))
    reusable_returns = len(re.findall(r"^\s*return\s+", code, re.MULTILINE))
    show_object = "show_object" in code or "fixture" in code or "result" in code or "chair" in code
    score = 0.35 * min(1.0, functions / 6.0)
    score += 0.20 * min(1.0, named_values / 8.0)
    score += 0.25 * min(1.0, reusable_returns / max(1, functions))
    score += 0.20 if show_object else 0.0
    return {
        "score": float(max(0.0, min(1.0, score))),
        "function_count": functions,
        "named_numeric_assignments": named_values,
        "return_count": reusable_returns,
        "has_final_object": show_object,
    }


def weighted_reference_score(comparison: dict[str, Any]) -> dict[str, Any]:
    def one(ref: dict[str, Any]) -> float:
        return float(0.25 * ref["bbox"] + 0.35 * ref["chamfer"] + 0.40 * ref["silhouette"]["score"])

    glb_score = one(comparison["glb_reference"])
    ideal_score = one(comparison["ideal_cadquery"]) if "ideal_cadquery" in comparison else glb_score
    return {
        "score": float(0.60 * ideal_score + 0.40 * glb_score),
        "ideal_score": float(ideal_score),
        "glb_score": float(glb_score),
    }


def verifier_markdown(result: dict[str, Any]) -> str:
    reward = result["reward"]
    lines = [
        f"# CadQuery Environment Report: {result['episode_id']} / {result['step_id']}",
        "",
        f"- Total reward: `{reward['total']:.3f}`",
        f"- Build: `{reward['build']:.3f}`",
        f"- Topology: `{reward['topology']:.3f}`",
        f"- Contact/gaps: `{reward.get('contact', 0.0):.3f}`",
        f"- Semantic parts: `{reward['semantic_parts']:.3f}`",
        f"- Reference similarity: `{reward['reference_similarity']:.3f}`",
        f"- Silhouette: `{reward['silhouette']:.3f}`",
        f"- Editability: `{reward['editability']:.3f}`",
        "",
        "## Topology",
        "",
        "```json",
        json.dumps(result["topology"], indent=2),
        "```",
        "",
        "## Notes",
    ]
    for note in result["notes"]:
        lines.append(f"- {note}")
    return "\n".join(lines) + "\n"


def evaluate_code(
    code: str,
    episode_id: str,
    step_id: str,
    task_prompt: str = "",
    reference_root: Path = REFERENCE_ROOT,
    reward_mode: str = "full",
    task_spec: dict[str, Any] | None = None,
) -> dict[str, Any]:
    if reward_mode not in {"fast", "full"}:
        raise ValueError(f"reward_mode must be fast or full, got {reward_mode}")
    using_default_reference = reference_root.resolve() == REFERENCE_ROOT.resolve()
    reference_summary = load_reference_summary(reference_root, auto_default=using_default_reference)
    reference_available = bool(reference_summary) and (reference_root / "glb_reference" / "metrics.json").exists()
    episode_dir = RUNS_ROOT / safe_slug(episode_id) / safe_slug(step_id)
    if episode_dir.exists():
        shutil.rmtree(episode_dir)
    episode_dir.mkdir(parents=True, exist_ok=True)
    (episode_dir / "candidate.py").write_text(code)
    if task_prompt:
        (episode_dir / "task_prompt.txt").write_text(task_prompt)
    if task_spec:
        (episode_dir / "task_spec.json").write_text(json.dumps(task_spec, indent=2))

    run = run_cadquery(code, episode_dir, "candidate")
    if not run["ok"]:
        error_detail = concise_error(run["stdout"], run["stderr"])
        notes = ["CadQuery code did not execute or did not export an STL."]
        if error_detail:
            notes.append(f"Build error: {error_detail}")
        result = {
            "episode_id": episode_id,
            "step_id": step_id,
            "ok": False,
            "artifacts_dir": str(episode_dir),
            "error": "CadQuery execution failed.",
            "stdout": run["stdout"],
            "stderr": run["stderr"],
            "elapsed_ms": run["elapsed_ms"],
            "reward": {
                "total": -1.0,
                "build": 0.0,
                "topology": 0.0,
                "semantic_parts": 0.0,
                "reference_similarity": 0.0,
                "silhouette": 0.0,
                "editability": 0.0,
                "efficiency": 0.0,
            },
            "notes": notes,
        }
        (episode_dir / "reward.json").write_text(json.dumps(result, indent=2))
        (episode_dir / "verifier_report.md").write_text(verifier_markdown({**result, "topology": {}}))
        return result

    stl_path = Path(run["output"]["stl_path"])
    mesh = mesh_from_file(stl_path)
    normalized = normalize_mesh(mesh)
    normalized_stl = episode_dir / "candidate_normalized.stl"
    normalized.export(normalized_stl)
    masks = save_silhouettes(normalized, episode_dir / "masks") if reward_mode == "full" else {}
    renders = save_color_renders(normalized, episode_dir / "renders") if reward_mode == "full" else {}

    topology = topology_metrics(normalized)
    topo_reward = topology_reward(topology)
    contact = contact_metrics(normalized)
    semantic = semantic_reward(code, normalized, task_spec)
    editability = editability_reward(code)
    if reference_available:
        comparison = compare_to_references(
            normalized,
            reference_root,
            mode=reward_mode,
            candidate_silhouettes=masks if reward_mode == "full" else None,
        )
        reference = weighted_reference_score(comparison)
        if "ideal_cadquery" in comparison:
            silhouette = {
                "score": float(
                    0.60 * comparison["ideal_cadquery"]["silhouette"]["score"]
                    + 0.40 * comparison["glb_reference"]["silhouette"]["score"]
                ),
                "ideal": comparison["ideal_cadquery"]["silhouette"],
                "glb": comparison["glb_reference"]["silhouette"],
            }
        else:
            silhouette = {
                "score": float(comparison["glb_reference"]["silhouette"]["score"]),
                "glb": comparison["glb_reference"]["silhouette"],
            }
    else:
        comparison = {}
        reference = {"score": 0.50, "ideal_score": 0.50, "glb_score": 0.50, "reference_available": False}
        silhouette = {"score": 0.50, "reference_available": False, "per_view": {}}
    build = 1.0
    efficiency = 1.0
    if reward_mode == "fast":
        # Fast mode is for dense intermediate RL feedback. Avoid over-weighting
        # bbox-only similarity because blocky impostors can game that signal.
        total = (
            0.22 * build
            + 0.17 * topo_reward["score"]
            + 0.12 * contact["score"]
            + 0.25 * semantic["score"]
            + 0.10 * reference["score"]
            + 0.10 * editability["score"]
            + 0.04 * efficiency
        )
    else:
        total = (
            0.18 * build
            + 0.17 * topo_reward["score"]
            + 0.10 * contact["score"]
            + 0.15 * semantic["score"]
            + 0.15 * reference["score"]
            + 0.10 * silhouette["score"]
            + 0.10 * editability["score"]
            + 0.05 * efficiency
        )
    notes = []
    if topology["components"] > 120:
        notes.append(f"Candidate has {topology['components']} mesh components; acceptable for assemblies only if the parts are intentional.")
    if not topology["watertight"]:
        notes.append("Candidate mesh is not fully watertight; this is tolerated for chair assemblies but should improve for monolithic parts.")
    if contact["large_gap_components"] > 0:
        notes.append(f"Candidate has {contact['large_gap_components']} significant separated components; smaller assembly gaps are okay, large gaps reduce contact reward.")
    if semantic["score"] < 0.45:
        target_name = task_spec.get("id", "target object") if task_spec else "Markus-chair"
        notes.append(f"Candidate is weak on {target_name} semantic hints; add/organize recognizable subassemblies in code.")
    if reference_available and reference["score"] < 0.35:
        notes.append("Candidate is still far from the ideal CadQuery/reference GLB shape.")
    if not reference_available:
        notes.append("No task-specific GLB reference is available yet; reward uses build, topology, contact, task semantics, bbox profile, and editability.")

    result = {
        "episode_id": episode_id,
        "step_id": step_id,
        "ok": True,
        "reward_mode": reward_mode,
        "artifacts_dir": str(episode_dir),
        "candidate_stl": str(stl_path),
        "candidate_normalized_stl": str(normalized_stl),
        "renders": renders,
        "masks": masks,
        "elapsed_ms": run["elapsed_ms"],
        "bbox": bbox_dict(mesh),
        "normalized_bbox": bbox_dict(normalized),
        "topology": topology,
        "contact": contact,
        "semantic_parts": semantic,
        "editability": editability,
        "reference_comparison": comparison,
        "reference_similarity": reference,
        "silhouette": silhouette,
        "reward": {
            "total": float(max(-1.0, min(1.0, total))),
            "build": build,
            "topology": topo_reward["score"],
            "contact": contact["score"],
            "semantic_parts": semantic["score"],
            "reference_similarity": reference["score"],
            "silhouette": silhouette["score"],
            "editability": editability["score"],
            "efficiency": efficiency,
        },
        "notes": notes or ["Candidate built and scored successfully."],
    }
    (episode_dir / "reward.json").write_text(json.dumps(result, indent=2))
    (episode_dir / "verifier_report.md").write_text(verifier_markdown(result))
    return result


def smoke(episodes: int, reward_mode: str = "full") -> dict[str, Any]:
    load_reference_summary()
    ideal_code = read_code(DEFAULT_IDEAL_CODE)
    simple_code = "\n".join(
        [
            "import cadquery as cq",
            "seat = cq.Workplane('XY').box(520, 480, 70).translate((0, 0, 450))",
            "back = cq.Workplane('XY').box(440, 45, 720).translate((0, -235, 760))",
            "column = cq.Workplane('XY').cylinder(360, 28).translate((0, 0, 210))",
            "base = cq.Workplane('XY').cylinder(55, 55).translate((0, 0, 55))",
            "fixture = seat.union(back).union(column).union(base).clean()",
        ]
    )
    results = []
    for i in range(max(1, episodes)):
        code = ideal_code if i == 0 else simple_code
        step = "ideal_reference" if i == 0 else f"simple_candidate_{i}"
        results.append(evaluate_code(code, "smoke", step, "Smoke-test CadQuery Markus chair environment.", reward_mode=reward_mode))
    summary = {
        "episodes": len(results),
        "ok": all(item["ok"] for item in results),
        "best_reward": max(item["reward"]["total"] for item in results),
        "mean_reward": float(np.mean([item["reward"]["total"] for item in results])),
        "results": [
            {
                "step_id": item["step_id"],
                "ok": item["ok"],
                "reward": item["reward"],
                "artifacts_dir": item["artifacts_dir"],
            }
            for item in results
        ],
    }
    (RUNS_ROOT / "smoke_summary.json").write_text(json.dumps(summary, indent=2))
    return summary


def main() -> None:
    parser = argparse.ArgumentParser(description="CadQuery CADForge reference preprocessing and reward environment.")
    sub = parser.add_subparsers(dest="command", required=True)

    prep = sub.add_parser("preprocess-reference")
    prep.add_argument("--glb", default=str(DEFAULT_GLB))
    prep.add_argument("--ideal-code", default=str(DEFAULT_IDEAL_CODE))
    prep.add_argument("--out-root", default=str(REFERENCE_ROOT))

    eval_parser = sub.add_parser("evaluate")
    eval_parser.add_argument("--code-file")
    eval_parser.add_argument("--episode-id", default="manual")
    eval_parser.add_argument("--step-id", default="step-0")
    eval_parser.add_argument("--task-prompt", default="")
    eval_parser.add_argument("--reward-mode", choices=["fast", "full"], default="full")
    eval_parser.add_argument("--task-spec", default="")
    eval_parser.add_argument("--reference-root", default=str(REFERENCE_ROOT))

    smoke_parser = sub.add_parser("smoke")
    smoke_parser.add_argument("--episodes", type=int, default=2)
    smoke_parser.add_argument("--reward-mode", choices=["fast", "full"], default="full")

    args = parser.parse_args()
    if args.command == "preprocess-reference":
        ideal_code = Path(args.ideal_code) if args.ideal_code else None
        result = preprocess_reference(Path(args.glb), ideal_code, Path(args.out_root))
    elif args.command == "evaluate":
        if args.code_file:
            code = read_code(Path(args.code_file))
        else:
            payload = json.loads(sys.stdin.read() or "{}")
            code = extract_code(str(payload.get("code", "")))
        result = evaluate_code(
            code,
            args.episode_id,
            args.step_id,
            args.task_prompt,
            reference_root=Path(args.reference_root),
            reward_mode=args.reward_mode,
            task_spec=read_task_spec(args.task_spec),
        )
    elif args.command == "smoke":
        result = smoke(args.episodes, reward_mode=args.reward_mode)
    else:
        raise AssertionError(args.command)
    print(json.dumps(result, indent=2))


if __name__ == "__main__":
    main()