from vortex.marker import Flags, ScanBoundary, VolumeBoundary
from vortex.scan import RasterScanConfig, RadialScanConfig, FreeformScanConfig, FreeformScan
from vortex_tools.scan import plot_annotated_waveforms_space

rsc = RasterScanConfig()
rsc.flags = Flags(0x1)

raster_segments = rsc.to_segments()

asc = RadialScanConfig()
asc.set_aiming()
# add an offset just to illustrate
asc.offset = (4, 0)
asc.flags = Flags(0x2)

aiming_segments = asc.to_segments()

pattern = []
idx = np.linspace(0, len(raster_segments), 10, dtype=int)
for (i, (a, b)) in enumerate(zip(idx[:-1], idx[1:])):
    if i > 0:
        markers = raster_segments[a].markers
        markers.insert(0, VolumeBoundary(0, 0, False))
        markers.insert(0, ScanBoundary(0, 0))

    pattern += raster_segments[a:b]
    pattern += aiming_segments

ffsc = FreeformScanConfig()
ffsc.pattern = pattern
ffsc.loop = True

scan = FreeformScan()
scan.initialize(ffsc)

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('Raster Scan and Offset Aiming Scan')
fig.tight_layout()