app/agents/visual_lesson_agent.py

from __future__ import annotations

import logging
import re
from typing import Any, Literal

from pydantic import BaseModel, Field

from app.agents.cerebras_client import CerebrasClient
from app.agents.senses_agent import VISION_MODEL_ID
from app.schemas.visual_lesson import (
    AttentionLessonDraft,
    EvidenceBundle,
    ExtractedChartDataset,
    PlanCritique,
    VisualizationRequest,
)


class CapabilityDecision(BaseModel):
    capability: Literal["attention", "protein", "nucleic_acid", "molecular", "thermodynamics", "differential_equation", "graph_algorithm", "array_algorithm", "composition", "legacy"]
    reasoning: str


class CompositionPlanDecision(BaseModel):
    kind: Literal["registered_geometry", "chart", "unsupported"]
    target: Literal["sphere", "ellipsoid", "prolate_spheroid", "torus", "mobius_strip", "line", "scatter", "bar", "heatmap", "force-graph", "unsupported"]
    panel_kinds: list[Literal["three_scene", "chart", "diagram", "equation", "annotation"]]
    interpretation: str


class CompositionPlanCritique(BaseModel):
    approved: bool
    feedback: list[str]


class VisualAmbiguityOption(BaseModel):
    option_id: str
    label: str
    description: str
    resolved_prompt: str


class VisualizationIntentDecision(BaseModel):
    specificity: Literal["exact", "open", "ambiguous"]
    interpretation: str
    requested_views: list[str] = Field(default_factory=list)
    hard_requirements: list[str] = Field(default_factory=list)
    requested_interactions: list[str] = Field(default_factory=list)
    clarification_question: str = ""
    clarification_reason: str = ""
    ambiguity_options: list[VisualAmbiguityOption] = Field(default_factory=list)


class SemanticNodePlan(BaseModel):
    node_id: str
    label: str
    role: Literal["input", "state", "process", "decision", "output", "data", "parameter", "concept"] = "concept"
    group: str = ""
    emphasis: Literal["primary", "secondary", "muted"] = "secondary"


class SemanticEdgePlan(BaseModel):
    source: str
    target: str
    label: str = ""
    direction: Literal["forward", "backward", "bidirectional", "none"] = "forward"
    style: Literal["solid", "dashed", "feedback"] = "solid"


class SemanticFacetPlan(BaseModel):
    facet_id: str
    title: str
    purpose: str
    layout: Literal["flow_horizontal", "flow_vertical", "recurrent_unrolled", "cycle", "layered"] = "flow_horizontal"
    priority: Literal["primary", "supporting"] = "supporting"
    nodes: list[SemanticNodePlan]
    edges: list[SemanticEdgePlan]


class SemanticCompositionPlan(BaseModel):
    title: str
    interpretation: str
    facets: list[SemanticFacetPlan]
    answer_markdown: str


_ATTENTION_TERMS = (
    "attention",
    "attend",
    "query key value",
    "qkv",
    "softmax weights",
    "transformer token",
)

