backend/app/protein/models.py

"""
Data models for the protein visualization API.
"""
from typing import Optional, List, Dict, Any
from pydantic import BaseModel, Field


class CrystalGrow(BaseModel):
    """Crystal growth conditions."""
    pH: Optional[float] = None
    pH_structured: Optional[float] = None
    pH_from_details: Optional[float] = None
    temp: Optional[float] = None
    pdbx_details: Optional[str] = None


class EntryInfo(BaseModel):
    """Entry-level information."""
    rcsb_id: Optional[str] = None
    title: Optional[str] = None
    initial_release_date: Optional[str] = None
    exptl_method: Optional[str] = None
    resolution_combined: Optional[List[float]] = None
    crystal_grow: Optional[CrystalGrow] = None


class PolymerEntity(BaseModel):
    """Polymer entity (sequence-level) information."""
    rcsb_id: Optional[str] = None
    entity_id: Optional[str] = None
    sequence_length: Optional[int] = None
    seq_can: Optional[str] = None
    reference_sequences: Optional[List[Dict[str, str]]] = None


class PolymerEntityInstance(BaseModel):
    """Polymer entity instance (chain-level) information."""
    rcsb_id: Optional[str] = None
    asym_id: Optional[str] = None
    auth_asym_id: Optional[str] = None
    binding_affinity: Optional[Any] = None


class NonpolymerEntity(BaseModel):
    """Nonpolymer entity (ligand/binder) information."""
    rcsb_id: Optional[str] = None
    name: Optional[str] = None
    formula_weight: Optional[float] = None
    smiles: Optional[str] = None
    inchi: Optional[str] = None


class ChainMetadata(BaseModel):
    """Complete chain metadata from RCSB."""
    pdb_id: str
    chain_id: str
    entry: EntryInfo
    polymer_entity: Optional[PolymerEntity] = None
    polymer_entity_instance: Optional[PolymerEntityInstance] = None
    nonpolymer_entities: Optional[List[NonpolymerEntity]] = None
    binding_status: Optional[str] = None
    chain_composition: Optional[str] = None
    cath_id: Optional[str] = None
    cath_superfamily: Optional[str] = None
    state_id: Optional[str] = None


class FilterParams(BaseModel):
    """Request model for filtering protein data."""
    chain_ids: Optional[List[str]] = Field(default=None, description="List of chain IDs to filter")
    pdb_id: Optional[str] = Field(default=None, description="PDB ID for single chain search")
    auth_asym_id: Optional[str] = Field(default=None, description="Chain ID for single chain search")
    rmsd_min: Optional[float] = Field(default=None, ge=0, description="Minimum RMSD value")
    rmsd_max: Optional[float] = Field(default=None, ge=0, description="Maximum RMSD value")
    tm_score_min: Optional[float] = Field(default=None, ge=0, le=1, description="Minimum TM-score")
    tm_score_max: Optional[float] = Field(default=None, ge=0, le=1, description="Maximum TM-score")
    length_min: Optional[int] = Field(default=None, ge=0, description="Minimum sequence length")
    length_max: Optional[int] = Field(default=None, ge=0, description="Maximum sequence length")
    include_duplicates: Optional[bool] = Field(default=None, description="Include duplicate structures")
    cluster_id: Optional[str] = Field(default=None, description="Structural similarity cluster ID")
    limit: Optional[int] = Field(default=1000, ge=1, le=10000, description="Maximum results to return")


class FilterOptions(BaseModel):
    """Available filter options and ranges."""
    chain_ids: List[str] = Field(description="Available chain IDs")
    rmsd_range: Dict[str, float] = Field(description="Min and max RMSD values")
    tm_score_range: Dict[str, float] = Field(description="Min and max TM-score values")
    length_range: Dict[str, int] = Field(description="Min and max sequence lengths")
    clusters: List[str] = Field(description="Available similarity clusters")
    total_records: int = Field(description="Total number of records")


class DataRecord(BaseModel):
    """Single protein comparison record."""
    pdb_id_a: str
    auth_asym_id_a: str
    pdb_id_b: str
    auth_asym_id_b: str
    tm_score: float
    rmsd: float
    structure_sim: Optional[float] = None
    sequence_identity: Optional[float] = None
    length_a: Optional[int] = None
    length_b: Optional[int] = None
    aligned_length: Optional[int] = None
    status: Optional[str] = None
    cluster_id: Optional[str] = None
    exptl_method: Optional[str] = None
    temp: Optional[float] = None
    pH: Optional[float] = None
    # Matched-chain (B) annotations from the node table
    binding_status: Optional[str] = None
    chain_composition: Optional[str] = None
    # Energy score deltas fetched directly from the edge table
    delta_rosetta: Optional[float] = None
    delta_foldx: Optional[float] = None
    delta_evoef2: Optional[float] = None
    delta_rw: Optional[float] = None
    delta_rw_plus: Optional[float] = None
    # Alternative state similarity fields from edge table
    state_id_b: Optional[str] = None
    state_fidelity: Optional[str] = None
    avg_sim: Optional[str] = None
    observation_fidelity: Optional[str] = None


class DataResponse(BaseModel):
    """Response model for filtered data."""
    data: List[DataRecord]
    total: int
    filtered: int


class SummaryStats(BaseModel):
    """Aggregate statistics for the dataset."""
    total_records: int
    unique_chains: int
    avg_rmsd: float
    avg_tm_score: float
    avg_sequence_length: float
    rmsd_distribution: Dict[str, int]
    tm_score_distribution: Dict[str, int]
    length_distribution: Dict[str, int]