import sys
import json
import tensorflow.keras as K
import tensorflow as tf
import os

# coding = utf-8
def main():
    modelfile=sys.argv[1]
    print(modelfile)
    jsonfilepath=sys.argv[2]
    modeltype=sys.argv[3]  #0: keras, 1:h5, 2:folder
    # pytorchmodel=sys.argv[4]
    print(jsonfilepath)
    print(modeltype)
    
    # if int(modeltype)==0:
    #     model=K.models.load_model(modelfile)
    #     model_config=model.to_json()
    if int(modeltype)==1:
        model=K.models.load_model(modelfile)
        model_config=model.to_json()
    if int(modeltype)==2:
        model=K.models.model_from_yaml(modelfile)
        model_config=model.to_json()
    # if int(modeltype)==3:
    #     import torch as tt
    #     import onnx
    #     from onnx2keras import onnx_to_keras
    #     if pytorchmodel == "_":
    #         raise Exception("Need to select a model type")
    #     torchmodel=tt.load(modelfile)
    #     dummyinput=tt.randn(1,3,224,224,requires_grad=True)
    #     temppathindex=jsonfilepath.rfind('/')
    #     temppath=jsonfilepath[:temppathindex+1]+'temponnx.onnx'
    #
    #     temponnx=tt.onnx.export(torchmodel,dummyinput,temppath,
    #                             verbose=True,input_names=['input'],output_names=['output'])
    #     onnx_model = onnx.load(temppath)
    #     model = onnx_to_keras(onnx_model, ['input'])
    #     model_config = model.to_json()

    print(model_config)
    y=json.loads(model_config)
    out_file=open(jsonfilepath,"w")
    json.dump(y,out_file,indent=6)


if __name__ == '__main__':
    main()
    

