import pandas as pd
import endaq

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

# Get a Longer Dataset with DatetimeIndex
engine = endaq.ide.get_primary_sensor_data('https://info.endaq.com/hubfs/data/Commute.ide', measurement_type='accel',
                                           time_mode='datetime')

# Compute PSD
psd = endaq.calc.psd.rolling_psd(engine, num_slices=500, add_resultant=True, octave_bins=12, fstart=4)

# Visualize as a Heatmap
heatmap2 = endaq.plot.plots.spectrum_over_time(psd, plot_type='Heatmap', var_to_process='Resultant', zsmooth='best',
                                               log_freq=True, log_val=True)
heatmap2.show()

# Visualize as an Animation
animation = endaq.plot.plots.spectrum_over_time(psd, plot_type='Animation', var_to_process='Resultant',
                                                log_freq=True, log_val=True)
animation.show()

# Use Rolling PSD to Calculate RMS in Certain Frequency Bins
rms = endaq.calc.psd.rolling_psd(engine, num_slices=500, scaling='rms', add_resultant=True,
                                 freq_splits=[1, 20, 60, 300, 3000])

# Plot the RMS per Frequency Bin Over Time
lines = endaq.plot.plots.spectrum_over_time(rms, plot_type='Lines', log_val=False, var_to_process='Resultant')
lines.show()

# Compute Pseudo Velocity at Specific Times
pvss = endaq.calc.shock.rolling_shock_spectrum(engine, slice_width=2.0, add_resultant=True,
                                               mode='pvss', init_freq=4, damp=0.05,
                                               index_values=pd.DatetimeIndex(['2016-08-02 12:07:15',
                                                                              '2016-08-02 12:08:01',
                                                                              '2016-08-02 12:10:06'], tz='UTC'))

# Visualize as a Waterfall
waterfall2 = endaq.plot.plots.spectrum_over_time(pvss, plot_type='Waterfall', log_freq=True, log_val=True,
                                                 var_to_process='Resultant', waterfall_line_sequence=False)
waterfall2.show()