#!/usr/bin/python
import os
def generateCPUFile(rootDir):
    cpuDir = rootDir
    cpuRegFile = os.path.join(cpuDir, "BF16OpRegister.cpp")
    fileNames = os.listdir(cpuDir)
    print(fileNames)
    if len(fileNames) <= 1:
        # Error dirs
        return
    funcNames = []
    for fi in fileNames:
        f = os.path.join(cpuDir, fi)
        if os.path.isdir(f):
            continue
        print(f)
        with open(f) as fileC:
            c = fileC.read().split('\n')
            c = list(filter(lambda l:l.find('REGISTER_BF16_OP_CREATOR')>=0, c))
            c = list(filter(lambda l:l.find('OpType')>=0, c))
            for l in c:
                l = l.split('(')[1]
                l = l.split(')')[0]
                l = l.replace(' ', '')
                l = l.split(',')
                funcName = '___' + l[0] + '__' + l[1] + '__'
                funcNames.append(funcName)
    with open(cpuRegFile, 'w') as f:
        f.write('// This file is generated by Shell for ops register\n')
        f.write('namespace MNN {\n')
        for l in funcNames:
            f.write("extern void " + l + '();\n')
        f.write('\n')
        f.write('void registerBF16Ops() {\n')
        for l in funcNames:
            f.write(l+'();\n')
        f.write("}\n}\n")

import sys
generateCPUFile(sys.argv[1])
