import numpy as np
import matplotlib.pyplot as plt
from matplotlib.gridspec import SubplotSpec, GridSpec

# Create the figure and subplots using subplots_mosaic
fig = plt.figure(constrained_layout=True)
spec = GridSpec(1, 2, figure=fig, width_ratios=[1, 1])

# Diagram A: Unit Circle
ax_A = fig.add_subplot(spec[0, 0])
t = np.linspace(0, 2*np.pi, 100)
e_1 = np.array([1, 0])
e_2 = np.array([0, 1])
unit_circle = np.cos(t)[:, np.newaxis] * e_1 + np.sin(t)[:, np.newaxis] * e_2
ax_A.plot(unit_circle[:, 0], unit_circle[:, 1])
ax_A.quiver(0, 0, e_1[0], e_1[1], angles='xy', scale_units='xy', scale=1, color='red', label=r'$e_1$')
ax_A.quiver(0, 0, e_2[0], e_2[1], angles='xy', scale_units='xy', scale=1, color='blue', label=r'$e_2$')
ax_A.grid()

ax_A.set_aspect('equal', adjustable='datalim')
ax_A.set_title('Unit Circle')
ax_A.legend()

# Diagram B: Ellipse
ax_B = fig.add_subplot(spec[0, 1])
v_1 = np.array([1.3, 0])
v_2 = np.array([2, 1])
ellipse = np.cos(t)[:, np.newaxis] * v_1 + np.sin(t)[:, np.newaxis] * v_2
ax_B.plot(ellipse[:, 0], ellipse[:, 1])
ax_B.quiver(0, 0, v_1[0], v_1[1], angles='xy', scale_units='xy', scale=1, color='red', label=r'$v_2$')
ax_B.quiver(0, 0, v_2[0], v_2[1], angles='xy', scale_units='xy', scale=1, color='blue', label=r'$v_2$')
ax_B.grid()

ax_B.set_aspect('equal', adjustable='datalim')
ax_B.set_title('Ellipse')
ax_B.legend()
plt.savefig('cholesky_ellipse.svg')
plt.show()