import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import brentq

def plot_affine_construction(curve_func, t_p, t_view_min, t_view_max, num_lines=30, line_spacing=0.002, search_width=0.3):
    t_dense = np.linspace(t_view_min, t_view_max, 600)
    curve_dense = np.array([curve_func(t) for t in t_dense])
    
    P = np.array(curve_func(t_p))
    
    dt = 1e-5
    P_plus = np.array(curve_func(t_p + dt))
    P_minus = np.array(curve_func(t_p - dt))
    
    T = P_plus - P_minus
    T_norm = np.linalg.norm(T)
    T_unit = T / T_norm
    
    Acc = (P_plus - 2*P + P_minus) / (dt**2)
    
    N = np.array([-T_unit[1], T_unit[0]])
    
    if np.dot(Acc, N) < 0:
        N = -N

    midpoints = [P]
    lines_to_plot = []
    
    def elevation_func(t, h):
        pt = np.array(curve_func(t))
        dist = np.dot(pt - P, N)
        return dist - h

    for i in range(1, num_lines + 1):
        h = i * line_spacing
        
        t_left = None
        for step in np.linspace(0.001, search_width, 50):
            val = elevation_func(t_p - step, h)
            if val > 0:
                t_left = t_p - step
                break
                
        t_right = None
        for step in np.linspace(0.001, search_width, 50):
            val = elevation_func(t_p + step, h)
            if val > 0:
                t_right = t_p + step
                break

        if t_left is not None and t_right is not None:
            try:
                t1 = brentq(elevation_func, t_left, t_p, args=(h,))
                t2 = brentq(elevation_func, t_p, t_right, args=(h,))
                
                p1 = np.array(curve_func(t1))
                p2 = np.array(curve_func(t2))
                
                lines_to_plot.append((p1, p2))
                midpoints.append((p1 + p2) / 2.0)
            except ValueError:
                break
        else:
            break

    midpoints = np.array(midpoints)

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

    ax.plot(curve_dense[:, 0], curve_dense[:, 1], color="#4a5a90", linewidth=2)
    ax.scatter(P[0], P[1], color="black", s=50, zorder=5)

    tan_len = search_width
    tan_start = P - tan_len * T_unit
    tan_end   = P + tan_len * T_unit
    ax.plot([tan_start[0], tan_end[0]], [tan_start[1], tan_end[1]], 'k--', alpha=0.4, linewidth=1)

    for (p1, p2) in lines_to_plot:
        ax.plot([p1[0], p2[0]], [p1[1], p2[1]], color="#938fba", linewidth=1, alpha=0.6)
        ax.scatter([p1[0], p2[0]], [p1[1], p2[1]], color="#938fba", s=10, alpha=0.6)

    if len(midpoints) > 1:
        ax.plot(midpoints[:, 0], midpoints[:, 1], color="#4a5a90", linewidth=2, marker='o', markersize=3)

    # ax.set_aspect("equal")
    
    all_pts = np.vstack((midpoints, *lines_to_plot)) if lines_to_plot else midpoints
    if len(all_pts) > 1:
        xmin, ymin = np.min(all_pts, axis=0)
        xmax, ymax = np.max(all_pts, axis=0)
        dx = xmax - xmin
        dy = ymax - ymin
        margin_x = max(dx * 0.2, 0.01)
        margin_y = max(dy * 0.2, 0.005)
        ax.set_xlim(xmin - margin_x, xmax + margin_x)
        ax.set_ylim(ymin - margin_y, ymax + margin_y)
        # ax.set_xlim(xmin , xmax )
        # ax.set_ylim(ymin , ymax )
    else:
        ax.set_xlim(-0.1, 0.1)
        ax.set_ylim(0, 0.1)
    return fig, ax

def specific_curve(t):
    x = t
    y = t**2 - 3 * t**3
    return np.array([x, y])

fig, ax = plot_affine_construction(specific_curve, t_p=0.0, t_view_min=-0.15, t_view_max=0.25, num_lines=50, line_spacing=0.0005, search_width=0.3)

ys = np.linspace(0, 0.02, 200)
xs = 3/2 * ys
ax.plot(xs, ys, linestyle='--', alpha=0.5, color="#938fba")
fig.savefig("affine_midpoints_cubic.svg")
fig.show()