import math
import torch
import torch.nn as nn
from ..build import build_nn, decode


class MBNetV2Backbone(nn.Module):
    """
    Wrapper of pytorch image models to use in bit_cv.
    https://github.com/rwightman/pytorch-image-models
    """
    def __init__(self, arch_def,
                 out_stage_ids=[3, 6, 8],
                 **kwargs):
        super(MBNetV2Backbone, self).__init__(**kwargs)
        assert arch_def is not None
        stages = []
        params_config = decode(arch_def)
        print(params_config)
        for s_id, layer_params in params_config.items():
            layers = []
            for k, v in layer_params.items():
                op_type = v.pop('op_type', None)
                layers.append(build_nn(op_type, v))
            stages.append(nn.Sequential(*layers))
        backbone = nn.Sequential(*stages)
        self.backbone = backbone
        self.out_stage_ids = out_stage_ids
        efficientnet_init_weights(self)

    def forward(self, x):
        out_features = []
        for id, module in enumerate(self.backbone):
            x = module(x)
            if id in self.out_stage_ids:
                out_features.append(x)
        return out_features


class FeatureHook(nn.Module):
    def __init__(self, model, out_indices, feature_info=None, feature_concat=False, flatten_sequential=False):
        super(FeatureHook, self).__init__()
        if feature_info is None:
            feature_info = model.feature_info
        remaining = {feature_info[i]['module']: feature_info[i] for i in out_indices}
        modules = _module_list(model, flatten_sequential=flatten_sequential)
        for new_name, old_name, module in modules:
            self.add_module(new_name, module)
            for fn, fm in module.named_modules(prefix=old_name):
                if fn in remaining:
                    hook_fn = self._collect_output_hook
                    hook_type = remaining[fn].get('hook_type', 'forward')
                    if hook_type == 'forward_pre':
                        fm.register_forward_pre_hook(hook_fn)
                    elif hook_type == 'forward':
                        fm.register_forward_hook(hook_fn)
                    else:
                        assert False, "Unsupported hook type"
                    del remaining[fn]
            if not remaining:
                break
        assert not remaining, 'Return layers ({}) are not present in model'.format(remaining)
        self.feature_concat = feature_concat
        self.feature_outputs = []

    def _collect_output_hook(self, *args):
        x = args[-1]
        if isinstance(x, (tuple, list)):
            if len(x) > 1 and self.feature_concat:
                x = torch.cat(1)
            else:
                x = x[0]
        self.feature_outputs.append(x)

    def forward(self, x):
        for module in self.children():
            x = module(x)
        output = self.feature_outputs
        self.feature_outputs = []
        return output


def _module_list(module, flatten_sequential=False):
    # a yield/iter would be better for this but wouldn't be compatible with torchscript
    ml = []
    for name, module in module.named_children():
        if flatten_sequential and isinstance(module, nn.Sequential):
            # first level of Sequential containers is flattened into containing model
            for child_name, child_module in module.named_children():
                combined = [name, child_name]
                ml.append(('_'.join(combined), '.'.join(combined), child_module))
        else:
            ml.append((name, name, module))
    return ml


def _init_weight_goog(m, n='', fix_group_fanout=True):
    """ Weight initialization as per Tensorflow official implementations.

    Args:
        m (nn.Module): module to init
        n (str): module name
        fix_group_fanout (bool): enable correct (matching Tensorflow TPU impl) fanout calculation w/ group convs

    Handles layers in EfficientNet, EfficientNet-CondConv, MixNet, MnasNet, MobileNetV3, etc:
    * https://github.com/tensorflow/tpu/blob/master/models/official/mnasnet/mnasnet_model.py
    * https://github.com/tensorflow/tpu/blob/master/models/official/efficientnet/efficientnet_model.py
    """
    if isinstance(m, nn.Conv2d):
        fan_out = m.kernel_size[0] * m.kernel_size[1] * m.out_channels
        if fix_group_fanout:
            fan_out //= m.groups
        nn.init.normal_(m.weight, 0, math.sqrt(2.0 / fan_out))
        if m.bias is not None:
            nn.init.zeros_(m.bias)
    elif isinstance(m, nn.BatchNorm2d):
        nn.init.ones_(m.weight)
        nn.init.zeros_(m.bias)
    elif isinstance(m, nn.Linear):
        fan_out = m.weight.size(0)  # fan-out
        fan_in = 0
        if 'routing_fn' in n:
            fan_in = m.weight.size(1)
        init_range = 1.0 / math.sqrt(fan_in + fan_out)
        nn.init.uniform_(m.weight, -init_range, init_range)
        nn.init.zeros_(m.bias)


def efficientnet_init_weights(model: nn.Module, init_fn=None):
    init_fn = init_fn or _init_weight_goog
    for n, m in model.named_modules():
        init_fn(m, n)


