from __future__ import annotations

from typing import Any, Literal

from pydantic import BaseModel, Field


class EvidenceBox(BaseModel):
    page: int = 0
    x: float = 0.0
    y: float = 0.0
    w: float = 0.0
    h: float = 0.0


class SelectionSnippet(BaseModel):
    page_number: int = 0
    text: str = ""
    boxes: list[EvidenceBox] = Field(default_factory=list)


class VisualizationRequest(BaseModel):
    request_id: str
    prompt: str
    intention: str = "learn"
    selection_text: str = ""
    surrounding_context: str = ""
    selection_snippets: list[SelectionSnippet] = Field(default_factory=list)
    selection_image_base64: str = ""
    active_document_ids: list[str] = Field(default_factory=list)
    # Hard access policy. False means orchestration must not enumerate, embed,
    # retrieve, or recall anything from the active project. Explicit selection
    # text/images remain direct user input and are not project retrieval.
    project_context_enabled: bool = True
    familiarity: str = "graduate"
    resolution_token: str = ""
    resolution_candidate_id: str = ""


class EvidenceSource(BaseModel):
    source_id: str
    origin: Literal["selection", "project", "web", "builtin", "database"]
    title: str
    url: str = ""
    document_id: str = ""
    page_number: int = 0
    anchor: str = ""
    excerpt: str = ""
    boxes: list[EvidenceBox] = Field(default_factory=list)
    authority: Literal["explicit", "project", "primary", "official", "secondary", "canonical", "curated", "unknown"] = "unknown"


class EvidenceClaim(BaseModel):
    claim_id: str
    text: str
    claim_type: Literal["source_fact", "derived_fact", "standard_definition", "illustrative_choice", "missing_fact"]
    support_level: Literal["direct", "derived", "canonical", "illustrative", "missing"]
    source_ids: list[str] = Field(default_factory=list)


class EvidenceBundle(BaseModel):
    sources: list[EvidenceSource] = Field(default_factory=list)
    claims: list[EvidenceClaim] = Field(default_factory=list)
    warnings: list[str] = Field(default_factory=list)
    live_research_used: bool = False


class AttentionLessonDraft(BaseModel):
    title: str
    interpretation: str
    requested_variant: str = "standard_self_attention"
    lesson_variant: Literal["standard_self_attention", "causal_self_attention", "generic_fallback"] = "standard_self_attention"
    variant_notice: str = ""
    tokens: list[str] = Field(default_factory=list)
    assumptions: list[str] = Field(default_factory=list)
    teaching_steps: list[str] = Field(default_factory=list)
    claim_ids: list[str] = Field(default_factory=list)


class PlanCritique(BaseModel):
    approved: bool
    feedback: list[str] = Field(default_factory=list)
    missing_claim_ids: list[str] = Field(default_factory=list)


class TensorInput(BaseModel):
    tensor_id: str
    label: str
    values: list[list[float]]


class TensorOperation(BaseModel):
    operation_id: str
    op: Literal[
        "identity",
        "matmul",
        "matmul_transpose_right",
        "divide_sqrt_dimension",
        "apply_causal_mask",
        "softmax_rows",
    ]
    inputs: list[str]
    output: str
    dimension: int = 0


class NumericAssertion(BaseModel):
    assertion_id: str
    kind: Literal["shape", "finite", "row_sum", "masked_zero", "matmul_close"]
    tensor_id: str
    tolerance: float = 1e-6
    passed: bool = False
    detail: str = ""


class PanelSpec(BaseModel):
    panel_id: str
    panel_type: Literal[
        "token_strip",
        "tensor_matrix",
        "attention_heatmap",
        "weighted_attention",
        "vector_aggregation",
        "equation",
        "explanation",
        "evidence",
    ]
    title: str
    bindings: list[str] = Field(default_factory=list)
    claim_ids: list[str] = Field(default_factory=list)
    order: int


class TimelineStep(BaseModel):
    step_id: str
    label: str
    explanation: str
    operation_ids: list[str] = Field(default_factory=list)
    active_panel_ids: list[str] = Field(default_factory=list)
    claim_ids: list[str] = Field(default_factory=list)


