from pydantic import BaseModel, Field
from typing import List, Optional


class PredictionRequest(BaseModel):
    id: int
    Age: int
    Sex: int
    Chest_pain_type: int = Field(alias="Chest pain type")
    BP: int
    Cholesterol: int
    FBS_over_120: int = Field(alias="FBS over 120")
    EKG_results: int = Field(alias="EKG results")
    Max_HR: int = Field(alias="Max HR")
    Exercise_angina: int = Field(alias="Exercise angina")
    ST_depression: float = Field(alias="ST depression")
    Slope_of_ST: int = Field(alias="Slope of ST")
    Number_of_vessels_fluro: int = Field(alias="Number of vessels fluro")
    Thallium: int

    class Config:
        populate_by_name = True


class BatchPredictionRequest(BaseModel):
    data: List[PredictionRequest]


class TrainingRequest(BaseModel):
    force_retrain: Optional[bool] = False