# Import necessary modules
from tensorflow import keras
from memryx import NeuralCompiler, Benchmark

# Step 1: Load the MobileNet model
# --------------------------------
# MobileNet is a lightweight neural network commonly used for mobile applications.
# We will use it as an example model to compile for the MemryX Accelerator (MXA).
mobilenet = keras.applications.MobileNet()

# Step 2: Initialize the Neural Compiler
# --------------------------------------
# NeuralCompiler is a class that compiles neural network models into a Dataflow Program (DFP).
# The DFP can then be executed on the MXA. Here, we initialize the compiler and specify verbose output.
nc = NeuralCompiler(models=mobilenet, verbose=1)

# Step 3: Compile the MobileNet model into a DFP
# ----------------------------------------------
# This step converts the MobileNet model into a DFP, which is optimized for MXA hardware.
# The 'run()' function compiles the model.
mobilenet_dfp = nc.run()

# Optional: Save the compiled DFP for future use
# ----------------------------------------------
# We can save the DFP to a file for reuse without recompiling the model.
mobilenet_dfp.write("mobilenet.dfp")


# Step 4: Initialize the Benchmark
# --------------------------------
# The Benchmark class is responsible for running inference on the compiled model (DFP).
# We use the compiled DFP (mobilenet_dfp) to initialize the benchmark.
benchmark = Benchmark(dfp=mobilenet_dfp)

# Step 5: Run the Benchmark with Random Data
# ------------------------------------------
# Here, we run inference with 1000 frames of random data to measure the performance of the MXA.
# The 'run()' function processes the frames, and we can retrieve the FPS (frames per second).
# The 'with' context is used to ensure proper resource management during benchmarking.
with benchmark as accl:
    # Run 1000 frames of random data
    _, _, fps = accl.run(frames=1000)
    # Print the FPS (Frames Per Second) performance
    print(f"FPS of MobileNet Accelerated on MXA: {fps:.2f}")