import json
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 (JSON Format)")

        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):
        self.paned_window = ttk.PanedWindow(self.root, orient=tk.HORIZONTAL)
        self.paned_window.pack(fill=tk.BOTH, expand=True)

        # --- Sidebar ---
        sidebar = ttk.Frame(self.paned_window, width=300, padding="10")
        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)

        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")
        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):
        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)

        # Updated to look for .json or .log files
        log_files = sorted(
            [
                f
                for f in self.results_dir.glob("*")
                if f.suffix in [".json", ".log", ".txt"]
            ]
        )
        self.file_listbox.delete(0, tk.END)
        for f in log_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:
                # Use Pandas' built-in fast JSON engine
                # 'lines=True' treats each line as a JSON object
                self.df = pd.read_json(self.current_file, lines=True)

                # Check if the columns are nested and flatten them if necessary
                # If read_json doesn't flatten automatically, we expand the dicts:
                for col in [
                    "angvel_target",
                    "angvel_current",
                    "mot_current",
                    "mot_target",
                    "rot_target",
                    "rot_current",
                ]:
                    if col in self.df.columns:
                        # This expands the dict column into new columns (col.x, col.y, col.z)
                        expanded = pd.json_normalize(self.df[col])
                        expanded.columns = [f"{col}.{c}" for c in expanded.columns]
                        self.df = pd.concat(
                            [self.df.drop(columns=[col]), expanded], axis=1
                        )

                self.update_plot()
            except Exception as e:
                messagebox.showerror("Error", f"Could not load JSON file: {e}")

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

        c = self.axis_var.get()  # current axis (x, y, or z)
        self.ax.clear()

        # Mapping new nested keys to the plot
        # The keys follow the pattern: category.axis
        self.ax.plot(
            self.df["time"], self.df[f"rot_target.{c}"], label=f"Rot Target {c}"
        )
        self.ax.plot(
            self.df["time"], self.df[f"rot_current.{c}"], label=f"Rot Current {c}"
        )

        self.ax.plot(
            self.df["time"], self.df[f"angvel_target.{c}"], label=f"AngVel Target {c}"
        )
        self.ax.plot(
            self.df["time"], self.df[f"angvel_current.{c}"], label=f"AngVel Current {c}"
        )

        self.ax.plot(
            self.df["time"],
            self.df[f"mot_target.{c}"],
            label=f"Mot Target {c}",
            linestyle="--",
        )
        self.ax.plot(
            self.df["time"],
            self.df[f"mot_current.{c}"],
            label=f"Mot Current {c}",
            alpha=0.7,
        )

        self.ax.set_xlabel("Time (s)")
        self.ax.set_ylabel("Value")
        self.ax.set_title(f"{self.current_file.name} — {c.upper()} Axis")
        self.ax.legend(loc="upper right", fontsize="small")
        self.ax.grid(True, alpha=0.3)

        self.canvas.draw()


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