import torch
from darknet2any.tool.torch_utils import convert2cpu


def parse_cfg(cfgfile):
  blocks = []
  fp = open(cfgfile, 'r')
  block = None
  line = fp.readline()
  while line != '':
    line = line.rstrip()
    if line == '' or line[0] == '#':
      line = fp.readline()
      continue
    elif line[0] == '[':
      if block:
        blocks.append(block)
      block = dict()
      block['type'] = line.lstrip('[').rstrip(']')
      # set default value
      if block['type'] == 'convolutional':
        block['batch_normalize'] = 0
    else:
      key, value = line.split('=')
      key = key.strip()
      if key == 'type':
        key = '_type'
      value = value.strip()
      block[key] = value
    line = fp.readline()

  if block:
    blocks.append(block)
  fp.close()
  return blocks


def print_cfg(blocks):
  print('layer   filters  size        input        output');
  prev_width = 416
  prev_height = 416
  prev_filters = 3
  out_filters = []
  out_widths = []
  out_heights = []
  ind = -2
  for block in blocks:
    ind = ind + 1
    if block['type'] == 'net':
      prev_width = int(block['width'])
      prev_height = int(block['height'])
      continue
    elif block['type'] == 'convolutional':
      filters = int(block['filters'])
      kernel_size = int(block['size'])
      stride = int(block['stride'])
      is_pad = int(block['pad'])
      pad = (kernel_size - 1) // 2 if is_pad else 0
      width = (prev_width + 2 * pad - kernel_size) // stride + 1
      height = (prev_height + 2 * pad - kernel_size) // stride + 1
      print('%5d %-6s %4d  %d x %d / %d   %3d x %3d x%4d   ->   %3d x %3d x%4d' % (
        ind, 'conv', filters, kernel_size, kernel_size, stride, prev_width, prev_height, prev_filters, width,
        height, filters))
      prev_width = width
      prev_height = height
      prev_filters = filters
      out_widths.append(prev_width)
      out_heights.append(prev_height)
      out_filters.append(prev_filters)
    elif block['type'] == 'maxpool':
      pool_size = int(block['size'])
      stride = int(block['stride'])
      width = prev_width // stride
      height = prev_height // stride
      print('%5d %-6s     %d x %d / %d   %3d x %3d x%4d   ->   %3d x %3d x%4d' % (
        ind, 'max', pool_size, pool_size, stride, prev_width, prev_height, prev_filters, width, height,
        filters))
      prev_width = width
      prev_height = height
      prev_filters = filters
      out_widths.append(prev_width)
      out_heights.append(prev_height)
      out_filters.append(prev_filters)
    elif block['type'] == 'avgpool':
      width = 1
      height = 1
      print('%5d %-6s           %3d x %3d x%4d   ->  %3d' % (
        ind, 'avg', prev_width, prev_height, prev_filters, prev_filters))
      prev_width = width
      prev_height = height
      prev_filters = filters
      out_widths.append(prev_width)
      out_heights.append(prev_height)
      out_filters.append(prev_filters)
    elif block['type'] == 'softmax':
      print('%5d %-6s                  ->  %3d' % (ind, 'softmax', prev_filters))
      out_widths.append(prev_width)
      out_heights.append(prev_height)
      out_filters.append(prev_filters)
    elif block['type'] == 'cost':
      print('%5d %-6s                   ->  %3d' % (ind, 'cost', prev_filters))
      out_widths.append(prev_width)
      out_heights.append(prev_height)
      out_filters.append(prev_filters)
    elif block['type'] == 'reorg':
      stride = int(block['stride'])
      filters = stride * stride * prev_filters
      width = prev_width // stride
      height = prev_height // stride
      print('%5d %-6s       / %d   %3d x %3d x%4d   ->   %3d x %3d x%4d' % (
        ind, 'reorg', stride, prev_width, prev_height, prev_filters, width, height, filters))
      prev_width = width
      prev_height = height
      prev_filters = filters
      out_widths.append(prev_width)
      out_heights.append(prev_height)
      out_filters.append(prev_filters)
    elif block['type'] == 'upsample':
      stride = int(block['stride'])
      filters = prev_filters
      width = prev_width * stride
      height = prev_height * stride
      print('%5d %-6s       * %d   %3d x %3d x%4d   ->   %3d x %3d x%4d' % (
        ind, 'upsample', stride, prev_width, prev_height, prev_filters, width, height, filters))
      prev_width = width
      prev_height = height
      prev_filters = filters
      out_widths.append(prev_width)
      out_heights.append(prev_height)
      out_filters.append(prev_filters)
    elif block['type'] == 'route':
      layers = block['layers'].split(',')
      layers = [int(i) if int(i) > 0 else int(i) + ind for i in layers]
      if len(layers) == 1:
        print('%5d %-6s %d' % (ind, 'route', layers[0]))
        prev_width = out_widths[layers[0]]
        prev_height = out_heights[layers[0]]
        prev_filters = out_filters[layers[0]]
      elif len(layers) == 2:
        print('%5d %-6s %d %d' % (ind, 'route', layers[0], layers[1]))
        prev_width = out_widths[layers[0]]
        prev_height = out_heights[layers[0]]
        assert (prev_width == out_widths[layers[1]])
        assert (prev_height == out_heights[layers[1]])
        prev_filters = out_filters[layers[0]] + out_filters[layers[1]]
      elif len(layers) == 4:
        print('%5d %-6s %d %d %d %d' % (ind, 'route', layers[0], layers[1], layers[2], layers[3]))
        prev_width = out_widths[layers[0]]
        prev_height = out_heights[layers[0]]
        assert (prev_width == out_widths[layers[1]] == out_widths[layers[2]] == out_widths[layers[3]])
        assert (prev_height == out_heights[layers[1]] == out_heights[layers[2]] == out_heights[layers[3]])
        prev_filters = out_filters[layers[0]] + out_filters[layers[1]] + out_filters[layers[2]] + out_filters[
          layers[3]]
      else:
        print("route error !!! {} {} {}".format(sys._getframe().f_code.co_filename,
                            sys._getframe().f_code.co_name, sys._getframe().f_lineno))

      out_widths.append(prev_width)
      out_heights.append(prev_height)
      out_filters.append(prev_filters)
    elif block['type'] in ['region', 'yolo']:
      print('%5d %-6s' % (ind, 'detection'))
      out_widths.append(prev_width)
      out_heights.append(prev_height)
      out_filters.append(prev_filters)
    elif block['type'] == 'shortcut':
      from_id = int(block['from'])
      from_id = from_id if from_id > 0 else from_id + ind
      print('%5d %-6s %d' % (ind, 'shortcut', from_id))
      prev_width = out_widths[from_id]
      prev_height = out_heights[from_id]
      prev_filters = out_filters[from_id]
      out_widths.append(prev_width)
      out_heights.append(prev_height)
      out_filters.append(prev_filters)
    elif block['type'] == 'sam':
      from_id = int(block['from'])
      from_id = from_id if from_id > 0 else from_id + ind
      print('%5d %-6s %d' % (ind, 'sam', from_id))
      prev_width = out_widths[from_id]
      prev_height = out_heights[from_id]
      prev_filters = out_filters[from_id]
      out_widths.append(prev_width)
      out_heights.append(prev_height)
      out_filters.append(prev_filters)
    elif block['type'] == 'connected':
      filters = int(block['output'])
      print('%5d %-6s              %d  ->  %3d' % (ind, 'connected', prev_filters, filters))
      prev_filters = filters
      out_widths.append(1)
      out_heights.append(1)
      out_filters.append(prev_filters)
    else:
      print('unknown type %s' % (block['type']))


