import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import quad
import matplotlib.cm as cm

def get_theta(t, x0):
    h = np.cos(np.pi * (t - 1))
    
    # Solve A*cos(theta) + B*sin(theta) = h
    A = t
    B = -x0
    R = np.sqrt(A**2 + B**2)
    delta = np.arctan2(B, A)  # Phase shift
    
    # Clip h/R to [-1, 1] to avoid numerical domain errors close to boundaries
    val = np.clip(h / R, -1.0, 1.0)
    theta = delta + np.arccos(val)
    return theta

def get_dtheta_dt(t, theta, x0):
    h_prime = -np.pi * np.sin(np.pi * (t - 1))
    numerator = np.cos(theta) - h_prime
    denominator = t * np.sin(theta) + x0 * np.cos(theta)
    return numerator / denominator

def get_envelope_point(t, x0):
    theta = get_theta(t, x0)
    dtheta = get_dtheta_dt(t, theta, x0)
    
    h = np.cos(np.pi * (t - 1))
    h_prime = -np.pi * np.sin(np.pi * (t - 1))
    
    # Derived from system of equations for envelope:
    # x*sin(theta) - y*cos(theta) = -h
    # x*cos(theta) + y*sin(theta) = -h' / theta'
    
    u = -h
    v = -h_prime / dtheta
    
    x = v * np.cos(theta) + u * np.sin(theta)
    y = v * np.sin(theta) - u * np.cos(theta)
    
    return x, y, theta

def plot_line(ax, x0, y0, theta, length_forward=2, length_backward=1, **kwargs):
    dx = np.cos(theta)
    dy = np.sin(theta)
    ax.plot(
        [x0 - length_backward*dx, x0 + length_forward*dx],
        [y0 - length_backward*dy, y0 + length_forward*dy],
        **kwargs
    )

def plotter(x0):
    # --- Configuration ---
    num_points = 500
    t_values = np.linspace(0, 1, num_points)
    
    # --- 1. Calculate Envelope (Caustic) ---
    env_xs = []
    env_ys = []
    thetas = []
    
    for t in t_values:
        ex, ey, th = get_envelope_point(t, x0)
        env_xs.append(ex)
        env_ys.append(ey)
        thetas.append(th)
    
    env_xs = np.array(env_xs)
    env_ys = np.array(env_ys)
    thetas = np.array(thetas)
    
    # --- 2. Calculate Arc Length of Envelope for Involutes ---
    d_env_x = np.gradient(env_xs, t_values)
    d_env_y = np.gradient(env_ys, t_values)
    ds_dt = np.sqrt(d_env_x**2 + d_env_y**2)
    
    # Integrate ds to get arc length s(t)
    s_values = np.zeros_like(t_values)
    s_values[1:] = np.cumsum(
        0.5 * (ds_dt[1:] + ds_dt[:-1]) * (t_values[1] - t_values[0])
    )
    
    # --- Plotting ---
    fig, ax = plt.subplots(figsize=(12, 12))
    
    # Plot Unit Circle X
    circle = plt.Circle((0, 0), 1, color='gray', fill=False, linestyle='--', alpha=0.5, linewidth=1.5)
    ax.add_patch(circle)
    
    # Plot Line Segment Y
    ax.plot([x0, x0], [0, 1], color='black', linewidth=3, label='Segment Y')
    
    # Plot the Family of Lines (Tangents to Caustic)
    num_lines = 50
    indices = np.linspace(0, num_points-1, num_lines, dtype=int)
    for i in indices:
        t = t_values[i]
        theta = thetas[i]
        color = cm.viridis(t)
        plot_line(
            ax, x0, t, theta,
            length_forward=5, length_backward=5,
            color=color, linewidth=0.5, alpha=0.6
        )
    
    # Plot the Envelope (Caustic)
    ax.plot(env_xs, env_ys, color='#4a5a90', linewidth=1.5, label='Caustic (Evolute)')
    
    # Plot Involutes (currently computed but not drawn)
    num_involutes = 15
    l_range = np.linspace(-1, 3, num_involutes)
    
    for l0 in l_range:
        inv_xs = env_xs + (l0 - s_values) * np.cos(thetas)
        inv_ys = env_ys + (l0 - s_values) * np.sin(thetas)
        # ax.plot(inv_xs, inv_ys, color='green', linewidth=0.8, alpha=0.8)
    
    # Setup Plot Limits and Style
    ax.set_title(f"Caustic for $x_0={x0:.02f}$")
    ax.set_xlim(-2.5, 3.5)
    ax.set_ylim(-2.5, 3.5)
    ax.legend(loc='upper right')
    ax.grid(True, alpha=0.3)
    fig.tight_layout()
    
    # Return the figure (and axes) so caller can do fig.savefig(...)
    return fig

import os
from tqdm import tqdm
movie_name = "monge_caustic"
dir_path = f"./{movie_name}"
tmax = 6
frames_per_t = 24
t_values = 1.5 + 0.5 * np.cos(np.linspace(0, 2 * np.pi, frames_per_t * tmax))

for N, t in tqdm(enumerate(t_values)):
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)

    plotter(t).savefig(f"{dir_path}/{N:03d}.png",bbox_inches='tight')
    plt.close()

import os
from tqdm import tqdm
movie_name = "monge_caustic"
dir_path = f"./{movie_name}"
tmax = 6
frames_per_t = 24
t_values = 1.5 + 0.5 * np.cos(np.linspace(0, 2 * np.pi, frames_per_t * tmax))

for N, t in tqdm(enumerate(t_values)):
    if not os.path.exists(dir_path):
        os.makedirs(dir_path)

    plotter(t).savefig(f"{dir_path}/{N:03d}.png",bbox_inches='tight')
    plt.close()