RustPhysicsMQ/analyze.py

import tkinter as tk
from pathlib import Path
from tkinter import messagebox, ttk

import matplotlib.pyplot as plt
import pandas as pd
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg, NavigationToolbar2Tk


class SimVisualizer:
    def __init__(self, root: tk.Tk):
        self.root = root
        self.root.title("Simulation Data Viewer")

        # --- Protocol Handler ---
        # This handles the "X" button on the window frame
        self.root.protocol("WM_DELETE_WINDOW", self.quit_app)

        self.results_dir = Path("results")
        self.df = None
        self.current_file = None

        self.setup_ui()
        self.refresh_file_list()

    def setup_ui(self):
        # 1. Create a PanedWindow (Horizontal orientation)
        self.paned_window = ttk.PanedWindow(self.root, orient=tk.HORIZONTAL)
        self.paned_window.pack(fill=tk.BOTH, expand=True)

        # --- Sidebar Frame ---
        # Note: We still define a width, but it's now the *initial* width
        sidebar = ttk.Frame(self.paned_window, width=300, padding="10")

        # 2. Add the sidebar to the PanedWindow
        # 'weight=0' keeps it from growing automatically when the window is resized
        self.paned_window.add(sidebar, weight=0)

        ttk.Label(sidebar, text="Result Files", font=("Helvetica", 12, "bold")).pack(
            pady=5
        )

        self.file_listbox = tk.Listbox(sidebar, height=15)
        self.file_listbox.pack(fill=tk.X, pady=5)
        self.file_listbox.bind("<<ListboxSelect>>", self.on_file_select)

        ttk.Separator(sidebar, orient="horizontal").pack(fill="x", pady=10)

        ttk.Label(sidebar, text="Select Axis", font=("Helvetica", 10, "bold")).pack(
            pady=5
        )
        self.axis_var = tk.StringVar(value="x")
        for axis in ["x", "y", "z"]:
            ttk.Radiobutton(
                sidebar,
                text=f"{axis.upper()} Axis",
                variable=self.axis_var,
                value=axis,
                command=self.update_plot,
            ).pack(anchor=tk.W)

        # --- Quit Button ---
        spacer = ttk.Label(sidebar, text="")
        spacer.pack(fill=tk.Y, expand=True)

        self.quit_button = ttk.Button(
            sidebar, text="Quit Program", command=self.quit_app
        )
        self.quit_button.pack(fill=tk.X, pady=10)

        # --- Plot Frame ---
        self.plot_frame = ttk.Frame(self.paned_window, padding="10")

        # 3. Add the plot frame to the PanedWindow
        # 'weight=1' ensures the plot area expands to take up all remaining space
        self.paned_window.add(self.plot_frame, weight=1)

        self.fig, self.ax = plt.subplots(figsize=(8, 6))
        self.canvas = FigureCanvasTkAgg(self.fig, master=self.plot_frame)
        self.canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)

        self.toolbar = NavigationToolbar2Tk(self.canvas, self.plot_frame)
        self.toolbar.update()

    def quit_app(self):
        """Cleanly close the plots and the application."""
        plt.close("all")
        self.root.destroy()
        self.root.quit()

    def refresh_file_list(self):
        if not self.results_dir.exists():
            self.results_dir.mkdir(parents=True)

        csv_files = sorted(self.results_dir.glob("*.csv"))
        self.file_listbox.delete(0, tk.END)
        for f in csv_files:
            self.file_listbox.insert(tk.END, f.name)

    def on_file_select(self, event):
        selection = self.file_listbox.curselection()
        if selection:
            filename = self.file_listbox.get(selection[0])
            self.current_file = self.results_dir / filename
            try:
                self.df = pd.read_csv(self.current_file)
                self.update_plot()
            except Exception as e:
                messagebox.showerror("Error", f"Could not load file: {e}")

    def update_plot(self):
        if self.df is None:
            return

        coord = self.axis_var.get()
        self.ax.clear()

        self.ax.plot(
            self.df["time"], self.df[f"target_{coord}"], label=f"target_{coord}"
        )
        self.ax.plot(
            self.df["time"], self.df[f"current_{coord}"], label=f"current_{coord}"
        )
        self.ax.plot(
            self.df["time"], self.df[f"mot_{coord}"].clip(-1, 1), label=f"mot_{coord}"
        )
        # self.ax.plot(
        #     self.df["time"],
        #     self.df[f"dmot_{coord}"].clip(-1, 1),
        #     label=f"desired mot_{coord}",
        #     linestyle="--",
        # )

        self.ax.set_xlabel("Time (s)")
        self.ax.set_ylabel("Value")
        self.ax.set_title(f"{self.current_file.name} — {coord.upper()} Axis")
        self.ax.legend(loc="upper right")
        self.ax.grid(True)

        ymin, ymax = self.ax.get_ylim()
        self.ax.set_ylim(min(ymin, -1.1), max(ymax, 1.1))
        self.canvas.draw()


if __name__ == "__main__":
    root = tk.Tk()
    app = SimVisualizer(root)
    root.mainloop()