# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

# -*- coding: utf-8 -*-
# This is a python script that generates operator wrappers such as FullyConnected,
# based on current libmxnet.dll. This script is written so that we don't need to
# write new operator wrappers when new ones are added to the library.

from ctypes import *
from ctypes.util import find_library
import os
import logging
import platform
import re
import sys
import tempfile
import filecmp
import shutil
import codecs

def gen_enum_value(value):
    return 'k' + value[0].upper() + value[1:]

class EnumType:
    name = ''
    enumValues = []
    def __init__(self, typeName = 'ElementWiseOpType', \
                 typeString = "{'avg', 'max', 'sum'}"):
        self.name = typeName
        if (typeString[0] == '{'):  # is a enum type
            isEnum = True
            # parse enum
            self.enumValues = typeString[typeString.find('{') + 1:typeString.find('}')].split(',')
            for i in range(0, len(self.enumValues)):
                self.enumValues[i] = self.enumValues[i].strip().strip("'")
        else:
            logging.warn("trying to parse none-enum type as enum: %s" % typeString)
    def GetDefinitionString(self, indent = 0):
        indentStr = ' ' * indent
        ret = indentStr + 'enum class %s {\n' % self.name
        for i in range(0, len(self.enumValues)):
            ret = ret + indentStr + '  %s = %d' % (gen_enum_value(self.enumValues[i]), i)
            if (i != len(self.enumValues) -1):
                ret = ret + ","
            ret = ret + "\n"
        ret = ret + "};\n"
        return ret
    def GetDefaultValueString(self, value = ''):
        return self.name + "::" + gen_enum_value(value)
    def GetEnumStringArray(self, indent = 0):
        indentStr = ' ' * indent
        ret = indentStr + 'static const char *%sValues[] = {\n' % self.name
        for i in range(0, len(self.enumValues)):
            ret = ret + indentStr + '  "%s"' % self.enumValues[i]
            if (i != len(self.enumValues) -1):
                ret = ret + ","
            ret = ret + "\n"
        ret = ret + indentStr + "};\n"
        return ret
    def GetConvertEnumVariableToString(self, variable=''):
        return "%sValues[int(%s)]" % (self.name, variable)


class Arg:
    typeDict = {'boolean':'bool',\
        'boolean or None':'dmlc::optional<bool>',\
        'Shape(tuple)':'Shape',\
        'Symbol':'Symbol',\
        'NDArray':'Symbol',\
        'NDArray-or-Symbol':'Symbol',\
        'Symbol[]':'const std::vector<Symbol>&',\
        'Symbol or Symbol[]':'const std::vector<Symbol>&',\
        'NDArray[]':'const std::vector<Symbol>&',\
        'caffe-layer-parameter':'::caffe::LayerParameter',\
        'NDArray-or-Symbol[]':'const std::vector<Symbol>&',\
        'float':'mx_float',\
        'real_t':'mx_float',\
        'int':'int',\
        'int (non-negative)': 'uint32_t',\
        'long (non-negative)': 'uint64_t',\
        'int or None':'dmlc::optional<int>',\
        'float or None':'dmlc::optional<float>',\
        'long':'int64_t',\
        'double':'double',\
        'double or None':'dmlc::optional<double>',\
        'Shape or None':'dmlc::optional<Shape>',\
        'string':'const std::string&',\
        'tuple of <float>':'nnvm::Tuple<mx_float>',\
        'tuple of <>':'mxnet::cpp::Shape',\
        '':'index_t'}
    name = ''
    type = ''
    description = ''
    isEnum = False
    enum = None
    hasDefault = False
    defaultString = ''
    def __init__(self, opName = '', argName = '', typeString = '', descString = ''):
        self.name = argName
        self.description = descString
        if (typeString[0] == '{'):  # is enum type
            self.isEnum = True
            self.enum = EnumType(self.ConstructEnumTypeName(opName, argName), typeString)
            self.type = self.enum.name
        else:
            try:
                self.type = self.typeDict[typeString.split(',')[0]]
            except:
                print('argument "%s" of operator "%s" has unknown type "%s"' % (argName, opName, typeString))
                pass
        if typeString.find('default=') != -1:
            self.hasDefault = True
            self.defaultString = typeString.split('default=')[1].strip().strip("'")
            if typeString.startswith('string'):
                self.defaultString = self.MakeCString(self.defaultString)
            elif self.isEnum:
                self.defaultString = self.enum.GetDefaultValueString(self.defaultString)
            elif self.defaultString == 'None':
                self.defaultString = self.type + '()'
            elif self.type == "bool":
                if self.defaultString == "1" or self.defaultString == "True":
                    self.defaultString = "true"
                else:
                    self.defaultString = "false"
            elif self.defaultString[0] == '(':
                self.defaultString = 'Shape' + self.defaultString
            elif self.defaultString[0] == '[':
                self.defaultString = 'Shape(' + self.defaultString[1:-1] + ")"
            elif self.type == 'dmlc::optional<int>':