class VisualLessonSpec(BaseModel):
    version: Literal["1.0"] = "1.0"
    compiler_version: str = "attention-1"
    capability: Literal["attention"] = "attention"
    project_id: str
    prompt: str
    title: str
    interpretation: str
    answer_markdown: str = ""
    requested_variant: str
    lesson_variant: Literal["standard_self_attention", "causal_self_attention", "generic_fallback"]
    variant_notice: str = ""
    tokens: list[str]
    assumptions: list[str] = Field(default_factory=list)
    evidence_sources: list[EvidenceSource] = Field(default_factory=list)
    evidence_claims: list[EvidenceClaim] = Field(default_factory=list)
    inputs: list[TensorInput] = Field(default_factory=list)
    operations: list[TensorOperation] = Field(default_factory=list)
    assertions: list[NumericAssertion] = Field(default_factory=list)
    panels: list[PanelSpec] = Field(default_factory=list)
    timeline: list[TimelineStep] = Field(default_factory=list)
    seed: int = 170603762
    created_at: float


class CompiledBranch(BaseModel):
    branch_id: Literal["unmasked", "causal"]
    tensors: dict[str, list[list[float]]]
    assertions: list[NumericAssertion] = Field(default_factory=list)


class CompiledLesson(BaseModel):
    branches: list[CompiledBranch]
    assertions_passed: bool


class VisualLessonPayload(BaseModel):
    capability: Literal["attention"] = "attention"
    spec: VisualLessonSpec
    compiled: CompiledLesson
    warnings: list[str] = Field(default_factory=list)


class ProteinIdentity(BaseModel):
    requested_label: str
    protein_name: str
    gene_name: str = ""
    organism: str = ""
    taxonomy_id: int = 0
    uniprot_accession: str = ""
    synonyms: list[str] = Field(default_factory=list)
    source_ids: list[str] = Field(default_factory=list)


class ProteinChain(BaseModel):
    chain_id: str
    auth_chain_id: str = ""
    entity_id: str = ""
    description: str = ""
    sequence_length: int = 0
    uniprot_accession: str = ""


class ProteinLigand(BaseModel):
    comp_id: str
    name: str = ""
    chain_id: str = ""
    instance_count: int = 1


class ProteinFeatureSelector(BaseModel):
    chain_id: str = ""
    auth_chain_id: str = ""
    numbering: Literal["label", "author", "uniprot", "component"] = "label"
    start: int = 0
    end: int = 0
    comp_id: str = ""


class ProteinFeature(BaseModel):
    feature_id: str
    kind: Literal["chain", "domain", "ligand", "binding_site", "residue", "mutation"]
    label: str
    selector: ProteinFeatureSelector
    color: str = "#D27A3A"
    claim_ids: list[str] = Field(default_factory=list)
    mapping_note: str = ""


class ProteinStructure(BaseModel):
    source: Literal["pdb", "alphafold"]
    structure_id: str
    pdb_id: str = ""
    assembly_id: str = ""
    coordinate_asset_id: str
    coordinate_sha256: str
    coordinate_format: Literal["bcif", "cif"]
    source_url: str
    experimental_method: str = ""
    resolution_angstrom: float = 0.0
    coverage_fraction: float = 0.0
    confidence_mode: Literal["none", "plddt"] = "none"
    chains: list[ProteinChain] = Field(default_factory=list)
    ligands: list[ProteinLigand] = Field(default_factory=list)


class ProteinViewSpec(BaseModel):
    representation: Literal["cartoon", "surface", "ball_and_stick"] = "cartoon"
    color_scheme: Literal["chain", "secondary_structure", "confidence"] = "chain"
    spin: bool = False
    background: Literal["cream", "white", "dark"] = "cream"
    visible_feature_ids: list[str] = Field(default_factory=list)


