# %%
"""
Physics-Based Roller Coaster Rail Generator

Creates smooth roller coaster rails by:
1. Detecting curvature discontinuities at segment junctions
2. Replacing them with G2-continuous quintic polynomial transitions
3. Computing physics-based binormals (hanging mass direction)
4. Generating parallel rails offset from centerline

The result is a track where a marble rolling along it will experience
smooth, continuous forces without abrupt banking changes.
"""

import numpy as np
from build123d import *  # type: ignore
from ocp_vscode import *  # type: ignore

# Physical constants
GRAVITY = 9.81  # m/s²


# ==============================================================================
# UTILITY FUNCTIONS
# ==============================================================================


def normalize(v):
    """
    Normalize a vector to unit length.

    Args:
        v: Vector (numpy array or array-like)

    Returns:
        Normalized vector, or zero vector if input has negligible magnitude
    """
    magnitude = np.linalg.norm(v)
    if magnitude < 1e-12:
        return np.zeros_like(v) if hasattr(v, "__len__") else 0
    return v / magnitude


def sample_wire(wire, steps):
    """
    Sample points along a build123d Wire at uniform parameter intervals.

    Args:
        wire: build123d Wire object
        steps: Number of points to sample

    Returns:
        numpy array of shape (steps, 3) containing [x, y, z] coordinates
    """
    pts = []
    for i in range(steps):
        u = i / (steps - 1)
        p = wire.position_at(u)
        pts.append([p.X, p.Y, p.Z])
    return np.array(pts)


def arc_length(points):
    """
    Compute cumulative arc length along a path.

    Args:
        points: Array of shape (n, 3) representing path points

    Returns:
        Array of shape (n,) with cumulative arc length at each point
    """
    n = len(points)
    s = np.zeros(n)
    for i in range(1, n):
        ds = np.linalg.norm(points[i] - points[i - 1])
        s[i] = s[i - 1] + ds
    return s


# ==============================================================================
# DIFFERENTIAL GEOMETRY
# ==============================================================================


def compute_geometry(points, s):
    """
    Compute differential geometry properties along a path.

    Uses central finite differences for numerical derivatives.

    Args:
        points: Array of shape (n, 3) representing path points
        s: Array of shape (n,) with arc length at each point

    Returns:
        T: Tangent vectors (n, 3) - normalized first derivative
        k: Curvature magnitudes (n,) - ||d²r/ds²||
        k_vector: Curvature vectors (n, 3) - d²r/ds²
        dk: Curvature derivative (n,) - dk/ds
    """
    n = len(points)

    # Tangent: dr/ds (first derivative w.r.t. arc length)
    T = np.zeros_like(points)
    for i in range(1, n - 1):
        ds = s[i + 1] - s[i - 1]
        if ds > 1e-12:
            T[i] = (points[i + 1] - points[i - 1]) / ds
    T[0], T[-1] = T[1], T[-2]
    T = np.array([normalize(t) for t in T])

    # Curvature vector: d²r/ds² (second derivative w.r.t. arc length)
    k_vector = np.zeros_like(points)
    for i in range(1, n - 1):
        ds1 = s[i] - s[i - 1]
        ds2 = s[i + 1] - s[i]
        if ds1 > 1e-12 and ds2 > 1e-12:
            k_vector[i] = (
                2
                * (
                    (points[i + 1] - points[i]) / ds2
                    - (points[i] - points[i - 1]) / ds1
                )
                / (ds1 + ds2)
            )
    k_vector[0], k_vector[-1] = k_vector[1], k_vector[-2]

    # Curvature magnitude
    k = np.linalg.norm(k_vector, axis=1)

    # Curvature derivative (rate of curvature change)
    dk = np.gradient(k, s)

    return T, k, k_vector, dk


# ==============================================================================
# JUNCTION DETECTION
# ==============================================================================


