from typing import Dict, Final, Optional, Tuple, Union

import torch
import torch.nn as nn
from transformers import CLIPVisionModelWithProjection, logging
from transformers.modeling_outputs import ImageClassifierOutputWithNoAttention

from .configuration_predictor import AestheticsPredictorConfig

logging.set_verbosity_error()

URLS: Final[Dict[str, str]] = {
    "openai/clip-vit-base-patch16": "https://github.com/LAION-AI/aesthetic-predictor/raw/main/sa_0_4_vit_b_16_linear.pth",
    "openai/clip-vit-base-patch32": "https://github.com/LAION-AI/aesthetic-predictor/raw/main/sa_0_4_vit_b_32_linear.pth",
    "openai/clip-vit-large-patch14": "https://github.com/LAION-AI/aesthetic-predictor/raw/main/sa_0_4_vit_l_14_linear.pth",
}


class AestheticsPredictorV1(CLIPVisionModelWithProjection):
    def __init__(self, config: AestheticsPredictorConfig) -> None:
        super().__init__(config)
        self.predictor = nn.Linear(config.projection_dim, 1)
        self.post_init()

    def forward(
        self,
        pixel_values: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, ImageClassifierOutputWithNoAttention]:
        return_dict = (
            return_dict if return_dict is not None else self.config.use_return_dict
        )

        outputs = super().forward(
            pixel_values=pixel_values,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        image_embeds = outputs[0]  # image_embeds
        image_embeds /= image_embeds.norm(dim=-1, keepdim=True)

        prediction = self.predictor(image_embeds)

        if not return_dict:
            return (None, prediction, image_embeds)

        return ImageClassifierOutputWithNoAttention(
            loss=None,
            logits=prediction,
            hidden_states=image_embeds,
        )


def convert_from_openai_clip(
    openai_model_name: str, config: Optional[AestheticsPredictorConfig] = None
) -> AestheticsPredictorV1:
    config = config or AestheticsPredictorConfig.from_pretrained(openai_model_name)
    model = AestheticsPredictorV1(config)

    clip_model = CLIPVisionModelWithProjection.from_pretrained(openai_model_name)
    model.load_state_dict(clip_model.state_dict(), strict=False)

    state_dict = torch.hub.load_state_dict_from_url(URLS[openai_model_name])
    model.predictor.load_state_dict(state_dict)
    model.eval()

    return model