class ProteinLessonSpec(BaseModel):
    version: Literal["1.0"] = "1.0"
    capability: Literal["protein"] = "protein"
    project_id: str
    prompt: str
    title: str
    interpretation: str
    answer_markdown: str = ""
    identity: ProteinIdentity
    structure: ProteinStructure
    features: list[ProteinFeature] = Field(default_factory=list)
    view: ProteinViewSpec = Field(default_factory=ProteinViewSpec)
    assumptions: list[str] = Field(default_factory=list)
    evidence_sources: list[EvidenceSource] = Field(default_factory=list)
    evidence_claims: list[EvidenceClaim] = Field(default_factory=list)
    timeline: list[TimelineStep] = Field(default_factory=list)
    created_at: float


class ProteinLessonPayload(BaseModel):
    capability: Literal["protein"] = "protein"
    spec: ProteinLessonSpec
    warnings: list[str] = Field(default_factory=list)


class ProteinLessonDraft(BaseModel):
    title: str
    interpretation: str
    assumptions: list[str] = Field(default_factory=list)
    teaching_steps: list[str] = Field(default_factory=list)
    feature_ids: list[str] = Field(default_factory=list)
    initial_representation: Literal["cartoon", "surface", "ball_and_stick"] = "cartoon"
    initial_color_scheme: Literal["chain", "secondary_structure", "confidence"] = "chain"


class ProteinResolutionCandidate(BaseModel):
    candidate_id: str
    protein_name: str
    gene_name: str = ""
    organism: str = ""
    uniprot_accession: str = ""
    structure_source: Literal["pdb", "alphafold", "unknown"] = "unknown"
    structure_id: str = ""
    reason: str
    source_url: str = ""


class ProteinResolutionPayload(BaseModel):
    resolution_token: str
    prompt: str
    candidates: list[ProteinResolutionCandidate]
    expires_in_seconds: int = 600


class NucleicFormParameters(BaseModel):
    form: Literal["a_dna", "b_dna", "z_dna", "rna"]
    handedness: Literal["right", "left"]
    bases_per_turn: float
    rise_angstrom: float
    radius_angstrom: float
    strand_offset_degrees: float = 144.0


class NucleicFeature(BaseModel):
    feature_id: str
    kind: Literal[
        "backbone", "base", "base_pair", "hydrogen_bond", "direction",
        "major_groove", "minor_groove", "sugar", "phosphate", "chain",
    ]
    label: str
    color: str
    claim_ids: list[str] = Field(default_factory=list)
    chain_id: str = ""
    start: int = 0
    end: int = 0
    comp_id: str = ""


class NucleicStructure(BaseModel):
    pdb_id: str
    assembly_id: str = ""
    title: str = ""
    coordinate_asset_id: str
    coordinate_sha256: str
    coordinate_format: Literal["bcif", "cif"]
    source_url: str
    experimental_method: str = ""
    resolution_angstrom: float = 0.0
    polymer_types: list[str] = Field(default_factory=list)
    chains: list[ProteinChain] = Field(default_factory=list)


class NucleicComposition(BaseModel):
    primary_view: Literal["helix_3d", "molecule_3d", "chemistry", "comparison", "structure"]
    visible_feature_ids: list[str] = Field(default_factory=list)
    caption: str = ""
    comparison_items: list[Literal["dna", "rna", "a_dna", "b_dna", "z_dna"]] = Field(default_factory=list)


class NucleobaseAtom(BaseModel):
    atom_id: int
    element: Literal["H", "C", "N", "O", "P", "S"]
    x: float
    y: float
    z: float


class NucleobaseBond(BaseModel):
    atom_a: int
    atom_b: int
    order: int = 1


class NucleobaseMolecule(BaseModel):
    name: Literal["adenine", "cytosine", "guanine", "thymine", "uracil"]
    symbol: Literal["A", "C", "G", "T", "U"]
    molecular_formula: str
    pubchem_cid: int
    source_url: str
    atoms: list[NucleobaseAtom]
    bonds: list[NucleobaseBond]