def detect_junctions(k, dk, threshold_percentile=88.0, min_separation=50):
    """
    Detect curvature discontinuities (segment junctions) in a path.

    Identifies points where curvature changes abruptly, indicating
    a junction between two spline segments with different curvatures.

    Args:
        k: Curvature magnitudes along path
        dk: Curvature derivatives (dk/ds)
        threshold_percentile: Percentile for severity threshold (85-95 typical)
        min_separation: Minimum point separation to prevent overlapping windows

    Returns:
        List of junction indices, sorted by position
    """
    # Compute severity metric: combines curvature change and jerk
    abs_dk = np.abs(dk)
    d2k = np.abs(np.gradient(dk))  # Curvature jerk
    severity = abs_dk + 2 * d2k
    threshold = np.percentile(severity, threshold_percentile)

    # Find local maxima above threshold
    junctions = []
    for i in range(10, len(severity) - 10):
        if severity[i] > threshold:
            if severity[i] > severity[i - 1] and severity[i] > severity[i + 1]:
                junctions.append((i, severity[i]))

    # Sort by severity (highest first) and enforce minimum separation
    if len(junctions) > 0:
        junctions.sort(key=lambda x: x[1], reverse=True)

        merged = [junctions[0][0]]  # Start with highest severity junction
        for j_idx, j_severity in junctions[1:]:
            # Only keep if far enough from all existing junctions
            if all(abs(j_idx - m) > min_separation for m in merged):
                merged.append(j_idx)

        junctions = sorted(merged)  # Sort by position along path

    return junctions


# ==============================================================================
# G2-CONTINUOUS TRANSITIONS
# ==============================================================================


def quintic_hermite_blend(t, p0, v0, a0, p1, v1, a1):
    """
    Quintic Hermite polynomial blending function.

    Provides G2 continuity (continuous position, tangent, and curvature).

    Args:
        t: Parameter in [0, 1]
        p0, p1: Start and end positions (scalars or vectors)
        v0, v1: Start and end velocities (tangent * length)
        a0, a1: Start and end accelerations (curvature * length²)

    Returns:
        Interpolated value at parameter t
    """
    # Precompute powers of t
    t2 = t * t
    t3 = t2 * t
    t4 = t3 * t
    t5 = t4 * t

    # Quintic Hermite basis functions
    # These ensure G2 continuity by matching position, first, and second derivatives
    h0 = 1 - 10 * t3 + 15 * t4 - 6 * t5
    h1 = t - 6 * t3 + 8 * t4 - 3 * t5
    h2 = 0.5 * t2 - 1.5 * t3 + 1.5 * t4 - 0.5 * t5
    h3 = 10 * t3 - 15 * t4 + 6 * t5
    h4 = -4 * t3 + 7 * t4 - 3 * t5
    h5 = 0.5 * t3 - t4 + 0.5 * t5

    return h0 * p0 + h1 * v0 + h2 * a0 + h3 * p1 + h4 * v1 + h5 * a1


def create_g2_transition(
    p_start, T_start, k_vector_start, p_end, T_end, k_vector_end, n_points=100
):
    """
    Create a G2-continuous transition curve between two points.

    Uses quintic Hermite interpolation to match position, tangent,
    and curvature at both endpoints.

    Args:
        p_start, p_end: Start and end positions (3D points)
        T_start, T_end: Start and end tangent vectors (normalized)
        k_vector_start, k_vector_end: Start and end curvature vectors
        n_points: Number of points in transition curve

    Returns:
        Array of shape (n_points, 3) forming the transition curve
    """
    L = np.linalg.norm(p_end - p_start)
    t = np.linspace(0, 1, n_points)

    # Scale derivatives appropriately
    v0 = T_start * L  # Velocity (tangent scaled by length)
    v1 = T_end * L
    a0 = k_vector_start * L * L  # Acceleration (curvature scaled by length²)
    a1 = k_vector_end * L * L

    # Generate transition points using quintic interpolation
    points = np.zeros((n_points, 3))
    for i, t_i in enumerate(t):
        points[i] = quintic_hermite_blend(t_i, p_start, v0, a0, p_end, v1, a1)

    return points


