from pathlib import Path

import matplotlib

matplotlib.use("Agg")

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.colors import Normalize


measurements = np.array(
    [
        [18, 20, 23, 26, 28],
        [17, 21, 24, 27, 30],
        [16, 19, 22, 25, 29],
        [15, 18, 21, 24, 27],
    ]
)
norm = Normalize(vmin=15, vmax=30)

fig, ax = plt.subplots(figsize=(6, 3.8), layout="constrained")
image = ax.imshow(measurements, cmap="viridis", norm=norm)

ax.set_title("Sensor temperature by rack")
ax.set_xlabel("Rack column")
ax.set_ylabel("Rack row")
ax.set_xticks(range(measurements.shape[1]))
ax.set_yticks(range(measurements.shape[0]))

colorbar = fig.colorbar(image, ax=ax)
colorbar.set_label("Temperature (deg C)")
colorbar.set_ticks([15, 20, 25, 30])

lower, upper = image.get_clim()
ticks = colorbar.get_ticks()
label = colorbar.ax.get_ylabel()

assert (lower, upper) == (15.0, 30.0)
np.testing.assert_allclose(ticks, [15, 20, 25, 30])
assert label == "Temperature (deg C)"

output = Path("colorbar-add.png")
fig.savefig(output, dpi=160)
plt.close(fig)

rendered = plt.imread(output)

print(f"colorbar limits: {lower:.1f} to {upper:.1f}")
print("colorbar ticks:", ", ".join(f"{tick:g}" for tick in ticks))
print(f"colorbar label: {label}")
print(f"saved: {output}")
print(f"image shape: {rendered.shape[0]} x {rendered.shape[1]} x {rendered.shape[2]}")