from pathlib import Path

import matplotlib
import matplotlib.pyplot as plt
import numpy as np
from matplotlib import colormaps

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

stations = ["North", "West", "Central", "South"]
days = ["Mon", "Tue", "Wed", "Thu", "Fri", "Sat", "Sun"]
output = Path("colormap-set.png")

cmap = colormaps["viridis"]

fig, ax = plt.subplots(figsize=(6.4, 3.8), layout="constrained")
image = ax.imshow(temperatures, cmap=cmap, vmin=15, vmax=31)
ax.set_xticks(range(len(days)), days)
ax.set_yticks(range(len(stations)), stations)
ax.set_xlabel("Day")
ax.set_title("Weekly station temperature")
colorbar = fig.colorbar(image, ax=ax, label="Temperature (deg C)")

fig.savefig(output, dpi=160)

print(f"matplotlib {matplotlib.__version__}")
print(f"colormap: {image.get_cmap().name}")
print(f"colorbar label: {colorbar.ax.get_ylabel()}")
print(f"saved: {output}")