from math import pi

from vortex.scan import FreeformScanConfig, FreeformScan, SequentialPattern
from vortex_tools.scan import plot_annotated_waveforms_space

(r, theta) = np.meshgrid(
    [2, 3, 4],
    np.linspace(0, 2*pi, 200),
    indexing='ij'
)

x = (r + 0.1*np.sin(r + 20*theta)) * np.sin(theta)
y = (r + 0.1*np.sin(r + 20*theta)) * np.cos(theta)

waypoints = np.stack((x, y), axis=-1)
pattern = SequentialPattern().to_pattern(waypoints)

cfg = FreeformScanConfig()
cfg.pattern = pattern
for limit in cfg.limits:
    limit.velocity *= 5
    limit.acceleration *= 10
cfg.bypass_limits_check = True
cfg.loop = True

scan = FreeformScan()
scan.initialize(cfg)

fg = plt.rcParams['lines.color']
fig, ax = plot_annotated_waveforms_space(scan.scan_buffer(), scan.scan_markers(), inactive_marker=None, scan_line=fg)
ax.set_title('Freeform Scan')
fig.tight_layout()