import matplotlib.pyplot as plt
import numpy as np
import astropy.units as u
from astropy.cosmology import Planck18

z = np.arange(0, 15, 0.01)

fig, ax1 = plt.subplots()
ax2 = ax1.twinx()

fontkwargs = {"fontweight": "bold", "fontsize": "large"}
color1 = "blue"
lines1 = ax1.plot(z, Planck18.comoving_distance(z).to(u.Glyr),
                  label="distance (left axis)", color=color1)
ax1.set_xlim(0, 15)
ax1.set_xlabel("redshift (z)", fontweight="bold")
ax1.set_xticklabels(ax1.get_xticks(), weight="bold")
ax1.set_ylabel("comoving distance (giga light years)", color=color1, **fontkwargs)
ax1.tick_params(axis='y', colors=color1)
ax1.set_yticklabels(ax1.get_yticks(), weight="bold")

color2 = "darkorange"
# put them on the same vertical scale:
ax2.set_ylim(ax1.get_ylim())
lines2 = ax2.plot(z, Planck18.lookback_time(z), "--",
                  label="time (right axis)", color=color2)
ax2.set_ylabel("lookback time (giga years)", color=color2, **fontkwargs)
ax2.tick_params(axis='y', colors=color2)
ax2.set_yticklabels(ax2.get_yticks(), weight="bold")

ax1.grid()

lines = lines1 + lines2
ax1.legend(lines, [line.get_label() for line in lines], fontsize="large")

plt.savefig("comoving_distance_and_time.png", dpi=300, bbox_inches="tight")
plt.show()
input()