app/schemas/composition.py

from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel, Field


RuntimeScalar = float | int | bool | str


class ParameterConstraints(BaseModel):
    minimum: float | None = None
    maximum: float | None = None
    exclusive_minimum: bool = False
    exclusive_maximum: bool = False
    allowed_values: list[str] = Field(default_factory=list)


class ParameterSpec(BaseModel):
    parameter_id: str
    label: str
    value_type: Literal["scalar", "integer", "boolean", "enum"]
    initial_value: RuntimeScalar
    constraints: ParameterConstraints = Field(default_factory=ParameterConstraints)
    recompute: Literal[
        "presentation-only", "derived-values", "geometry-transform", "chart-recompile",
        "field-resample", "remesh", "full-composition",
    ] = "derived-values"
    persistence: Literal["ephemeral", "session", "lesson"] = "session"


class ExpressionInstruction(BaseModel):
    instruction_id: str
    op: Literal[
        "constant", "variable", "negate", "add", "subtract", "multiply", "divide", "power",
        "sin", "cos", "tan", "sqrt", "abs", "exp", "log", "min", "max",
    ]
    args: list[str] = Field(default_factory=list)
    value: float | None = None
    name: str = ""


class ExpressionProgram(BaseModel):
    version: Literal["1.0"] = "1.0"
    instructions: list[ExpressionInstruction]
    result_id: str


class DerivedValueSpec(BaseModel):
    value_id: str
    label: str
    program: ExpressionProgram
    unit: str = ""
    display_precision: int = 3


class DatasetSpec(BaseModel):
    dataset_id: str
    columns: list[str]
    rows: list[dict[str, RuntimeScalar | None]]


class DataTransformSpec(BaseModel):
    transform_id: str
    kind: Literal["filter", "derive", "aggregate", "bin", "normalize"]
    input_field: str = ""
    output_field: str = ""
    operator: Literal["eq", "neq", "gt", "gte", "lt", "lte", "add", "subtract", "multiply", "divide", "mean", "sum", "count", "min", "max", "zscore", "minmax"] = "eq"
    operand: RuntimeScalar | None = None
    group_by: list[str] = Field(default_factory=list)
    bins: int = 10


class ChartEncoding(BaseModel):
    channel: Literal["x", "y", "color", "size", "label", "source", "target", "weight"]
    field: str
    value_type: Literal["quantitative", "ordinal", "nominal", "temporal"] = "quantitative"


class ChartSpec(BaseModel):
    chart_id: str
    family: Literal[
        "line", "scatter", "area", "bar", "histogram", "box", "violin", "heatmap", "contour",
        "error-band", "force-graph", "tree", "sankey", "hierarchy", "chord", "treemap", "packed-circles",
    ]
    dataset_id: str
    transforms: list[DataTransformSpec] = Field(default_factory=list)
    encodings: list[ChartEncoding]
    x_label: str = ""
    y_label: str = ""
    renderer_preference: Literal["automatic", "plotly", "d3", "svg"] = "automatic"
    interactions: list[Literal["zoom", "hover", "select-datum", "brush-range", "linked-selection"]] = Field(default_factory=list)


