src/statement_to_tikz/formalize.py

"""LLM formalizer: natural language geometry → GeometryIR."""

from __future__ import annotations

import json
import os
from typing import Any

import httpx

from statement_to_tikz.ir import GeometryIR

SYSTEM_PROMPT = """You convert olympiad-style plane geometry statements into a JSON GeometryIR.

Rules:
- Output ONLY valid JSON matching the schema. No markdown.
- List every named point in "points".
- Put metric/incidence relations in "constraints" (not coordinates).
- Do NOT invent numeric coordinates; optional "hint" only if the statement clearly implies a position.
- Use segments for edges to draw; set draw=true on segments that should appear.
- Use circles with center+through or through (3 points) for circumcircles.
- Constraint types (field "type"): equal_length, length, equal_angle, angle_measure,
  perpendicular, parallel, on_line, on_circle, midpoint, collinear, intersection.
- For equal_length: segments a1-a2 and b1-b2 have equal length.
- For midpoint: point is midpoint of a-b.
- For perpendicular: segment a1-a2 ⊥ b1-b2.
- For intersection: point lies on both lines ab and cd.
- Always include labels for all points (or leave labels empty to auto-fill).
- Keep the original statement in "statement".
"""

SCHEMA_HINT = {
    "statement": "string",
    "points": [{"name": "A", "hint": None}],
    "segments": [{"a": "A", "b": "B", "draw": True}],
    "lines": [{"a": "A", "b": "B", "draw": False}],
    "circles": [
        {
            "id": "omega",
            "center": "O",
            "radius": None,
            "through": ["A", "B", "C"],
            "draw": True,
        }
    ],
    "angle_marks": [
        {"vertex": "A", "a": "B", "b": "C", "right_angle": False, "label": None}
    ],
    "labels": [{"point": "A", "text": "A", "position": "above right"}],
    "constraints": [
        {"type": "equal_length", "a1": "A", "a2": "B", "b1": "A", "b2": "C"},
        {"type": "midpoint", "point": "D", "a": "B", "b": "C"},
        {"type": "perpendicular", "a1": "A", "a2": "D", "b1": "B", "b2": "C"},
    ],
}


def _extract_json(text: str) -> dict[str, Any]:
    text = text.strip()
    if text.startswith("```"):
        lines = text.split("\n")
        lines = lines[1:]
        if lines and lines[-1].strip().startswith("```"):
            lines = lines[:-1]
        text = "\n".join(lines).strip()
    return json.loads(text)


def formalize_offline(statement: str) -> GeometryIR | None:
    """Tiny pattern-based formalizer for demos/tests without an API key.

    Recognizes a few common phrasings; returns None if not matched.
    """
    s = " ".join(statement.lower().split())

    # Isosceles triangle with midpoint median
    if (
        "triangle abc" in s
        and ("ab = ac" in s or "ab=ac" in s)
        and "midpoint" in s
        and ("d" in s)
    ):
        return GeometryIR.model_validate(
            {
                "statement": statement,
                "points": [
                    {"name": "A", "hint": [0, 1.2]},
                    {"name": "B", "hint": [-1, 0]},
                    {"name": "C", "hint": [1, 0]},
                    {"name": "D"},
                ],
                "segments": [
                    {"a": "A", "b": "B"},
                    {"a": "B", "b": "C"},
                    {"a": "C", "b": "A"},
                    {"a": "A", "b": "D"},
                ],
                "constraints": [
                    {"type": "equal_length", "a1": "A", "a2": "B", "b1": "A", "b2": "C"},
                    {"type": "midpoint", "point": "D", "a": "B", "b": "C"},
                    {"type": "length", "a": "B", "b": "C", "value": 2.0},
                    {"type": "length", "a": "A", "b": "B", "value": 1.5},
                ],
                "labels": [],
            }
        )

    if "equilateral" in s and "triangle abc" in s:
        return GeometryIR.model_validate(
            {
                "statement": statement,
                "points": [
                    {"name": "A", "hint": [0, 0]},
                    {"name": "B", "hint": [1, 0]},
                    {"name": "C", "hint": [0.5, 0.8]},
                ],
                "segments": [
                    {"a": "A", "b": "B"},
                    {"a": "B", "b": "C"},
                    {"a": "C", "b": "A"},
                ],
                "constraints": [
                    {"type": "equal_length", "a1": "A", "a2": "B", "b1": "B", "b2": "C"},
                    {"type": "equal_length", "a1": "B", "a2": "C", "b1": "C", "b2": "A"},
                    {"type": "length", "a": "A", "b": "B", "value": 1.0},
                ],
                "labels": [],
            }
        )

    if "draw triangle abc" in s or s.strip() == "triangle abc":
        return GeometryIR.model_validate(
            {
                "statement": statement,
                "points": [
                    {"name": "A", "hint": [0, 1]},
                    {"name": "B", "hint": [-1, 0]},
                    {"name": "C", "hint": [1, 0]},
                ],
                "segments": [
                    {"a": "A", "b": "B"},
                    {"a": "B", "b": "C"},
                    {"a": "C", "b": "A"},
                ],
                "constraints": [],
                "labels": [],
            }
        )

    return None


def formalize_with_llm(
    statement: str,
    *,
    api_key: str | None = None,
    base_url: str | None = None,
    model: str | None = None,
    timeout: float = 60.0,
) -> GeometryIR:
    api_key = api_key or os.environ.get("OPENAI_API_KEY")
    if not api_key:
        raise RuntimeError(
            "OPENAI_API_KEY not set; use --ir, offline patterns, or set the key"
        )
    base_url = (base_url or os.environ.get("OPENAI_BASE_URL") or "https://api.openai.com/v1").rstrip(
        "/"
    )
    model = model or os.environ.get("STT_MODEL") or "gpt-4o-mini"

    user = (
        f"Statement:\n{statement}\n\n"
        f"JSON shape example (illustrative):\n{json.dumps(SCHEMA_HINT, indent=2)}"
    )

    payload = {
        "model": model,
        "temperature": 0,
        "messages": [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user},
        ],
        "response_format": {"type": "json_object"},
    }

    url = f"{base_url}/chat/completions"
    headers = {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}

    last_err: Exception | None = None
    for _ in range(2):
        try:
            with httpx.Client(timeout=timeout) as client:
                resp = client.post(url, headers=headers, json=payload)
                resp.raise_for_status()
                data = resp.json()
            content = data["choices"][0]["message"]["content"]
            obj = _extract_json(content)
            if "statement" not in obj:
                obj["statement"] = statement
            ir = GeometryIR.model_validate(obj)
            ir.ensure_labels()
            return ir
        except Exception as exc:  # noqa: BLE001
            last_err = exc
            continue
    raise RuntimeError(f"formalizer failed after retry: {last_err}")


def formalize(statement: str, *, allow_offline: bool = True, **kwargs: Any) -> GeometryIR:
    """Formalize statement via LLM, with optional offline pattern fallback."""
    if allow_offline:
        offline = formalize_offline(statement)
        if offline is not None:
            offline.ensure_labels()
            return offline
    return formalize_with_llm(statement, **kwargs)