path: root/src/gap_pipeline/surface.py

"""Surface-renaming generation, application, and deterministic validation."""

from __future__ import annotations

import asyncio
import hashlib
import re
from dataclasses import dataclass
from typing import Literal

from .clients import JsonLLM
from .models import CanonicalItem, ModelCallRecord
from .prompts import SURFACE_NAME_SYSTEM, surface_name_user
from .store import RunStore


SurfaceFamily = Literal[
    "descriptive_long",
    "descriptive_long_confusing",
    "descriptive_long_misleading",
    "garbled_string",
]

SURFACE_FAMILIES: tuple[SurfaceFamily, ...] = (
    "descriptive_long",
    "descriptive_long_confusing",
    "descriptive_long_misleading",
    "garbled_string",
)

IDENTIFIER_RE = re.compile(r"^[A-Za-z][A-Za-z0-9]{2,63}$")
MATH_SPAN_RE = re.compile(
    r"(?P<dollar>\${1,2}.*?\${1,2})"
    r"|(?P<paren>\\\(.*?\\\))"
    r"|(?P<bracket>\\\[.*?\\\])"
    r"|(?P<env>\\begin\{(?P<envname>[A-Za-z*]+)\}.*?\\end\{(?P=envname)\})",
    flags=re.DOTALL,
)


@dataclass(frozen=True)
class SurfaceVariant:
    family: SurfaceFamily
    rename_map: dict[str, str]
    problem: str
    solution: str

    def as_release_payload(self) -> dict[str, object]:
        return {
            "map": dict(self.rename_map),
            "question": self.problem,
            "solution": self.solution,
        }


def validate_rename_map(
    rename_map: dict[str, str],
    *,
    existing_identifiers: list[str],
    scientific_constants: list[str],
) -> None:
    if not rename_map:
        raise ValueError("surface rename map cannot be empty")
    if len(rename_map.values()) != len(set(rename_map.values())):
        raise ValueError("surface replacements must be one-to-one")
    originals = set(existing_identifiers) | set(scientific_constants)
    for old, new in rename_map.items():
        if old not in existing_identifiers:
            raise ValueError(f"rename source {old!r} is not in the symbol inventory")
        if new in originals:
            raise ValueError(f"replacement {new!r} collides with an existing identifier")
        if not IDENTIFIER_RE.fullmatch(new):
            raise ValueError(f"invalid replacement identifier {new!r}")
        if new.startswith("\\"):
            raise ValueError("replacement may not be a LaTeX command")


def _replace_in_span(span: str, rename_map: dict[str, str]) -> str:
    result = span
    for old in sorted(rename_map, key=lambda value: (-len(value), value)):
        # ASCII identifier boundaries avoid r -> radius changing \sqrt or prose.
        pattern = re.compile(
            rf"(?<![A-Za-z0-9]){re.escape(old)}(?![A-Za-z0-9])"
        )
        result = pattern.sub(lambda _: rename_map[old], result)
    return result


def apply_rename_map(text: str, rename_map: dict[str, str]) -> str:
    """Rename identifiers only inside explicit LaTeX math spans.

    This deliberately avoids the old prototype's ``a`` -> identifier mutation
    in ordinary English prose. PutnamGAP source records use explicit math
    delimiters for mathematical identifiers.
    """

    output: list[str] = []
    cursor = 0
    for match in MATH_SPAN_RE.finditer(text):
        output.append(text[cursor : match.start()])
        output.append(_replace_in_span(match.group(0), rename_map))
        cursor = match.end()
    output.append(text[cursor:])
    return "".join(output)


def create_surface_variant(
    item: CanonicalItem,
    family: SurfaceFamily,
    rename_map: dict[str, str],
) -> SurfaceVariant:
    existing = item.variables + item.parameters
    validate_rename_map(
        rename_map,
        existing_identifiers=existing,
        scientific_constants=item.scientific_constants,
    )
    return SurfaceVariant(
        family=family,
        rename_map=dict(rename_map),
        problem=apply_rename_map(item.problem, rename_map),
        solution=apply_rename_map(item.solution, rename_map),
    )


class SurfacePipeline:
    def __init__(self, proposer: JsonLLM, store: RunStore) -> None:
        self.proposer = proposer
        self.store = store

    async def propose_map(
        self,
        item: CanonicalItem,
        family: SurfaceFamily,
    ) -> dict[str, str]:
        symbols = [(value, "free variable") for value in item.variables] + [
            (value, "fixed parameter") for value in item.parameters
        ]
        forbidden = list(
            dict.fromkeys(
                item.variables + item.parameters + item.scientific_constants
            )
        )

        async def propose(symbol: str, role: str) -> tuple[str, str]:
            request_id = f"{item.item_id}.surface.{family}.{symbol}"
            user_prompt = surface_name_user(
                symbol=symbol,
                role=role,
                family=family,
                context=(item.problem + "\n" + item.solution)[:12_000],
                forbidden=forbidden,
            )
            response = await self.proposer.generate_json(
                system_prompt=SURFACE_NAME_SYSTEM,
                user_prompt=user_prompt,
                request_id=request_id,
            )
            self.store.write_call(
                ModelCallRecord(
                    request_id=request_id,
                    model=response.model,
                    system_prompt=SURFACE_NAME_SYSTEM,
                    user_prompt=user_prompt,
                    response_data=response.data,
                    raw_text=response.raw_text,
                    provider_response_id=response.response_id,
                    usage=response.usage,
                )
            )
            replacement = str(response.data["replacement"])
            return symbol, replacement

        proposals = await asyncio.gather(
            *(propose(symbol, role) for symbol, role in symbols)
        )
        rename_map = dict(proposals)
        validate_rename_map(
            rename_map,
            existing_identifiers=item.variables + item.parameters,
            scientific_constants=item.scientific_constants,
        )
        self.store.write_stage(
            f"surface_{family}_map",
            {"rename_map": rename_map},
            request_id=None,
        )
        return rename_map

    async def run_family(
        self,
        item: CanonicalItem,
        family: SurfaceFamily,
    ) -> SurfaceVariant:
        rename_map = await self.propose_map(item, family)
        variant = create_surface_variant(item, family, rename_map)
        self.store.write_stage(
            f"surface_{family}_variant",
            variant.as_release_payload(),
            request_id=None,
        )
        return variant

    async def run_all(self, item: CanonicalItem) -> dict[SurfaceFamily, SurfaceVariant]:
        # Run families sequentially so the immutable call log stays easy to
        # inspect; symbol proposals within each family remain concurrent.
        output: dict[SurfaceFamily, SurfaceVariant] = {}
        for family in SURFACE_FAMILIES:
            output[family] = await self.run_family(item, family)
        return output


def deterministic_garbled_name(item_id: str, symbol: str, length: int = 12) -> str:
    """Stable offline GS name for smoke tests; production follows the proposer."""

    digest = hashlib.sha256(f"{item_id}:{symbol}".encode()).hexdigest()
    return "v" + digest[: max(3, length - 1)]