from dataclasses import dataclass
from transformers import GPT2Config, CLIPVisionConfig

PREFIX_MAP = {
    "openai/clip-vit-base-patch32": 50,
    "openai/clip-vit-base-patch16": 197,
    "openai/clip-vit-large-patch14": 257,
    "openai/clip-vit-large-patch14-336": 577
}

TEXT_HIDDEN_SIZE_MAP = {
    "gpt2": 768,
    "gpt2-medium": 768,
    "gpt2-large": 1280,
    "gpt2-xl": 1600
}

IMAGE_HIDDEN_SIZE_MAP = {
    "openai/clip-vit-base-patch32": 768,
    "openai/clip-vit-base-patch16": 768,
    "openai/clip-vit-large-patch14": 768,
    "openai/clip-vit-large-patch14-336": 768
}


@dataclass
class CLIPGPT2Config:
    image_model: str = "openai/clip-vit-base-patch32"
    freeze_image_model: bool = True
    text_model: str = "gpt2-large"
    freeze_text_model: bool = True
    linear_mapping_type: int = "linear"
    add_image_token: bool = True
    freeze_ln: bool = False
    image_from_pretrained: bool = True
    text_from_pretrained: bool = True

    def __post_init__(self):
        self.prefix_length = PREFIX_MAP[self.image_model]
        self.image_hidden_size = IMAGE_HIDDEN_SIZE_MAP[self.image_model]
        self.text_hidden_size = TEXT_HIDDEN_SIZE_MAP[self.text_model]
        self.image_resize = 224 if "336" not in self.image_model else 336
        self.text_config = GPT2Config.from_pretrained(self.text_model)
        self.image_config = CLIPVisionConfig.from_pretrained(self.image_model)
        self.vocab_size = self.text_config.vocab_size + self.add_image_token