deepcad-canonical-equivalence

#!/usr/bin/env python3
"""DeepCAD canonical-equivalence solver.

Reads a pairs JSON, decides for each pair whether program_a and program_b
describe the same canonical solid, and writes predictions.json.

Approach: build a normalized, label-invariant signature for each program and
compare the two signatures for equality.
"""

from __future__ import annotations

import ast
import json
import sys
from pathlib import Path

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

META_KINDS = {"constraint", "dimension", "note", "metadata", "view", "comment"}
PROFILE_KINDS = {"rect", "circle", "slot"}
PROFILE_SKIP_FIELDS = {"id", "cmd", "sketch", "construction"}

TOL_DECIMALS = 6


# ---------------------------------------------------------------------------
# Expression evaluation
# ---------------------------------------------------------------------------


class Evaluator:
    """Evaluate scalar arithmetic expressions with parameter references."""

    def __init__(self, params):
        # params: name -> raw value (str expression or number)
        self.params = params
        self.cache = {}
        self.inprogress = set()

    def eval_name(self, name):
        if name in self.cache:
            return self.cache[name]
        if name not in self.params:
            raise KeyError(name)
        if name in self.inprogress:
            raise ValueError("parameter cycle: %s" % name)
        self.inprogress.add(name)
        try:
            val = self.eval_value(self.params[name])
        finally:
            self.inprogress.discard(name)
        self.cache[name] = val
        return val

    def eval_value(self, v):
        if isinstance(v, bool):
            raise ValueError("bool is not numeric")
        if isinstance(v, (int, float)):
            return float(v)
        if isinstance(v, str):
            tree = ast.parse(v, mode="eval")
            return self.eval_node(tree.body)
        raise ValueError("not a scalar expression: %r" % (v,))

    def eval_node(self, node):
        if isinstance(node, ast.Constant):
            if isinstance(node.value, bool) or not isinstance(node.value, (int, float)):
                raise ValueError("non-numeric constant")
            return float(node.value)
        if isinstance(node, ast.Num):  # pragma: no cover (py<3.8)
            return float(node.n)
        if isinstance(node, ast.BinOp):
            left = self.eval_node(node.left)
            right = self.eval_node(node.right)
            op = node.op
            if isinstance(op, ast.Add):
                return left + right
            if isinstance(op, ast.Sub):
                return left - right
            if isinstance(op, ast.Mult):
                return left * right
            if isinstance(op, ast.Div):
                return left / right
            raise ValueError("unsupported binary operator")
        if isinstance(node, ast.UnaryOp):
            operand = self.eval_node(node.operand)
            if isinstance(node.op, ast.USub):
                return -operand
            if isinstance(node.op, ast.UAdd):
                return operand
            raise ValueError("unsupported unary operator")
        if isinstance(node, ast.Name):
            return self.eval_name(node.id)
        raise ValueError("unsupported expression node: %r" % (node,))


# ---------------------------------------------------------------------------
# Numeric helpers
# ---------------------------------------------------------------------------


def round_num(x):
    if x is None:
        return None
    r = round(float(x), TOL_DECIMALS)
    if r == 0:
        r = 0.0
    return r


def ser(obj):
    """Canonical JSON string of a jsonable structure (for multiset sorting)."""
    return json.dumps(obj, sort_keys=True, ensure_ascii=False)


def canon_field(v, ev):
    """Canonicalize a geometric field value, evaluating numeric expressions."""
    if isinstance(v, bool):
        return v
    if isinstance(v, (int, float)):
        return round_num(v)
    if isinstance(v, str):
        try:
            return round_num(ev.eval_value(v))
        except Exception:
            return v
    if isinstance(v, list):
        return [canon_field(x, ev) for x in v]
    if isinstance(v, dict):
        return {k: canon_field(x, ev) for k, x in v.items()}
    return v


def normalize_plane(p):
    if isinstance(p, str):
        return p.upper()
    return p


# ---------------------------------------------------------------------------
# Signature construction
# ---------------------------------------------------------------------------


def canon_direction(d, ev):
    if isinstance(d, str):
        return d.lower()
    if isinstance(d, bool):
        return d
    if isinstance(d, (int, float)):
        return round_num(d)
    if isinstance(d, list):
        return [canon_field(x, ev) for x in d]
    return d


