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


class ProjectTrainIn(BaseModel):
    project_id: Optional[int] = Field(..., description="项目id")
    weights_id: Optional[str] = Field(None, description="权重文件")
    epochs: Optional[int] = Field(50, description="训练轮数")
    patience: Optional[int] = Field(20, description="早停的耐心值")


class ProjectTrainOut(BaseModel):
    """项目训练版本信息表"""
    id: Optional[int] = Field(None, description="训练id")
    train_version: Optional[str] = Field(None, description="训练版本号")
    weights_name: Optional[str] = Field(None, description="权重名称")
    epochs: Optional[int] = Field(None, description="训练轮数")
    patience: Optional[int] = Field(None, description="早停的耐心值")
    create_time: Optional[datetime] = Field(None, description="训练时间")

    class Config:
        orm_mode = True
        json_encoders = {
            datetime: lambda v: v.strftime("%Y-%m-%d %H:%M:%S")
        }