class SemanticDiagramNode(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 SemanticDiagramEdge(BaseModel):
    source: str
    target: str
    label: str = ""
    direction: Literal["forward", "backward", "bidirectional", "none"] = "forward"
    style: Literal["solid", "dashed", "feedback"] = "solid"


class SemanticDiagramSpec(BaseModel):
    diagram_id: str
    title: str = ""
    layout: Literal["flow_horizontal", "flow_vertical", "recurrent_unrolled", "cycle", "layered"] = "flow_horizontal"
    nodes: list[SemanticDiagramNode]
    edges: list[SemanticDiagramEdge]
    interactions: list[Literal["zoom", "pan", "select-node", "focus", "step"]] = Field(default_factory=lambda: ["zoom", "pan", "select-node"])


class Bounds3D(BaseModel):
    x: tuple[float, float]
    y: tuple[float, float]
    z: tuple[float, float]


class Resolution3D(BaseModel):
    preview: int = 32
    refined: int = 64
    high_quality: int = 96


class GeometryPresentation(BaseModel):
    color: str = "#2f80c9"
    opacity_parameter_id: str = "opacity"
    wireframe_parameter_id: str = "wireframe"
    show_axes_parameter_id: str = "show_axes"
    show_grid: bool = True
    show_bounding_box: bool = False


class RegisteredGeometrySpec(BaseModel):
    kind: Literal["registered"] = "registered"
    operator: Literal["sphere", "ellipsoid", "prolate_spheroid", "torus", "mobius_strip"]
    parameter_bindings: dict[str, str]
    presentation: GeometryPresentation = Field(default_factory=GeometryPresentation)


class RegisteredTopologySpec(BaseModel):
    kind: Literal["registered_topology"] = "registered_topology"
    operator: Literal["mobius_strip"]
    parameter_bindings: dict[str, str]
    presentation: GeometryPresentation = Field(default_factory=GeometryPresentation)


class ParametricGeometrySpec(BaseModel):
    kind: Literal["parametric_surface"] = "parametric_surface"
    x_program: ExpressionProgram
    y_program: ExpressionProgram
    z_program: ExpressionProgram
    u_domain: tuple[float, float]
    v_domain: tuple[float, float]
    u_segments: int = 64
    v_segments: int = 32
    presentation: GeometryPresentation = Field(default_factory=GeometryPresentation)


class ImplicitGeometrySpec(BaseModel):
    kind: Literal["implicit_surface"] = "implicit_surface"
    expression: ExpressionProgram
    iso_value_parameter_id: str = "iso_value"
    bounds: Bounds3D
    resolution: Resolution3D = Field(default_factory=Resolution3D)
    interpolation: Literal["linear"] = "linear"
    normal_method: Literal["gradient", "face-average"] = "gradient"
    deduplicate_vertices: bool = True
    presentation: GeometryPresentation = Field(default_factory=GeometryPresentation)


class VectorFieldGeometrySpec(BaseModel):
    kind: Literal["vector_field"] = "vector_field"
    x_program: ExpressionProgram
    y_program: ExpressionProgram
    z_program: ExpressionProgram
    bounds: Bounds3D
    samples_per_axis: int = 12
    presentation: GeometryPresentation = Field(default_factory=GeometryPresentation)


GeometrySpec = RegisteredGeometrySpec | RegisteredTopologySpec | ParametricGeometrySpec | ImplicitGeometrySpec | VectorFieldGeometrySpec


class ThreeScenePanelSpec(BaseModel):
    panel_id: str
    kind: Literal["three_scene"] = "three_scene"
    geometry: GeometrySpec
    bindings: list[str] = Field(default_factory=list)
    order: int = 0


class ChartPanelSpec(BaseModel):
    panel_id: str
    kind: Literal["chart"] = "chart"
    chart: ChartSpec
    bindings: list[str] = Field(default_factory=list)
    order: int = 0


class DiagramPanelSpec(BaseModel):
    panel_id: str
    kind: Literal["diagram"] = "diagram"
    diagram_type: Literal["ellipse_cross_section", "matrix", "algorithm", "annotated_axes", "semantic"]
    title: str = ""
    semantic: SemanticDiagramSpec | None = None
    bindings: list[str] = Field(default_factory=list)
    order: int = 0


class MolecularPanelSpec(BaseModel):
    panel_id: str
    kind: Literal["molecular"] = "molecular"
    asset_id: str
    coordinate_format: Literal["cif", "bcif"]
    bindings: list[str] = Field(default_factory=list)
    order: int = 0


class EquationPanelSpec(BaseModel):
    panel_id: str
    kind: Literal["equation"] = "equation"
    equations: list[str]
    bindings: list[str] = Field(default_factory=list)
    order: int = 0


class AnnotationPanelSpec(BaseModel):
    panel_id: str
    kind: Literal["annotation"] = "annotation"
    text: str
    bindings: list[str] = Field(default_factory=list)
    order: int = 0


CompositionPanelSpec = ThreeScenePanelSpec | ChartPanelSpec | DiagramPanelSpec | MolecularPanelSpec | EquationPanelSpec | AnnotationPanelSpec


class InteractionSpec(BaseModel):
    interaction_id: str
    kind: Literal["parameter", "selection", "timeline", "visibility"]
    source_panel_id: str = ""
    state_id: str
    target_panel_ids: list[str]


class SharedStateSpec(BaseModel):
    state_id: str
    value_type: Literal["scalar", "integer", "boolean", "enum", "selection", "range", "time"]
    initial_value: Any = None


class TimelineFrameSpec(BaseModel):
    frame_id: str
    label: str
    state_changes: dict[str, RuntimeScalar] = Field(default_factory=dict)


class CompositionTimelineSpec(BaseModel):
    timeline_id: str
    interval_ms: int = 1200
    loop: bool = False
    frames: list[TimelineFrameSpec]


class MathematicalProperty(BaseModel):
    name: str
    value: Any = None
    status: Literal["registered", "computed", "inferred", "unknown"]
    method: str = ""
    caveat: str = ""


class CompositionAssertion(BaseModel):
    assertion_id: str
    kind: Literal["finite", "constraint", "derived_close", "dataset_shape", "mesh_budget"]
    target_id: str
    expected: float | None = None
    tolerance: float = 1e-6


class ProvenanceSpec(BaseModel):
    interpretation: str
    operator_sources: list[str] = Field(default_factory=list)
    evidence_claim_ids: list[str] = Field(default_factory=list)
    research_used: bool = False


class VisualizationIntentRecord(BaseModel):
    specificity: Literal["exact", "open", "ambiguous"] = "open"
    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)


