import torch

from typing import Tuple
from dataclasses import dataclass
from transformers import PretrainedConfig, PreTrainedModel

from .csd import CSD
from .config import CSDConfig


@dataclass
class CSDOutput:
    image_embeds: torch.Tensor
    style_embeds: torch.Tensor
    content_embeds: torch.Tensor


class CSDModel(PreTrainedModel):
    config_class = CSDConfig

    def __init__(self, config: CSDConfig) -> None:
        super(CSDModel, self).__init__(config)

        self.model = CSD(
            vit_input_resolution=config.vit_input_resolution,
            vit_patch_size=config.vit_patch_size,
            vit_width=config.vit_width,
            vit_layers=config.vit_layers,
            vit_heads=config.vit_heads,
            vit_output_dim=config.vit_output_dim,
        )

    @torch.inference_mode()
    def forward(self, pixel_values: torch.Tensor, **kwargs) -> CSDOutput:
        image_embeds, style_embeds, content_embeds = self.model(pixel_values)
        return CSDOutput(image_embeds=image_embeds, style_embeds=style_embeds, content_embeds=content_embeds)