_PROTEIN_TERMS = (
    "protein", "uniprot", "alphafold", "amino acid", "residue", "binding site",
    "enzyme structure", "molecular structure", "ligand", "hemoglobin", "haemoglobin",
)
_NUCLEIC_TERMS = (
    "dna", "rna", "nucleic acid", "nucleotide", "base pair", "double helix",
    "a-dna", "b-dna", "z-dna", "ribose", "deoxyribose", "uracil", "thymine",
)
_THERMODYNAMICS_TERMS = (
    "piston", "ideal gas", "isothermal", "adiabatic", "isobaric", "isochoric",
    "pressure volume", "p-v diagram", "pv diagram", "gas compression", "gas expansion",
    "thermodynamic process", "boundary work", "first law of thermodynamics",
)
_DIFFERENTIAL_EQUATION_TERMS = (
    "differential equation", "direction field", "slope field", "phase portrait", "logistic growth",
    "carrying capacity", "exponential growth", "first-order ode", "first order ode", "damped oscillator",
    "harmonic oscillator", "initial value problem", "y' =", "y'=", "dy/dt",
    "pendulum", "bob and string", "swinging bob",
)
_GRAPH_ALGORITHM_TERMS = (
    "pathfinding", "path finding", "breadth-first search", "depth-first search", "dijkstra", "a-star",
    "astar", "graph traversal", "shortest path", "bfs", "dfs", "weighted graph",
)
_ARRAY_ALGORITHM_TERMS = (
    "bubble sort", "insertion sort", "selection sort", "merge sort", "quicksort", "quick sort",
    "binary search", "linear search", "sort this array", "sort [", "visualize sorting", "sorting algorithm", "search this array",
)
_COMPOSITION_TERMS = (
    "prolate spheroid", "prolate ellipsoid", "ellipsoid", "spheroid", "sphere", "torus", "toroid",
    "mobius strip", "möbius strip", "spherical surface", "line chart", "scatter plot", "bar chart",
    "heatmap", "force-directed graph", "force graph", "parabola", "quadratic curve", "quadratic function",
    "recurrent neural network", "rnns", "rnn ",
    "generative adversarial", "adversarial neural", "adversarial network", "gan architecture", "neural network architecture",
)
_UNSUPPORTED_NUCLEIC_TERMS = (
    "predict folding", "predict fold", "fold this rna", "sequence alignment", "align sequence",
    "rna folding", "dna folding", "secondary structure prediction", "molecular dynamics", "dock ", "docking",
    "electron density", "genome browser", "crispr",
)
_PDB_PATTERN = re.compile(r"(?i)\b(?=[0-9a-z]{4}\b)(?=[0-9a-z]*[a-z])[0-9][a-z0-9]{3}\b")
_UNIPROT_PATTERN = re.compile(r"(?i)\b(?:[opq][0-9][a-z0-9]{3}[0-9]|[a-nr-z][0-9](?:[a-z][a-z0-9]{2}[0-9]){1,2})\b")

_UNSUPPORTED_VARIANTS = {
    "cross-attention": "cross-attention",
    "cross attention": "cross-attention",
    "sparse attention": "sparse attention",
    "grouped-query": "grouped-query attention",
    "grouped query": "grouped-query attention",
    "multi-head": "multi-head comparison",
    "multi head": "multi-head comparison",
}


_logger = logging.getLogger(__name__)