def canon_profile(p, sketches, ev):
    kind = p.get("cmd")
    plane = sketches.get(p.get("sketch"))
    fields = {}
    for k, v in p.items():
        if k in PROFILE_SKIP_FIELDS:
            continue
        if k == "angle":
            try:
                a = ev.eval_value(v) % 180.0
                fields[k] = round_num(a)
                continue
            except Exception:
                pass
        fields[k] = canon_field(v, ev)
    return ["profile", kind, plane, fields]


def extrude_profile_ids(ex):
    if "profiles" in ex and isinstance(ex["profiles"], list):
        return list(ex["profiles"])
    if "profile" in ex:
        prof = ex["profile"]
        if isinstance(prof, list):
            return list(prof)
        return [prof]
    return []


def feature_sig(ex, profiles, sketches, ev):
    op = str(ex.get("operation", "new")).lower()
    extent = str(ex.get("extent", "one_side")).lower()
    direction = canon_direction(ex.get("direction", "normal"), ev)

    depth = None
    if "depth" in ex:
        try:
            depth = round_num(ev.eval_value(ex["depth"]))
        except Exception:
            depth = canon_field(ex["depth"], ev)

    ids = extrude_profile_ids(ex)
    prof_forms = []
    for pid in ids:
        if pid in profiles:
            prof_forms.append(ser(canon_profile(profiles[pid], sketches, ev)))
    prof_forms.sort()

    return [op, extent, direction, depth, prof_forms]


def partition_label(ex):
    if "body" in ex:
        return ("L", ex["body"])
    if "channel" in ex:
        return ("L", ex["channel"])
    return None


def build_signature(program):
    if not isinstance(program, list):
        program = []

    params = {}
    sketches = {}
    profiles = {}
    extrudes = []
    opaque = []

    for cmd in program:
        if not isinstance(cmd, dict):
            opaque.append(cmd)
            continue
        if cmd.get("construction") is True:
            continue
        kind = cmd.get("cmd")
        if kind in META_KINDS:
            continue
        if kind == "param":
            params[cmd.get("name")] = cmd.get("expr")
        elif kind == "sketch":
            sketches[cmd.get("id")] = normalize_plane(cmd.get("plane"))
        elif kind in PROFILE_KINDS:
            profiles[cmd.get("id")] = cmd
        elif kind == "extrude":
            extrudes.append(cmd)
        else:
            payload = {k: v for k, v in cmd.items() if k != "id"}
            opaque.append(payload)

    ev = Evaluator(params)

    default_group = []
    label_groups = {}
    for ex in extrudes:
        feat = feature_sig(ex, profiles, sketches, ev)
        lab = partition_label(ex)
        if lab is None:
            default_group.append(feat)
        else:
            label_groups.setdefault(ser(list(lab)), []).append(feat)

    partitions = []
    if default_group:
        partitions.append(["DEFAULT", sorted(ser(f) for f in default_group)])
    for _key, feats in label_groups.items():
        partitions.append(["EXPLICIT", sorted(ser(f) for f in feats)])

    partitions_sig = sorted(ser(p) for p in partitions)
    opaque_sig = sorted(ser(o) for o in opaque)

    return ser(["SIG", partitions_sig, opaque_sig])


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------


def main() -> int:
    input_path = Path(sys.argv[1]) if len(sys.argv) > 1 else Path("/workspace/data/public_pairs.json")
    output_path = Path(sys.argv[2]) if len(sys.argv) > 2 else Path("/workspace/predictions.json")

    payload = json.loads(input_path.read_text(encoding="utf-8"))
    predictions = []
    for pair in payload.get("pairs", []):
        pid = pair.get("pair_id")
        try:
            sig_a = build_signature(pair.get("program_a", []))
            sig_b = build_signature(pair.get("program_b", []))
            equivalent = sig_a == sig_b
        except Exception:
            sig_a = sig_b = None
            equivalent = False
        predictions.append(
            {
                "pair_id": pid,
                "equivalent": bool(equivalent),
                "signature_a": sig_a,
                "signature_b": sig_b,
            }
        )

    output_path.parent.mkdir(parents=True, exist_ok=True)
    output_path.write_text(
        json.dumps({"predictions": predictions}, indent=2) + "\n",
        encoding="utf-8",
    )
    return 0


if __name__ == "__main__":
    raise SystemExit(main())