# Interactive plot of Bezier interpolation.
# Interactively set polyline vertices or read vertices from a file.
# Buttons to subdivide polygonal line.
# Button to save vertices (points) to a file.
# Created: 1/24/2026 - R. Wenger
# Edited: 1/26/2026 - Zhe Yuan


import sys
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.widgets import Button, TextBox

def ibezier(inputA):
    # Global variables
    global vX, vY, plot_center, text_message, fig, ax, save_message # fix: fig should be global
    global jv, iv_selected, numv
    global default_output_filename, filename_textbox
    global cid_bpress, cid_pick
    global save_message_xloc
    global num_sample
    global curve_degree, is_subdivided # add for bezier

    # Compute a point on the Bezier curve at parameter t using de Casteljau's algorithm.
    # points: list of tuples/lists [(x0,y0), (x1,y1), ...]
    # Returns: (x, y) coordinates on the curve
    def deCasteljau(points, t):
        current_points = np.array(points)
        n = len(current_points) - 1 # degree
        
        # Iterate r from 1 to n
        for r in range(1, n + 1):
            # Calculate next level of points
            # P_i^r = (1-t)*P_i^{r-1} + t*P_{i+1}^{r-1}
            new_points = []
            for i in range(len(current_points) - 1):
                x = (1 - t) * current_points[i][0] + t * current_points[i+1][0]
                y = (1 - t) * current_points[i][1] + t * current_points[i+1][1]
                new_points.append([x, y])
            current_points = np.array(new_points)
            
        return (current_points[0][0], current_points[0][1])

    # Perform subdivision at t=0.5
    # Input: List of control points for ONE Bezier curve segment.
    # Output: A tuple (left_points, right_points) where they share the middle point.
    def subdivide_points_algorithm(segment_x, segment_y):
        # Combine to (x,y)
        points = list(zip(segment_x, segment_y))
        n = len(points) - 1
        # structure: level 0 has n+1 points, level n has 1 point
        levels = [points]
        
        # Run de Casteljau for t=0.5 and store all intermediate points
        current_points = points
        for r in range(1, n + 1):
            next_level = []
            for i in range(len(current_points) - 1):
                px = 0.5 * current_points[i][0] + 0.5 * current_points[i+1][0]
                py = 0.5 * current_points[i][1] + 0.5 * current_points[i+1][1]
                next_level.append((px, py))
            levels.append(next_level)
            current_points = next_level
            
        # Left curve control points are the first point of each level (P0^0, P0^1, ..., P0^n)
        left_x = [level[0][0] for level in levels]
        left_y = [level[0][1] for level in levels]
        
        # Right curve cpts are the last point of each level, reversed (P0^n, P1^{n-1}, ..., Pn^0)
        right_x = [level[-1][0] for level in levels]
        right_y = [level[-1][1] for level in levels]
        right_x.reverse()
        right_y.reverse()
        
        return (left_x, left_y), (right_x, right_y)

    # Draw the Bezier curve
    def drawBezierCurve(vX, vY):
        global num_sample, curve_degree
        # Number of points per segment = degree + 1
        nump = len(vX)
        num_segments = (nump - 1) // curve_degree
        for k in range(num_segments):
            start_idx = k * curve_degree
            end_idx = start_idx + curve_degree + 1
            segX = vX[start_idx:end_idx]
            segY = vY[start_idx:end_idx]
            points = list(zip(segX, segY))
            pX = []
            pY = []
            steps = np.linspace(0, 1, num_sample)
            for t in steps:
                (cx, cy) = deCasteljau(points, t)
                pX.append(cx)
                pY.append(cy)
            # Draw Bezier curve
            plt.plot(pX, pY, color="red", linewidth=1.5)
        # Draw control polygon
        plt.plot(vX, vY, '--', color='gray', linewidth=0.5)
        # Draw control points
        plt.plot(vX, vY, 'ob', picker=True, pickradius=5)
        return

    # Redraw the entire plot
    def redrawPlot(vX, vY):
        global text_message, ax
        
        plt.subplot(111)

        # Store current axes limits
        xL, xR = ax.get_xlim()
        yB, yT = ax.get_ylim()

        # Bug: Clearing the axes resets the navigation toolbar so you can't return
        #  to previous states.
        # In particular, you can't return to the starting (0,0), (1,1) configuration.
        plt.cla()
        text_message=None

        # Restore axes limits
        plt.xlim([xL,xR])
        plt.ylim([yB,yT])

        drawBezierCurve(vX, vY)

        outputPlotMessage("Drag and move control points")
        clearSaveMessage()
        plt.draw()
        return

    # Output a message above the plot
    def outputPlotMessage(s):
        global text_message, ax

        text_yloc=1.02

        if text_message is None:
            text_message=plt.text(0.0, text_yloc, s, transform=ax.transAxes)
        else:
            text_message.set_text(s)

        return

    # Output the message: Select vertex str(j) location
    def outputSelectVertexLocation(j):
        s = "Select vertex " + str(j) + " location"
        outputPlotMessage(s)
        return

    # Set a message related to saving control points to file
    def outputSaveMessage(s):
        global save_message, save_message_xloc
        text_yloc=1.02
        if save_message is not None:
            save_message.remove()

        save_message=plt.text(save_message_xloc-0.14, text_yloc, s, transform=ax.transAxes)
        return

    # Set the save message to ''
    def clearSaveMessage():
        global save_message

        if save_message is not None:
            save_message.remove()
            save_message = None
        return
    
    # CallbackS
    
    # Subdivide curve
    def subdivideCallback(event):
        global vX, vY, curve_degree, is_subdivided, smooth_button
        
        if len(vX) == 0: return

        tempX = []
        tempY = []
        
        # We must subdivide EVERY segment in the current curve
        num_segments = (len(vX) - 1) // curve_degree
        
        for k in range(num_segments):
            start_idx = k * curve_degree
            end_idx = start_idx + curve_degree + 1
            
            segX = vX[start_idx:end_idx]
            segY = vY[start_idx:end_idx]
            
            # Get divided points
            (lx, ly), (rx, ry) = subdivide_points_algorithm(segX, segY)
            
            if k == 0:
                tempX.extend(lx[:-1]) # Add all Left except last (shared)
                tempX.append(lx[-1])  # Add shared point
                tempX.extend(rx[1:])  # Add all Right except first (shared)
                tempY.extend(ly[:-1])
                tempY.append(ly[-1])
                tempY.extend(ry[1:])
            else:
                # For subsequent segments, the starting point is already in the list
                # lx[0] should equal tempX[-1]
                tempX.extend(lx[1:-1]) # Skip first, add middle
                tempX.append(lx[-1])   # Add shared middle
                tempX.extend(rx[1:])   # Skip first (shared), add rest
                
                tempY.extend(ly[1:-1])
                tempY.append(ly[-1])
                tempY.extend(ry[1:])

        vX = tempX
        vY = tempY
        
        # Set flag to enable Smooth button
        is_subdivided = True
        # update bottom to distinguish enabled or not
        smooth_button.color = 'white' 
        smooth_button.hovercolor = 'green'
        
        redrawPlot(vX, vY)
        
        return

    # Smooth button callback
    def smoothCallback(event):
        global vX, vY, curve_degree, is_subdivided
        
        if not is_subdivided:
            outputPlotMessage("Curve must be subdivided first.")
            plt.draw()
            return
            
        # Iterate through join points.
        nump = len(vX)
        
        changes = {} # Map index to new (x,y)
        
        for i in range(1, nump - 1):
            if i % curve_degree == 0:
                # This is a junction point between two segments
                # q_i = (q_{i-1} + q_{i+1}) / 2
                nx = (vX[i-1] + vX[i+1]) / 2.0
                ny = (vY[i-1] + vY[i+1]) / 2.0
                changes[i] = (nx, ny)
                
        # Apply changes
        for i, (nx, ny) in changes.items():
            vX[i] = nx
            vY[i] = ny
            
        redrawPlot(vX, vY)
        outputPlotMessage("Curve smoothed at junctions.")
        # Don't disable smooth button again
        plt.draw()
        return

    # Save control points to file
    def saveCallback(event):
        global vX, vY, filename_textbox, curve_degree
        
        if len(vX) == 0: return
        
        filename = filename_textbox.text
        
        if (filename == '' or filename.isspace()):
            # Should pop up an error window/message, but the GUI code would become
            #   even more complicated.
            outputSaveMessage('Illegal filename. Reset filename to default. No file saved.')
            filename_textbox.set_val(default_output_filename)
            plt.draw()
        else:
            # Prepare data for BEZIER format
            # Header:
            # BEZIER
            # # comments
            # {ndim} {nump} {ndegree}
            # coords...
            
            try:
                with open(filename, 'w') as f:
                    f.write("BEZIER\n")
                    f.write(f"# Saved from ibezier.py\n")
                    # ndim=2, nump=len(vX), ndegree=curve_degree
                    f.write(f"2 {len(vX)} {curve_degree}\n")
                    for i in range(len(vX)):
                        f.write(f"{vX[i]} {vY[i]}\n")
                        
                outputSaveMessage('Saved to ' + filename)
            except Exception as e:
                outputSaveMessage('Error saving file.')
                print(e)
                
            plt.draw()
        return

    # Create TextBox containing output filename and
    #   output message text box
    def createSaveTextBoxes(default_filename):
        global filename_textbox, save_message_xloc, fig

        # *** DEBUG ***
        print("Creating save text box.")
        
        ax_filename_textbox = fig.add_axes([save_message_xloc, 0.94, 0.3, 0.03])
        filename_textbox = TextBox(ax_filename_textbox, 'Output filename:', default_filename)
        plt.draw()

        return

    # Create buttons
    def createButtons(left_pos, button_width, button_height):
        global ax_subdivide_button, subdivide_button
        global ax_smooth_button, smooth_button
        global ax_save_button, save_button

        # Colors
        inactive_color = 'lightgray' 
        
        # Subdivide
        ax_subdivide_button = plt.axes([left_pos, 0.8, button_width, button_height])
        subdivide_button = Button(ax_subdivide_button, "Subdivide")
        subdivide_button.on_clicked(subdivideCallback)
        subdivide_button.color = inactive_color
        subdivide_button.hovercolor = inactive_color

        # Smooth
        ax_smooth_button = plt.axes([left_pos, 0.7, button_width, button_height])
        smooth_button = Button(ax_smooth_button, "Smooth")
        smooth_button.on_clicked(smoothCallback)
        smooth_button.color = inactive_color 
        smooth_button.hovercolor = inactive_color

        # Save
        ax_save_button = plt.axes([left_pos, 0.6, button_width, button_height])
        save_button = Button(ax_save_button, "Save")
        save_button.on_clicked(saveCallback)
        save_button.color = inactive_color
        save_button.hovercolor = inactive_color

        createSaveTextBoxes(default_output_filename)
        return

    # not enable all buttons as smooth needs subdivision first
    def enableButtons():
        global subdivide_button, smooth_button, save_button

        subdivide_button.color = 'white'
        subdivide_button.hovercolor = 'green'
        save_button.color = 'white'
        save_button.hovercolor = 'green'
        plt.draw()
        return
    
    # Add a vertex at (event.xdata, event.ydata)
    def addVertexCallback(event):
        global jv, vX, vY, numv, cid_bpress, fig, curve_degree

        plt.subplot(111)
        x=event.xdata
        y=event.ydata

        # Check that x and y are defined
        if (x is None) or (y is None): return

        # get (again) plot lower left (LL) and upper right (UR)
        # (Coordinates can change if window is moved from one screen
        #  to another.)
        axLL = ax.transData.transform((0,0));
        axUR = ax.transData.transform((1,1));

        
        # check that event.x and event.y are within draw region
        if ((event.x < axLL[0]) or (event.x > axUR[0])): return
        if ((event.y < axLL[1]) or (event.y > axUR[1])): return
        
        if (jv < numv):
            vX.append(x)
            vY.append(y)
            plt.plot(x, y, 'ob')
            jv =jv+1
            
            if (jv < numv):
                outputSelectVertexLocation(jv)
                plt.draw()
            else:
                # Disconnect button press callback before creating pick event
                fig.canvas.mpl_disconnect(cid_bpress)
                # Determine degree (n points => degree n-1)
                curve_degree = numv - 1
                drawBezierEnableButtonsPickEvent(vX, vY)
                enableButtons()
        return

    # Pick and move a vertex
    def pickPointCallback(event):
        global iv_selected, cid_moveCallback

        if (len(event.ind) < 1): return

        if (event.ind[0] >= 0) and (event.ind[0] < len(vX)):
            iv_selected = event.ind[0]
            cid_moveCallback = \
                fig.canvas.mpl_connect('button_release_event', moveVertexCallback)

        return

    # Move selected vertex to location (event.xdata, event.ydata)
    def moveVertexCallback(event):
        global cid_moveCallback, iv_selected, vX, vY
        global points, polyline

        plt.subplot(111)
        x=event.xdata
        y=event.ydata

        # Check that x and y are defined
        if (x is None) or (y is None): return

        # get (again) plot lower left (LL) and upper right (UR)
        # (Coordinates can change if window is moved from one screen
        #  to another.)
        axLL = ax.transData.transform((0,0));
        axUR = ax.transData.transform((1,1));
        
        # check that event.x and event.y are within draw region
        if ((event.x < axLL[0]) or (event.x > axUR[0])): return
        if ((event.y < axLL[1]) or (event.y > axUR[1])): return

        vX[iv_selected] = x
        vY[iv_selected] = y
        redrawPlot(vX, vY)

        fig.canvas.mpl_disconnect(cid_moveCallback)
        
        # reset iv_selected
        iv_selected = -1
        return  

    # After control points are created or read, draw bezier
    #   curve and enable buttons and pick event
    def drawBezierEnableButtonsPickEvent(vX, vY):
        global cid_pick, fig
        
        plt.subplot(111)
        outputPlotMessage("Select and move vertices")
        redrawPlot(vX, vY)
        
        cid_pick = fig.canvas.mpl_connect('pick_event', pickPointCallback)
        enableButtons()
        return

    # load vX and vY from a textfile
    # edit for BEZIER format
    def loadTextFile(filename):
        global vX, vY, jv, numv, curve_degree
        
        try:
            with open(filename, 'r') as f:
                lines = f.readlines()
                
            # Strip whitespace and remove empty lines
            lines = [l.strip() for l in lines if l.strip()]
            
            if not lines:
                raise ValueError("Empty file")
                
            if lines[0] != "BEZIER":
                raise ValueError("File must start with BEZIER")
                
            # Find the line with dimension/nump/degree
            data_start_idx = 1
            header_vals = []
            
            while data_start_idx < len(lines):
                line = lines[data_start_idx]
                if line.startswith('#'):
                    data_start_idx += 1
                    continue
                else:
                    # Found the header line
                    parts = line.split()
                    if len(parts) < 3:
                        raise ValueError("Invalid header line format")
                    header_vals = [int(p) for p in parts]
                    data_start_idx += 1
                    break
                    
            if not header_vals:
                raise ValueError("Could not find data dimensions")
                
            ndim = header_vals[0]
            nump = header_vals[1]
            ndegree = header_vals[2]
            
            # Read points
            tempX = []
            tempY = []
            
            for i in range(data_start_idx, len(lines)):
                line = lines[i]
                if line.startswith('#'): continue
                parts = line.split()
                if len(parts) >= 2:
                    tempX.append(float(parts[0]))
                    tempY.append(float(parts[1]))
                    
            if len(tempX) != nump:
                print(f"Warning: Expected {nump} points, found {len(tempX)}")
                
            vX = tempX
            vY = tempY
            numv = len(vX)
            jv = numv
            curve_degree = ndegree
            
        except Exception as e:
            print(f"Error loading file: {e}")
            sys.exit(1)
            
        return


    # Initialize global variables
    text_message = None
    save_message = None
    filename_textbox = None
    save_message_xloc = 0.5
    iv_selected = -1
    jv = 0
    plot_center = [ 0.5, 0.5 ]
    default_output_filename = 'controlpts.txt'
    num_sample = 100
    is_subdivided = False

    
    # If inputA is an integer, set number of vertices to inputA
    # If inputA is a string, read control points from inputA
    if isinstance(inputA, int):
        numv = inputA
        jv = 0
        vX = []
        vY = []
    elif isinstance(inputA, str):
        try:
            loadTextFile(inputA)
        except Exception as e:
            print(f"Error loading file: {e}")
            print('Unable to open filename: ', inputA)
            print('Check that file exists and that you have read permission.')
            print('Exiting.')
            return
    else:
        print('Illegal input: ', inputA)
        print('Exiting.')
        return

    if (numv < 2):
        print('Illegal number of vertices: ', numv)
        print('Requires at least two vertices.')
        print('Exiting')
        return
    
    fig, ax = plt.subplots(1, 1, figsize=(12, 8))
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.0])
    plt.subplots_adjust(right=0.9)
    
    outputSelectVertexLocation(0)

    # get plot lower left (LL) and upper right (UR)
    axLL = ax.transData.transform((0,0));
    axUR = ax.transData.transform((1,1));

    createButtons(0.91, 0.08, 0.05)

    if (jv < numv):
        # User interactively selects vertices
        cid_bpress = fig.canvas.mpl_connect('button_press_event', addVertexCallback)
    else:
        # Vertices already read in from input file
        drawBezierEnableButtonsPickEvent(vX, vY)
        # Enable smooth for already loaded curve if cpts - 1 is multiple of degree
        # and cpts - 1 != degree
        if (len(vX) - 1) % curve_degree == 0 and curve_degree != len(vX) - 1:
            # Set flag to enable Smooth button
            is_subdivided = True
            # update bottom to distinguish enabled or not
            smooth_button.color = 'white' 
            smooth_button.hovercolor = 'green'

    plt.show()

    return

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python ibezier.py <num_points> OR <filename>")
        sys.exit()
        
    arg = sys.argv[1]
    try:
        val = int(arg)
        ibezier(val)
    except ValueError:
        ibezier(arg)
    except Exception as e:
        print(f"Error: {e}")
        sys.exit()