# Source code for archai.common.model_summary

    # Copyright (c) Microsoft Corporation.
    # Licensed under the MIT license.

    from typing import Iterable, Mapping, Sized, Sequence
    import math

    import torch
    import torch.nn as nn

    from collections import OrderedDict
    import numpy as np

    from numbers import Number

    def summary(model, input_size):
        result, params_info = summary_string(model, input_size)
        print(result)

        return params_info

    def is_scaler(o):
        return isinstance(o, Number) or isinstance(o, str) or o is None

    def get_tensor_stat(tensor):
        assert isinstance(tensor, torch.Tensor)

        # some pytorch low-level memory management constant
        # the minimal allocate memory size (Byte)
        PYTORCH_MIN_ALLOCATE = 2 ** 9
        # the minimal cache memory size (Byte)
        PYTORCH_MIN_CACHE = 2 ** 20

        numel = tensor.numel()
        element_size = tensor.element_size()
        fact_numel = tensor.storage().size()
        fact_memory_size = fact_numel * element_size
        # since pytorch allocate at least 512 Bytes for any tensor, round
        # up to a multiple of 512
        memory_size = math.ceil(fact_memory_size / PYTORCH_MIN_ALLOCATE) \
                * PYTORCH_MIN_ALLOCATE

        # tensor.storage should be the actual object related to memory
        # allocation
        data_ptr = tensor.storage().data_ptr()
        size = tuple(tensor.size())
        # torch scalar has empty size
        if not size:
            size = (1,)

        return ([size], numel, memory_size)

    def get_all_tensor_stats(o):
        if is_scaler(o):
            return ([[]], 0, 0)
        elif isinstance(o, torch.Tensor):
            return get_tensor_stat(o)
        elif isinstance(o, Mapping):
            return get_all_tensor_stats(o.values())
        elif isinstance(o, Iterable): # tuple, list, maps
            stats = [[]], 0, 0
            for oi in o:
                tz = get_all_tensor_stats(oi)
                stats = tuple(x+y for x,y in zip(stats, tz))
            return stats
        elif hasattr(o, '__dict__'):
            return get_all_tensor_stats(o.__dict__)
        else:
            return ([[]], 0, 0)

    def get_shape(o):
        if is_scaler(o):
            return str(o)
        elif hasattr(o, 'shape'):
            return f'shape{o.shape}'
        elif hasattr(o, 'size'):
            return f'size{o.size()}'
        elif isinstance(o, Sequence):
            if len(o)==0:
                return 'seq[]'
            elif is_scaler(o[0]):
                return f'seq[{len(o)}]'
            return f'seq{[get_shape(oi) for oi in o]}'
        elif isinstance(o, Mapping):
            if len(o)==0:
                return 'map[]'
            elif is_scaler(next(o)):
                return f'map[{len(o)}]'
            arr = [(get_shape(ki), get_shape(vi)) for ki, vi in o]
            return f'map{arr}'
        else:
            return 'N/A'

    def summary_string(model, input_size, dtype=torch.float32):
        summary_str = ''

        # create properties
        summary = OrderedDict()
        hooks = []

        def register_hook(module):
            def hook(module, input, output):
                class_name = str(module.__class__).split(".")[-1].split("'")[0]
                module_idx = len(summary)
                m_key = "%s-%i" % (class_name, module_idx + 1)

                summary[m_key] = OrderedDict()
                summary[m_key]["input"] = get_all_tensor_stats(input)
                summary[m_key]["output"] = get_all_tensor_stats(output)

                params = 0
                if hasattr(module, "weight") and hasattr(module.weight, "size"):
                    params += torch.prod(torch.LongTensor(list(module.weight.size()))).item()
                    summary[m_key]["trainable"] = module.weight.requires_grad
                if hasattr(module, "bias") and hasattr(module.bias, "size"):
                    params += torch.prod(torch.LongTensor(list(module.bias.size()))).item()
                summary[m_key]["nb_params"] = params

            if (
                not isinstance(module, nn.Sequential)
                and not isinstance(module, nn.ModuleList)
            ):
                hooks.append(module.register_forward_hook(hook))

        # batch_size of 2 for batchnorm
        x = torch.rand(input_size, dtype=dtype,
                       device=next(model.parameters()).device)

        # register hook
        model.apply(register_hook)

        # make a forward pass
        # print(x.shape)
        model(x)

        # remove these hooks
        for h in hooks:
            h.remove()

        summary_str += "----------------------------------------------------------------" + "\n"
        line_new = "{:>20}  {:>25} {:>15}".format(
            "Layer (type)", "Output (elments, mem)", "Param #")
        summary_str += line_new + "\n"
        summary_str += "================================================================" + "\n"
        total_params = 0
        total_input = get_tensor_stat(x)
        total_output = ([[], 0, 0])
        trainable_params = 0
        for layer in summary:
            # input_shape, output_shape, trainable, nb_params
            line_new = "{:>20}  {:>25} {:>15}".format(
                layer,
                str(summary[layer]["output"][1:]),
                "{0:,}".format(summary[layer]["nb_params"]),
            )
            total_params += summary[layer]["nb_params"]

            total_output = tuple(x+y for x,y in zip(total_output, summary[layer]["output"]))
            if "trainable" in summary[layer]:
                if summary[layer]["trainable"] == True:
                    trainable_params += summary[layer]["nb_params"]
            summary_str += line_new + "\n"

        total_numel = total_params + total_output[1] + total_input[1]

        summary_str += "================================================================" + "\n"
        summary_str += "Total params: {0:,}".format(total_params) + "\n"
        summary_str += "Trainable params: {0:,}".format(trainable_params) + "\n"
        summary_str += "Non-trainable params: {0:,}".format(total_params -
                                                            trainable_params) + "\n"
        summary_str += "----------------------------------------------------------------" + "\n"
        summary_str += f"Input Elments: {total_input[1]:.4e}\n"
        summary_str += f"Input Mem: {total_input[2]:.4e}\n"
        summary_str += f"Layer Output Elements: {total_output[1]:.4e}\n"
        summary_str += f"Layer Output Mem: {total_output[2]:.4e}\n"
        summary_str += f"Params {total_params:.4e}\n"
        summary_str += f"Total Elements {total_numel:.4e}\n"
        summary_str += "----------------------------------------------------------------" + "\n"
        # return summary
        return summary_str, (total_params, trainable_params)