import numpy as np
import matplotlib.pyplot as plt

# General parabola in R^2: gamma(t) = p + u t + 0.5 a t^2
p = np.array([0.0, 0.0])
u = np.array([1.0, -0.2])
a = np.array([0.2, 2.0])

def gamma(t):
    return p + u * t + 0.5 * a * t**2

def gamma_dot(t):
    return u + a * t

def tangent_intersection(ti, tj):
    Pi = gamma(ti)
    Pj = gamma(tj)
    Vi = gamma_dot(ti)
    Vj = gamma_dot(tj)
    M = np.column_stack((Vi, -Vj))
    s = np.linalg.solve(M, Pj - Pi)
    return Pi + s[0] * Vi

def tangent_segment(t, s_min=-1.5, s_max=1.5):
    s_vals = np.linspace(s_min, s_max, 2)
    base = gamma(t).reshape(2, 1)
    dir_vec = gamma_dot(t).reshape(2, 1)
    pts = base + dir_vec * s_vals
    return pts[0], pts[1]

# Tangency parameters (arbitrary)
t0, t1, t2 = -1.0, 0.1, 1.2
P0, P1, P2 = gamma(t0), gamma(t1), gamma(t2)
I01 = tangent_intersection(t0, t1)
I12 = tangent_intersection(t1, t2)
I02 = tangent_intersection(t0, t2)

fig, ax = plt.subplots(figsize=(6, 4))

# Parabola plot
ts = np.linspace(t0 - 0.7, t2 + 0.7, 400)
curve = np.array([gamma(t) for t in ts])
ax.plot(curve[:, 0], curve[:, 1], linewidth=2)

# Tangent lines
for t in [t0, t1, t2]:
    xL, yL = tangent_segment(t)
    ax.plot(xL, yL, linestyle="--", linewidth=1)

# Triangles
tri01 = np.vstack([P0, P1, I01])
tri12 = np.vstack([P1, P2, I12])
tri02 = np.vstack([P0, P2, I02])

ax.fill(tri01[:, 0], tri01[:, 1], color="#4a5a90", alpha=0.75)
ax.fill(tri12[:, 0], tri12[:, 1], color="#4a5a90", alpha=0.75)
ax.plot(tri02[[0, 1, 2, 0], 0], tri02[[0, 1, 2, 0], 1], linewidth=1.25)

# Tangency points
ax.scatter([P0[0], P1[0], P2[0]], [P0[1], P1[1], P2[1]], s=25)

# Equation text
ax.text(
    0.5,
    0.64,
    r"$A_{01}^{1/3} + A_{12}^{1/3} = A_{02}^{1/3}$",
    transform=ax.transAxes,
    ha="center",
    va="bottom"
)

ax.set_aspect("equal", adjustable="box")
ax.axis("off")

# Tight bounds
all_pts = np.vstack([curve, P0, P1, P2, I01, I12, I02])
pad = 0.15 * (all_pts[:, 0].max() - all_pts[:, 0].min())
ax.set_xlim(all_pts[:, 0].min() - pad, all_pts[:, 0].max() + pad)
pad_y = 0.15 * (all_pts[:, 1].max() - all_pts[:, 1].min())
ax.set_ylim(all_pts[:, 1].min() - pad_y, all_pts[:, 1].max() + pad_y)

fig.savefig("affine_arclength_parabola.svg")
fig.show()