#
#  Copyright (c) 2025
#  Minh NGUYEN <vnguyen9@lakeheadu.ca>
#
import torch
import timm
import logging

from typing import TYPE_CHECKING, TypeVar, Generic, Optional, Mapping, Any
from typing_extensions import override
from torch import Tensor
from transformers import PreTrainedModel, PretrainedConfig
from timm.models.resnet import BasicBlock, Bottleneck, ResNet

from .configuration import ResnetConfig

if TYPE_CHECKING:
    ...

logger = logging.getLogger(__name__)
BLOCK_MAPPING = {"basic": BasicBlock, "bottleneck": Bottleneck}

ConfigClass = TypeVar("ConfigClass", bound=PretrainedConfig)


class ResnetModel(PreTrainedModel, Generic[ConfigClass]):
    """The custom resnet class."""

    config_class: ConfigClass = ResnetConfig
    """"""

    def __init__(self, config: ConfigClass):
        super().__init__(config=config)

        # reassign to reference custom config
        self.config: ResnetConfig = config

        block_layer = BLOCK_MAPPING[self.config.block_type]

        self.model = ResNet(
            block=block_layer,
            layers=self.config.layers,
            num_classes=self.config.num_classes,
            in_chans=self.config.input_channels,
            cardinality=self.config.cardinality,
            base_width=self.config.base_width,
            stem_width=self.config.stem_width,
            stem_type=self.config.stem_type,
            avg_down=self.config.avg_down,
        )

    def load_state_dict_from_source(self, source: Optional[str] = None):
        source = source or self.config.pretrained_state_source

        if source == "timm":
            pretrained_model = timm.create_model("resnet50d", pretrained=True)
            self.load_state_dict(
                state_dict=pretrained_model.state_dict(),
                assign=False,
                strict=False
            )

            del pretrained_model
        else:
            logger.warning(f"Now, only support loading pretrained weights from 'timm', but got {source!r}.")

    @override
    def load_state_dict(
        self, state_dict: Mapping[str, Any],
        strict: bool = True,
        assign: bool = False
    ):
        self.model.load_state_dict(state_dict, strict=strict, assign=assign)

    def forward(self, tensor: Tensor) -> Tensor:
        return self.model.forward(tensor)


class ResnetModelForImageClassification(ResnetModel):
    config_class = ResnetConfig

    def __init__(self, config: ResnetConfig):
        super().__init__(config)

        block_layer = BLOCK_MAPPING[config.block_type]
        self.model = ResNet(
            block_layer,
            config.layers,
            num_classes=config.num_classes,
            in_chans=config.input_channels,
            cardinality=config.cardinality,
            base_width=config.base_width,
            stem_width=config.stem_width,
            stem_type=config.stem_type,
            avg_down=config.avg_down,
        )

    @override
    def forward(self, tensor: Tensor, labels=None) -> dict[str, Tensor]:
        logits = self.model(tensor)
        if labels is not None:
            loss = torch.nn.functional.cross_entropy(logits, labels)
            return {"loss": loss, "logits": logits}

        return {"logits": logits}