class NucleicAcidSpec(BaseModel):
    version: Literal["1.0"] = "1.0"
    compiler_version: str = "nucleic-1"
    capability: Literal["nucleic_acid"] = "nucleic_acid"
    project_id: str
    prompt: str
    answer_markdown: str = ""
    mode: Literal["concept", "sequence", "comparison", "structure"]
    molecule: Literal["dna", "rna", "dna_rna", "dna_forms"]
    sequence: str
    complement: str = ""
    focus_base: Literal["A", "C", "G", "T", "U"] | None = None
    nucleobase_molecule: NucleobaseMolecule | None = None
    sequence_is_illustrative: bool = False
    forms: list[NucleicFormParameters] = Field(default_factory=list)
    structure: NucleicStructure | None = None
    features: list[NucleicFeature] = Field(default_factory=list)
    composition: NucleicComposition
    assumptions: list[str] = Field(default_factory=list)
    evidence_sources: list[EvidenceSource] = Field(default_factory=list)
    evidence_claims: list[EvidenceClaim] = Field(default_factory=list)
    created_at: float


class NucleicPoint(BaseModel):
    index: int
    base: str
    complement: str = ""
    strand_a: list[float]
    strand_b: list[float] = Field(default_factory=list)


class CompiledNucleicForm(BaseModel):
    form: Literal["a_dna", "b_dna", "z_dna", "rna"]
    points: list[NucleicPoint]


class CompiledNucleicGeometry(BaseModel):
    forms: list[CompiledNucleicForm]
    assertions_passed: bool = True


class NucleicAcidPayload(BaseModel):
    capability: Literal["nucleic_acid"] = "nucleic_acid"
    spec: NucleicAcidSpec
    compiled: CompiledNucleicGeometry
    warnings: list[str] = Field(default_factory=list)


class NucleicVisualizationDraft(BaseModel):
    primary_view: Literal["helix_3d", "molecule_3d", "chemistry", "comparison", "structure"]
    visible_feature_ids: list[str] = Field(default_factory=list)
    caption: str = Field(default="", max_length=180)


class ThermodynamicInputs(BaseModel):
    moles: float
    initial_pressure_pa: float
    initial_temperature_k: float
    initial_volume_m3: float
    gamma: float
    gas_constant: float = 8.314462618


class ThermodynamicSample(BaseModel):
    sample_index: int
    control_value: float
    pressure_pa: float
    volume_m3: float
    temperature_k: float
    heat_j: float
    work_by_j: float
    delta_internal_energy_j: float


class ThermodynamicProcessBranch(BaseModel):
    process_id: Literal["isothermal", "adiabatic", "isobaric", "isochoric"]
    label: str
    control_label: str
    invariant_latex: str
    equation_latex: str
    claim_ids: list[str] = Field(default_factory=list)
    samples: list[ThermodynamicSample]


class ThermodynamicsSpec(BaseModel):
    version: Literal["1.0"] = "1.0"
    compiler_version: str = "thermodynamics-1"
    capability: Literal["thermodynamics"] = "thermodynamics"
    project_id: str
    prompt: str
    title: str = "Ideal-gas piston"
    answer_markdown: str = ""
    primary_process: Literal["isothermal", "adiabatic", "isobaric", "isochoric"]
    initial_sample_index: int = 25
    inputs: ThermodynamicInputs
    process_ids: list[Literal["isothermal", "adiabatic", "isobaric", "isochoric"]] = Field(
        default_factory=lambda: ["isothermal", "adiabatic", "isobaric", "isochoric"]
    )
    assumptions: list[str] = Field(default_factory=list)
    evidence_sources: list[EvidenceSource] = Field(default_factory=list)
    evidence_claims: list[EvidenceClaim] = Field(default_factory=list)
    seed: int = 314159
    created_at: float


class CompiledThermodynamics(BaseModel):
    branches: list[ThermodynamicProcessBranch]
    assertions_passed: bool = True


class ThermodynamicsPayload(BaseModel):
    capability: Literal["thermodynamics"] = "thermodynamics"
    spec: ThermodynamicsSpec
    compiled: CompiledThermodynamics
    warnings: list[str] = Field(default_factory=list)


class DifferentialEquationParameters(BaseModel):
    growth_rate: float = 1.0
    carrying_capacity: float = 100.0
    linear_coefficient: float = 1.0
    forcing: float = 0.0
    initial_value: float = 1.0
    initial_velocity: float = 0.0
    natural_frequency: float = 1.0
    damping_ratio: float = 0.15


