from pymodhook import *
init_hook()
hook_modules("numpy", "matplotlib.pyplot", for_=["__main__"]) # Record calls to numpy and matplotlib
enable_hook()
import numpy as np
import matplotlib.pyplot as plt
arr = np.array(range(1,11))
arr_squared = arr ** 2
mean = np.mean(arr)
std_dev = np.std(arr)
print(mean, std_dev)

plt.plot(arr, arr_squared)
plt.show()

# Display the recorded code
print(f"Raw call trace:\n{get_code()}\n")
print(f"Optimized code:\n{get_optimized_code()}")