import numpy as np
import matplotlib.pyplot as plt

# Points
origin = np.array([0.0, 0.0])
y0 = np.array([3.0, -0.9])
v = np.array([4.0, 1.5])
v_minus_y0 = v - y0

# Parallelogram vertices
parallelogram = np.array([origin, y0, v, v_minus_y0, origin])

# Superellipse parameters
n = 4          # exponent
pow_factor = 2.0 / n   # 0.5

t = np.linspace(0, 2*np.pi, 500)

def superellipse(center, through, n):
    dx, dy = through - center
    R = (abs(dx)**n + abs(dy)**n)**(1.0/n)
    x = center[0] + R * np.sign(np.cos(t)) * np.abs(np.cos(t))**pow_factor
    y = center[1] + R * np.sign(np.sin(t)) * np.abs(np.sin(t))**pow_factor
    return x, y

# Superellipses
x_super1, y_super1 = superellipse(v, y0, n)          # around v passing y0
x_super2, y_super2 = superellipse(v_minus_y0, origin, n)  # around v-y0 passing origin

# Plot
fig, ax = plt.subplots(figsize=(6, 6))

# Line representing Y (slope -0.3)
x_line = np.linspace(-1, 6, 200)
ax.plot(x_line, -0.3 * x_line, linestyle='--', linewidth=1)

# Parallelogram
ax.plot(parallelogram[:, 0], parallelogram[:, 1], linewidth=1)

# Line from 0 to v
ax.plot([origin[0], v[0]], [origin[1], v[1]], linewidth=1)

# Superellipses
ax.plot(x_super1, y_super1, linewidth=1)
ax.plot(x_super2, y_super2, linewidth=1)

# Points
ax.scatter(*origin, s=40)
ax.scatter(*y0, s=40)
ax.scatter(*v, s=40)
ax.scatter(*v_minus_y0, s=40)

# Labels
ax.text(*origin, r'$0$', fontsize=12, verticalalignment='top', horizontalalignment='right')
ax.text(*y0, r'$y_0$', fontsize=12, verticalalignment='bottom', horizontalalignment='right')
ax.text(*v, r'$v$', fontsize=12, verticalalignment='bottom', horizontalalignment='left')
ax.text(*v_minus_y0, r'$v-y_0$', fontsize=12, verticalalignment='bottom', horizontalalignment='left')

ax.set_aspect('equal')
ax.axis('off')
plt.tight_layout()
plt.savefig('Riesz_lemma.svg')
plt.show()