def replace_junctions_with_g2_transitions(points, s, junctions, window_length=0.025):
    """
    Replace detected junctions with G2-continuous transition curves.

    For each junction, replaces a window of points around it with a smooth
    quintic transition that matches position, tangent, and curvature at
    both window boundaries.

    Args:
        points: Original path points (n, 3)
        s: Arc length array (n,)
        junctions: List of junction indices
        window_length: Arc length of transition window (in meters)

    Returns:
        Modified path points with junctions replaced by smooth transitions
    """
    if len(junctions) == 0:
        return points

    T, k, k_vector, dk = compute_geometry(points, s)
    pts_new = points.copy()

    # Calculate all window boundaries
    windows = []
    for junction_idx in junctions:
        s_junction = s[junction_idx]
        half_window = window_length / 2

        start_idx = np.searchsorted(s, s_junction - half_window)
        end_idx = np.searchsorted(s, s_junction + half_window)

        start_idx = max(5, start_idx)
        end_idx = min(len(points) - 5, end_idx)

        if end_idx - start_idx >= 10:
            windows.append((junction_idx, start_idx, end_idx))

    # Check for overlapping windows
    print(f"\nReplacing {len(windows)} junctions:")
    overlaps_found = False

    for i, (j1, s1, e1) in enumerate(windows):
        for j2, s2, e2 in windows[i + 1 :]:
            if not (e1 < s2 or e2 < s1):  # Overlap detected
                print(
                    f"  ⚠️  WARNING: Junction {i} [{s1}:{e1}] overlaps with "
                    f"junction {i + 1} [{s2}:{e2}]"
                )
                overlaps_found = True

    if overlaps_found:
        print(
            "  💡 Suggestion: increase junction_threshold or decrease junction_window_length"
        )

    # Replace junctions (skip overlapping ones)
    replaced_ranges = []

    for junction_idx, start_idx, end_idx in windows:
        # Check for overlap with previously replaced ranges
        overlap = False
        for r_start, r_end in replaced_ranges:
            if not (end_idx < r_start or r_end < start_idx):
                overlap = True
                print(
                    f"  ⏭️  Skipping junction at idx {junction_idx} "
                    f"(overlaps with previous replacement)"
                )
                break

        if overlap:
            continue

        # Extract boundary conditions
        p_start = pts_new[start_idx]
        T_start = T[start_idx]
        k_vec_start = k_vector[start_idx]

        p_end = pts_new[end_idx]
        T_end = T[end_idx]
        k_vec_end = k_vector[end_idx]

        # Generate G2 transition
        n_transition = end_idx - start_idx + 1
        transition = create_g2_transition(
            p_start,
            T_start,
            k_vec_start,
            p_end,
            T_end,
            k_vec_end,
            n_points=n_transition,
        )

        # Replace points
        pts_new[start_idx : end_idx + 1] = transition
        replaced_ranges.append((start_idx, end_idx))

        # Print diagnostic info
        s_junction = s[junction_idx]
        arc_length_window = s[end_idx] - s[start_idx]
        print(
            f"  ✓ Junction at s={s_junction:.3f}m (idx {junction_idx}): "
            f"window [{start_idx}:{end_idx}] ({n_transition} pts, "
            f"{arc_length_window * 1000:.1f}mm)"
        )

    return pts_new


# ==============================================================================
# PHYSICS SIMULATION
# ==============================================================================


