import numpy as np
import matplotlib.pyplot as plt

# Generic helper: build triangles from a param curve and its tangent
def build_triangles(curve_func, tangent_func, t_samples, eps=1e-8):
    pts = curve_func(t_samples)
    tans = tangent_func(t_samples)
    triangles = []
    for i in range(len(t_samples) - 1):
        Pn = pts[i]
        Pn1 = pts[i + 1]
        Tn = tans[i]
        Tn1 = tans[i + 1]

        A = np.column_stack((Tn, -Tn1))
        b = Pn1 - Pn
        det = np.linalg.det(A)
        if abs(det) < eps:
            continue
        s_u = np.linalg.solve(A, b)
        s = s_u[0]
        I = Pn + s * Tn
        triangles.append((Pn, Pn1, I))
    return pts, triangles

# Specific immersed cusp: y^2 = x^3 + x^2, parameterized by line y = t x
def cusp_curve(t):
    x = t**2 - 1
    y = t * (t**2 - 1)
    return np.vstack((x, y)).T

def cusp_tangent(t):
    dx = 2 * t
    dy = 3 * t**2 - 1
    return np.vstack((dx, dy)).T

t_min, t_max, n_samples = -1.2, 1.2, 9
t_values = np.linspace(t_min, t_max, n_samples)

points, triangles = build_triangles(cusp_curve, cusp_tangent, t_values)

t_dense = np.linspace(t_min, t_max, 400)
curve_dense = cusp_curve(t_dense)

fig, ax = plt.subplots(figsize=(5, 5))

# Plot curve
ax.plot(curve_dense[:, 0], curve_dense[:, 1], linewidth=2)

# Plot filled triangles
for Pn, Pn1, I in triangles:
    tri = np.vstack((Pn, Pn1, I))
    ax.fill(tri[:, 0], tri[:, 1], color="#4a5a90", alpha=0.7)

# Sample points
ax.scatter(points[:, 0], points[:, 1], s=15)

# Auto range from curve + triangles, tight around loop
all_xy = [curve_dense]
for Pn, Pn1, I in triangles:
    all_xy.append(np.vstack((Pn, Pn1, I)))
all_xy = np.vstack(all_xy)
xmin, ymin = np.min(all_xy, axis=0)
xmax, ymax = np.max(all_xy, axis=0)
dx, dy = xmax - xmin, ymax - ymin
margin_x = 0.15 * dx
margin_y = 0.15 * dy
ax.set_xlim(xmin - margin_x, xmax + margin_x)
ax.set_ylim(ymin - margin_y, ymax + margin_y)

ax.set_aspect("equal", adjustable="box")
ax.axis("off")
fig.savefig("affine_arclength_elliptic_curve.svg")
fig.show()