class VisualLessonAgent:
    def __init__(self, client: CerebrasClient | None = None) -> None:
        self._client = client or CerebrasClient()

    @staticmethod
    def _deterministic_route(prompt: str, selection_text: str) -> Literal["attention", "protein", "nucleic_acid", "molecular", "thermodynamics", "differential_equation", "graph_algorithm", "array_algorithm", "composition", "legacy"]:
        text = re.sub(r"\s+", " ", f"{prompt} {selection_text}").lower()
        if any(term in text for term in _UNSUPPORTED_NUCLEIC_TERMS):
            return "legacy"
        if _PDB_PATTERN.search(text):
            return "molecular"
        if any(term in text for term in _ATTENTION_TERMS):
            return "attention"
        if any(term in text for term in _NUCLEIC_TERMS):
            return "nucleic_acid"
        if any(term in text for term in _THERMODYNAMICS_TERMS):
            return "thermodynamics"
        if any(term in text for term in _DIFFERENTIAL_EQUATION_TERMS):
            return "differential_equation"
        if any(term in text for term in _GRAPH_ALGORITHM_TERMS):
            return "graph_algorithm"
        if any(term in text for term in _ARRAY_ALGORITHM_TERMS):
            return "array_algorithm"
        if any(term in text for term in _COMPOSITION_TERMS):
            return "composition"
        if "=" in text and any(variable in text for variable in ("x", "y", "z")):
            return "composition"
        if any(term in text for term in _PROTEIN_TERMS) or _UNIPROT_PATTERN.search(text):
            return "protein"
        return "legacy"

    def route(self, prompt: str, selection_text: str) -> Literal["attention", "protein", "nucleic_acid", "molecular", "thermodynamics", "differential_equation", "graph_algorithm", "array_algorithm", "composition", "legacy"]:
        deterministic = self._deterministic_route(prompt, selection_text)
        if deterministic != "legacy":
            return deterministic
        messages = [
            {
                "role": "system",
                "content": (
                    "Route an explicit visualization request. Choose `attention` only for transformer/self-attention, "
                    "query-key-value, attention matrices, attention weights, or token-attention mechanisms. Choose `protein` "
                    "for protein names, PDB/UniProt/AlphaFold identifiers, macromolecular protein structures, chains, ligands, "
                    "binding sites, or amino-acid residue requests. Choose `nucleic_acid` for DNA, RNA, nucleotides, base pairing, "
                    "helix forms, and nucleic-acid chemistry. Choose `molecular` for an exact PDB identifier whose polymer type must "
                    "be resolved officially. Choose `thermodynamics` for piston/cylinder ideal-gas processes, P–V work, isothermal, "
                    "adiabatic, isobaric, or isochoric process requests. Choose `differential_equation` for logistic or exponential "
                    "growth, linear first-order ODEs, direction fields, and damped-oscillator phase portraits. Choose "
                    "`graph_algorithm` for BFS, DFS, Dijkstra, A*, graph traversal, and shortest-path visualizations. Choose "
                    "`array_algorithm` for sorting arrays and linear or binary search visualizations. Choose "
                    "`composition` for registered mathematical geometry, explicit implicit-surface equations, ordinary charts, "
                    "heatmaps, and force-directed data layouts. "
                    "`legacy` for every other domain. Do not infer that "
                    "all AI topics are attention or that all biology topics are proteins."
                ),
            },
            {"role": "user", "content": f"REQUEST:\n{prompt}\n\nSELECTED CONTEXT:\n{selection_text[:1200]}"},
        ]
        try:
            decision = self._client.structured_complete(messages, CapabilityDecision, reasoning_effort="low")
            return decision.capability
        except Exception:
            return self._deterministic_route(prompt, selection_text)

    def plan_composition(self, prompt: str) -> tuple[CompositionPlanDecision, list[str]]:
        system = (
            "Select a verified registered visualization operator for the request. You may output only the listed target IDs and panel kinds. "
            "Never emit formulas, numbers, HTML, JavaScript, JSX, GLSL, Three.js calls, Plotly traces, D3 code, or renderer instructions. "
            "Use unsupported when no listed operator honestly answers the prompt."
        )
        try:
            plan = self._client.structured_complete([{"role":"system","content":system},{"role":"user","content":prompt}], CompositionPlanDecision, reasoning_effort="medium")
        except Exception:
            return CompositionPlanDecision(kind="unsupported", target="unsupported", panel_kinds=[], interpretation="No registered operator resolved the request."), ["Composition planning was unavailable."]
        try:
            critique = self._client.structured_complete([
                {"role":"system","content":"Check that the selected registered operator directly answers the prompt, uses only allowed IDs, and does not imply unsupported facts. Return concise feedback."},
                {"role":"user","content":f"PROMPT:\n{prompt}\n\nPLAN:\n{plan.model_dump_json()}"},
            ], CompositionPlanCritique, reasoning_effort="low")
            if not critique.approved:
                revised_prompt = f"{prompt}\n\nCritic feedback:\n" + "\n".join(f"- {item}" for item in critique.feedback)
                plan = self._client.structured_complete([{"role":"system","content":system},{"role":"user","content":revised_prompt}], CompositionPlanDecision, reasoning_effort="medium")
        except Exception:
            return plan, ["Composition plan review was unavailable; the schema-valid plan was validated deterministically."]
        return plan, []

    def resolve_visual_intent(self, prompt: str, selection_text: str = "") -> VisualizationIntentDecision:
        lowered = re.sub(r"\s+", " ", prompt).lower()
        if ("adversarial neural" in lowered or "adversarial network" in lowered) and not any(
            term in lowered for term in ("gan", "generator", "discriminator", "adversarial example", "attack", "perturbation", "adversarial training")
        ):
            return VisualizationIntentDecision(
                specificity="ambiguous",
                interpretation="The phrase can refer to distinct mechanisms with incompatible visual structures.",
                clarification_question="Which adversarial-neural-network concept should I visualize?",
                clarification_reason="A GAN, an adversarial attack, and adversarial training require different components and interactions.",
                ambiguity_options=[
                    VisualAmbiguityOption(option_id="gan", label="GAN competition", description="Generator and discriminator training against one another.", resolved_prompt="Visualize a GAN generator-discriminator competition with data flow and training feedback."),
                    VisualAmbiguityOption(option_id="attack", label="Adversarial attack", description="A perturbation changes a classifier's prediction.", resolved_prompt="Visualize an adversarial-example attack against a neural-network classifier."),
                    VisualAmbiguityOption(option_id="training", label="Adversarial training", description="Clean and attacked examples jointly train a robust model.", resolved_prompt="Visualize the adversarial-training pipeline for a neural-network classifier."),
                ],
            )
        system = (
            "Interpret a visualization command, without designing renderer code. Classify it as exact when the user specifies concrete views, "
            "components, interactions, counts, parameters, labels, or animation behavior; open when the scientific subject is clear but useful views "
            "must be selected; ambiguous only when different meanings would change the scientific subject. Broad but coherent requests should be open, "
            "not ambiguous, because they can use several coordinated views. Extract every explicit visualization instruction as a short hard requirement. "
            "If the request asks for a 2D plot, 3D scene, animation, simulation, or spatial environment but does not identify the scientific subject "
            "or mathematical object to render, classify it as ambiguous and ask for that missing subject. For ambiguity return concrete options when "
            "they are genuinely available; an empty option list is allowed when free-text refinement is the honest path. Never ask about cosmetic choices."
        )
        user_content = (
            f"REQUEST:\n{prompt}\n\nSELECTED CONTEXT:\n{selection_text[:1600]}" if selection_text else prompt
        )
        try:
            decision = self._client.structured_complete(
                [{"role": "system", "content": system}, {"role": "user", "content": user_content}],
                VisualizationIntentDecision,
                reasoning_effort="low",
            )
            return self._normalize_visual_intent(prompt, selection_text, decision)
        except Exception:
            words = len(re.findall(r"[A-Za-z0-9]+", f"{prompt} {selection_text}"))
            return self._normalize_visual_intent(prompt, selection_text, VisualizationIntentDecision(
                specificity="exact" if words >= 12 else "open",
                interpretation=prompt.strip(),
                requested_views=["primary mechanism"],
                hard_requirements=[prompt.strip()] if words >= 12 else [],
                requested_interactions=["zoom", "pan", "select"],
            ))

    @staticmethod
    def _normalize_visual_intent(prompt: str, selection_text: str, decision: VisualizationIntentDecision) -> VisualizationIntentDecision:
        """Keep modality-only requests out of the semantic-diagram renderer."""
        lowered = re.sub(r"\s+", " ", f"{prompt} {selection_text}").lower()
        asks_for_mode = any(term in lowered for term in (
            "3d", "3-d", "three dimensional", "three-dimensional", "spatial environment",
            "2d", "2-d", "two dimensional", "two-dimensional", "animation", "simulation",
        ))
        known_subject = any(term in lowered for term in (
            *_ATTENTION_TERMS, *_PROTEIN_TERMS, *_NUCLEIC_TERMS, *_THERMODYNAMICS_TERMS,
            *_DIFFERENTIAL_EQUATION_TERMS, *_GRAPH_ALGORITHM_TERMS, *_ARRAY_ALGORITHM_TERMS,
            *_COMPOSITION_TERMS, "equation", "surface", "plot", "chart", "graph",
        ))
        if asks_for_mode and not known_subject:
            return VisualizationIntentDecision(
                specificity="ambiguous",
                interpretation=decision.interpretation or prompt.strip(),
                clarification_question="What scientific system or mathematical object should the visual represent?",
                clarification_reason="The requested rendering mode is clear, but a 2D/3D scene or simulation needs a subject before it can be compiled faithfully.",
            )
        return decision

    def extract_chart_dataset(
        self, prompt: str, context_text: str = "", image_base64: str = ""
    ) -> ExtractedChartDataset:
        """Read a chart-ready dataset strictly from selected context (table text and/or a region image).

        Never invents numbers -> has_data=False when the context does not contain a real
        plottable dataset, so the caller can fall back to the illustrative/schematic path.
        """
        system = (
            "Extract a chart-ready dataset strictly from the provided context. Each row is one "
            "labeled data point: label is the category/series/axis name exactly as written (e.g. a "
            "method name, a year, a group), value is its associated number. Transcribe only values "
            "that are actually visible in the context -> never invent, estimate, or average unstated "
            "numbers. If the request asks to compare a specific column or metric, prefer rows for "
            "that metric; otherwise pick the single most prominent numeric column. If the context "
            "contains no real numeric dataset that answers the request, set has_data to false and "
            "leave rows empty.\n\n"
            "Also choose chart_family, the visual best suited to what the labels actually are -> not "
            "a default. Use bar whenever labels are named, unordered categories being compared against "
            "each other (methods, models, subjects, groups, conditions) -> this is the common case for "
            "a results/comparison table and must not become a line chart. Use line only when labels are "
            "a genuine ordered sequence (time, epoch, sequential step) where a trend or trajectory is "
            "the point. Use scatter for two independent continuous numeric measurements with no "
            "inherent order. Use heatmap only for a two-dimensional category-by-category matrix of "
            "values. If the request explicitly names a chart type, follow it unless the data makes "
            "that literally impossible."
        )
        text_part = f"REQUEST:\n{prompt}\n\nCONTEXT:\n{context_text[:4000]}" if context_text else f"REQUEST:\n{prompt}"
        if image_base64:
            user_content: Any = [
                {"type": "image_url", "image_url": {"url": f"data:image/png;base64,{image_base64}"}},
                {"type": "text", "text": text_part},
            ]
        else:
            user_content = text_part
        try:
            return self._client.structured_complete(
                [{"role": "system", "content": system}, {"role": "user", "content": user_content}],
                ExtractedChartDataset,
                model=VISION_MODEL_ID if image_base64 else None,
                reasoning_effort="low",
            )
        except Exception:
            return ExtractedChartDataset(has_data=False)

    def plan_semantic_composition(self, prompt: str, intent: VisualizationIntentDecision) -> SemanticCompositionPlan:
        system = (
            "Create a visual-first scientific composition from bounded semantic diagrams. Return 1-4 complementary facets that share one interpretation. "
            "Use multiple facets for an open concept (for example architecture plus an unrolled process plus a state-flow view). Do not make a tutorial, "
            "timeline of teaching steps, or text-only facet. Nodes and edges must encode the requested mechanism. Node IDs must be unique within each facet; "
            "every edge endpoint must exist. Keep each facet under 18 nodes. Do not emit coordinates, formulas, HTML, JavaScript, SVG, D3, Three.js, or callbacks. "
            "Do not claim observed data or paper-specific facts. The answer_markdown must briefly explain what was rendered and any illustrative assumptions."
        )
        try:
            plan = self._client.structured_complete(
                [
                    {"role": "system", "content": system},
                    {"role": "user", "content": f"REQUEST:\n{prompt}\n\nINTENT:\n{intent.model_dump_json()}"},
                ],
                SemanticCompositionPlan,
                reasoning_effort="medium",
            )
        except Exception:
            _logger.exception("Semantic composition planning failed for prompt=%r; using deterministic fallback plan", prompt)
            lowered = prompt.lower()
            if "rnn" in lowered or "recurrent neural" in lowered:
                facets = [
                    SemanticFacetPlan(
                        facet_id="recurrent-cell", title="Recurrent state update", purpose="Show the recurrent computation", layout="flow_horizontal", priority="primary",
                        nodes=[
                            SemanticNodePlan(node_id="x_t", label="input x(t)", role="input"),
                            SemanticNodePlan(node_id="h_prev", label="state h(t−1)", role="state"),
                            SemanticNodePlan(node_id="cell", label="recurrent cell", role="process", emphasis="primary"),
                            SemanticNodePlan(node_id="h_t", label="state h(t)", role="state"),
                            SemanticNodePlan(node_id="y_t", label="output y(t)", role="output"),
                        ],
                        edges=[
                            SemanticEdgePlan(source="x_t", target="cell", label="input"), SemanticEdgePlan(source="h_prev", target="cell", label="memory"),
                            SemanticEdgePlan(source="cell", target="h_t", label="update"), SemanticEdgePlan(source="h_t", target="y_t", label="readout"),
                        ],
                    ),
                    SemanticFacetPlan(
                        facet_id="unrolled-sequence", title="Unrolled through time", purpose="Expose state propagation across timesteps", layout="recurrent_unrolled", priority="supporting",
                        nodes=[SemanticNodePlan(node_id=f"x{i}", label=f"x({i})", role="input", group=f"t{i}") for i in range(1,4)] + [SemanticNodePlan(node_id=f"h{i}", label=f"h({i})", role="state", group=f"t{i}", emphasis="primary") for i in range(1,4)],
                        edges=[SemanticEdgePlan(source=f"x{i}", target=f"h{i}", label="input") for i in range(1,4)] + [SemanticEdgePlan(source=f"h{i}", target=f"h{i+1}", label="state") for i in range(1,3)],
                    ),
                ]
            else:
                facets = [SemanticFacetPlan(
                    facet_id="mechanism", title=intent.interpretation[:80] or "Requested mechanism", purpose="Render the requested components and relationships", priority="primary",
                    nodes=[SemanticNodePlan(node_id="input", label="input", role="input"), SemanticNodePlan(node_id="mechanism", label=(intent.interpretation or prompt)[:40], role="process", emphasis="primary"), SemanticNodePlan(node_id="output", label="output", role="output")],
                    edges=[SemanticEdgePlan(source="input", target="mechanism"), SemanticEdgePlan(source="mechanism", target="output")],
                )]
            plan = SemanticCompositionPlan(
                title=(intent.interpretation or prompt).strip().capitalize()[:90],
                interpretation=intent.interpretation or prompt,
                facets=facets,
                answer_markdown="I rendered the requested mechanism as coordinated schematic views. The layout is illustrative; the labeled relationships carry the scientific meaning.",
            )
        if not 1 <= len(plan.facets) <= 4:
            raise ValueError("A semantic composition must contain one to four visual facets")
        return plan

    @staticmethod
    def _default_draft(request: VisualizationRequest) -> AttentionLessonDraft:
        lowered = f"{request.prompt} {request.selection_text}".lower()
        causal = any(term in lowered for term in ("causal", "autoregressive", "masked attention"))
        return AttentionLessonDraft(
            title="Inside scaled dot-product attention",
            interpretation="Follow four tokens through query, key, and value projection, scoring, scaling, masking, softmax, and value aggregation.",
            requested_variant="causal_self_attention" if causal else "standard_self_attention",
            lesson_variant="causal_self_attention" if causal else "standard_self_attention",
            tokens=["Research", "Mate", "reads", "papers"],
            assumptions=[
                "The tensors are small illustrative teaching values, not weights extracted from a trained model.",
                "The lesson shows one attention computation rather than a multi-head comparison.",
            ],
            teaching_steps=[
                "Start with four tokens and small illustrative embeddings.",
                "Project each embedding into query, key, and value vectors.",
                "Compare every query with every key using dot products.",
                "Scale logits by the square root of the key dimension.",
                "Apply the optional causal mask before normalization.",
                "Normalize each row into attention weights with softmax.",
                "Use those weights to combine the value vectors into outputs.",
            ],
            claim_ids=[
                "qkv-projections", "query-key-products", "scale-key-dimension",
                "row-softmax", "value-aggregation", "causal-mask", "illustrative-values",
            ],
        )

    @staticmethod
    def _evidence_text(bundle: EvidenceBundle) -> str:
        sources_by_id = {source.source_id: source for source in bundle.sources}
        lines: list[str] = []
        for claim in bundle.claims:
            source_titles = [sources_by_id[source_id].title for source_id in claim.source_ids if source_id in sources_by_id]
            lines.append(
                f"- [{claim.claim_id}] {claim.text} | support={claim.support_level} | sources={', '.join(source_titles) or 'none'}"
            )
        return "\n".join(lines)[:9000]

    def _plan(
        self,
        request: VisualizationRequest,
        bundle: EvidenceBundle,
        profile_context: str,
        revision_feedback: list[str] | None = None,
    ) -> AttentionLessonDraft:
        feedback = "\n".join(f"- {item}" for item in (revision_feedback or []))
        messages = [
            {
                "role": "system",
                "content": (
                    "Plan a compact Attention Laboratory lesson. Use only claim IDs present in the evidence ledger. "
                    "Never emit tensors, renderer code, or trained-model weights. Return exactly four short token labels "
                    "and exactly seven teaching steps. Supported variants are ordinary self-attention and causal "
                    "self-attention. Unsupported variants must be described as a generic fallback rather than silently "
                    "approximated. Keep the interpretation readable in a narrow research-paper side panel."
                ),
            },
            {
                "role": "user",
                "content": (
                    f"REQUEST:\n{request.prompt}\n\nSELECTED CONTEXT:\n{request.selection_text[:1600]}\n\n"
                    f"FAMILIARITY: {request.familiarity}\n\nLEARNER/STYLE CONTEXT:\n{profile_context[:1600]}\n\n"
                    f"EVIDENCE LEDGER:\n{self._evidence_text(bundle)}"
                    + (f"\n\nREVISION FEEDBACK:\n{feedback}" if feedback else "")
                ),
            },
        ]
        return self._client.structured_complete(messages, AttentionLessonDraft, reasoning_effort="high")

    def _critic(self, draft: AttentionLessonDraft, bundle: EvidenceBundle) -> PlanCritique:
        messages = [
            {
                "role": "system",
                "content": (
                    "Review an Attention Laboratory plan for evidence linkage, explicit illustrative assumptions, "
                    "seven-step sequencing, side-panel density, and honest unsupported-variant disclosure. Be strict "
                    "but concise. Do not review renderer code or recalculate tensors."
                ),
            },
            {
                "role": "user",
                "content": f"PLAN:\n{draft.model_dump_json()}\n\nAVAILABLE CLAIMS:\n{self._evidence_text(bundle)}",
            },
        ]
        return self._client.structured_complete(messages, PlanCritique, reasoning_effort="medium")

    @staticmethod
    def _enforce_variant(request: VisualizationRequest, draft: AttentionLessonDraft) -> AttentionLessonDraft:
        lowered = f"{request.prompt} {request.selection_text}".lower()
        unsupported = next((label for term, label in _UNSUPPORTED_VARIANTS.items() if term in lowered), "")
        if unsupported:
            return draft.model_copy(update={
                "requested_variant": unsupported,
                "lesson_variant": "generic_fallback",
                "variant_notice": (
                    f"The request refers to {unsupported}, which is not rendered in this first capability. "
                    "This lesson shows standard scaled dot-product self-attention and does not claim to reproduce that variant."
                ),
            })
        if any(term in lowered for term in ("causal", "autoregressive", "masked attention")):
            return draft.model_copy(update={
                "requested_variant": "causal_self_attention",
                "lesson_variant": "causal_self_attention",
            })
        return draft.model_copy(update={
            "requested_variant": "standard_self_attention",
            "lesson_variant": "standard_self_attention",
        })

    def plan_with_review(
        self,
        request: VisualizationRequest,
        bundle: EvidenceBundle,
        profile_context: str = "",
    ) -> tuple[AttentionLessonDraft, list[str]]:
        warnings: list[str] = []
        try:
            draft = self._plan(request, bundle, profile_context)
        except Exception:
            return self._enforce_variant(request, self._default_draft(request)), [
                "The lesson planner was unavailable; ResearchMate used the deterministic attention lesson plan."
            ]

        draft = self._enforce_variant(request, draft)
        try:
            critique = self._critic(draft, bundle)
            if not critique.approved:
                revised = self._plan(request, bundle, profile_context, critique.feedback)
                draft = self._enforce_variant(request, revised)
                final_critique = self._critic(draft, bundle)
                if not final_critique.approved:
                    warnings.append("The plan critic remained critical after the single allowed revision; the validated revision was used.")
        except Exception:
            warnings.append("Plan review was unavailable; the schema-valid lesson plan was compiled directly.")
        return draft, warnings