def compute_velocity_profile(points, mass, v0, friction_coeff, g_scaled):
    """
    Compute velocity along the path using energy conservation.

    Energy equation:
    KE_new = KE_old + PE_lost - friction_work
    (1/2)mv² = (1/2)mv₀² + mg*Δh - μ*mg*Δs

    Args:
        points: Path points (n, 3)
        mass: Mass of rolling object (kg)
        v0: Initial velocity (m/s)
        friction_coeff: Rolling resistance coefficient (dimensionless)
        g_scaled: Gravitational acceleration (m/s²)

    Returns:
        Array of velocities at each point (m/s)
    """
    n = len(points)
    v = np.zeros(n)
    v[0] = v0

    for i in range(1, n):
        # Distance traveled
        d = np.linalg.norm(points[i] - points[i - 1])

        # Height change (positive when going downhill)
        dh = points[i - 1][2] - points[i][2]

        # Energy lost to friction
        friction_work = friction_coeff * mass * g_scaled * d

        # Apply energy conservation
        v2 = v[i - 1] ** 2 + 2 * g_scaled * dh - 2 * friction_work / mass
        v[i] = np.sqrt(max(v2, 0))

    return v


def compute_hanging_binormals(points, velocities, T, k, k_vector, g_scaled):
    """
    Compute binormals representing the direction a hanging mass would point.

    Physics model: In the reference frame of the cart, a hanging mass experiences
    apparent forces from gravity and centrifugal effects. The binormal points in
    the direction the mass hangs.

    F_apparent = F_gravity + F_centrifugal
    where F_centrifugal points outward (away from curve center)

    Args:
        points: Path points (n, 3)
        velocities: Velocity at each point (n,)
        T: Tangent vectors (n, 3)
        k: Curvature magnitudes (n,)
        k_vector: Curvature vectors (n, 3) - point toward curve center
        g_scaled: Gravitational acceleration (m/s²)

    Returns:
        Binormal vectors (n, 3) - direction of hanging mass
    """
    n = len(points)
    gravity = np.array([0, 0, -g_scaled])
    B = np.zeros((n, 3))

    # First point: establish reference orientation
    T0 = normalize(T[0])

    if k[0] < 1e-9:
        # Straight section: mass hangs straight down
        apparent_force = gravity
    else:
        # Curved section: add centrifugal force
        R = 1 / k[0]
        N_inward = normalize(k_vector[0])  # Points toward curve center
        centrifugal_mag = velocities[0] ** 2 / R
        F_centrifugal = centrifugal_mag * (-N_inward)  # Points outward
        apparent_force = gravity + F_centrifugal

    # Project onto plane perpendicular to tangent
    apparent_force_perp = apparent_force - np.dot(apparent_force, T0) * T0
    B[0] = normalize(apparent_force_perp)

    # Ensure first binormal points generally downward
    if B[0][2] > 0:
        B[0] = -B[0]

    # Compute remaining binormals with continuity enforcement
    for i in range(1, n):
        T_i = normalize(T[i])

        if k[i] < 1e-9:
            apparent_force = gravity
        else:
            R = 1 / k[i]
            N_inward = normalize(k_vector[i])
            centrifugal_mag = velocities[i] ** 2 / R
            F_centrifugal = centrifugal_mag * (-N_inward)
            apparent_force = gravity + F_centrifugal

        apparent_force_perp = apparent_force - np.dot(apparent_force, T_i) * T_i
        b = normalize(apparent_force_perp)

        # Enforce continuity: prevent 180° flips
        if np.dot(b, B[i - 1]) < 0:
            b = -b

        B[i] = b

    return B


def smooth_binormals(B, tangents, iterations=5):
    """
    Apply light smoothing to binormals to remove high-frequency jitter.

    Uses weighted averaging while maintaining orthogonality to tangent.

    Args:
        B: Binormal vectors (n, 3)
        tangents: Tangent vectors (n, 3)
        iterations: Number of smoothing passes

    Returns:
        Smoothed binormal vectors (n, 3)
    """
    B_smooth = B.copy()
    n = len(B)

    for _ in range(iterations):
        B_new = B_smooth.copy()
        for i in range(2, n - 2):
            # Weighted average (favor current value)
            avg = (B_smooth[i - 1] + 2 * B_smooth[i] + B_smooth[i + 1]) / 4

            # Ensure orthogonality to tangent
            T = normalize(tangents[i])
            avg_perp = avg - np.dot(avg, T) * T
            B_new[i] = normalize(avg_perp)

        B_smooth = B_new

    return B_smooth


