from vortex import Range
from vortex.scan import RasterScan, RasterScanConfig

cfg = RasterScanConfig()
cfg.volume_extent = Range(-0.1, 0.1)
cfg.bscan_extent = Range(-0.05, 0.05)
cfg.samples_per_segment = 20
cfg.segments_per_volume = 10
cfg.loop = True

scan = RasterScan()
scan.initialize(cfg)

fig, ax = plt.subplots()

for segment in cfg.to_waypoints():
    ax.plot(segment[:, 0], segment[:, 1], 'x')

path = scan.scan_buffer()
ax.plot(path[:, 0], path[:, 1], 'w-', zorder=-1)

ax.set_title('Raster Scan - Unidirectional')
ax.set_xlabel('x (au)')
ax.set_ylabel('y (au)')
ax.axis('equal')

cfg.bidirectional_segments = True

scan = RasterScan()
scan.initialize(cfg)

fig, ax = plt.subplots()

for segment in cfg.to_waypoints():
    ax.plot(segment[:, 0], segment[:, 1], 'x')

path = scan.scan_buffer()
ax.plot(path[:, 0], path[:, 1], 'w-', zorder=-1)

ax.set_title('Raster Scan - Bidirectional')
ax.set_xlabel('x (au)')
ax.set_ylabel('y (au)')
ax.axis('equal')