import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import pickle

"""
Script originally written by Ben Hokanson-Fasig
"""

########################## Generate the data to plot ##########################
with open('ray_data.pkl', 'rb') as f:
    data = pickle.load(f)
    ray_paths = data['ray_paths']
    ray_tofs = data['ray_tofs']



##################### Plot ray paths with times of flight #####################
# Create two sub-figure axes, one for the plot and one for the colorbar
fig, axs = plt.subplots(ncols=2, figsize=(8, 5),
                        gridspec_kw={'width_ratios': [7.5, 0.5]})
# Set the colormap to be used
cmap = plt.cm.get_cmap('viridis_r')
# Define the maximum and minimum times of flight to use in the colorbar
tof_min = 2900e-9
tof_max = 3500e-9
# Plot each path, giving it a color from the colormap which corresponds to its
# time of flight
for path, tof in zip(ray_paths, ray_tofs):
    axs[0].plot(path[0], path[2],
                color=cmap((tof-tof_min)/(tof_max-tof_min)))
# Draw the colorbar using the defined tof range from above
cb = mpl.colorbar.ColorbarBase(axs[-1], cmap=cmap,
                               boundaries=np.linspace(tof_min, tof_max, 101)*1e9)
cb.set_label("Time of Flight [ns]", fontsize=12,
             rotation=270, verticalalignment='bottom')
# Give the plot a title and axis labels
axs[0].set_title("Radio Ray Paths to a 200 m Antenna", fontsize=14)
axs[0].set_xlabel("x [m]", fontsize=12)
axs[0].set_ylabel("z [m]", fontsize=12)
# Save the plot as a PDF so it scales nicely in presentations
plt.savefig('ray_tofs.pdf')