# ==============================================================================
# MAIN PIPELINE
# ==============================================================================


def build_physics_based_roller_coaster_rails(
    wire,
    steps=5000,
    rail_spacing=5.0,
    mass=0.005,
    initial_velocity=0.01,
    friction_coeff=0.05,
    junction_window_length=0.030,
    junction_threshold=88.0,
    binormal_smooth_iterations=5,
    downsample_factor=10,
    unit_scale=0.001,
):
    """
    Generate physics-based roller coaster rails from a build123d Wire.

    Creates two parallel rails that maintain proper banking angles based on
    the physics of a rolling mass. Ensures smooth transitions at segment
    junctions using G2-continuous quintic polynomial blending.

    Args:
        wire: build123d Wire defining the centerline path
        steps: Number of points to sample from wire
        rail_spacing: Distance between rails (in build123d units)
        mass: Mass of rolling object (kg)
        initial_velocity: Starting velocity (m/s)
        friction_coeff: Rolling resistance coefficient
        junction_window_length: Arc length of G2 transition windows (m)
        junction_threshold: Percentile threshold for junction detection (85-95)
        binormal_smooth_iterations: Number of binormal smoothing passes
        downsample_factor: Factor to downsample output rails (1 = no downsampling)
        unit_scale: Conversion from build123d units to meters (0.001 for mm)

    Returns:
        rail_1, rail_2: Two build123d Splines representing the rails
    """
    # Sample wire and convert to meters
    pts = sample_wire(wire, steps)
    pts_m = pts * unit_scale
    s = arc_length(pts_m)

    print(f"Original track: {len(pts)} points, {s[-1]:.3f} m length")

    # Compute initial geometry
    T, k, k_vector, dk = compute_geometry(pts_m, s)

    # Detect segment junctions
    avg_point_spacing = s[-1] / len(pts_m)
    min_separation_points = int((junction_window_length * 1.2) / avg_point_spacing)

    junctions = detect_junctions(
        k,
        dk,
        threshold_percentile=junction_threshold,
        min_separation=min_separation_points,
    )

    print(f"Detected {len(junctions)} segment junctions")

    # Replace junctions with G2 transitions
    if len(junctions) > 0:
        pts_m = replace_junctions_with_g2_transitions(
            pts_m, s, junctions, window_length=junction_window_length
        )
        s = arc_length(pts_m)
        print(f"After junction replacement: {s[-1]:.3f} m length")

    # Recompute geometry after modifications
    T, k, k_vector, dk = compute_geometry(pts_m, s)

    # Compute velocity profile
    velocities = compute_velocity_profile(
        pts_m, mass, initial_velocity, friction_coeff, GRAVITY
    )

    # Compute physics-based binormals
    B = compute_hanging_binormals(pts_m, velocities, T, k, k_vector, GRAVITY)

    # Optional light smoothing
    if binormal_smooth_iterations > 0:
        B = smooth_binormals(B, T, iterations=binormal_smooth_iterations)

    # Print diagnostics
    print_diagnostics(pts_m, velocities, k, B)

    # Convert back to build123d units
    pts_out = pts_m / unit_scale

    # Generate rail positions using cross product
    # Rails are perpendicular to both tangent and binormal
    rail_1_pts = [
        Vector(*(p + rail_spacing * np.cross(b, t))) for p, b, t in zip(pts_out, B, T)
    ]
    rail_2_pts = [
        Vector(*(p - rail_spacing * np.cross(b, t))) for p, b, t in zip(pts_out, B, T)
    ]

    # Downsample for final splines
    rail_1_pts = [p for i, p in enumerate(rail_1_pts) if i % downsample_factor == 0]
    rail_2_pts = [p for i, p in enumerate(rail_2_pts) if i % downsample_factor == 0]

    return Spline(*rail_1_pts), Spline(*rail_2_pts)


