import matplotlib.pyplot as plt
from io import StringIO
from matplotlib.patches import Rectangle

# LoRA on one weight matrix: the pre-trained W (d x d) stays frozen; only the correction B·A is
# trained, with B of shape d x r and A of shape r x d. For d = 4096 and r = 8 that is
# 2 x 4096 x 8 = 65,536 trained numbers against 16,777,216 in W. At true scale B and A would be
# 8/4096 of W's side (a fraction of a pixel), so their thickness is drawn THICK times larger.
BLUE, ORANGE, GREEN = '#58a6ff', '#f0883e', '#3fb950'
d, r = 4096, 8
THICK = 40
S = 4.0                                  # side of W in the drawing
t = S * r / d * THICK                    # drawn thickness of B and A

plt.rcParams.update({'text.color': 'gray'})
fig, ax = plt.subplots(figsize=(9.6, 4.6))

# W, frozen
ax.add_patch(Rectangle((0, 0), S, S, fc=BLUE, alpha=0.18, ec=BLUE, lw=1.4))
ax.text(S / 2, S / 2 + 0.25, 'W', ha='center', va='center', fontsize=26, color=BLUE, fontweight='bold')
ax.text(S / 2, S / 2 - 0.45, f'{d:,} × {d:,}', ha='center', va='center', fontsize=13)
ax.text(S / 2, -0.45, 'pre-trained, frozen', ha='center', fontsize=13, color=BLUE)
ax.text(S / 2, -0.95, f'{d * d:,} numbers', ha='center', fontsize=13, color=BLUE, fontweight='bold')

ax.text(S + 0.55, S / 2, '+', ha='center', va='center', fontsize=30)

# B (d x r): as tall as W, thin
xb = S + 1.1
ax.add_patch(Rectangle((xb, 0), t, S, fc=GREEN, alpha=0.55, ec=GREEN, lw=1.4))
ax.text(xb + t / 2, S + 0.3, 'B', ha='center', fontsize=20, color=GREEN, fontweight='bold')
ax.text(xb - 0.12, 0.35, f'{d:,} × {r}', ha='right', va='bottom', fontsize=12, rotation=90)

ax.plot(xb + t + 0.5, S - t / 2, 'o', color='gray', ms=5)                    # the product B·A

# A (r x d): as wide as W, flat
xa = xb + t + 1.0
ax.add_patch(Rectangle((xa, S - t), S, t, fc=GREEN, alpha=0.55, ec=GREEN, lw=1.4))
ax.text(xa + S / 2, S + 0.3, 'A', ha='center', fontsize=20, color=GREEN, fontweight='bold')
ax.text(xa + S / 2, S - t - 0.4, f'{r} × {d:,}', ha='center', fontsize=12)

n_lora = d * r + r * d
cx = xa + S / 2
ax.text(cx, 1.75, 'trained', ha='center', fontsize=13, color=GREEN)
ax.text(cx + 0.2, 1.15, f'{d:,}·{r} + {r}·{d:,} = {n_lora:,} numbers', ha='center', fontsize=12.5, color=GREEN, fontweight='bold')
ax.text(cx, 0.55, f'{n_lora / (d * d):.2%} of W; the product B·A has rank ≤ {r}', ha='center', fontsize=12.5)
ax.text(cx, -0.45, 'B starts at zero, so training starts', ha='center', fontsize=12, style='italic')
ax.text(cx, -0.95, 'exactly at the pre-trained model', ha='center', fontsize=12, style='italic')

ax.text(cx, -1.7, f'output = W·x + B·(A·x)', ha='center', fontsize=12.5)
ax.text(S / 2, -1.7, f'B and A drawn {THICK}× thicker than to scale', ha='center', fontsize=11, style='italic')

ax.set_xlim(-0.3, xa + S + 0.3)
ax.set_ylim(-2.0, S + 0.8)
ax.set_aspect('equal')
ax.axis('off')
plt.tight_layout()

buffer = StringIO()
plt.savefig(buffer, format="svg", transparent=True)
print(buffer.getvalue())

plt.close()