# Interactive plot of Bspline 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: 2/24/2026 - Zhe Yuan


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

def ibspline(inputA, inputB=None):
    # 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, knot_vector, add_cp_mode # added for Bspline

    # de Boor's algorithm to evaluate a point on a B-spline.
    # Returns (x, y) coordinate at parameter t
    def deBoor(points, knots, degree, t):
        m = len(points)
        # Find knot span k such that knots[k] <= t < knots[k+1]
        k = degree - 1
        while k < m -2 and knots[k + 1] <= t:
            k += 1
        if k == m-2 and t == knots[m-1]:
            pass # k reach max
            
        # Copy relevant control points: P_{k-degree} to P_k
        d = [list(points[j]) for j in range(k +1 - degree, k + 2)]
        
        # Apply de Boor's algorithm
        for r in range(1, degree + 1):
            for j in range(degree, r - 1, -1):
                idx_i = j + k - degree
                denom = knots[idx_i + degree - r + 1] - knots[idx_i]
                
                # Prevent division by zero if knots are duplicated
                alpha = 0.0
                if denom > 1e-8:
                    alpha = (t - knots[idx_i]) / denom
                    
                d[j][0] = (1.0 - alpha) * d[j-1][0] + alpha * d[j][0]
                d[j][1] = (1.0 - alpha) * d[j-1][1] + alpha * d[j][1]
                
        return (d[degree][0], d[degree][1])

    # Insert a knot using Boehm's algorithm given the selected control point index h
    def applyBoehm(h):
        global vX, vY, knot_vector, curve_degree
        n = curve_degree
        m = len(vX)
        
        # !If h < n, do nothing, print error message
        if h < n:
            outputPlotMessage(f"Cannot add here. Please select index h >= {n}.")
            plt.draw()
            return
            
        t_h = knot_vector[h-1]
        t_h1 = knot_vector[h]
        t_hat = (t_h + t_h1) / 2.0
        
        points = list(zip(vX, vY))
        new_points = []
        
        for i in range(m + 1):
            if i <= h - n:
                new_points.append(points[i])
            elif i >= h + 1:
                new_points.append(points[i - 1])
            else:
                denom = knot_vector[i+n-1] - knot_vector[i-1]
                alpha = (t_hat - knot_vector[i-1]) / denom if denom > 1e-8 else 0.0
                px = (1 - alpha) * points[i-1][0] + alpha * points[i][0]
                py = (1 - alpha) * points[i-1][1] + alpha * points[i][1]
                new_points.append((px, py))
                
        # Update globals
        vX = [p[0] for p in new_points]
        vY = [p[1] for p in new_points]
        
        # Update knot vector: insert t_hat
        knot_vector.insert(h, t_hat)
        
        outputPlotMessage(f"Added Control Point. (m={len(vX)})")
        redrawPlot(vX, vY)

    # Draw the B-Spline curve
    def drawBSplineCurve(vX, vY):
        global num_sample, curve_degree, knot_vector
        m = len(vX)
        n = curve_degree
        
        points = list(zip(vX, vY))
        pX, pY = [], []
        
        # Curve is defined on [t_n, t_m]
        t_start = knot_vector[n-1]
        t_end = knot_vector[m-1]
        
        steps = np.linspace(t_start, t_end, num_sample)
        for t in steps:
            cx, cy = deBoor(points, knot_vector, n, t)
            pX.append(cx)
            pY.append(cy)
            
        # Draw B-Spline 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])

        drawBSplineCurve(vX, vY) # change to bspline

        outputPlotMessage("Drag control points, or add control point then select")
        # change message, we now use bottom to add cp.
        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
    
    # callback of add cp
    def addCPButtonCallback(event):
        global add_cp_mode
        if len(vX) == 0: return
        add_cp_mode = True
        outputPlotMessage("Select an existing control point p_h (h >= n) to split")
        plt.draw()
        return
    
    # update pick point callback to handle mova & add cp
    def pickPointCallback(event):
        global iv_selected, cid_moveCallback, add_cp_mode
        if (len(event.ind) < 1): return

        if add_cp_mode:
            h = event.ind[0]
            applyBoehm(h)
            add_cp_mode = False # Reset state
            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

    # Save control points to file
    def saveCallback(event):
        global vX, vY, filename_textbox, curve_degree,  knot_vector # add knot
        
        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 Bspline format
  
            try:
                with open(filename, 'w') as f:
                    f.write("BSPLINE\n")
                    f.write(f"# Saved from ibspline.py\n")
                    f.write(f"2 {len(vX)} {curve_degree}\n")
                    # Write knot vector as space-separated string
                    knots_str = " ".join([str(k) for k in knot_vector])
                    f.write(f"{knots_str}\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_add_button, add_button
        global ax_save_button, save_button

        # Colors
        inactive_color = 'lightgray' 
        
        # ADd cp
        ax_add_button = plt.axes([left_pos, 0.8, button_width, button_height])
        add_button = Button(ax_add_button, "Add CP")
        add_button.on_clicked(addCPButtonCallback)
        add_button.color = inactive_color
        add_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 add_button, save_button

        add_button.color = 'white'
        add_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, knot_vector # add knot_vector

        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)
                # init knot vector
                knot_vector = list(range(1, numv + curve_degree))
                drawBezierEnableButtonsPickEvent(vX, vY)
                enableButtons()
        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, knot_vector # add knot
        
        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] != "BSPLINE":
                raise ValueError("File must start with BSPLINE")
                
            # 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 = lines[data_start_idx].split()
                    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 knot
            for i in range(data_start_idx, len(lines)):
                if lines[i].startswith('#'):
                    data_start_idx += 1
                    continue
                else:
                    knot_line = lines[data_start_idx]
                    knot_vector = [float(k) for k in knot_line.split()]
                    data_start_idx += 1
                    break
            
            
            tempX, tempY = [], []
            for i in range(data_start_idx, len(lines)):
                if lines[i].startswith('#'): continue
                parts = lines[i].split()
                if len(parts) >= 2:
                    tempX.append(float(parts[0]))
                    tempY.append(float(parts[1]))
                    
            vX, vY = tempX, 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 = 'bspline_pts.txt'
    num_sample = 100
    add_cp_mode = False # State to track if user is adding a control point
    knot_vector = []


    # If inputA is a string, read control points from inputA
    if isinstance(inputA, int) and isinstance(inputB, int):
        numv = inputA
        curve_degree = inputB
        if numv < curve_degree + 1:
            print(f"Error: m ({numv}) must be >= n+1 ({curve_degree+1})")
            return
        jv = 0
        vX, vY = [], []
    elif isinstance(inputA, str):
        # file logic remains same
        loadTextFile(inputA)
    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

    plt.show()

    return

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Usage: python ibspline.py <m> <n> OR <filename>")
        sys.exit()
        
    arg1 = sys.argv[1]
    if len(sys.argv) == 3:
        try:
            m = int(arg1)
            n = int(sys.argv[2])
            ibspline(m, n)
        except ValueError:
            print("m and n must be integers.")
    else:
        ibspline(arg1)