class RequirementCoverage(BaseModel):
    requirement: str
    status: Literal["satisfied", "unsupported", "conflicting"]
    panel_ids: list[str] = Field(default_factory=list)
    explanation: str = ""


class CompositionSpec(BaseModel):
    version: Literal["1.0"] = "1.0"
    schema_version: Literal["1.0"] = "1.0"
    compiler_version: str = "composition-1"
    operator_registry_version: str = "geometry-1"
    capability: Literal["composition"] = "composition"
    project_id: str
    prompt: str
    title: str
    answer_markdown: str = ""
    parameters: list[ParameterSpec]
    derived_values: list[DerivedValueSpec] = Field(default_factory=list)
    datasets: list[DatasetSpec] = Field(default_factory=list)
    panels: list[CompositionPanelSpec]
    interactions: list[InteractionSpec] = Field(default_factory=list)
    shared_state: list[SharedStateSpec] = Field(default_factory=list)
    timeline: CompositionTimelineSpec | None = None
    assumptions: list[str] = Field(default_factory=list)
    provenance: ProvenanceSpec
    assertions: list[CompositionAssertion] = Field(default_factory=list)
    mathematical_properties: list[MathematicalProperty] = Field(default_factory=list)
    request_intent: VisualizationIntentRecord | None = None
    requirement_coverage: list[RequirementCoverage] = Field(default_factory=list)
    created_at: float


class RuntimeCheckpoint(BaseModel):
    composition_id: str
    checkpoint_revision: int = 0
    parameter_revision: int = 0
    parameter_values: dict[str, RuntimeScalar] = Field(default_factory=dict)
    compiler_versions: dict[str, str] = Field(default_factory=dict)


class CheckpointUpdate(BaseModel):
    base_checkpoint_revision: int
    parameter_revision: int
    parameter_values: dict[str, RuntimeScalar]
    compiler_versions: dict[str, str] = Field(default_factory=dict)


class CompositionPayload(BaseModel):
    capability: Literal["composition"] = "composition"
    spec: CompositionSpec
    checkpoint: RuntimeCheckpoint | None = None
    warnings: list[str] = Field(default_factory=list)