def print_diagnostics(points, velocities, k, B):
    """Print diagnostic information about the track."""
    print("\n=== Ride Diagnostics ===")
    print(f"Height range: {np.min(points[:, 2]):.3f} - {np.max(points[:, 2]):.3f} m")
    print(f"Velocity: {np.min(velocities):.3f} - {np.max(velocities):.3f} m/s")

    # G-forces (filter unrealistic tight curves)
    g_forces = []
    max_g = 0
    max_g_idx = 0

    for i in range(len(points)):
        if k[i] > 1e-6:
            R = 1 / k[i]
            if R > 0.002:  # 2mm minimum radius
                centripetal_accel = velocities[i] ** 2 / R
                g_force = centripetal_accel / GRAVITY
                g_forces.append(g_force)
                if g_force > max_g:
                    max_g = g_force
                    max_g_idx = i

    if g_forces:
        print(f"G-forces: {np.min(g_forces):.2f}g - {np.max(g_forces):.2f}g")
        print(
            f"  Max G at idx {max_g_idx}: "
            f"R={1 / k[max_g_idx] * 1000:.2f}mm, "
            f"v={velocities[max_g_idx]:.3f}m/s"
        )

    # Stall check
    stall_idx = np.where(velocities < 0.01)[0]
    if len(stall_idx) == 0:
        print("✓ No stalls")
    else:
        print(f"⚠️  Stalls at {100 * stall_idx[0] / len(points):.1f}%")

    # Binormal quality metrics
    binormal_changes = np.array(
        [np.linalg.norm(B[i] - B[i - 1]) for i in range(1, len(B))]
    )
    print(
        f"\nBinormal variation: "
        f"max={np.max(binormal_changes):.3f}, "
        f"avg={np.mean(binormal_changes):.4f}"
    )

    down_pointing = np.sum(B[:, 2] < 0)
    print(f"Downward binormal: {100 * down_pointing / len(B):.1f}%")


# ==============================================================================
# EXAMPLE USAGE
# ==============================================================================

if __name__ == "__main__":
    # Define track segments
    ln1 = Line((-20, 0), (0, 0))

    powerup = Spline(
        (39, 0, 71),
        (40, 0, 70),
        (100, 0, 0),
        tangents=((1, 0, 0), (1, 0, 0)),
        tangent_scalars=(0.5, 2),
    )

    corner = RadiusArc(powerup @ 1, (100, 60, 0), -30)
    screw = Helix(75, 155, 15, center=(75, 40, 15), direction=(-1, 0, 0))
    sp1 = Spline(corner @ 1, screw @ 0, tangents=(corner % 1, screw % 0))
    sp2 = Spline(screw @ 1, (-100, 30, 10), ln1 @ 0, tangents=(screw % 1, powerup % 0))

    # Assemble wire
    track_centerline = Wire([powerup, corner, sp1, screw, sp2, ln1])

    # Generate physics-based rails
    rail_1, rail_2 = build_physics_based_roller_coaster_rails(
        track_centerline,
        steps=1000,
        rail_spacing=1.0,  # 1mm spacing
        mass=0.002,  # 2 grams
        initial_velocity=0.01,  # Nearly at rest
        friction_coeff=0.05,
        junction_window_length=0.08,  # 80mm transition windows
        junction_threshold=88,
        binormal_smooth_iterations=100,
        downsample_factor=10,
        unit_scale=0.001,  # mm to m
    )

    # Create rail geometry
    sk1 = Location(rail_1 ^ 0) * Circle(0.4)
    rail_1_solid = sweep(sections=sk1.face(), path=rail_1, binormal=rail_2)

    sk2 = Location(rail_2 ^ 0) * Circle(0.4)
    rail_2_solid = sweep(sections=sk2.face(), path=rail_2, binormal=rail_1)

    show(rail_1_solid, rail_2_solid)