from matplotlib import pyplot as plt

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

cfg = RasterScanConfig()
cfg.volume_extent = Range.symmetric(1)
cfg.bscan_extent = Range.symmetric(1)
cfg.samples_per_segment = 20
cfg.segments_per_volume = 10
for limit in cfg.limits:
    limit.velocity *= 10
    limit.acceleration *= 40
cfg.loop = True

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

cfgs.append(cfg.copy())
names.append('Raster Scan - Unidirectional')

cfgs.append(cfg.copy())
cfgs[-1].bidirectional_segments = True
names.append('Raster Scan - Bidirectional')

for (name, cfg, ax) in zip(names, cfgs, axs):
    scan = RasterScan()
    scan.initialize(cfg)

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

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

    ax.set_xlabel('x (au)')
    ax.set_ylabel('y (au)')
    ax.set_title(name)