import tkinter as tk
from tkinter import ttk, messagebox, filedialog
import serial
import serial.tools.list_ports
import threading
import time
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator, AutoMinorLocator
from matplotlib.backends.backend_tkagg import FigureCanvasTkAgg
from scipy.interpolate import make_interp_spline

TUBE_MODES = {
    "Diode (Dual)": 1,
    "Triode (Dual)": 2,
    "Pentode": 3,
    "Pentode G2 Curves": 4,
    "Quick Test Mode": 5      
}

class TubeTracerApp:
    def __init__(self, root):
        self.root = root
        self.root.title("Valve Analyser MkII (valvewizard.co.uk)")
        self.root.geometry("1350x850") 
        self.root.configure(bg="#1e1e24")
        
        self.serial_port = None
        self.is_running = False
        self.qt_running = False  
        
        self.curves_ch1 = {}  
        self.curves_ch2 = {}  
        self.last_update_time = 0  
        
        self.entry_font = ("Segoe UI", 12)
        
        self.setup_styles()
        self.create_widgets()
        self.on_mode_change()
        self.on_qt_type_change()
        self.toggle_qt_dht_fields()
        
    def setup_styles(self):
        style = ttk.Style()
        style.theme_use('clam')
        
        style.configure(".", background="#1e1e24", foreground="#ffffff", font=("Segoe UI", 12))
        
        style.configure("TLabelframe", background="#2a2a35", bordercolor="#3f3f52", borderwidth=1)
        style.configure("TLabelframe.Label", background="#2a2a35", foreground="#00adb5", font=("Segoe UI", 13, "bold"))
        
        style.configure("TLabel", background="#2a2a35", foreground="#eeeeee", font=("Segoe UI", 12))
        
        style.configure("TEntry", fieldbackground="#1e1e24", foreground="#ffffff", bordercolor="#3f3f52")
        
        style.configure("TCombobox", fieldbackground="#2a2a35", foreground="#ffffff", bordercolor="#3f3f52", arrowcolor="#888888", font=("Segoe UI", 11))
        style.map("TCombobox", fieldbackground=[("readonly", "#2a2a35")], foreground=[("readonly", "#ffffff")])
        
        style.configure("Start.TButton", background="#00bbf9", foreground="#ffffff", font=("Segoe UI", 11, "bold"), borderwidth=0)
        style.map("Start.TButton", background=[("active", "#0096c7"), ("disabled", "#3f3f52")])
        
        style.configure("Stop.TButton", background="#f15bb5", foreground="#ffffff", font=("Segoe UI", 11, "bold"), borderwidth=0)
        style.map("Stop.TButton", background=[("active", "#d90429"), ("disabled", "#3f3f52")])
        
        style.configure("Save.TButton", background="#00adb5", foreground="#ffffff", font=("Segoe UI", 11, "bold"), borderwidth=0)
        style.map("Save.TButton", background=[("active", "#007a87")])

        style.configure("Scan.TButton", background="#3f3f52", foreground="#ffffff", font=("Segoe UI", 10, "bold"), borderwidth=0)
        style.map("Scan.TButton", background=[("active", "#575770")])

        style.configure("TNotebook", background="#1e1e24", borderwidth=0)
        style.configure("TNotebook.Tab", background="#2a2a35", foreground="#eeeeee", padding=[12, 5], font=("Segoe UI", 10, "bold"))
        style.map("TNotebook.Tab", background=[("selected", "#00adb5")], foreground=[("selected", "#ffffff")])

        style.configure("QT.TButton", background="#ff9f1c", foreground="#ffffff", font=("Segoe UI", 11, "bold"), borderwidth=0)
        style.map("QT.TButton", background=[("active", "#e67e22"), ("disabled", "#3f3f52")])

    def create_widgets(self):
        self.notebook = ttk.Notebook(self.root)
        self.notebook.pack(side=tk.LEFT, fill=tk.Y, padx=15, pady=15)
        
        self.sweep_tab = ttk.Frame(self.notebook)
        self.qt_tab = ttk.Frame(self.notebook)
        
        self.notebook.add(self.sweep_tab, text="  Curve Tracer  ")
        self.notebook.add(self.qt_tab, text="  Quick Single Tester  ")
        
        config_frame = ttk.LabelFrame(self.sweep_tab, text=" Curve Tracer ", padding=12)
        config_frame.pack(fill=tk.BOTH, expand=True)
        
        ttk.Label(config_frame, text="Serial Port:").pack(anchor=tk.W, pady=(0, 2))
        
        port_frame = ttk.Frame(config_frame)
        port_frame.pack(anchor=tk.W, pady=(0, 10), fill=tk.X)
        
        self.port_combobox = ttk.Combobox(port_frame, values=self.get_serial_ports(), width=16, state="readonly")
        self.port_combobox.pack(side=tk.LEFT, padx=(0, 5))
        if self.port_combobox['values']:
            self.port_combobox.current(0)
            
        self.scan_btn = ttk.Button(port_frame, text="🔄 Scan", command=self.refresh_serial_ports, style="Scan.TButton", width=7)
        self.scan_btn.pack(side=tk.LEFT)
            
        ttk.Label(config_frame, text="Test Mode:").pack(anchor=tk.W, pady=(0, 2))
        self.mode_combobox = ttk.Combobox(config_frame, values=list(TUBE_MODES.keys())[:4], width=25, state="readonly")
        self.mode_combobox.pack(anchor=tk.W, pady=(0, 10))
        self.mode_combobox.current(1)
        self.mode_combobox.bind("<<ComboboxSelected>>", self.on_mode_change)
        
        ttk.Label(config_frame, text="Max Anode Voltage (V): [Up to 500V]").pack(anchor=tk.W, pady=(0, 2))
        self.va_entry = ttk.Entry(config_frame, width=27, font=self.entry_font)
        self.va_entry.insert(0, "250")
        self.va_entry.pack(anchor=tk.W, pady=(0, 10))
        
        self.vg_label = ttk.Label(config_frame, text="Min Grid Bias 1 (-V): [Down to -60V]")
        self.vg_label.pack(anchor=tk.W, pady=(0, 2))
        self.vg_entry = ttk.Entry(config_frame, width=27, font=self.entry_font)
        self.vg_entry.insert(0, "-4.00")
        self.vg_entry.pack(anchor=tk.W, pady=(0, 10))
        
        self.step_label = ttk.Label(config_frame, text="Grid Step Size (V):")
        self.step_label.pack(anchor=tk.W, pady=(0, 2))
        self.step_entry = ttk.Entry(config_frame, width=27, font=self.entry_font)
        self.step_entry.insert(0, "0.50")
        self.step_entry.pack(anchor=tk.W, pady=(0, 10))

        self.screen_label = ttk.Label(config_frame, text="Screen Volts (V): [Up to 500V]")
        self.screen_label.pack(anchor=tk.W, pady=(0, 2))
        self.screen_entry = ttk.Entry(config_frame, width=27, font=self.entry_font)
        self.screen_entry.insert(0, "200")
        self.screen_entry.pack(anchor=tk.W, pady=(0, 10))

        self.ul_label = ttk.Label(config_frame, text="UL Percent (%):")
        self.ul_label.pack(anchor=tk.W, pady=(0, 2))
        self.ul_entry = ttk.Entry(config_frame, width=27, font=self.entry_font)
        self.ul_entry.insert(0, "0")
        self.ul_entry.pack(anchor=tk.W, pady=(0, 10))

        self.vg2_label = ttk.Label(config_frame, text="Second Grid Bias (-V): [Down to -60V]")
        self.vg2_label.pack(anchor=tk.W, pady=(0, 2))
        self.vg2_entry = ttk.Entry(config_frame, width=27, font=self.entry_font)
        self.vg2_entry.insert(0, "0.00")
        self.vg2_entry.pack(anchor=tk.W, pady=(0, 12))

        ttk.Label(config_frame, text="Max Current Limit (mA) [Optional]:").pack(anchor=tk.W, pady=(0, 2))
        self.max_ia_entry = ttk.Entry(config_frame, width=27, font=self.entry_font)
        self.max_ia_entry.insert(0, "")
        self.max_ia_entry.pack(anchor=tk.W, pady=(0, 10))

        self.dht_var = tk.BooleanVar(value=False)
        self.dht_check = ttk.Checkbutton(config_frame, text="Directly Heated Tube (DHT)", variable=self.dht_var, command=self.toggle_dht_fields)
        self.dht_check.pack(anchor=tk.W, pady=(0, 10))
        
        self.v_filament_label = ttk.Label(config_frame, text="Filament Volts (V):")
        self.v_filament_entry = ttk.Entry(config_frame, width=27, font=self.entry_font)
        self.v_filament_entry.insert(0, "5.0")
        
        self.start_btn = ttk.Button(config_frame, text="▶  Start Sweep", command=self.start_sweep, style="Start.TButton", width=25)
        self.start_btn.pack(anchor=tk.W, pady=(0, 6))
        
        self.stop_btn = ttk.Button(config_frame, text="■  Stop / Reset", command=self.stop_sweep, state=tk.DISABLED, style="Stop.TButton", width=25)
        self.stop_btn.pack(anchor=tk.W, pady=(0, 6))
        
        self.save_btn = ttk.Button(config_frame, text="💾  Save Current Plot", command=self.save_plot, style="Save.TButton", width=25)
        self.save_btn.pack(anchor=tk.W, pady=(0, 10))
        
        self.sep_line = ttk.Separator(config_frame, orient='horizontal')
        self.sep_line.pack(fill='x', pady=5)
        
        self.status_label = ttk.Label(config_frame, text="Status: Idle", font=("Segoe UI", 12, "bold"), foreground="#888888")
        self.status_label.pack(anchor=tk.W, pady=5)
        
        # ---------------------------------------------------------
        # Quick Single Tester
        # ---------------------------------------------------------
        qt_frame = ttk.LabelFrame(self.qt_tab, text=" Quick Single Tester ", padding=12)
        qt_frame.pack(fill=tk.BOTH, expand=True)
        
        ttk.Label(qt_frame, text="Tube Type:").pack(anchor=tk.W, pady=(0, 2))
        self.qt_type_var = tk.StringVar(value="Triode")
        self.qt_type_combo = ttk.Combobox(qt_frame, textvariable=self.qt_type_var, values=["Triode", "Pentode"], state="readonly", width=25)
        self.qt_type_combo.pack(anchor=tk.W, pady=(0, 10))
        self.qt_type_combo.bind("<<ComboboxSelected>>", lambda e: self.on_qt_type_change())
        
        ttk.Label(qt_frame, text="Target Vg1 (-V):").pack(anchor=tk.W, pady=(0, 2))
        self.qt_vg_entry = ttk.Entry(qt_frame, width=27, font=self.entry_font)
        self.qt_vg_entry.insert(0, "2.0")
        self.qt_vg_entry.pack(anchor=tk.W, pady=(0, 10))
        
        ttk.Label(qt_frame, text="Target Va (V):").pack(anchor=tk.W, pady=(0, 2))
        self.qt_va_entry = ttk.Entry(qt_frame, width=27, font=self.entry_font)
        self.qt_va_entry.insert(0, "250")
        self.qt_va_entry.pack(anchor=tk.W, pady=(0, 10))
        
        self.qt_vsc_lbl = ttk.Label(qt_frame, text="Screen Vg2 (V):")
        self.qt_vsc_entry = ttk.Entry(qt_frame, width=27, font=self.entry_font)
        self.qt_vsc_entry.insert(0, "150")
        
        self.qt_dht_var = tk.BooleanVar(value=False)
        self.qt_dht_check = ttk.Checkbutton(qt_frame, text="Directly Heated Tube (DHT)", variable=self.qt_dht_var, command=self.toggle_qt_dht_fields)
        self.qt_dht_check.pack(anchor=tk.W, pady=(5, 5))
        
        self.qt_vf_lbl = ttk.Label(qt_frame, text="Filament Volts (V):")
        self.qt_vf_entry = ttk.Entry(qt_frame, width=27, font=self.entry_font)
        self.qt_vf_entry.insert(0, "5.0")
        
        self.results_box = ttk.LabelFrame(qt_frame, text=" Test Results: ", padding=8)
        self.results_box.pack(fill=tk.X, pady=(10, 15))
        
        ttk.Label(self.results_box, text="Param", font=("Segoe UI", 12, "bold"), foreground="#00adb5").grid(row=0, column=0, padx=5, pady=4, sticky="w")
        self.ch1_title = ttk.Label(self.results_box, text="CH1 (Triode1)", font=("Segoe UI", 12, "bold"), foreground="#00adb5")
        self.ch1_title.grid(row=0, column=1, padx=12, pady=4, sticky="w")
        self.ch2_title = ttk.Label(self.results_box, text="CH2 (Triode2)", font=("Segoe UI", 12, "bold"), foreground="#00adb5")
        self.ch2_title.grid(row=0, column=2, padx=12, pady=4, sticky="w")
        
        ttk.Label(self.results_box, text="Vg1:").grid(row=1, column=0, padx=5, pady=3, sticky="w")
        self.ch1_vg_lbl = ttk.Label(self.results_box, text="-", font=("Segoe UI", 12, "bold"))
        self.ch1_vg_lbl.grid(row=1, column=1, padx=12, pady=3, sticky="w")
        self.ch2_vg_lbl = ttk.Label(self.results_box, text="-", font=("Segoe UI", 12, "bold"))
        self.ch2_vg_lbl.grid(row=1, column=2, padx=12, pady=3, sticky="w")
        
        ttk.Label(self.results_box, text="Va:").grid(row=2, column=0, padx=5, pady=3, sticky="w")
        self.ch1_va_lbl = ttk.Label(self.results_box, text="-", font=("Segoe UI", 12, "bold"))
        self.ch1_va_lbl.grid(row=2, column=1, padx=12, pady=3, sticky="w")
        self.ch2_va_lbl = ttk.Label(self.results_box, text="-", font=("Segoe UI", 12, "bold"))
        self.ch2_va_lbl.grid(row=2, column=2, padx=12, pady=3, sticky="w")
        
        ttk.Label(self.results_box, text="Ia:").grid(row=3, column=0, padx=5, pady=3, sticky="w")
        self.ch1_ia_lbl = ttk.Label(self.results_box, text="-", font=("Segoe UI", 12, "bold"))
        self.ch1_ia_lbl.grid(row=3, column=1, padx=12, pady=3, sticky="w")
        self.ch2_ia_lbl = ttk.Label(self.results_box, text="-", font=("Segoe UI", 12, "bold"))
        self.ch2_ia_lbl.grid(row=3, column=2, padx=12, pady=3, sticky="w")

        self.vsc_result_lbl = ttk.Label(self.results_box, text="Vsc:")
        self.ch1_vsc_lbl = ttk.Label(self.results_box, text="-", font=("Segoe UI", 12, "bold"))
        
        self.isc_result_lbl = ttk.Label(self.results_box, text="Isc:")
        self.ch1_isc_lbl = ttk.Label(self.results_box, text="-", font=("Segoe UI", 12, "bold"))
        
        self.gm_result_lbl = ttk.Label(self.results_box, text="gm:")
        self.gm_result_lbl.grid(row=6, column=0, padx=5, pady=3, sticky="w")
        self.ch1_gm_lbl = ttk.Label(self.results_box, text="-", font=("Segoe UI", 12, "bold"))
        self.ch1_gm_lbl.grid(row=6, column=1, padx=12, pady=3, sticky="w")
        self.ch2_gm_lbl = ttk.Label(self.results_box, text="-", font=("Segoe UI", 12, "bold"))
        self.ch2_gm_lbl.grid(row=6, column=2, padx=12, pady=3, sticky="w")
        
        self.ra_result_lbl = ttk.Label(self.results_box, text="ra:")
        self.ra_result_lbl.grid(row=7, column=0, padx=5, pady=3, sticky="w")
        self.ch1_ra_lbl = ttk.Label(self.results_box, text="-", font=("Segoe UI", 12, "bold"))
        self.ch1_ra_lbl.grid(row=7, column=1, padx=12, pady=3, sticky="w")
        self.ch2_ra_lbl = ttk.Label(self.results_box, text="-", font=("Segoe UI", 12, "bold"))
        self.ch2_ra_lbl.grid(row=7, column=2, padx=12, pady=3, sticky="w")
        
        self.mu_result_lbl = ttk.Label(self.results_box, text="mu:")
        self.mu_result_lbl.grid(row=8, column=0, padx=5, pady=3, sticky="w")
        self.ch1_mu_lbl = ttk.Label(self.results_box, text="-", font=("Segoe UI", 12, "bold"))
        self.ch1_mu_lbl.grid(row=8, column=1, padx=12, pady=3, sticky="w")
        self.ch2_mu_lbl = ttk.Label(self.results_box, text="-", font=("Segoe UI", 12, "bold"))
        self.ch2_mu_lbl.grid(row=8, column=2, padx=12, pady=3, sticky="w")
        
        self.qt_start_btn = ttk.Button(qt_frame, text="⚡  Run Quick Test", command=self.execute_quick_test, style="QT.TButton", width=25)
        self.qt_start_btn.pack(anchor=tk.W, pady=(5, 4))
        
        self.qt_stop_btn = ttk.Button(qt_frame, text="■  Stop / Reset ", command=self.stop_quick_test, style="Stop.TButton", width=25)
        self.qt_stop_btn.pack(anchor=tk.W, pady=(0, 5))
        self.qt_stop_btn.config(state=tk.DISABLED)

        plot_frame = ttk.Frame(self.root, padding=10)
        plot_frame.pack(side=tk.RIGHT, fill=tk.BOTH, expand=True, padx=(0, 15), pady=15)
        
        plt.style.use('default') 
        self.fig, self.ax = plt.subplots(figsize=(7, 5), facecolor="#ffffff")
        self.setup_plot_axes()
        
        self.canvas = FigureCanvasTkAgg(self.fig, master=plot_frame)
        self.canvas.get_tk_widget().pack(fill=tk.BOTH, expand=True)

    def on_mode_change(self, event=None):
        selected_mode = self.mode_combobox.get()
        
        self.vg_entry.config(state=tk.NORMAL)
        self.step_entry.config(state=tk.NORMAL)
        self.screen_entry.config(state=tk.NORMAL)
        self.ul_entry.config(state=tk.NORMAL)
        self.vg2_entry.config(state=tk.NORMAL)
        self.step_label.config(text="Grid Step Size (V):")
        
        if selected_mode == "Diode (Dual)":
            self.vg_entry.config(state=tk.DISABLED)
            self.step_entry.config(state=tk.DISABLED)
            self.screen_entry.config(state=tk.DISABLED)
            self.ul_entry.config(state=tk.DISABLED)
            self.vg2_entry.config(state=tk.DISABLED)
        elif selected_mode == "Triode (Dual)":
            self.screen_entry.config(state=tk.DISABLED)
            self.ul_entry.config(state=tk.DISABLED)
            self.vg2_entry.config(state=tk.DISABLED)
        elif selected_mode == "Pentode":
            pass
        elif selected_mode == "Pentode G2 Curves":
            self.vg_entry.config(state=tk.DISABLED)
            self.ul_entry.config(state=tk.DISABLED)
            self.vg2_entry.config(state=tk.DISABLED)
            self.step_label.config(text="G2 Step Size (V):")
            
        self.update_plot()

    def on_qt_type_change(self):
        if self.qt_type_var.get() == "Triode":
            self.qt_vsc_lbl.pack_forget()
            self.qt_vsc_entry.pack_forget()
            
            self.vsc_result_lbl.grid_forget()
            self.ch1_vsc_lbl.grid_forget()
            self.isc_result_lbl.grid_forget()
            self.ch1_isc_lbl.grid_forget()
            
            self.ch1_title.config(text="CH1 (Triode1)")
            self.ch2_title.grid(row=0, column=2, padx=12, pady=4, sticky="w")
            self.ch2_vg_lbl.grid(row=1, column=2, padx=12, pady=3, sticky="w")
            self.ch2_va_lbl.grid(row=2, column=2, padx=12, pady=3, sticky="w")
            self.ch2_ia_lbl.grid(row=3, column=2, padx=12, pady=3, sticky="w")
            self.ch2_gm_lbl.grid(row=6, column=2, padx=12, pady=3, sticky="w")
            self.ch2_ra_lbl.grid(row=7, column=2, padx=12, pady=3, sticky="w")
            self.ch2_mu_lbl.grid(row=8, column=2, padx=12, pady=3, sticky="w")
        else:
            self.qt_vsc_lbl.pack(before=self.qt_dht_check, anchor=tk.W, pady=(0, 2))
            self.qt_vsc_entry.pack(before=self.qt_dht_check, anchor=tk.W, pady=(0, 10))
            
            self.vsc_result_lbl.grid(row=4, column=0, padx=5, pady=3, sticky="w")
            self.ch1_vsc_lbl.grid(row=4, column=1, padx=12, pady=3, sticky="w")
            self.isc_result_lbl.grid(row=5, column=0, padx=5, pady=3, sticky="w")
            self.ch1_isc_lbl.grid(row=5, column=1, padx=12, pady=3, sticky="w")
            
            self.ch1_title.config(text="Result")
            self.ch2_title.grid_forget()
            self.ch2_vg_lbl.grid_forget()
            self.ch2_va_lbl.grid_forget()
            self.ch2_ia_lbl.grid_forget()
            self.ch2_gm_lbl.grid_forget()
            self.ch2_ra_lbl.grid_forget()
            self.ch2_mu_lbl.grid_forget()

    def toggle_dht_fields(self):
        if self.dht_var.get():
            self.start_btn.pack_forget()
            self.stop_btn.pack_forget()
            self.save_btn.pack_forget()
            self.sep_line.pack_forget()
            self.status_label.pack_forget()
            
            self.v_filament_label.pack(anchor=tk.W, pady=(0, 2))
            self.v_filament_entry.pack(anchor=tk.W, pady=(0, 10))
            
            self.start_btn.pack(anchor=tk.W, pady=(0, 6))
            self.stop_btn.pack(anchor=tk.W, pady=(0, 6))
            self.save_btn.pack(anchor=tk.W, pady=(0, 10))
            self.sep_line.pack(fill='x', pady=5)
            self.status_label.pack(anchor=tk.W, pady=5)
        else:
            self.v_filament_label.pack_forget()
            self.v_filament_entry.pack_forget()

    def toggle_qt_dht_fields(self):
        if self.qt_dht_var.get():
            self.qt_vf_lbl.pack(before=self.results_box, anchor=tk.W, pady=(0, 2))
            self.qt_vf_entry.pack(before=self.results_box, anchor=tk.W, pady=(0, 10))
        else:
            self.qt_vf_lbl.pack_forget()
            self.qt_vf_entry.pack_forget()

    def setup_plot_axes(self):
        self.ax.set_facecolor("#ffffff")
        current_mode = self.mode_combobox.get() if hasattr(self, 'mode_combobox') else "Tube"
        self.ax.set_title(f"Tube Characteristic Curves ({current_mode})", fontsize=13, fontweight='bold', pad=15, color="#1e1e24")
        self.ax.set_xlabel("Anode Voltage Va (V)", fontsize=11, labelpad=8, color="#1e1e24")
        self.ax.set_ylabel("Anode Current Ia (mA)", fontsize=11, labelpad=8, color="#1e1e24")
        
        self.ax.minorticks_on()

        self.ax.xaxis.set_major_locator(MultipleLocator(50)) 
        self.ax.xaxis.set_minor_locator(MultipleLocator(10))

        self.ax.yaxis.set_major_locator(plt.MaxNLocator(nbins='auto', steps=[1, 5, 10]))
        self.ax.yaxis.set_minor_locator(AutoMinorLocator())            
        self.ax.grid(True, which='major', linestyle='--', alpha=0.6, color="#cccccc")
        self.ax.grid(True, which='minor', linestyle=':', alpha=0.3, color="#aaaaaa")
    
        self.ax.tick_params(colors='#1e1e24', labelsize=10)
        self.ax.tick_params(which='minor', bottom=False, left=False)
        self.ax.grid(True, linestyle='--', alpha=0.5, color="#cccccc")
        self.ax.tick_params(colors='#1e1e24', labelsize=10)

    def get_serial_ports(self):
        return [port.device for port in serial.tools.list_ports.comports()]

    def refresh_serial_ports(self):
        ports = self.get_serial_ports()
        self.port_combobox['values'] = ports
        if ports:
            self.port_combobox.current(0)
            self.status_label.config(text="Status: Ports refreshed", foreground="#00adb5")
        else:
            self.port_combobox.set('')
            self.status_label.config(text="Status: No ports found", foreground="#f15bb5")

    def start_sweep(self):
        port = self.port_combobox.get()
        if not port:
            messagebox.showerror("Error", "Please select a valid Serial Port!")
            return
            
        try:
            self.serial_port = serial.Serial(port, 57600, timeout=1)
            self.is_running = True
            
            self.start_btn.config(state=tk.DISABLED)
            self.stop_btn.config(state=tk.NORMAL)
            self.status_label.config(text="Status: Connecting...", foreground="#ffb703")
            
            self.curves_ch1.clear()
            self.curves_ch2.clear()
            self.ax.clear()
            self.setup_plot_axes()
            self.canvas.draw()
            self.last_update_time = time.time()
            
            threading.Thread(target=self.serial_handler_thread, daemon=True).start()
            
        except Exception as e:
            messagebox.showerror("Connection Error", f"Could not open serial port: {e}")
            self.stop_sweep()

    def serial_handler_thread(self):
        time.sleep(2.5)
        buffer = ""
        init_time = time.time()
        while self.is_running and (time.time() - init_time < 5):
            if self.serial_port.in_waiting:
                buffer += self.serial_port.read(self.serial_port.in_waiting).decode('utf-8', errors='ignore')
                if "Select test" in buffer or "1 =" in buffer:
                    break
        
        mode_idx = TUBE_MODES[self.mode_combobox.get()]
        
        try:
            va_max = float(self.va_entry.get().strip())
            vg_min = float(self.vg_entry.get().replace('-', '').strip())
            step_size = float(self.step_entry.get().strip())
            screen_v = float(self.screen_entry.get().strip())
            ul_pct = float(self.ul_entry.get().strip())
            vg2_bias = float(self.vg2_entry.get().replace('-', '').strip())
            max_ia_input = self.max_ia_entry.get().strip()
            max_ia_limit = float(max_ia_input) if max_ia_input else None
            max_current_to_send = f"{int(float(max_ia_input))}\n" if max_ia_input else "0\n"
        except ValueError:
            messagebox.showerror("Error", "Please enter valid numeric values.")
            self.stop_sweep()
            return

        v_offset = 0.0
        if self.dht_var.get():
            try:
                v_offset = float(self.v_filament_entry.get()) / 2.0
            except ValueError:
                pass
        
        try:
            self.serial_port.reset_input_buffer()
            
            fmt = lambda v: f"{float(v):g}"
            
            self.serial_port.write(f"{mode_idx}\n".encode())
            time.sleep(0.25)
            
            self.serial_port.write(f"{fmt(va_max)}\n".encode())
            time.sleep(0.25)
            
            if mode_idx == 2:  # Triode

                self.serial_port.write(f"{fmt(abs(vg_min))}\n".encode())
                time.sleep(0.25)
                
                self.serial_port.write(f"{fmt(step_size)}\n".encode())
                time.sleep(0.25)
                
            elif mode_idx == 3:  # Pentode
                self.serial_port.write(f"{fmt(abs(vg_min))}\n".encode())
                time.sleep(0.25)
                self.serial_port.write(f"{fmt(step_size)}\n".encode())
                time.sleep(0.25)
                self.serial_port.write(f"{fmt(screen_v)}\n".encode())
                time.sleep(0.25)
                self.serial_port.write(f"{fmt(ul_pct)}\n".encode())
                time.sleep(0.25)
                self.serial_port.write(f"{fmt(abs(vg2_bias))}\n".encode())
                time.sleep(0.25)
                
            elif mode_idx == 4:  
                self.serial_port.write(f"{fmt(screen_v)}\n".encode())
                time.sleep(0.25)
                self.serial_port.write(f"{fmt(step_size)}\n".encode())  
                time.sleep(0.25)
                
            self.serial_port.write(max_current_to_send.encode())
            time.sleep(0.25)

            self.serial_port.write(b"y\n")
            time.sleep(0.1)
            
            self.root.after(0, lambda: self.status_label.config(text="Status: Sweeping...", foreground="#00bbf9"))
        except Exception as e:
            print(f"Error sending config: {e}")
            self.stop_sweep()
            return

        while self.is_running:
            try:
                if self.serial_port.in_waiting:
                    line = self.serial_port.readline().decode('utf-8', errors='ignore').strip()
                    if not line:
                        continue
                    
                    print(f"Sweep Output Line: {line}")
                    
                    if "error" in line.lower():
                        self.root.after(0, lambda msg=line: messagebox.showerror("Arduino Hardware Error", msg))
                        self.root.after(0, self.stop_sweep)
                        break
                        
                    if "Complete" in line or "End" in line or "Select test" in line:
                        self.root.after(0, lambda: self.status_label.config(text="Status: Test Complete", foreground="#0096c7"))
                        break
                        
                    parts = [p.strip() for p in line.split(',') if p.strip()]
                    if any("Vb" in p or "HV" in p or "Ia" in p for p in parts):
                        continue
                        
                    if len(parts) >= 6:
                        try:
                            v_grid1 = float(parts[0])
                            v_anode1 = float(parts[1])
                            i_anode1 = float(parts[2])
                            
                            v_grid2 = float(parts[3])
                            v_anode2 = float(parts[4])
                            i_anode2 = float(parts[5])

                            if self.dht_var.get():
                                v_grid1 -= v_offset  
                                v_grid2 -= v_offset
                                v_anode1 = max(0.0, v_anode1 - v_offset)
                                v_anode2 = max(0.0, v_anode2 - v_offset)
                            

                            if mode_idx == 4:
                                step = step_size if step_size > 0 else 10.0
                                v_key = float(round(v_anode2 / step) * step)
                                
                                if v_key not in self.curves_ch1:
                                    self.curves_ch1[v_key] = ([], [])
                                self.curves_ch1[v_key][0].append(v_anode1)
                                self.curves_ch1[v_key][1].append(i_anode1)
                            else:

                                v_grid1 = round(v_grid1, 2)
                                
                                if v_grid1 not in self.curves_ch1:
                                    self.curves_ch1[v_grid1] = ([], [])
                                self.curves_ch1[v_grid1][0].append(v_anode1)
                                self.curves_ch1[v_grid1][1].append(i_anode1)
                            
                                if v_grid1 not in self.curves_ch2:
                                    self.curves_ch2[v_grid1] = ([], [])
                                self.curves_ch2[v_grid2][0].append(v_anode2)
                                self.curves_ch2[v_grid1][1].append(i_anode2)

                            current_time = time.time()
                            if current_time - self.last_update_time > 0.15:  
                                self.last_update_time = current_time
                                self.root.after(0, self.update_plot)
                            
                        except ValueError:
                            pass
            except Exception as e:
                break
                
        self.root.after(0, self.update_plot)
        if self.is_running:
            self.is_running = False
            if self.serial_port and self.serial_port.is_open:
                try: self.serial_port.close()
                except: pass
            self.start_btn.config(state=tk.NORMAL)
            self.stop_btn.config(state=tk.DISABLED)

    def draw_smooth_curve(self, x_data, y_data):
        x = np.array(x_data, dtype=float)
        y = np.array(y_data, dtype=float)
        idx = np.argsort(x)
        x = x[idx]; y = y[idx]
        _, unique_indices = np.unique(x, return_index=True)
        x = x[unique_indices]; y = y[unique_indices]
        
        if len(x) < 2: 
            return None, None
            
        if len(x) >= 4:
            try:
                x_new = np.linspace(x.min(), x.max(), 300)
                spl = make_interp_spline(x, y, k=3)
                y_new = spl(x_new)
                y_new = np.clip(y_new, 0, None)
                return x_new, y_new
            except Exception:
                pass
                
        return x, y

    def format_vg(self, v_val):
        try:
            val = float(v_val)
            rounded = round(val, 1)
            if rounded.is_integer():
                return f"{int(rounded)}V"
            return f"{rounded}V"
        except (ValueError, TypeError):
            return f"{v_val}V"

    def update_plot(self):
        self.ax.clear()
        self.setup_plot_axes()
        max_valid_ia = 0.1 
        
        current_mode = self.mode_combobox.get()
        
        max_ia_input = self.max_ia_entry.get().strip()
        try:
            max_ia_limit = float(max_ia_input) if max_ia_input else None
            if max_ia_limit is not None and max_ia_limit <= 0:
                max_ia_limit = None
        except ValueError:
            max_ia_limit = None

        all_vgs = sorted(list(self.curves_ch1.keys()), key=float, reverse=True)
        color_cycle = plt.rcParams['axes.prop_cycle'].by_key()['color']
        vg_color_map = {vg: color_cycle[i % len(color_cycle)] for i, vg in enumerate(all_vgs)}
        
        for v_grid, (va, ia) in self.curves_ch1.items():
            if len(va) > 1:
                if max(ia) > max_valid_ia: max_valid_ia = max(max_valid_ia, max(ia))
                color = vg_color_map.get(v_grid)
                
                x_curve, y_curve = self.draw_smooth_curve(va, ia)
                
                if x_curve is not None and len(x_curve) > 0:
                    if max_ia_limit is not None:
                        valid_indices = np.where(y_curve <= max_ia_limit)[0]
                        if len(valid_indices) > 0:
                            last_idx = valid_indices[-1]
                            x_plot = x_curve[:last_idx + 1]
                            y_plot = y_curve[:last_idx + 1]
                        else:
                            x_plot, y_plot = x_curve[:2], y_curve[:2]
                    else:
                        x_plot, y_plot = x_curve, y_curve

                    if current_mode == "Pentode G2 Curves":
                        legend_label = f"Vg2={v_grid}V"
                    elif current_mode == "Pentode":
                        legend_label = f"Ia (Vg1={v_grid}V)"
                    else:
                        legend_label = f"Ch1 Vg={v_grid}V"

                    self.ax.plot(x_plot, y_plot, linestyle='-', linewidth=2.5, color=color, label=legend_label)

                    x_end = x_plot[-1]
                    y_end = y_plot[-1]
                    
                    dx = x_plot[-1] - x_plot[-2] if len(x_plot) > 1 else 1.0
                    dy = y_plot[-1] - y_plot[-2] if len(x_plot) > 1 else 0.0
                    angle = np.degrees(np.arctan2(dy, dx)) * 0.7
                    
                    v_formatted = self.format_vg(v_grid) if hasattr(self, 'format_vg') else f"{v_grid}V"
                    if current_mode == "Pentode G2 Curves":
                        vg_str = f"Vg2={v_formatted}"
                    else:
                        vg_str = v_formatted

                    if max_ia_limit and y_end >= (max_ia_limit - 0.5):
                        y_text_pos = max_ia_limit + (max_ia_limit * 0.015)
                        self.ax.text(x_end, y_text_pos, vg_str, 
                                    fontsize=9, color=color, fontweight='bold',
                                    rotation=angle, rotation_mode='anchor', 
                                    verticalalignment='bottom', 
                                    horizontalalignment='center')
                    else:
                        self.ax.text(x_end + 2.0, y_end, vg_str, 
                                    fontsize=9, color=color, fontweight='bold',
                                    rotation=0, rotation_mode='anchor', 
                                    verticalalignment='center', 
                                    horizontalalignment='left')

        for v_grid, (va, _) in self.curves_ch1.items():
            if v_grid in self.curves_ch2:
                ig2 = self.curves_ch2[v_grid][1]
                
                min_len = min(len(va), len(ig2))
                if min_len > 1:
                    va_plot = va[:min_len]
                    ig2_plot = ig2[:min_len]
                    
                    if max(ig2_plot) > max_valid_ia:
                        max_valid_ia = max(max_valid_ia, max(ig2_plot))
                    
                    color = vg_color_map.get(v_grid)
                    
                    x_curve, y_curve = self.draw_smooth_curve(va_plot, ig2_plot)
                    
                    if x_curve is not None and len(x_curve) > 0:
                        if max_ia_limit is not None:
                            valid_indices = np.where(y_curve <= max_ia_limit)[0]
                            if len(valid_indices) > 0:
                                last_idx = valid_indices[-1]
                                x_plot = x_curve[:last_idx + 1]
                                y_plot = y_curve[:last_idx + 1]
                            else:
                                x_plot, y_plot = x_curve[:2], y_curve[:2]
                        else:
                            x_plot, y_plot = x_curve, y_curve

                        if current_mode == "Pentode":
                            lbl = f"Ig2 (Vg1={v_grid}V)"
                        else:
                            lbl = f"Ch2 Vg={v_grid}V"

                        self.ax.plot(x_plot, y_plot, linestyle='--', linewidth=1.5, alpha=0.8, color=color, label=lbl)
        
        handles, labels = self.ax.get_legend_handles_labels()
        if handles:
            self.ax.legend(handles, labels, loc="upper left", bbox_to_anchor=(1.02, 1), fontsize='small', facecolor="#ffffff", edgecolor="#cccccc")
        
        if max_ia_limit:
            self.ax.set_ylim(bottom=0, top=max_ia_limit * 1.08)
        else:
            if max_valid_ia > 0.1:
                self.ax.set_ylim(bottom=0, top=max_valid_ia * 1.15)
            else:
                self.ax.set_ylim(bottom=0)

        self.ax.set_xlim(left=0)
        self.fig.tight_layout() 
        self.canvas.draw()
        self.root.update_idletasks()

    def save_plot(self):
        file_path = filedialog.asksaveasfilename(defaultextension=".png", filetypes=[("PNG Image", "*.png"), ("All Files", "*.*")])
        if file_path:
            try:
                self.fig.savefig(file_path, dpi=300, bbox_inches='tight', facecolor="#ffffff")
                messagebox.showinfo("Success", "Plot saved successfully!")
            except Exception as e: messagebox.showerror("Error", f"Could not save plot: {e}")

    def stop_sweep(self):
        self.is_running = False
        if self.serial_port and self.serial_port.is_open:
            try:
                self.serial_port.dtr = False
                self.serial_port.rts = False
                time.sleep(0.05)
                self.serial_port.dtr = True
                self.serial_port.rts = True
                time.sleep(0.05)
                self.serial_port.close()
            except: 
                pass
        self.start_btn.config(state=tk.NORMAL)
        self.stop_btn.config(state=tk.DISABLED)
        if "Complete" not in self.status_label.cget("text") and "Error" not in self.status_label.cget("text"):
            self.status_label.config(text="Status: Stopped & Reset", foreground="#f15bb5")

    def execute_quick_test(self):
        if self.qt_running:
            return
            
        port = self.port_combobox.get()
        if not port:
            messagebox.showerror("Error", "Please select a valid Serial Port!")
            return
            
        try:
            vg_input = float(self.qt_vg_entry.get().replace('-', '').strip())
            va_input = float(self.qt_va_entry.get().replace('-', '').strip())
            vsc_input = float(self.qt_vsc_entry.get().replace('-', '').strip()) if self.qt_type_var.get() == "Pentode" else 0.0
            tube_type_num = 1 if self.qt_type_var.get() == "Triode" else 2
        except ValueError:
            messagebox.showerror("Error", "Please enter valid parameters for Quick Test.")
            return

        is_dht = self.qt_dht_var.get()
        v_offset = 0.0
        if is_dht:
            try:
                v_offset = float(self.qt_vf_entry.get().replace('-', '').strip()) / 2.0
            except ValueError:
                v_offset = 0.0

        vg_to_send = max(0.0, vg_input - v_offset)
        va_to_send = va_input + v_offset
        vsc_to_send = vsc_input + v_offset

        try:
            self.serial_port = serial.Serial(port, 57600, timeout=1)
            self.qt_running = True
            self.qt_start_btn.config(state=tk.DISABLED)
            self.qt_stop_btn.config(state=tk.NORMAL)
            self.status_label.config(text="Status: Connecting QT...", foreground="#ffb703")
        except Exception as e:
            messagebox.showerror("Connection Error", f"Could not open serial port: {e}")
            return

        for lbl in [self.ch1_vg_lbl, self.ch1_va_lbl, self.ch1_ia_lbl, self.ch1_gm_lbl, self.ch1_ra_lbl, self.ch1_mu_lbl,
                    self.ch2_vg_lbl, self.ch2_va_lbl, self.ch2_ia_lbl, self.ch2_gm_lbl, self.ch2_ra_lbl, self.ch2_mu_lbl,
                    self.ch1_vsc_lbl, self.ch1_isc_lbl]:
            lbl.config(text="-")

        def run_qt():
            time.sleep(2.5)
            buffer = ""
            init_time = time.time()
            
            while self.qt_running and (time.time() - init_time < 5):
                if self.serial_port.in_waiting:
                    buffer += self.serial_port.read(self.serial_port.in_waiting).decode('utf-8', errors='ignore')
                    if "Select test" in buffer or "1 =" in buffer:
                        break
                        
            try:
                self.serial_port.reset_input_buffer()
                self.serial_port.write(b"5\n")
                time.sleep(0.4)
                
                self.serial_port.write(f"{tube_type_num}\n".encode())
                time.sleep(0.4)
                
                self.serial_port.write(f"{vg_to_send:.2f}\n".encode())
                time.sleep(0.4)
                
                self.serial_port.write(f"{int(va_to_send)}\n".encode())
                time.sleep(0.4)
                
                if tube_type_num == 2:
                    self.serial_port.write(f"{int(vsc_to_send)}\n".encode())
                    time.sleep(0.4)
                    
                self.serial_port.write(b"y\n")
                self.root.after(0, lambda: self.status_label.config(text="Status: Quick Testing...", foreground="#00bbf9"))
                
                start_time = time.time()
                while self.qt_running and (time.time() - start_time < 15):
                    if self.serial_port.in_waiting:
                        line = self.serial_port.readline().decode('utf-8', errors='ignore').strip()
                        if not line: continue
                        
                        print(f"QT Output Line: {line}")
                        
                        if "error" in line.lower():
                            self.root.after(0, lambda msg=line: messagebox.showerror("Arduino Hardware Error", msg))
                            self.root.after(0, self.stop_quick_test)
                            break
                        
                        if "," in line and not any(k in line for k in ["Vg1", "Va1", "gm1", "Vb", "Vg(V)", "Va(V)"]):
                            parts = [p.strip() for p in line.split(',') if p.strip()]
                            
                            if tube_type_num == 1 and len(parts) >= 12:
                                try:
                                    v_g1 = float(parts[0])
                                    v_a1 = float(parts[1])
                                    v_g2 = float(parts[6])
                                    v_a2 = float(parts[7])
                                    
                                    if is_dht:
                                        v_g1 = -(abs(v_g1) + v_offset)
                                        v_g2 = -(abs(v_g2) + v_offset)
                                        v_a1 = v_a1 - v_offset
                                        v_a2 = v_a2 - v_offset
                                        
                                    val_vg1, val_va1 = v_g1, v_a1
                                    val_vg2, val_va2 = v_g2, v_a2
                                    
                                    self.root.after(0, lambda v=val_vg1: self.ch1_vg_lbl.config(text=f"{v:.2f} V"))
                                    self.root.after(0, lambda v=val_va1: self.ch1_va_lbl.config(text=f"{v:.1f} V"))
                                    self.root.after(0, lambda v=val_vg2: self.ch2_vg_lbl.config(text=f"{v:.2f} V"))
                                    self.root.after(0, lambda v=val_va2: self.ch2_va_lbl.config(text=f"{v:.1f} V"))

                                    self.root.after(0, lambda: self.ch1_ia_lbl.config(text=f"{parts[2]} mA"))
                                    self.root.after(0, lambda: self.ch1_gm_lbl.config(text=f"{parts[3]} mA/V"))
                                    self.root.after(0, lambda: self.ch1_ra_lbl.config(text=f"{parts[4]} kΩ"))
                                    self.root.after(0, lambda: self.ch1_mu_lbl.config(text=parts[5]))
                                    
                                    self.root.after(0, lambda: self.ch2_ia_lbl.config(text=f"{parts[8]} mA"))
                                    self.root.after(0, lambda: self.ch2_gm_lbl.config(text=f"{parts[9]} mA/V"))
                                    self.root.after(0, lambda: self.ch2_ra_lbl.config(text=f"{parts[10]} kΩ"))
                                    self.root.after(0, lambda: self.ch2_mu_lbl.config(text=parts[11]))
                                    self.root.after(0, lambda: self.status_label.config(text="Status: QT Finished", foreground="#0096c7"))
                                    break
                                except ValueError:
                                    pass
                                
                            elif tube_type_num == 2 and len(parts) >= 8:
                                try:
                                    v_g1 = float(parts[0])
                                    v_a1 = float(parts[1])
                                    try:
                                        v_sc1 = float(parts[3])
                                    except ValueError:
                                        v_sc1 = 0.0
                                    
                                    if is_dht:
                                        v_g1 = -(abs(v_g1) + v_offset)
                                        v_a1 = v_a1 - v_offset
                                        v_sc1 = v_sc1 - v_offset
                                        
                                    val_vg1, val_va1, val_vsc1 = v_g1, v_a1, v_sc1
                                    
                                    self.root.after(0, lambda v=val_vg1: self.ch1_vg_lbl.config(text=f"{v:.2f} V"))
                                    self.root.after(0, lambda v=val_va1: self.ch1_va_lbl.config(text=f"{v:.1f} V"))
                                    self.root.after(0, lambda v=val_vsc1: self.ch1_vsc_lbl.config(text=f"{v:.1f} V"))

                                    self.root.after(0, lambda: self.ch1_ia_lbl.config(text=f"{parts[2]} mA"))
                                    self.root.after(0, lambda: self.ch1_isc_lbl.config(text=f"{parts[4]} mA"))
                                    self.root.after(0, lambda: self.ch1_gm_lbl.config(text=f"{parts[5]} mA/V"))
                                    self.root.after(0, lambda: self.ch1_ra_lbl.config(text=f"{parts[6]} kΩ"))
                                    self.root.after(0, lambda: self.ch1_mu_lbl.config(text=parts[7]))
                                    self.root.after(0, lambda: self.status_label.config(text="Status: QT Finished", foreground="#0096c7"))
                                    break
                                except ValueError:
                                    pass
            except Exception as e:
                print(f"Error in QT serial thread: {e}")
            finally:
                self.cleanup_qt_connection()

        threading.Thread(target=run_qt, daemon=True).start()

    def stop_quick_test(self):
        self.qt_running = False
        self.cleanup_qt_connection()
        if "Finished" not in self.status_label.cget("text"):
            self.status_label.config(text="Status: QT Stopped & Reset", foreground="#f15bb5")

    def cleanup_qt_connection(self):
        if self.serial_port and self.serial_port.is_open:
            try:
                self.serial_port.dtr = False
                self.serial_port.rts = False
                time.sleep(0.05)
                self.serial_port.dtr = True
                self.serial_port.rts = True
                time.sleep(0.05)
                self.serial_port.close()
            except: 
                pass
        self.qt_running = False
        self.root.after(0, lambda: self.qt_start_btn.config(state=tk.NORMAL))
        self.root.after(0, lambda: self.qt_stop_btn.config(state=tk.DISABLED))

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