from matplotlib import pyplot as plt

from vortex import Range
from vortex.scan import RepeatedRasterScan, RepeatedRasterScanConfig
from vortex_tools.scan import plot_annotated_waveforms_space

fg = plt.rcParams['lines.color']
fig, axs = plt.subplots(2, 2, sharex=True, sharey=True, constrained_layout=True, subplot_kw=dict(adjustable='box', aspect='equal'))
cfgs = []
names = []

cfg = RepeatedRasterScanConfig()
cfg.segment_extent = Range.symmetric(1)
cfg.volume_extent = Range.symmetric(2)
cfg.segments_per_volume = 6
cfg.samples_per_segment = 50
for limit in cfg.limits:
    limit.acceleration *= 5
cfg.loop = True

names.append('Standard')
cfgs.append(cfg.copy())
cfgs[-1].repeat_count = 1
cfgs[-1].repeat_period = 1

names.append('Repeat')
cfgs.append(cfg.copy())

names.append('Even Repeat Period + Bidirectional')
cfgs.append(cfg.copy())
cfgs[-1].repeat_period = 2
cfgs[-1].bidirectional_segments = True

names.append('Odd Repeat Period + Bidirectional')
cfgs.append(cfg.copy())
cfgs[-1].repeat_period = 3
cfgs[-1].bidirectional_segments = True

for (name, cfg, ax) in zip(names, cfgs, axs.flat):
    scan = RepeatedRasterScan()
    scan.initialize(cfg)
    plot_annotated_waveforms_space(scan.scan_buffer(), scan.scan_markers(), inactive_marker=None, scan_line=fg, axes=ax)
    ax.set_title(name)