class DifferentialEquationSample(BaseModel):
    sample_index: int
    t: float
    primary: float
    derivative: float
    secondary: float = 0.0


class DirectionFieldPoint(BaseModel):
    t: float
    state: float
    slope: float


class DifferentialEquationFeature(BaseModel):
    feature_id: str
    label: str
    t: float = 0.0
    value: float
    kind: Literal["equilibrium", "inflection", "initial_state"]


class DifferentialEquationSpec(BaseModel):
    version: Literal["1.0"] = "1.0"
    compiler_version: str = "differential-equation-1"
    capability: Literal["differential_equation"] = "differential_equation"
    project_id: str
    prompt: str
    title: str = "Differential equation"
    answer_markdown: str = ""
    family: Literal["logistic", "linear_first_order", "damped_oscillator"]
    interpretation: str
    equation_latex: str
    parameters: DifferentialEquationParameters
    t_start: float = 0.0
    t_end: float = 10.0
    sample_count: int = 401
    initial_sample_index: int = 0
    enabled_views: list[Literal["trajectory", "direction_field", "phase_portrait"]] = Field(default_factory=list)
    assumptions: list[str] = Field(default_factory=list)
    evidence_sources: list[EvidenceSource] = Field(default_factory=list)
    evidence_claims: list[EvidenceClaim] = Field(default_factory=list)
    created_at: float


class CompiledDifferentialEquation(BaseModel):
    samples: list[DifferentialEquationSample]
    direction_field: list[DirectionFieldPoint] = Field(default_factory=list)
    features: list[DifferentialEquationFeature] = Field(default_factory=list)
    assertions_passed: bool = True


class DifferentialEquationPayload(BaseModel):
    capability: Literal["differential_equation"] = "differential_equation"
    spec: DifferentialEquationSpec
    compiled: CompiledDifferentialEquation
    warnings: list[str] = Field(default_factory=list)


class GraphNodeSpec(BaseModel):
    node_id: str
    label: str
    x: float
    y: float


class GraphEdgeSpec(BaseModel):
    edge_id: str
    source: str
    target: str
    weight: float = 1.0


class GraphSearchStep(BaseModel):
    step_index: int
    current_node: str = ""
    frontier: list[str] = Field(default_factory=list)
    visited: list[str] = Field(default_factory=list)
    distances: dict[str, float] = Field(default_factory=dict)
    parent_edge_ids: list[str] = Field(default_factory=list)
    active_edge_ids: list[str] = Field(default_factory=list)
    final_path: list[str] = Field(default_factory=list)
    description: str = ""


class GraphSearchBranch(BaseModel):
    algorithm: Literal["bfs", "dfs", "dijkstra", "astar"]
    start_node: str
    target_node: str
    found: bool
    path: list[str] = Field(default_factory=list)
    path_cost: float = 0.0
    steps: list[GraphSearchStep]


class GraphAlgorithmSpec(BaseModel):
    version: Literal["1.0"] = "1.0"
    compiler_version: str = "graph-algorithm-1"
    capability: Literal["graph_algorithm"] = "graph_algorithm"
    project_id: str
    prompt: str
    title: str = "Graph pathfinding"
    answer_markdown: str = ""
    directed: bool = False
    nodes: list[GraphNodeSpec]
    edges: list[GraphEdgeSpec]
    primary_algorithm: Literal["bfs", "dfs", "dijkstra", "astar"] = "dijkstra"
    initial_start_node: str
    initial_target_node: str
    graph_is_illustrative: bool = True
    assumptions: list[str] = Field(default_factory=list)
    evidence_sources: list[EvidenceSource] = Field(default_factory=list)
    evidence_claims: list[EvidenceClaim] = Field(default_factory=list)
    created_at: float


class CompiledGraphAlgorithm(BaseModel):
    branches: list[GraphSearchBranch]
    assertions_passed: bool = True


class GraphAlgorithmPayload(BaseModel):
    capability: Literal["graph_algorithm"] = "graph_algorithm"
    spec: GraphAlgorithmSpec
    compiled: CompiledGraphAlgorithm
    warnings: list[str] = Field(default_factory=list)


