import pandas as pd
import endaq

# Set Theme
endaq.plot.utilities.set_theme()

# Get Vibration Data
df_vibe = pd.read_csv('https://info.endaq.com/hubfs/data/motorcycle-vibration-moving-frequency.csv', index_col=0)

# Calculate a Rolling FFT
fft = endaq.calc.fft.rolling_fft(df_vibe, num_slices=200, add_resultant=True)

# Visualize the Rolling FFT as a Heatmap
heatmap = endaq.plot.plots.spectrum_over_time(fft, plot_type='Heatmap', freq_max=200, var_to_process='Resultant')
heatmap.show()

# Plot the Peak Frequency vs Time
peak = endaq.plot.plots.spectrum_over_time(fft, plot_type='Peak', freq_max=200, var_to_process='Resultant')
peak.show()

# Visualize as a Surface Plot
surface = endaq.plot.plots.spectrum_over_time(fft, plot_type='Surface', freq_max=200, var_to_process='Resultant')
surface.show()

# Visualize as a Waterfall
waterfall = endaq.plot.plots.spectrum_over_time(fft, plot_type='Waterfall', freq_max=200, var_to_process='Resultant')
waterfall.show()