backend/models.py

from __future__ import annotations

from typing import Literal

from pydantic import BaseModel, Field


RelName = str
Quantifier = Literal["any", "all", "none", "exactly"]
ColorBy = Literal["relation", "family", "confidence", "depth", "macroarea"]
ViewName = Literal["tree", "radial", "map", "table", "stats"]


class WalsClause(BaseModel):
    feature_id: str
    values: list[int] = Field(default_factory=list)


class NodePredicate(BaseModel):
    languages: list[str] = Field(default_factory=list)
    families: list[str] = Field(default_factory=list)
    macroareas: list[str] = Field(default_factory=list)
    statuses: list[str] = Field(default_factory=list)
    term_contains: str = ""
    term_regex: str | None = None
    require_coords: bool = False
    min_lat: float | None = None
    max_lat: float | None = None
    min_lon: float | None = None
    max_lon: float | None = None
    wals: list[WalsClause] = Field(default_factory=list)
    phonemes_have: list[str] = Field(default_factory=list)
    phonemes_lack: list[str] = Field(default_factory=list)
    require_tone: bool | None = None
    # Popularity within language (modern lects). Zipf ~3 rare … ~7 very common.
    min_zipf: float | None = Field(default=None, ge=0.0, le=8.0)
    # Keep leaves whose corpus rank is <= N (1 = most frequent). Orders of magnitude presets in UI.
    max_rank: int | None = Field(default=None, ge=1, le=5_000_000)
    keep_unknown: bool = False


class EdgeFilter(BaseModel):
    relations: list[str] = Field(default_factory=list)
    min_confidence: float = 0.7
    max_depth: int = Field(default=3, ge=1, le=8)
    max_visit: int = Field(default=50_000, ge=100, le=200_000)
    # Also walk toward etymons (forward CSR) and place them opposite descendants.
    include_ancestors: bool = True
    max_ancestor_depth: int = Field(default=4, ge=0, le=8)


class PathFilter(BaseModel):
    node: NodePredicate = Field(default_factory=NodePredicate)
    quantifier: Quantifier = "any"
    exactly_k: int = 1
    relations_any: list[str] = Field(default_factory=list)
    relations_none: list[str] = Field(default_factory=list)
    apply_to_root: bool = False


class TreeQuery(BaseModel):
    term: str
    lang: str
    # Editorial etymology id / etymology_number key; omit to auto-resolve.
    ety: str | None = None
    edges: EdgeFilter = Field(default_factory=EdgeFilter)
    leaf: NodePredicate = Field(default_factory=NodePredicate)
    path: PathFilter = Field(default_factory=PathFilter)
    expand: list[str] = Field(default_factory=list)
    # High default = do not collapse children into cluster stubs.
    cluster_threshold: int = Field(default=1_000_000, ge=1, le=1_000_000)
    max_payload: int = Field(default=50_000, ge=50, le=100_000)
    color_by: ColorBy = "relation"


class SuggestHit(BaseModel):
    term: str
    lang: str
    ety: str = ""
    lang_display: str | None = None
    family_name: str | None = None
    child_count: int = 0
    id: int