import re
import os

print("Generating new layers.py from the following files\n")
files = list(filter(lambda x: x[-10:] == '_layers.py', os.listdir("onnx_caffe")))
print(files)

layers_file = open(os.path.join('onnx_caffe', 'layers.py'), 'w')
layers_file.write("# This file is generated by generate_layers.py\n")
doc_file = open("Operators.md", 'w')
doc_file.write("# Supported Layer List\n\n")
doc_file.write("**This file is generated by generate_layers.py. Please do not edit manually.**\n\n")
doc_file.write("## List\n\n")
count = 0
for fname in files:
    with open(os.path.join('onnx_caffe', fname), 'rb') as f:
        content = f.read()
        layers = re.findall('class\ \S+\(Layer\)' , str(content))
        layers = list(map(lambda x: x[6:-7], layers))
        count += len(layers)
        #print(layers)
        for layer in layers:
            layers_file.write('from {} import {}\n'.format('.' + fname[:-3], layer))
            doc_file.write('* {}\n'.format(layer))

layers_file.close()
doc_file.close()
print("Totally {} kinds of layers generated.".format(count))