import tensorflowjs as tfjs
import json

# After saving keras model with ZeroPadding3D as model.json  Modify with this code:


def fixjson_file(model_file, verbose=False, scube=38):
    # Open json file to modify
    with open(model_file) as f:
        model_dict = json.load(f)

    # Convert
    for layer in model_dict['modelTopology']['model_config']['config']['layers']:
        if layer['class_name'] == "ZeroPadding3D":
            model_dict['modelTopology']['model_config']['config']['layers'].remove(layer)

    prev_layer_name = ""
    for layer in model_dict['modelTopology']['model_config']['config']['layers']:
        if layer['class_name'] == "InputLayer":
            layer["config"]["batch_input_shape"] = [None, scube, scube, scube, 1]

        if layer['class_name'] == "Conv3D":
            layer["config"]["padding"] = "same"
            layer["config"]["data_format"] = "channels_last"
            layer['inbound_nodes'][0][0][0] = prev_layer_name

        prev_layer_name = layer["config"]["name"]

    if verbose:
        # Verification
        for layer in model_dict['modelTopology']['model_config']['config']['layers']:
            print(layer)
            print("-------------------------------------------------------")

    # Save model.json file
    with open(model_file, 'w') as fp:
        json.dump(model_dict, fp)