import numpy as np
import matplotlib.pyplot as plt

def plot_affine_normals(curve_func, t0, t1, t_samples, d,
                        n_curve=400,
                        eps=1e-8):
    t_samples = np.asarray(t_samples, dtype=float)

    # --- dense curve for backbone and scaling ---
    t_dense = np.linspace(t0, t1, n_curve)
    curve_dense = np.asarray(curve_func(t_dense))
    if curve_dense.ndim == 1:
        curve_dense = curve_dense.reshape(-1, 2)

    # --- sample points for lines and normals ---
    pts = np.asarray(curve_func(t_samples))
    if pts.ndim == 1:
        pts = pts.reshape(-1, 2)

    # numerical derivatives wrt t for affine normal
    x = pts[:, 0]
    y = pts[:, 1]

    dx = np.gradient(x, t_samples, edge_order=2)
    dy = np.gradient(y, t_samples, edge_order=2)
    ddx = np.gradient(dx, t_samples, edge_order=2)
    ddy = np.gradient(dy, t_samples, edge_order=2)
    dddx = np.gradient(ddx, t_samples, edge_order=2)
    dddy = np.gradient(ddy, t_samples, edge_order=2)

    # tangent (Euclidean)
    T = np.vstack((dx, dy)).T
    T_norm = np.linalg.norm(T, axis=1, keepdims=True)
    T_norm[T_norm < eps] = 1.0
    T_unit = T / T_norm

    # affine normal direction (unnormalized)
    # Δ = [γ', γ''], Δ' = d/dt Δ
    Delta = dx * ddy - dy * ddx
    Delta_p = np.gradient(Delta, t_samples, edge_order=2)

    N_aff = np.vstack(
        (Delta * ddx - (1.0 / 3.0) * Delta_p * dx,
         Delta * ddy - (1.0 / 3.0) * Delta_p * dy)
    ).T

    N_aff_norm = np.linalg.norm(N_aff, axis=1, keepdims=True)
    mask_aff = N_aff_norm[:, 0] > eps
    N_aff_unit = np.zeros_like(N_aff)
    N_aff_unit[mask_aff] = N_aff[mask_aff] / N_aff_norm[mask_aff]

    # Euclidean normal for "inward" offset
    N_e = np.vstack((-T_unit[:, 1], T_unit[:, 0])).T

    # geometric scale for line lengths
    all_xy = curve_dense
    xmin, ymin = np.min(all_xy, axis=0)
    xmax, ymax = np.max(all_xy, axis=0)
    dx_box, dy_box = xmax - xmin, ymax - ymin
    scale = max(dx_box, dy_box)
    if scale < eps:
        scale = 1.0

    line_half_len = 0.18 * scale
    normal_half_len = 0.22 * scale

    # center for "inward" direction
    center = np.mean(all_xy, axis=0)

    # sign for inward offset
    to_center = center[None, :] - pts
    dot_nc = np.sum(N_e * to_center, axis=1)
    sign = np.sign(dot_nc)
    sign[sign == 0] = 1.0
    sign = sign[:, None]

    offset1 = sign * d * N_e
    offset2 = sign * (2 * d) * N_e

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

    # main curve
    ax.plot(curve_dense[:, 0], curve_dense[:, 1], linewidth=2, color="#4a5a90")

    # tangents and parallels
    for P, Tu, off1, off2 in zip(pts, T_unit, offset1, offset2):
        # base tangent
        P0 = P - line_half_len * Tu
        P1 = P + line_half_len * Tu
        ax.plot([P0[0], P1[0]], [P0[1], P1[1]], linewidth=1.0, alpha=0.9, color="#938fba")

        # offset by d
        Q = P + off1
        Q0 = Q - line_half_len * Tu
        Q1 = Q + line_half_len * Tu
        ax.plot([Q0[0], Q1[0]], [Q0[1], Q1[1]], linewidth=1.0, alpha=0.5, color="#938fba")

        # offset by 2d
        R = P + off2
        R0 = R - line_half_len * Tu
        R1 = R + line_half_len * Tu
        ax.plot([R0[0], R1[0]], [R0[1], R1[1]], linewidth=1.0, alpha=0.3, color="#938fba")

    # affine normal lines
    for P, Na, ok in zip(pts, N_aff_unit, mask_aff):
        if not ok:
            continue
        A0 = P - 0 * normal_half_len * Na
        A1 = P + 6 * normal_half_len * Na
        ax.plot([A0[0], A1[0]], [A0[1], A1[1]], linewidth=1.0, alpha=0.8, color="#4f0808")

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

    # auto range with margin
    all_xy = np.vstack((all_xy, pts))
    xmin, ymin = np.min(all_xy, axis=0)
    xmax, ymax = np.max(all_xy, axis=0)
    dx_box, dy_box = xmax - xmin, ymax - ymin
    margin_x = 0.15 * dx_box if dx_box > 0 else 1.0
    margin_y = 0.15 * dy_box if dy_box > 0 else 1.0
    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")

    return fig, ax

def three_lobes(t, epsilon=0.02):
    # r = 1 + epsilon * cos(3 t)
    r = 1.0 + epsilon * np.cos(3 * t)
    x = r * np.cos(t)
    y = r * np.sin(t)
    return np.vstack((x, y)).T

fig, ax = plot_affine_normals(three_lobes, 0, 2*np.pi, np.linspace(0, 2*np.pi, 57), 0.02)
fig.savefig("affine_focal_set_three_lobes.svg")
fig.show()