import copy
import torch.nn as nn


def as_list(data):
    if isinstance(data, list):
        return data
    else:
        return [data]

class NameManager(object):
    def __init__(self):
        self._counter = {}

    def get(self, hint):
        if hint not in self._counter:
            self._counter[hint] = 0
        name = '%s%d' % (hint, self._counter[hint])
        self._counter[hint] += 1
        return name


class ConvBlock(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1,
                 groups=1, bias=None, op_list=('conv', ), op_params=None, **kwargs):
        super(ConvBlock, self).__init__()
        op_params = copy.deepcopy(op_params) if op_params is not None else {}
        if bias is None:
            bias = 'norm' not in op_list
        self.grad_backward = kwargs.pop('grad_backward', True)
        conv_kwargs = {
            'in_channels': in_channels,
            'out_channels': out_channels,
            'kernel_size': kernel_size,
            'stride': stride,
            'padding': padding,
            'dilation': dilation,
            'groups': groups,
            'bias': bias,
        }
        conv_kwargs.update(kwargs)
        name_manager = NameManager()
        self.op_names = []
        for op_type in op_list:
            op_name = name_manager.get(op_type)
            if op_type == 'conv':
                op = nn.Conv2d(**conv_kwargs)
                in_channels = out_channels
            elif op_type == 'norm':
                op = nn.BatchNorm2d(in_channels)
            elif op_type == 'act':
                op = nn.ReLU()
            else:
                assert False
            self.add_module(op_name, op)
            self.op_names.append(op_name)

    def forward(self, x):
        if not self.grad_backward:
            x = x.detach()
        for op_name in self.op_names:
            x = getattr(self, op_name)(x)
        return x


class Compose(nn.Module):
    def __init__(self, op_params_list, in_channels=0):
        super(Compose, self).__init__()
        name_manager = NameManager()
        cur_channels = in_channels
        self.op_names = []
        for op_params in op_params_list:
            # get op_type
            op_type = op_params.pop('op_type', None)
            if op_type is None:
                if op_params['type'] == 'ConvBlock':
                    op_type = 'conv'
            # get op
            if op_type == 'conv':
                op_params = _process_conv_params(op_params, cur_channels)
                for j, op_params_j in enumerate(op_params):
                    op_name_j = op_params_j.pop('op_name', None)
                    if op_name_j is None:
                        op_name_j = name_manager.get(op_type)
                    op_params_j.pop('type')
                    op_j = ConvBlock(**op_params_j)
                    self.add_module(op_name_j, op_j)
                    self.op_names.append(op_name_j)
                cur_channels = op_params[-1]['out_channels']

    def forward(self, x):
        for op_name in self.op_names:
            x = getattr(self, op_name)(x)
        return x


def _process_conv_params(conv_params, in_channels=0):
    assert 'in_channels' not in conv_params
    conv_params = {key: as_list(val) for key, val in conv_params.items()}
    num_output = len(conv_params['out_channels'])
    conv_params_list = []
    for i in range(num_output):
        conv_params_i = {key: val[i % len(val)] for key, val in conv_params.items()}
        conv_params_i['in_channels'] = in_channels
        conv_params_list.append(conv_params_i)
        in_channels = conv_params_i['out_channels']
    return conv_params_list


def _process_linear_params(linear_params, in_features=0):
    assert 'in_features' not in linear_params
    linear_params = {key: as_list(val) for key, val in linear_params.items()}
    num_output = len(linear_params['out_features'])
    linear_params_list = []
    for i in range(num_output):
        linear_params_i = {key: val[i % len(val)] for key, val in linear_params.items()}
        linear_params_i['in_features'] = in_features
        linear_params_list.append(linear_params_i)
        in_features = linear_params_i['out_features']
    return linear_params_list
