import pathlib

import pydantic

MAX_DOWNLOAD_TIME = 0.2

IMAGE_DOWNLOAD_PATH = pathlib.Path("/tmp/images")


class DataConfig(pydantic.BaseModel):
    buffer_size: int = 1000
    data_len: int = 100
    train_len: int = 90
    small_dataset: str = "laion/220k-gpt4vision-captions-from-livis"
    large_dataset: str = "laion/laion400m"
    dataset: str = small_dataset


class ModelConfig(pydantic.BaseModel):
    text_model: str = "microsoft/xtremedistil-l6-h256-uncased"  # 51 mb
    vision_model: str = "edgenext_small"  # 20 mb
    projection_layers: int = 3
    embed_dim: int = 256
    transformer_embed_dim: int = 768
    max_len: int = 128  # 77
    cls_type: bool = True
    freeze_vision_base: bool = False
    freeze_text_base: bool = False


class TrainerConfig(pydantic.BaseModel):
    epochs: int = 20
    batch_size: int = 256
    learning_rate: float = 5e-4
    accumulate_grad_batches: int = 1
    temperature: float = 1.0
    vision_freeze_layers: int = 2
    lambda_1: float = 1.0
    lambda_2: float = 1.0

    val_check_interval: int = 1000

    run_openai_clip: bool = False

    _model_config: ModelConfig = ModelConfig()
    _data_config: DataConfig = DataConfig()