def load_conv(buf, start, conv_model):
  num_w = conv_model.weight.numel()
  num_b = conv_model.bias.numel()
  conv_model.bias.data.copy_(torch.from_numpy(buf[start:start + num_b]));
  start = start + num_b
  conv_model.weight.data.copy_(torch.from_numpy(buf[start:start + num_w]).reshape(conv_model.weight.data.shape));
  start = start + num_w
  return start


def save_conv(fp, conv_model):
  if conv_model.bias.is_cuda:
    convert2cpu(conv_model.bias.data).numpy().tofile(fp)
    convert2cpu(conv_model.weight.data).numpy().tofile(fp)
  else:
    conv_model.bias.data.numpy().tofile(fp)
    conv_model.weight.data.numpy().tofile(fp)


def load_conv_bn(buf, start, conv_model, bn_model):
  num_w = conv_model.weight.numel()
  num_b = bn_model.bias.numel()
  bn_model.bias.data.copy_(torch.from_numpy(buf[start:start + num_b]));
  start = start + num_b
  bn_model.weight.data.copy_(torch.from_numpy(buf[start:start + num_b]));
  start = start + num_b
  bn_model.running_mean.copy_(torch.from_numpy(buf[start:start + num_b]));
  start = start + num_b
  bn_model.running_var.copy_(torch.from_numpy(buf[start:start + num_b]));
  start = start + num_b
  conv_model.weight.data.copy_(torch.from_numpy(buf[start:start + num_w]).reshape(conv_model.weight.data.shape));
  start = start + num_w
  return start


def save_conv_bn(fp, conv_model, bn_model):
  if bn_model.bias.is_cuda:
    convert2cpu(bn_model.bias.data).numpy().tofile(fp)
    convert2cpu(bn_model.weight.data).numpy().tofile(fp)
    convert2cpu(bn_model.running_mean).numpy().tofile(fp)
    convert2cpu(bn_model.running_var).numpy().tofile(fp)
    convert2cpu(conv_model.weight.data).numpy().tofile(fp)
  else:
    bn_model.bias.data.numpy().tofile(fp)
    bn_model.weight.data.numpy().tofile(fp)
    bn_model.running_mean.numpy().tofile(fp)
    bn_model.running_var.numpy().tofile(fp)
    conv_model.weight.data.numpy().tofile(fp)


def load_fc(buf, start, fc_model):
  num_w = fc_model.weight.numel()
  num_b = fc_model.bias.numel()
  fc_model.bias.data.copy_(torch.from_numpy(buf[start:start + num_b]));
  start = start + num_b
  fc_model.weight.data.copy_(torch.from_numpy(buf[start:start + num_w]));
  start = start + num_w
  return start


def save_fc(fp, fc_model):
  fc_model.bias.data.numpy().tofile(fp)
  fc_model.weight.data.numpy().tofile(fp)


if __name__ == '__main__':
  import sys

  blocks = parse_cfg('cfg/yolo.cfg')
  if len(sys.argv) == 2:
    blocks = parse_cfg(sys.argv[1])
  print_cfg(blocks)