class ArrayAlgorithmStep(BaseModel):
    step_index: int
    values: list[float]
    compared_indices: list[int] = Field(default_factory=list)
    active_indices: list[int] = Field(default_factory=list)
    sorted_indices: list[int] = Field(default_factory=list)
    pivot_index: int = -1
    found_index: int = -1
    range_start: int = -1
    range_end: int = -1
    operation: Literal["initial", "compare", "swap", "write", "partition", "found", "complete", "not_found"] = "initial"
    description: str = ""
    comparisons: int = 0
    writes: int = 0


class ArrayAlgorithmBranch(BaseModel):
    algorithm: Literal["bubble_sort", "insertion_sort", "selection_sort", "merge_sort", "quick_sort", "linear_search", "binary_search"]
    steps: list[ArrayAlgorithmStep]
    final_values: list[float]
    found_index: int = -1
    comparisons: int = 0
    writes: int = 0


class ArrayAlgorithmSpec(BaseModel):
    version: Literal["1.0"] = "1.0"
    compiler_version: str = "array-algorithm-1"
    capability: Literal["array_algorithm"] = "array_algorithm"
    project_id: str
    prompt: str
    title: str = "Array algorithm"
    answer_markdown: str = ""
    input_values: list[float]
    primary_algorithm: Literal["bubble_sort", "insertion_sort", "selection_sort", "merge_sort", "quick_sort", "linear_search", "binary_search"]
    enabled_algorithms: list[Literal["bubble_sort", "insertion_sort", "selection_sort", "merge_sort", "quick_sort", "linear_search", "binary_search"]]
    search_target: float = 0.0
    has_search_target: bool = False
    input_is_illustrative: bool = True
    assumptions: list[str] = Field(default_factory=list)
    evidence_sources: list[EvidenceSource] = Field(default_factory=list)
    evidence_claims: list[EvidenceClaim] = Field(default_factory=list)
    created_at: float


class CompiledArrayAlgorithm(BaseModel):
    branches: list[ArrayAlgorithmBranch]
    assertions_passed: bool = True


class ArrayAlgorithmPayload(BaseModel):
    capability: Literal["array_algorithm"] = "array_algorithm"
    spec: ArrayAlgorithmSpec
    compiled: CompiledArrayAlgorithm
    warnings: list[str] = Field(default_factory=list)


class MolecularResolutionCandidate(BaseModel):
    candidate_id: Literal["protein", "nucleic_acid", "whole_assembly"]
    label: str
    description: str
    pdb_id: str


class MolecularResolutionPayload(BaseModel):
    resolution_token: str
    prompt: str
    candidates: list[MolecularResolutionCandidate]
    expires_in_seconds: int = 600


class VisualizationClarificationCandidate(BaseModel):
    candidate_id: str
    label: str
    description: str = ""


class ExtractedChartRow(BaseModel):
    label: str
    value: float


class ExtractedChartDataset(BaseModel):
    has_data: bool
    chart_family: Literal["bar", "line", "scatter", "heatmap"] = "bar"
    x_label: str = ""
    y_label: str = ""
    rows: list[ExtractedChartRow] = Field(default_factory=list)


class VisualizationClarificationPayload(BaseModel):
    resolution_token: str
    prompt: str
    question: str
    reason: str = ""
    candidates: list[VisualizationClarificationCandidate]
    allow_free_text: bool = True
    expires_in_seconds: int = 600


from app.schemas.composition import CompositionPayload


LessonPayload = VisualLessonPayload | ProteinLessonPayload | NucleicAcidPayload | ThermodynamicsPayload | DifferentialEquationPayload | GraphAlgorithmPayload | ArrayAlgorithmPayload | CompositionPayload
ResolutionPayload = ProteinResolutionPayload | MolecularResolutionPayload | VisualizationClarificationPayload


class VisualizationReadyPayload(BaseModel):
    request_id: str
    topic: str
    result_kind: Literal["lesson", "legacy", "resolution_required", "clarification_required"]
    lesson_id: str = ""
    lesson: LessonPayload | None = None
    resolution: ResolutionPayload | None = None
    visual: dict[str, Any] | None = None
    warning: str = ""
    answer_markdown: str = ""