# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/04b_classification.model.meta_arch.common.ipynb (unless otherwise specified).

__all__ = ['GeneralizedImageClassifier']

# Cell
import logging
from collections import namedtuple
from typing import *

import torch
from omegaconf import DictConfig, OmegaConf
from pytorch_lightning.core.memory import get_human_readable_count
from torch.nn import Module

from ..backbones import ImageClassificationBackbone
from ..build import build_backbone, build_head
from ..heads import ImageClassificationHead
from ....core_classes import BasicModule
from ....utils.shape_spec import ShapeSpec

_logger = logging.getLogger(__name__)

# Cell
class GeneralizedImageClassifier(BasicModule):
    """
    A General Image Classifier. Any models that contains the following 2 components:
    1. Feature extractor (aka backbone)
    2. Image Classification head (Pooling + Classifier)
    """

    _hypers = namedtuple("hypers", field_names=["lr", "wd"])

    def __init__(
        self,
        backbone: ImageClassificationBackbone,
        head: ImageClassificationHead,
    ):
        """
        Arguments:
        1. `backbone`: a `ImageClassificationBackbone` module, must follow gale's backbone interface
        2. `head`: a head containg the classifier. and the pooling layer, must be an instance of
        `ImageClassificationHead`.
        """
        super(GeneralizedImageClassifier, self).__init__()
        self.backbone = backbone
        assert isinstance(backbone, ImageClassificationBackbone)
        self.head = head
        assert isinstance(head, ImageClassificationHead)

    def forward(self, batched_inputs: torch.Tensor) -> torch.Tensor:
        """
        Runs the batched_inputs through `backbone` followed by the `head`.
        Returns a Tensor which contains the logits for the batched_inputs.
        """
        # forward pass through the backbone
        out = self.backbone(batched_inputs)
        # pass through the classification layer
        out = self.head(out)
        return out

    @classmethod
    def from_config_dict(cls, cfg: DictConfig):
        """
        Instantiate the Meta Architecture from gale config
        """
        if not hasattr(cfg.model, "backbone"):
            raise ValueError("Configuration for model backbone not found")

        if not hasattr(cfg.model, "head"):
            raise ValueError("Configuration for model head not found")

        input_shape = ShapeSpec(cfg.input.channels, cfg.input.height, cfg.input.width)
        _logger.debug(f"Inputs: {input_shape}")

        backbone = build_backbone(cfg, input_shape=input_shape)
        param_count = get_human_readable_count(
            sum([m.numel() for m in backbone.parameters()])
        )
        _logger.debug(
            "Backbone {} created, param count: {}.".format(
                cfg.model.backbone.name, param_count
            )
        )

        head = build_head(cfg, backbone.output_shape())
        param_count = get_human_readable_count(
            sum([m.numel() for m in head.parameters()])
        )
        _logger.debug(
            "Head {} created, param count: {}.".format(cfg.model.head.name, param_count)
        )

        kwds = {"backbone": backbone, "head": head}

        instance = cls(**kwds)
        instance.input_shape = input_shape

        param_count = get_human_readable_count(
            sum([m.numel() for m in instance.parameters()])
        )
        _logger.info("Model created, param count: {}.".format(param_count))

        return instance

    def build_param_dicts(self):
        """
        Builds up the Paramters dicts for optimization
        """
        backbone_params = self.backbone.build_param_dicts()
        head_params = self.head.build_param_dicts()
        return backbone_params + head_params

    @property
    def hypers(self) -> Tuple:
        """
        Returns list of parameters like `lr` and `wd`
        for each param group
        """
        lrs = []
        wds = []

        for p in self.build_param_dicts():
            lrs.append(p["lr"])
            wds.append(p["weight_decay"])
        return self._hypers(lrs, wds)