from pydantic import BaseModel, Field


class TrainRequest(BaseModel):
    epochs:     int   = Field(10,    description="학습 에폭 수")
    lr:         float = Field(1e-5,  description="학습률")
    batch_size: int   = Field(4,     description="배치 사이즈")
    ewc_lambda: float = Field(1000., description="EWC 패널티 강도 (클수록 기존 지식 보호)")
    strategy:   str   = Field("full",description="auto | full | ewc | replay")


class TrainStatus(BaseModel):
    running:  bool
    epoch:    int
    total:    int
    loss:     float
    val_acc:  float = 0.0
    message:  str
    batch:         int = 0    # 현재 배치
    total_batches: int = 0    # 전체 배치 수

class PredictResult(BaseModel):
    filename:       str
    label:          str | None = None
    confidence:     float | None = None
    all_probs:      dict | None = None
    key_tokens:     list[str] | None = None
    ocr_word_count: int | None = None
    total_pages:    int | None = None
    per_page:       list | None = None
    error:          str | None = None


class DatasetStatus(BaseModel):
    total_samples: int
    by_label:      dict
    label2id:      dict