# Copyright 2024 The SoftCon Authors and The HuggingFace Inc. team.
"""Self-contained SoftCon model and config for trust_remote_code loading."""

from typing import Optional

import torch
import torch.nn as nn
from torchvision import models

from transformers.configuration_utils import PretrainedConfig as PreTrainedConfig
from transformers.modeling_outputs import BaseModelOutputWithPooling, ImageClassifierOutput
from transformers.modeling_utils import PreTrainedModel
from transformers.models.dinov2.configuration_dinov2 import Dinov2Config
from transformers.models.dinov2.modeling_dinov2 import Dinov2Model
from transformers.processing_utils import Unpack
from transformers.utils import TransformersKwargs, logging


logger = logging.get_logger(__name__)


class SoftConConfig(PreTrainedConfig):
    model_type = "softcon"

    def __init__(
        self,
        backbone="resnet50",
        num_channels=13,
        modality="s2c",
        hidden_size=2048,
        image_size=224,
        patch_size=14,
        init_values=1e-5,
        num_hidden_layers=None,
        num_attention_heads=None,
        num_register_tokens=0,
        block_chunks=0,
        num_labels=0,
        **kwargs,
    ):
        super().__init__(**kwargs)

        self.backbone = backbone
        self.num_channels = num_channels
        self.modality = modality
        self.hidden_size = hidden_size
        self.image_size = image_size
        self.patch_size = patch_size
        self.init_values = init_values
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.num_register_tokens = num_register_tokens
        self.block_chunks = block_chunks
        self.num_labels = num_labels

    def get_dinov2_config(self) -> Dinov2Config:
        if self.backbone not in {"vit_small", "vit_base"}:
            raise ValueError(f"Backbone '{self.backbone}' is not a ViT encoder.")
        if self.num_hidden_layers is None or self.num_attention_heads is None:
            raise ValueError(
                "ViT models require `num_hidden_layers` and `num_attention_heads` in the model config."
            )

        return Dinov2Config(
            hidden_size=self.hidden_size,
            num_hidden_layers=self.num_hidden_layers,
            num_attention_heads=self.num_attention_heads,
            image_size=self.image_size,
            patch_size=self.patch_size,
            num_channels=self.num_channels,
            layerscale_value=self.init_values,
            use_mask_token=False,
        )


class SoftConPreTrainedModel(PreTrainedModel):
    config: SoftConConfig
    base_model_prefix = "softcon"
    main_input_name = "pixel_values"
    input_modalities = ("image",)
    supports_gradient_checkpointing = False


def _build_resnet_encoder(config: SoftConConfig) -> nn.Module:
    backbone = models.resnet50(weights=None)
    if config.num_channels != 3:
        backbone.conv1 = nn.Conv2d(config.num_channels, 64, kernel_size=7, stride=2, padding=3, bias=False)
    backbone.fc = nn.Identity()
    return backbone


def _build_vit_encoder(config: SoftConConfig) -> Dinov2Model:
    return Dinov2Model(config.get_dinov2_config())


def _build_encoder(config: SoftConConfig) -> nn.Module:
    if config.backbone == "resnet50":
        return _build_resnet_encoder(config)
    if config.backbone in {"vit_small", "vit_base"}:
        return _build_vit_encoder(config)
    raise ValueError(f"Unsupported backbone '{config.backbone}'")


class SoftConModel(SoftConPreTrainedModel):
    def __init__(self, config: SoftConConfig):
        super().__init__(config)
        self.encoder = _build_encoder(config)
        self.post_init()

    def _forward_resnet(self, pixel_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        x = self.encoder.conv1(pixel_values)
        x = self.encoder.bn1(x)
        x = self.encoder.relu(x)
        x = self.encoder.maxpool(x)
        x = self.encoder.layer1(x)
        x = self.encoder.layer2(x)
        x = self.encoder.layer3(x)
        x = self.encoder.layer4(x)
        last_hidden_state = x.flatten(2).transpose(1, 2)
        pooler_output = self.encoder.avgpool(x).flatten(1)
        return last_hidden_state, pooler_output

    def _forward_vit(self, pixel_values: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        outputs = self.encoder(pixel_values=pixel_values, return_dict=True)
        return outputs.last_hidden_state, outputs.pooler_output

    def forward(
        self,
        pixel_values: Optional[torch.Tensor] = None,
        return_dict: Optional[bool] = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> BaseModelOutputWithPooling:
        if pixel_values is None:
            raise ValueError("You must specify `pixel_values`")

        pixel_values = pixel_values.to(dtype=self.dtype)
        if return_dict is None:
            return_dict = self.config.use_return_dict

        if self.config.backbone == "resnet50":
            last_hidden_state, pooler_output = self._forward_resnet(pixel_values)
        else:
            last_hidden_state, pooler_output = self._forward_vit(pixel_values)

        if not return_dict:
            return (last_hidden_state, pooler_output)

        return BaseModelOutputWithPooling(last_hidden_state=last_hidden_state, pooler_output=pooler_output)


class SoftConForImageClassification(SoftConPreTrainedModel):
    def __init__(self, config: SoftConConfig):
        super().__init__(config)
        self.softcon = SoftConModel(config)
        self.classifier = (
            nn.Linear(config.hidden_size, config.num_labels) if config.num_labels > 0 else nn.Identity()
        )
        self.post_init()

    def forward(
        self,
        pixel_values: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        return_dict: Optional[bool] = None,
        **kwargs: Unpack[TransformersKwargs],
    ) -> ImageClassifierOutput:
        outputs = self.softcon(pixel_values=pixel_values, return_dict=True, **kwargs)
        logits = self.classifier(outputs.pooler_output)
        loss = None
        if labels is not None:
            loss = self.loss_function(labels, logits, self.config, **kwargs)
        if not return_dict:
            output = (logits,) + outputs[1:]
            return ((loss,) + output) if loss is not None else output
        return ImageClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


__all__ = [
    "SoftConConfig",
    "SoftConForImageClassification",
    "SoftConModel",
    "SoftConPreTrainedModel",
]