""" layer_info.py """
from typing import Any, Dict, Iterable, List, Optional, Sequence, Type, Union

import torch
from torch import nn
from torch.jit import ScriptModule

DETECTED_INPUT_OUTPUT_TYPES = Union[
    Sequence[Any], Dict[Any, torch.Tensor], torch.Tensor
]


class LayerInfo:
    """Class that holds information about a layer module."""

    def __init__(
        self,
        var_name: str,
        module: nn.Module,
        depth: int,
        depth_index: Optional[int] = None,
        parent_info: Optional["LayerInfo"] = None,
    ) -> None:
        # Identifying information
        self.layer_id = id(module)
        self.module = module
        self.class_name = (
            str(module.original_name)
            if isinstance(module, ScriptModule)
            else module.__class__.__name__
        )
        # {layer name: {row_name: row_value}}
        self.inner_layers: Dict[str, Dict[str, Any]] = {}
        self.depth = depth
        self.depth_index = depth_index
        self.executed = False
        self.parent_info = parent_info
        self.var_name = var_name
        self.is_leaf_layer = not any(self.module.children())

        # Statistics
        self.trainable = True
        self.is_recursive = False
        self.input_size: List[int] = []
        self.output_size: List[int] = []
        self.kernel_size: List[int] = []
        self.num_params = 0
        self.macs = 0
        self.quantized = False

    def __repr__(self) -> str:
        return f"{self.class_name}: {self.depth}"

    @staticmethod
    def calculate_size(
        inputs: DETECTED_INPUT_OUTPUT_TYPES, batch_dim: Optional[int]
    ) -> List[int]:
        """Set input_size or output_size using the model's inputs."""

        def nested_list_size(inputs: Sequence[Any]) -> List[int]:
            """Flattens nested list size."""
            if hasattr(inputs, "tensors"):
                return nested_list_size(inputs.tensors)  # type: ignore[attr-defined]
            if (
                isinstance(inputs, torch.Tensor)
                or not hasattr(inputs, "__getitem__")
                or not inputs
            ):
                return []
            if isinstance(inputs[0], dict):
                return nested_list_size(list(inputs[0].items()))
            if hasattr(inputs[0], "size") and callable(inputs[0].size):
                return list(inputs[0].size())
            if isinstance(inputs, (list, tuple)):
                return nested_list_size(inputs[0])
            return []

        size = []
        # pack_padded_seq and pad_packed_seq store feature into data attribute
        if isinstance(inputs, (list, tuple)) and inputs and hasattr(inputs[0], "data"):
            size = list(inputs[0].data.size())
            if batch_dim is not None:
                size = size[:batch_dim] + [1] + size[batch_dim + 1 :]

        elif isinstance(inputs, dict):
            # TODO avoid overwriting the previous size every time?
            for _, output in inputs.items():
                size = list(output.size())
                if batch_dim is not None:
                    size = [size[:batch_dim] + [1] + size[batch_dim + 1 :]]

        elif isinstance(inputs, torch.Tensor):
            size = list(inputs.size())
            if batch_dim is not None and batch_dim < len(size):
                size[batch_dim] = 1

        elif isinstance(inputs, (list, tuple)):
            size = nested_list_size(inputs)

        else:
            raise TypeError(
                "Model contains a layer with an unsupported input or output type: "
                f"{inputs}, type: {type(inputs)}"
            )

        return size

    def get_layer_name(self, show_var_name: bool, show_depth: bool) -> str:
        layer_name = self.class_name
        if show_var_name and self.var_name:
            layer_name += f" ({self.var_name})"
        if show_depth and self.depth > 0:
            layer_name += f": {self.depth}"
            if self.depth_index is not None:
                layer_name += f"-{self.depth_index}"
        return layer_name
    
    def determine_if_quantized(self, quantization_base_class: Optional[Type] = None) -> None:
        """check if module has a quantization function as a named parameter
        which determines this module as a possible quantized layer
        """
        if quantization_base_class:
            attribute_types = [type(getattr(self.module, attribute)) for attribute in dir(self.module)]
            self.quantized = any([issubclass(attribute_type, quantization_base_class) for attribute_type in attribute_types])
        else:
            self.quantized = False

    def calculate_num_params(self) -> None:
        """
        Set num_params, trainable, inner_layers, and kernel_size
        using the module's parameters.
        """
        name = ""
        for name, param in self.module.named_parameters():
            self.num_params += param.nelement()
            self.trainable &= param.requires_grad

            ksize = list(param.size())
            if name == "weight":
                # to make [in_shape, out_shape, ksize, ksize]
                if len(ksize) > 1:
                    ksize[0], ksize[1] = ksize[1], ksize[0]
                self.kernel_size = ksize

            # RNN modules have inner weights such as weight_ih_l0
            self.inner_layers[name] = {
                "kernel_size": str(ksize),
                "num_params": f"├─{param.nelement():,}",
            }
        if self.inner_layers:
            self.inner_layers[name][
                "num_params"
            ] = f"└─{self.inner_layers[name]['num_params'][2:]}"

    def calculate_macs(self) -> None:
        """
        Set MACs using the module's parameters and layer's output size, which is
        used for computing number of operations for Conv layers.

        Please note: Returned MACs is the number of MACs for the full tensor,
        i.e., taking the batch-dimension into account.
        """
        for name, param in self.module.named_parameters():
            if name in ("weight", "bias"):
                # ignore C when calculating Mult-Adds in ConvNd
                if "Conv" in self.class_name:
                    self.macs += int(
                        param.nelement()
                        * prod(self.output_size[:1] + self.output_size[2:])
                    )
                else:
                    self.macs += self.output_size[0] * param.nelement()
            # RNN modules have inner weights such as weight_ih_l0
            elif "weight" in name or "bias" in name:
                self.macs += prod(self.output_size[:2]) * param.nelement()

    def check_recursive(self, summary_list: List["LayerInfo"]) -> None:
        """
        If the current module is already-used, mark as (recursive).
        Must check before adding line to the summary.
        """
        if any(self.module.named_parameters()):
            for other_layer in summary_list:
                if self.layer_id == other_layer.layer_id:
                    self.is_recursive = True

    def macs_to_str(
        self, reached_max_depth: bool, children_layers: List["LayerInfo"]
    ) -> str:
        """Convert MACs to string."""
        if self.macs <= 0:
            return "--"
        if self.is_leaf_layer or self.quantized:
            return f"{self.macs:,}"
        if reached_max_depth:
            sum_child_macs = sum(
                child.macs for child in children_layers if child.is_leaf_layer or child.quantized
            )
            return f"{sum_child_macs:,}"
        return "--"

    def num_params_to_str(self, reached_max_depth: bool) -> str:
        """Convert num_params to string."""
        if self.is_recursive:
            return "(recursive)"
        if self.num_params > 0 and (reached_max_depth or self.is_leaf_layer or self.quantized ):
            param_count_str = f"{self.num_params:,}"
            return param_count_str if self.trainable else f"({param_count_str})"
        return "--"


def prod(num_list: Union[Iterable[int], torch.Size]) -> int:
    result = 1
    if isinstance(num_list, Iterable):
        for item in num_list:
            result *= prod(item) if isinstance(item, Iterable) else item
    return result
