from pathlib import Path

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


output = Path("error-bar-add.png")

days = np.array([1, 2, 3, 4, 5])
mean_latency = np.array([42, 39, 36, 38, 35])
latency_error = np.array([2.5, 2.0, 1.8, 2.2, 1.6])

batch_size = np.array([20, 40, 60, 80, 100])
throughput = np.array([120, 168, 205, 236, 251])
batch_error = np.array([3, 4, 5, 5, 6])
lower_error = np.array([9, 12, 15, 13, 14])
upper_error = np.array([14, 16, 20, 18, 21])
asymmetric_error = np.vstack([lower_error, upper_error])

fig, (ax_left, ax_right) = plt.subplots(
    ncols=2,
    figsize=(8.2, 3.6),
    layout="constrained",
)

symmetric = ax_left.errorbar(
    days,
    mean_latency,
    yerr=latency_error,
    fmt="o-",
    capsize=4,
    elinewidth=1.4,
    ecolor="0.25",
    label="Mean latency",
)
ax_left.set_title("Symmetric y error")
ax_left.set_xlabel("Test day")
ax_left.set_ylabel("Latency (ms)")
ax_left.grid(True, axis="y", alpha=0.25)
ax_left.legend()

asymmetric = ax_right.errorbar(
    batch_size,
    throughput,
    xerr=batch_error,
    yerr=asymmetric_error,
    fmt="s-",
    capsize=4,
    elinewidth=1.4,
    ecolor="0.25",
    color="tab:green",
    label="Throughput",
)
ax_right.set_title("Asymmetric y and x error")
ax_right.set_xlabel("Batch size")
ax_right.set_ylabel("Rows per second")
ax_right.grid(True, axis="y", alpha=0.25)
ax_right.legend()

fig.savefig(output, dpi=160)
plt.close(fig)

symmetric_caplines = symmetric.lines[1]
asymmetric_caplines = asymmetric.lines[1]

print(f"matplotlib {matplotlib.__version__}")
print(f"symmetric points: {len(days)}")
print(f"symmetric cap artists: {len(symmetric_caplines)}")
print(f"asymmetric shape: {asymmetric_error.shape}")
print(f"x error points: {len(batch_error)}")
print(f"asymmetric cap artists: {len(asymmetric_caplines)}")
print(f"saved: {output}")
print(f"bytes: {output.stat().st_size}")