from pathlib import Path

import matplotlib

matplotlib.use("Agg")

import matplotlib.image as mpimg
import pandas as pd


sales = pd.DataFrame(
    {
        "quarter": ["2026 Q1", "2026 Q2", "2026 Q3", "2026 Q4"],
        "hardware": [18000, 21500, 23100, 26000],
        "software": [9500, 12000, 14800, 17100],
    }
)

plot_data = sales.set_index("quarter")[["hardware", "software"]]

ax = plot_data.plot(
    kind="bar",
    figsize=(7, 4),
    rot=0,
    title="Quarterly revenue",
)
ax.set_xlabel("Quarter")
ax.set_ylabel("Revenue (USD)")
ax.legend(title="Segment")

figure = ax.get_figure()
figure.set_layout_engine("constrained")

series = ax.get_legend_handles_labels()[1]
bar_count = len(ax.patches)

assert plot_data.index.name == "quarter"
assert series == ["hardware", "software"]
assert bar_count == len(plot_data) * len(series)

output = Path("sales-by-quarter.png")
figure.savefig(output, dpi=150)

saved_image = mpimg.imread(output)
height, width = saved_image.shape[:2]

assert (width, height) == (1050, 600)

print(f"index: {plot_data.index.name}")
print(f"series: {', '.join(series)}")
print(f"bars: {bar_count}")
print(f"image: {width} x {height}")
print(f"saved: {output}")