#!/usr/bin/python3
import numpy as np
import matplotlib
matplotlib.use("TkAgg")
import matplotlib.pyplot as plt

import tkinter as tk
from tkinter import ttk, filedialog, messagebox
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg

def on_exit():
    plt.close('all')   # closes all matplotlib figures
    root.quit()        # stops the Tk event loop
    root.destroy()     # destroys the window cleanly


# ---------------------------------------------------------
# Compute pattern and draw plots
# ---------------------------------------------------------
def compute_and_plot(sw, sh, p, lr, pnts, plot_frame):
    # Clear old plot widgets
    for widget in plot_frame.winfo_children():
        widget.destroy()

    # -----------------------------
    # Original math
    # -----------------------------
    nh = int(np.floor((sh - 2 * lr) / p) + 1)
    nv = int(np.floor((sw - 2 * lr) / p) + 1)
    cx = np.sqrt(lr**2 - (p / 2)**2)
    lw = sw / 2 - lr - cx
    lv = sh / 2 - lr - cx

    tstart = np.arcsin(p / (2 * lr))
    tend = 2 * np.pi - tstart
    t = np.linspace(tstart, tend, pnts + 1)

    loopx = -lr * np.cos(t)
    loopy = -lr * np.sin(t)

    lwh = (nh - 1) * p / 2
    hx = np.concatenate([
        [-lw],
        lw + loopx + cx,
        [-lw],
        -lw - loopx - cx
    ])
    hy = np.concatenate([
        [-lwh],
        -lwh + loopy + p / 2,
        [-lwh + p],
        -lwh + loopy + 3 * p / 2
    ])

    lvh = (nv - 1) * p / 2
    vx = np.concatenate([
        [-lvh],
        -lvh + loopy + p / 2,
        [-lvh + p],
        -lvh + loopy + 3 * p / 2
    ])
    vy = np.concatenate([
        [lv],
        -lv - loopx - cx,
        [lv - cx],
        lv + loopx + cx
    ])

    nah = len(hx)
    nav = len(vx)
    cnth = nh // 2
    cntv = nv // 2

    bigx = -99 * np.ones(cnth * nah + cntv * nav)
    bigy = -99 * np.ones(cnth * nah + cntv * nav)

    for l in range(cnth):
        bigx[l * nah:(l + 1) * nah] = hx
        bigy[l * nah:(l + 1) * nah] = hy + l * p * 2

    start = cnth * nah
    for l in range(cntv):
        bigx[start + l * nav:start + (l + 1) * nav] = vx + 2 * p * l
        bigy[start + l * nav:start + (l + 1) * nav] = vy

    # -----------------------------
    # Plotting
    # -----------------------------
    fig, axs = plt.subplots(2, 2, figsize=(8, 8))

    axs[0, 0].plot(hx, hy)
    axs[0, 0].set_title("Horizontal Loop Segment")

    axs[0, 1].plot(vx, vy)
    axs[0, 1].set_title("Vertical Loop Segment")

    axs[1, 0].plot(bigx[:cnth * nah], bigy[:cnth * nah])
    axs[1, 0].set_title("Horizontal Repeated Pattern")

    axs[1, 1].plot(bigx, bigy)
    axs[1, 1].set_title("Full Pattern")

    for ax in axs.flat:
        ax.set_aspect("equal")
        ax.grid(True)

    fig.tight_layout()

    canvas = FigureCanvasTkAgg(fig, master=plot_frame)
    canvas.draw()
    canvas.get_tk_widget().pack(fill="both", expand=True)

    return bigx, bigy


# ---------------------------------------------------------
# Save G-code
# ---------------------------------------------------------
def save_gcode(bigx, bigy):
    if bigx is None:
        messagebox.showerror("Error", "No pattern generated yet.")
        return

    dbigxy = np.vstack((np.diff(bigx), np.diff(bigy)))

    filename = filedialog.asksaveasfilename(
        defaultextension=".gcode",
        filetypes=[("G-code files", "*.gcode")]
    )
    if not filename:
        return

    with open(filename, "w") as f:
        f.write("G90\nG91\nG1 F500\n")
        f.write(f"G1 X{bigx[0]:f} Y{bigy[0]:f}\n")
        for dx, dy in zip(dbigxy[0], dbigxy[1]):
            f.write(f"G1 X{dx:f} Y{dy:f}\n")
        f.write(f"G1 X{-bigx[-1]:f} Y{-bigy[-1]:f}\n")

    messagebox.showinfo("Saved", f"G-code saved to:\n{filename}")


# ---------------------------------------------------------
# GUI Setup
# ---------------------------------------------------------
root = tk.Tk()
root.title("Scaffold Pattern Generator")
root.geometry("1000x800")

root.protocol("WM_DELETE_WINDOW", on_exit)

# Top frame: inputs + buttons
top_frame = ttk.Frame(root)
top_frame.pack(side="top", fill="x", padx=10, pady=10)

input_frame = ttk.Frame(top_frame)
input_frame.pack(side="left", fill="x", expand=False)

btn_frame = ttk.Frame(top_frame)
btn_frame.pack(side="right", padx=10)

# Bottom frame: plots
plot_frame = ttk.Frame(root)
plot_frame.pack(side="bottom", fill="both", expand=True)

# Input fields
fields = {
    "sw (scaffold width)": 100,
    "sh (scaffold height)": 100,
    "p (pitch)": 0.3,
    "lr (loop radius)": 1.0,
    "pnts (# loop points)": 4
}

entries = {}

for label, default in fields.items():
    row = ttk.Frame(input_frame)
    row.pack(fill="x", pady=2)
    ttk.Label(row, text=label, width=25).pack(side="left")
    ent = ttk.Entry(row)
    ent.insert(0, str(default))
    ent.pack(side="left")
    entries[label] = ent

bigx = None
bigy = None


# ---------------------------------------------------------
# Button callbacks
# ---------------------------------------------------------
def on_plot():
    global bigx, bigy
    try:
        sw = float(entries["sw (scaffold width)"].get())
        sh = float(entries["sh (scaffold height)"].get())
        p = float(entries["p (pitch)"].get())
        lr = float(entries["lr (loop radius)"].get())
        pnts = int(entries["pnts (# loop points)"].get())
    except ValueError:
        messagebox.showerror("Error", "Invalid numeric input.")
        return

    bigx, bigy = compute_and_plot(sw, sh, p, lr, pnts, plot_frame)


def on_save():
    save_gcode(bigx, bigy)

# Buttons
ttk.Button(btn_frame, text="Plot", command=on_plot).pack(side="top", pady=5)
ttk.Button(btn_frame, text="Save G-code", command=on_save).pack(side="top", pady=5)
ttk.Button(btn_frame, text="Exit", command=on_exit).pack(side="top", pady=5)

root.mainloop()