import numpy
import gradio
import huggingface_hub
import json

class NumpyArrayEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, numpy.ndarray):
            return obj.tolist()
        return JSONEncoder.default(self, obj)

analysis_network = huggingface_hub.from_pretrained_keras("cmudrc/wave-energy-analysis")
synthesis_network = huggingface_hub.from_pretrained_keras("cmudrc/wave-energy-synthesis")
   
with gradio.Blocks() as demo:
    geometry = gradio.Textbox(label="geometry")     
    spectrum = gradio.Textbox(label="spectrum")

    analyze_it = gradio.Button("Analyze")
    synthesize_it = gradio.Button("Synthesize")

    analyze_it.click(fn=lambda x: json.dumps(analysis_network.predict(numpy.asarray([json.loads(x)])), cls=NumpyArrayEncoder), inputs=[geometry], outputs=[spectrum], api_name="analyze")
    synthesize_it.click(fn=lambda x: json.dumps(synthesis_network.predict(numpy.asarray(json.loads(x))), cls=NumpyArrayEncoder), inputs=[spectrum], outputs=[geometry], api_name="synthesize")
    
demo.launch(debug=True)