request.available_collections("EFI_IDM", details=False)

print(request.available_measurements("EFI_IDM"))

idm_vars = [
    # Positional information in geodetic (GD) and geocentric (GC) frames
    #  redundant with VirES variables Latitude, Longitude, Radius (in geocentric frame)
    "Latitude_GD", "Longitude_GD", "Height_GD", "Radius_GC",
    # Quasi-dipole magnetic latitude and local time
    #  redundant with VirES auxiliaries, QDLat & MLT
    "Latitude_QD", "MLT_QD",
    # Satellite velocity in NEC frame
    "V_sat_nec",
    # Estimated ion effective mass, uncertainty, and validity flags
    "M_i_eff", "M_i_eff_err", "M_i_eff_Flags",
    # Effective masses from ruhlik et al. (2015) topside empirical model
    "M_i_eff_tbt_model",
    # Along-track ion drift velocity, uncertainty, validity flags, and velocity without detrending
    "V_i", "V_i_err", "V_i_Flags", "V_i_raw",
    # Ion density, uncertainty, and validity flags
    "N_i", "N_i_err", "N_i_Flags",
    # Modified-OML faceplate area, and Langmuir spherical probe radius
    "A_fp", "R_p",
    # Electron temperature, and spacecraft floating potential
    "T_e", "Phi_sc",
]

start = dt.datetime(2016, 1, 2)
end = dt.datetime(2016, 1, 3)

request = SwarmRequest(SERVER_URL)
request.set_collection("SW_PREL_EFIAIDM_2_")
request.set_products(
    measurements=idm_vars,
    auxiliaries=["OrbitNumber", "OrbitDirection", "MLT"]
)
data = request.get_between(start, end)

df = data.as_dataframe()
df.head()

ds = data.as_xarray()
ds

fig, axes = plt.subplots(nrows=4, sharex=True, figsize=(10, 7))
axes_r = [ax.twinx() for ax in axes]
# Plot quantity (left axis) and error (right axis) for each quantity
ds.plot.scatter(x="Timestamp", y="M_i_eff", ax=axes[0], s=1, linewidths=0)
ds.plot.scatter(x="Timestamp", y="M_i_eff_err", ax=axes_r[0], s=0.1, color="tab:orange")
ds.plot.scatter(x="Timestamp", y="V_i", ax=axes[1], s=1, linewidths=0)
ds.plot.scatter(x="Timestamp", y="V_i_err", ax=axes_r[1], s=0.1, color="tab:orange")
ds.plot.scatter(x="Timestamp", y="N_i", ax=axes[2], s=1, linewidths=0)
ds.plot.scatter(x="Timestamp", y="N_i_err", ax=axes_r[2], s=0.1, color="tab:orange")
ds.plot.scatter(x="Timestamp", y="T_e", ax=axes[3], s=1, linewidths=0)
fig.subplots_adjust(hspace=0)

# Add legend to identify each side
blue = mpl.patches.Patch(color="tab:blue", label="Quantities")
orange = mpl.patches.Patch(color="tab:orange", label="Errors")
axes[0].legend(handles=[blue, orange])

# # Generate additional ticklabels for x-axis
# Use time xticks to get dataset vars at those xticks
locx = axes[-1].get_xticks()
times = mpl.dates.num2date(locx)
times = [t.replace(tzinfo=None) for t in times]
_ds_xticks = ds.reindex({"Timestamp": times}, method="nearest")
# Build ticklabels from dataset vars
xticklabels = np.stack([
    _ds_xticks["Timestamp"].dt.strftime("%H:%M").values,
    np.round(_ds_xticks["Latitude"].values, 2).astype(str),
    np.round(_ds_xticks["Longitude"].values, 2).astype(str),
])
xticklabels = ["\n".join(row) for row in xticklabels.T]
# Add labels to first xtick
_xt0 = xticklabels[0].split("\n")
xticklabels[0] = f"Time:  {_xt0[0]}\nLat:  {_xt0[1]}\nLon: {_xt0[2]}"
axes[-1].set_xticks(axes[-1].get_xticks())
axes[-1].set_xticklabels(xticklabels)
axes[-1].set_xlabel("")
# Adjust title
sources = "\n".join([i for i in ds.attrs["Sources"] if "IDM" in i])
title = "".join([
    f"Swarm {ds['Spacecraft'].data[0]} ion effective mass and drift, ",
    ds["Timestamp"].dt.date.data[0].isoformat(),
    f"\n{sources}"
])
fig.suptitle(title);