'''
@author: insightface
'''

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import sys
import os
import json
import argparse
import numpy as np
import mxnet as mx


def is_no_bias(attr):
  ret = False
  if 'no_bias' in attr and (attr['no_bias']==True or attr['no_bias']=='True'):
    ret = True
  return ret

def count_fc_flops(input_filter, output_filter, attr):
  #print(input_filter, output_filter ,attr)
  ret = 2*input_filter*output_filter
  if is_no_bias(attr):
    ret -= output_filter
  return int(ret)


def count_conv_flops(input_shape, output_shape, attr):
  kernel = attr['kernel'][1:-1].split(',')
  kernel = [int(x) for x in kernel]

  #print('kernel', kernel)
  if is_no_bias(attr):
    ret = (2*input_shape[1]*kernel[0]*kernel[1]-1)*output_shape[2]*output_shape[3]*output_shape[1]
  else:
    ret = 2*input_shape[1]*kernel[0]*kernel[1]*output_shape[2]*output_shape[3]*output_shape[1]
  num_group = 1
  if 'num_group' in attr:
    num_group = int(attr['num_group'])
  ret /= num_group
  return int(ret)


def count_flops(sym, **data_shapes):
  all_layers = sym.get_internals()
  #print(all_layers)
  arg_shapes, out_shapes, aux_shapes = all_layers.infer_shape(**data_shapes)
  out_shape_dict = dict(zip(all_layers.list_outputs(), out_shapes))

  nodes = json.loads(sym.tojson())['nodes']
  nodeid_shape = {}
  for nodeid, node in enumerate(nodes):
    name = node['name']
    layer_name = name+"_output"
    if layer_name in out_shape_dict:
      nodeid_shape[nodeid] = out_shape_dict[layer_name]
  #print(nodeid_shape)
  FLOPs = 0
  for nodeid, node in enumerate(nodes):
    flops = 0
    if node['op']=='Convolution':
      output_shape = nodeid_shape[nodeid]
      name = node['name']
      attr = node['attrs']
      input_nodeid = node['inputs'][0][0]
      input_shape = nodeid_shape[input_nodeid]
      flops = count_conv_flops(input_shape, output_shape, attr)
    elif node['op']=='FullyConnected':
      attr = node['attrs']
      output_shape = nodeid_shape[nodeid]
      input_nodeid = node['inputs'][0][0]
      input_shape = nodeid_shape[input_nodeid]
      output_filter = output_shape[1]
      input_filter = input_shape[1]*input_shape[2]*input_shape[3]
      #assert len(input_shape)==4 and input_shape[2]==1 and input_shape[3]==1
      flops = count_fc_flops(input_filter, output_filter, attr)
    #print(node, flops)
    FLOPs += flops

  return FLOPs

def flops_str(FLOPs):
  preset = [ (1e12, 'T'), (1e9, 'G'), (1e6, 'M'), (1e3, 'K') ]

  for p in preset:
    if FLOPs//p[0]>0:
      N = FLOPs/p[0]
      ret = "%.1f%s"%(N, p[1])
      return ret
  ret = "%.1f"%(FLOPs)
  return ret

if __name__ == '__main__':
  parser = argparse.ArgumentParser(description='flops counter')
  # general
  #parser.add_argument('--model', default='../models2/y2-arcface-retinat1/model,1', help='path to load model.')
  #parser.add_argument('--model', default='../models2/r100fc-arcface-retinaa/model,1', help='path to load model.')
  parser.add_argument('--model', default='../models2/r50fc-arcface-emore/model,1', help='path to load model.')
  args = parser.parse_args()
  _vec = args.model.split(',')
  assert len(_vec)==2
  prefix = _vec[0]
  epoch = int(_vec[1])
  print('loading',prefix, epoch)
  sym, arg_params, aux_params = mx.model.load_checkpoint(prefix, epoch)
  all_layers = sym.get_internals()
  sym = all_layers['fc1_output']
  FLOPs = count_flops(sym, data=(1,3,112,112))
  print('FLOPs:', FLOPs)