"""Pydantic request/response models for the NephroScreen API.

Every clinical field is optional: missing values are imputed by the same
train-fitted pipeline used in training, so the form stays usable even when a
patient's full lab panel is not available.
"""
from typing import Literal, Optional

from pydantic import BaseModel, Field


class PatientInput(BaseModel):
    # Numeric labs
    age: Optional[float] = Field(None, ge=0, le=120, description="Age (years)")
    bp: Optional[float] = Field(None, ge=0, description="Blood pressure (mm/Hg)")
    sg: Optional[float] = Field(None, description="Specific gravity")
    al: Optional[float] = Field(None, ge=0, le=5, description="Albumin (0-5)")
    su: Optional[float] = Field(None, ge=0, le=5, description="Sugar (0-5)")
    bgr: Optional[float] = Field(None, ge=0, description="Blood glucose random (mgs/dl)")
    bu: Optional[float] = Field(None, ge=0, description="Blood urea (mgs/dl)")
    sc: Optional[float] = Field(None, ge=0, description="Serum creatinine (mgs/dl)")
    sod: Optional[float] = Field(None, description="Sodium (mEq/L)")
    pot: Optional[float] = Field(None, description="Potassium (mEq/L)")
    hemo: Optional[float] = Field(None, ge=0, description="Hemoglobin (gms)")
    pcv: Optional[float] = Field(None, ge=0, description="Packed cell volume")
    wbcc: Optional[float] = Field(None, ge=0, description="White blood cell count (cells/cmm)")
    rbcc: Optional[float] = Field(None, ge=0, description="Red blood cell count (millions/cmm)")
    # Categorical indicators
    rbc: Optional[Literal["normal", "abnormal"]] = None
    pc: Optional[Literal["normal", "abnormal"]] = None
    pcc: Optional[Literal["present", "notpresent"]] = None
    ba: Optional[Literal["present", "notpresent"]] = None
    htn: Optional[Literal["yes", "no"]] = None
    dm: Optional[Literal["yes", "no"]] = None
    cad: Optional[Literal["yes", "no"]] = None
    appet: Optional[Literal["good", "poor"]] = None
    pe: Optional[Literal["yes", "no"]] = None
    ane: Optional[Literal["yes", "no"]] = None

    model_config = {
        "json_schema_extra": {
            "example": {
                "age": 62, "bp": 80, "sg": 1.01, "al": 3, "su": 0,
                "bgr": 148, "bu": 86, "sc": 3.2, "sod": 135, "pot": 4.6,
                "hemo": 9.5, "pcv": 28, "wbcc": 9800, "rbcc": 3.4,
                "rbc": "abnormal", "pc": "abnormal", "pcc": "present", "ba": "notpresent",
                "htn": "yes", "dm": "yes", "cad": "no", "appet": "poor",
                "pe": "yes", "ane": "yes",
            }
        }
    }


class Indicator(BaseModel):
    feature: str
    label: str
    value: float
    normal_range: str
    flag: Literal["low", "high", "normal"]


class PredictionResponse(BaseModel):
    prediction: Literal["CKD", "Not CKD"]
    probability: float = Field(..., description="Model probability of CKD (0-1)")
    risk_band: Literal["Low", "Moderate", "High"]
    threshold: float = Field(..., description="Recall-tuned decision threshold used")
    key_indicators: list[Indicator]
    disclaimer: str