import os
from pathlib import Path
import numpy as np
import pandas as pd
import cv2 as cv
import json
from numbers import Number
from PySide6 import QtWidgets
from PySide6.QtWidgets import QMessageBox
from PySide6.QtCore import Signal, QThread, Qt
from PySide6.QtGui import QIcon, QPixmap
import matplotlib.pyplot as plt
from matplotlib.colors import rgb_to_hsv
from matplotlib.widgets import RectangleSelector
from arcjetCV.gui.arcjetCV_gui import Ui_MainWindow
from arcjetCV.utils.video import Video, VideoMeta
from arcjetCV.utils.processor import ArcjetProcessor, ProcessorWorker
from arcjetCV.utils.utils import (
    splitfn,
    getOutlierMask,
    annotateImage,
    annotate_image_with_frame_number,
)
from arcjetCV.calibration.calibration_controller import CalibrationController
from PySide6.QtWidgets import QFileDialog
import json


class MainWindow(QtWidgets.QMainWindow):
    frame_processed = Signal()

    def __init__(self):
        super().__init__()

        # Initialize the user interface
        self.ui = Ui_MainWindow()
        self.ui.setupUi(self)

        self.testing = False

        # Set the application icon
        self.logo_white_path = os.path.join(
            Path(__file__).parent.absolute(), "logo/arcjetCV_logo_white.png"
        )
        self.logo_path = os.path.join(
            Path(__file__).parent.absolute(), "logo/arcjetCV_logo_.png"
        )
        self.setWindowIcon(QIcon(self.logo_path))

        # Load and process the application logo
        self.rgb_frame = cv.imread(self.logo_white_path)
        self.rgb_frame = cv.cvtColor(self.rgb_frame, cv.COLOR_BGR2RGB)

        # Initialize frame and plotting windows
        self.calibrated = None
        self.calibration_data = None
        self.pixels_per_mm = None
        self._plot_ref = None
        self._tplot_ref = None
        self._brightness_ref = None
        self.show_img()

        # Initialize line start and stop positions
        self.start_line = None
        self.stop_line = None

        # Flags for video loading and processing
        self.VIDEO_LOADED = False
        self.NEW_VIDEO = False

        # Initialize folder/file properties
        self.folder = None
        self.path = None
        self.filename = None
        self.ext = None
        self.video = None
        self.videometa = None

        # Processor objects
        self.processor = None

        # Data structures
        self.raw_outputs = []
        self.time_series = None
        self.PLOTKEYS = []
        self.fit_dict = None
        self._plot_ref = None

        # Connect interface buttons to their respective functions
        self.ui.pushButton_process.clicked.connect(self.process_all)
        self.ui.pushButton_loadVideo.clicked.connect(self.load_video)
        self.ui.pushButton_loadCalibration.clicked.connect(self.load_calibration)
        self.ui.pushButton_export_csv.clicked.connect(self.export_to_csv)
        self.ui.pushButton_fitData.clicked.connect(self.fit_data)
        self.ui.pushButton_LoadFiles.clicked.connect(self.load_outputs)
        self.ui.pushButton_PlotData.clicked.connect(self.plot_outputs)
        self.ui.checkBox_display_shock2.stateChanged.connect(self.plot_outputs)
        self.ui.checkBox_m95_radius.stateChanged.connect(self.plot_outputs)
        self.ui.checkBox_m50_radius.stateChanged.connect(self.plot_outputs)
        self.ui.checkBox_model_center.stateChanged.connect(self.plot_outputs)
        self.ui.checkBox_50_radius.stateChanged.connect(self.plot_outputs)
        self.ui.checkBox_95_radius.stateChanged.connect(self.plot_outputs)
        self.ui.checkBox_shock_area.stateChanged.connect(self.plot_outputs)
        self.ui.checkBox_model_rad.stateChanged.connect(self.plot_outputs)
        self.ui.checkBox_shock_center.stateChanged.connect(self.plot_outputs)
        self.ui.checkBox_shockmodel.stateChanged.connect(self.plot_outputs)
        self.ui.checkBox_ypos.stateChanged.connect(self.plot_outputs)
        self.ui.checkBox_kalman_filter.stateChanged.connect(self.plot_outputs)
        self.ui.comboBox_units.setCurrentText("[mm]")
        self._set_tooltips()

        # init_plot_brightness
        self.ui.Window3.canvas.axes.clear()
        self.ui.Window3.canvas.axes.set_yticklabels([])
        self.ui.Window3.canvas.axes.get_xaxis().set_visible(False)
        self.ui.Window3.canvas.axes.get_yaxis().set_visible(False)
        if self.ui.Window3.canvas.toolbar is not None:
            self.ui.Window3.canvas.toolbar.setVisible(False)
        self.ui.Window3.canvas.draw()

        # Set some UI elements to be initially disabled
        self.ui.spinBox_FrameIndex.setDisabled(True)
        self.ui.comboBox_flowDirection.setDisabled(True)
        self.ui.comboBox_filterType.setDisabled(True)
        self.ui.FilterTabs.setDisabled(True)
        self.ui.comboBox_filterType.setDisabled(True)
        self.ui.comboBox_filterType.setDisabled(True)

        # Show the main window
        self.show()

    def _set_tooltips(self):
        """Set hover help for key UI controls."""
        tooltips = {
            "pushButton_loadCalibration": "Load a calibration file to convert pixels to physical units.",
            "pushButton_loadVideo": "Open a video file for processing.",
            "pushButton_save_frame": "Save the currently displayed frame as an image.",
            "pushButton_process": "Process the selected frame range and extract time-series outputs.",
            "pushButton_LoadFiles": "Load previously exported output files for plotting.",
            "pushButton_export_csv": "Export current results to a CSV file.",
            "pushButton_PlotData": "Plot selected quantities from loaded or processed data.",
            "pushButton_fitData": "Fit the selected data over the chosen time interval.",
            "applyCrop": "Apply the current crop bounds to the video processing region.",
            "spinBox_FrameIndex": "Current frame index shown in the preview.",
            "comboBox_flowDirection": "Direction of the flow used by the detector.",
            "comboBox_filterType": "Segmentation method: HSV, GRAY, CNN, or AutoHSV.",
            "checkBox_display_shock": "Enable/disable shock detection overlay on the preview frame.",
            "checkBox_crop": "Show/hide crop rectangle and use crop bounds during processing.",
            "checkBox_annotate": "Overlay computed values and labels on the preview frame.",
            "spinBox_crop_xmin": "Left crop boundary (pixels).",
            "spinBox_crop_xmax": "Right crop boundary (pixels).",
            "spinBox_crop_ymin": "Top crop boundary (pixels).",
            "spinBox_crop_ymax": "Bottom crop boundary (pixels).",
            "spinBox_FirstGoodFrame": "First frame included in batch processing.",
            "spinBox_LastGoodFrame": "Last frame included in batch processing.",
            "spinBox_frame_skips": "Process every Nth frame to speed up analysis.",
            "checkBox_writeVideo": "Write an output video with overlays while processing.",
            "comboBox_units": "Unit system used in plots and derived quantities.",
            "doubleSpinBox_fps": "Frame rate used for the time axis (frames per second).",
            "spinBox_mask_frames": "Minimum number of valid points required for plotting/fitting masks.",
            "checkBox_display_shock2": "Include shock position in the time-series plot.",
            "checkBox_kalman_filter": "Apply Kalman smoothing to plotted signals.",
            "comboBox_fit_type": "Model used for fitting (currently linear).",
            "doubleSpinBox_fit_start_time": "Fit window start time (seconds).",
            "doubleSpinBox_fit_last_time": "Fit window end time (seconds).",
            "checkBox_m95_radius": "Plot model 95% radius.",
            "checkBox_m50_radius": "Plot model 50% radius.",
            "checkBox_model_center": "Plot model center position.",
            "checkBox_50_radius": "Plot measured 50% radius.",
            "checkBox_95_radius": "Plot measured 95% radius.",
            "checkBox_shock_area": "Plot shock area.",
            "checkBox_model_rad": "Plot model radius.",
            "checkBox_shock_center": "Plot shock center position.",
            "checkBox_shockmodel": "Plot shock model result.",
            "checkBox_ypos": "Plot Y-position metric.",
        }

        for name, tip in tooltips.items():
            widget = getattr(self.ui, name, None)
            if widget is not None:
                widget.setToolTip(tip)
                widget.setStatusTip(tip)

    def _set_calibration_resolution_label(self):
        """Display loaded calibration resolution as mm/px and px/mm in the UI."""
        label = getattr(self.ui, "label_calibrationResolution", None)
        if label is None:
            return

        ppm = getattr(self, "pixels_per_mm", None)
        if ppm in (None, 0):
            label.setText("Resolution: None")
            return

        mm_per_px = 1.0 / ppm
        label.setText(f"Resolution: {mm_per_px:.6f} mm/px | {ppm:.6f} px/mm")

    @staticmethod
    def _flow_axis_sign(flow_direction):
        """
        Return a sign that makes downstream motion positive across flow directions.
        """
        flow = str(flow_direction).strip().lower()
        return -1.0 if flow in {"left", "up"} else 1.0

    def closeEvent(self, event):
        print("🔴 Closing application...")
        try:
            if hasattr(self, "worker") and self.worker is not None:
                print("🛑 Attempting to stop worker...")
                self.worker.stop()
                self.worker.deleteLater()
                self.worker = None

            if hasattr(self, "thread") and isinstance(self.thread, QThread):
                if self.thread.isRunning():
                    print("🛑 Waiting for thread to finish...")
                    self.thread.quit()
                    self.thread.wait(5000)
                    print("✅ Thread fully stopped.")
                self.thread.deleteLater()
                self.thread = None

        except Exception as e:
            print(f"[ERROR] during closeEvent: {e}")

        event.accept()


    def show_img(self):
        # if self.calibrated == True:
        #     self.rgb_frame = CalibrationController.apply_calibration(
        #         self.rgb_frame, self.calibration_data
        #     )
        #     print("Calibrated")
        if self._plot_ref is None:
            # Create a new plot reference and define the cursor data format
            self._plot_ref = self.ui.Window0.canvas.axes.imshow(
                self.rgb_frame, aspect="equal"
            )
            self._plot_ref.axes.get_xaxis().set_visible(False)
            self._plot_ref.axes.get_yaxis().set_visible(False)

            # Customize axis spines color to match the background (hiding the border)
            self._plot_ref.axes.spines["top"].set_color("none")
            self._plot_ref.axes.spines["bottom"].set_color("none")
            self._plot_ref.axes.spines["left"].set_color("none")
            self._plot_ref.axes.spines["right"].set_color("none")

            def format_cursor_data(data):
                try:
                    data[0]
                except (TypeError, IndexError):
                    data = [data]
                data_str = "RGB = [" + ", ".join(
                    "{:0.3g}".format(item) for item in data if isinstance(item, Number)
                )
                hsv = rgb_to_hsv(np.array(data) / 255.0)
                hsv_str = (
                    "HSV = ["
                    + "%i " % (hsv[0] * 180)
                    + ","
                    + "%i " % (hsv[1] * 255)
                    + ","
                    + "%i " % (hsv[2] * 255)
                    + "]"
                )
                return data_str + "] " + hsv_str

            self._plot_ref.format_cursor_data = format_cursor_data
            self.rect_selector = RectangleSelector(self._plot_ref.axes, self.onselect)
        else:
            self._plot_ref.set_data(self.rgb_frame)

        self.ui.Window0.canvas.draw()

    # def onselect(self, eclick, erelease):
    #     x1, y1 = eclick.xdata, eclick.ydata
    #     x2, y2 = erelease.xdata, erelease.ydata
    #     self.ui.spinBox_crop_xmin.setValue(x1)
    #     self.ui.spinBox_crop_xmax.setValue(x2)
    #     self.ui.spinBox_crop_ymin.setValue(y1)
    #     self.ui.spinBox_crop_ymax.setValue(y2)
    #     self.update_crop()

    def onselect(self, eclick, erelease):
        x1, y1 = eclick.xdata, eclick.ydata
        x2, y2 = erelease.xdata, erelease.ydata

        # If image was rotated 90° clockwise, reverse it
        # Convert (x, y) to pre-rotation coordinates
        print(self.calibrated)
        # if self.calibrated:
        #     print("good crop")
        #     h, w = self.rgb_frame.shape[:2]

        #     def unrotate_point(x, y):
        #         # Reverse 90° clockwise = rotate 90° counter-clockwise
        #         return y, w - x

        #     x1, y1 = unrotate_point(x1, y1)
        #     x2, y2 = unrotate_point(x2, y2)

        #     # Optional: undo homography if you want original pixel space
        #     H_inv = np.array(self.calibration_data["homography_inverse"])
        #     pts = np.array([[x1, y1], [x2, y2]], dtype=np.float32).reshape(-1, 1, 2)
        #     pts_original = cv.perspectiveTransform(pts, H_inv).reshape(-1, 2)
        #     x1, y1 = pts_original[0]
        #     x2, y2 = pts_original[1]

        # Store corrected coordinates
        self.ui.spinBox_crop_xmin.setValue(int(min(x1, x2)))
        self.ui.spinBox_crop_xmax.setValue(int(max(x1, x2)))
        self.ui.spinBox_crop_ymin.setValue(int(min(y1, y2)))
        self.ui.spinBox_crop_ymax.setValue(int(max(y1, y2)))
        self.update_crop()

    def update_crop(self):
        # Validate bounding box inputs here
        if (
            self.ui.spinBox_crop_xmax.value() - self.ui.spinBox_crop_xmin.value() >= 40
            and self.ui.spinBox_crop_ymax.value() - self.ui.spinBox_crop_ymin.value()
            >= 40
        ):
            self.videometa["CROP_XMIN"] = self.ui.spinBox_crop_xmin.value()
            self.videometa["CROP_XMAX"] = self.ui.spinBox_crop_xmax.value()
            self.videometa["CROP_YMIN"] = self.ui.spinBox_crop_ymin.value()
            self.videometa["CROP_YMAX"] = self.ui.spinBox_crop_ymax.value()
            self.videometa.write()
            self.update_frame_index()
        else:
            pass

    def brightness_click_slot(self, x, y):
        try:
            if x > 0 and x < self.video.nframes - 1:
                self.ui.spinBox_FrameIndex.setValue(int(x))
        except:
            pass

    def update_frame_index(self):

        if self.NEW_VIDEO:
            self.ui.spinBox_crop_xmin.setValue(self.videometa["CROP_XMIN"])
            self.ui.spinBox_crop_xmax.setValue(self.videometa["CROP_XMAX"])
            self.ui.spinBox_crop_ymin.setValue(self.videometa["CROP_YMIN"])
            self.ui.spinBox_crop_ymax.setValue(self.videometa["CROP_YMAX"])

        self.ui.maxHue.setDisabled(False)
        self.ui.minHue.setDisabled(False)
        self.ui.minIntensity.setDisabled(False)
        self.ui.maxIntensity.setDisabled(False)
        self.ui.minSaturation.setDisabled(False)
        self.ui.maxSaturation.setDisabled(False)
        self.ui.maxHue_2.setDisabled(False)
        self.ui.minHue_2.setDisabled(False)
        self.ui.minIntensity_2.setDisabled(False)
        self.ui.maxIntensity_2.setDisabled(False)
        self.ui.minSaturation_2.setDisabled(False)
        self.ui.maxSaturation_2.setDisabled(False)

        # disabled the tabs with CNN
        if (
            self.ui.comboBox_filterType.currentText() == "CNN"
            or self.ui.comboBox_filterType.currentText() == "AutoHSV"
        ):
            self.ui.tab_3.setDisabled(True)
            self.ui.tab_4.setDisabled(True)
        elif self.ui.comboBox_filterType.currentText() == "GRAY":
            self.ui.tab_3.setDisabled(False)
            self.ui.tab_4.setDisabled(False)
            self.ui.maxHue.setDisabled(True)
            self.ui.minHue.setDisabled(True)
            self.ui.minIntensity.setDisabled(False)
            self.ui.maxIntensity.setDisabled(False)
            self.ui.minSaturation.setDisabled(True)
            self.ui.maxSaturation.setDisabled(True)
            self.ui.maxHue_2.setDisabled(True)
            self.ui.minHue_2.setDisabled(True)
            self.ui.minIntensity_2.setDisabled(False)
            self.ui.maxIntensity_2.setDisabled(False)
            self.ui.minSaturation_2.setDisabled(True)
            self.ui.maxSaturation_2.setDisabled(True)
        else:
            self.ui.tab_3.setDisabled(False)
            self.ui.tab_4.setDisabled(False)

        frame_index = self.ui.spinBox_FrameIndex.value()
        self.rgb_frame = self.video.get_frame(frame_index)
        if self.rgb_frame is None:
            return
        if self.calibrated == True:
            self.rgb_frame = CalibrationController.apply_calibration(
                self.rgb_frame, self.calibration_data
            )
            print("Calibrated")

        self.processor.flow_dir = self.ui.comboBox_flowDirection.currentText()

        crop_range = [
            [self.ui.spinBox_crop_ymin.value(), self.ui.spinBox_crop_ymax.value()],
            [self.ui.spinBox_crop_xmin.value(), self.ui.spinBox_crop_xmax.value()],
        ]

        # Apply crop params
        if self.processor.crop != crop_range:
            self.processor.crop = crop_range
            if (
                self.VIDEO_LOADED
            ):  # To avoid crash (can be applied to all the function update_frame)
                self.videometa["CROP_YMIN"] = self.ui.spinBox_crop_ymin.value()
                self.videometa["CROP_XMIN"] = self.ui.spinBox_crop_xmin.value()
                self.videometa["CROP_YMAX"] = self.ui.spinBox_crop_ymax.value()
                self.videometa["CROP_XMAX"] = self.ui.spinBox_crop_xmax.value()
                self.videometa.write()

        # Process frame
        contour_dict, argdict = self.processor.process(
            self.rgb_frame, self.grab_ui_values()
        )

        # Draw contours
        width = self.rgb_frame.shape[1]
        thickness = max(
            1, width // 500
        )  # Scale thickness with width, minimum thickness of 1
        for key in contour_dict.keys():
            if key == "MODEL":
                cv.drawContours(
                    self.rgb_frame, contour_dict[key], -1, (0, 255, 0), thickness
                )
            elif key == "SHOCK" and self.ui.checkBox_display_shock.isChecked():
                cv.drawContours(
                    self.rgb_frame, contour_dict[key], -1, (255, 0, 0), thickness
                )

        # Draw annotations
        annotate_image_with_frame_number(self.rgb_frame, frame_index)
        if self.ui.checkBox_annotate.isChecked():
            annotateImage(self.rgb_frame, argdict, top=True, left=True)

        if self.ui.checkBox_crop.isChecked():
            start_point = (self.processor.crop[1][0], self.processor.crop[0][0])
            end_point = (self.processor.crop[1][1], self.processor.crop[0][1])
            cv.rectangle(
                self.rgb_frame, start_point, end_point, (255, 255, 255), thickness
            )

        self.frame_processed.emit()
        self.plot_location()

    def save_frame(self):
        frame_index = self.ui.spinBox_FrameIndex.value()
        frame = self.video.get_frame(frame_index)
        # Convert from BGR to RGB
        image = cv.cvtColor(frame, cv.COLOR_RGB2BGR)

        file_dialog = QtWidgets.QFileDialog()
        file_dialog.setAcceptMode(QtWidgets.QFileDialog.AcceptSave)
        file_dialog.setNameFilter("Images (*.png *.jpg *.bmp *.tif)")
        try:
            if file_dialog.exec():
                fpath = file_dialog.selectedFiles()[0]
                cv.imwrite(fpath, image)
                print("------SAVED FRAME %i-----" % frame_index)
                self.ui.basebar.setText("------SAVED FRAME %i-----" % frame_index)
        except:
            self.arcjetcv_message_box("Warning", "Could not save frame")

    def connect_elements(self):
        self.ui.spinBox_FrameIndex.valueChanged.connect(self.update_frame_index)
        self.frame_processed.connect(self.show_img)
        self.ui.checkBox_display_shock.stateChanged.connect(self.update_frame_index)
        self.ui.maxHue.valueChanged.connect(self.update_frame_index)
        self.ui.minHue.valueChanged.connect(self.update_frame_index)
        self.ui.minIntensity.valueChanged.connect(self.update_frame_index)
        self.ui.maxIntensity.valueChanged.connect(self.update_frame_index)
        self.ui.minSaturation.valueChanged.connect(self.update_frame_index)
        self.ui.maxSaturation.valueChanged.connect(self.update_frame_index)
        self.ui.maxHue_2.valueChanged.connect(self.update_frame_index)
        self.ui.minHue_2.valueChanged.connect(self.update_frame_index)
        self.ui.minIntensity_2.valueChanged.connect(self.update_frame_index)
        self.ui.maxIntensity_2.valueChanged.connect(self.update_frame_index)
        self.ui.minSaturation_2.valueChanged.connect(self.update_frame_index)
        self.ui.maxSaturation_2.valueChanged.connect(self.update_frame_index)
        self.ui.applyCrop.clicked.connect(self.update_crop)
        self.ui.pushButton_save_frame.clicked.connect(self.save_frame)
        self.ui.checkBox_crop.stateChanged.connect(self.update_frame_index)
        if (
            self.ui.spinBox_crop_xmin.value() < self.ui.spinBox_crop_xmax.value()
            and self.ui.spinBox_crop_ymin.value() < self.ui.spinBox_crop_ymax.value()
        ):
            self.ui.applyCrop.clicked.connect(self.update_frame_index)
        self.ui.comboBox_filterType.currentTextChanged.connect(self.update_frame_index)
        self.ui.comboBox_flowDirection.currentTextChanged.connect(
            self.update_frame_index
        )

    def load_video(self):
        print("------New Video-----")
        dialog = QtWidgets.QFileDialog()
        options = QtWidgets.QFileDialog.Options()
        pathmask = dialog.getOpenFileName(
            None,
            "Select Video",
            "",
            "Video Files (*.mp4 *.avi *.mov *.m4v);;All Files (*)",
            options=options,
        )
        self.NEW_VIDEO = True

        self.path = pathmask[0]
        self._plot_ref = None
        if self.path != "":
            self.folder, self.filename, self.ext = splitfn(self.path)
            loading_dialog = None

            if not self.testing:
                loading_dialog = QtWidgets.QProgressDialog(
                    "Loading video...", "Cancel", 0, 100, self
                )
                loading_dialog.setWindowTitle("Loading Video")
                loading_dialog.setWindowModality(Qt.WindowModal)
                loading_dialog.setMinimumDuration(0)
                loading_dialog.setAutoClose(False)
                loading_dialog.setValue(0)
                loading_dialog.show()

            def update_loading_progress(value, message):
                if loading_dialog is None:
                    return
                loading_dialog.setValue(int(max(0, min(100, value))))
                if message:
                    loading_dialog.setLabelText(message)
                QtWidgets.QApplication.processEvents()
                if loading_dialog.wasCanceled():
                    raise RuntimeError("Video loading canceled by user.")

            # Old matplotlib artists become invalid after axes.clear(); drop refs.
            self.start_line = None
            self.stop_line = None
            self._tplot_ref = None

            self.ui.spinBox_FrameIndex.setDisabled(False)
            self.ui.comboBox_flowDirection.setDisabled(False)
            self.ui.comboBox_filterType.setDisabled(False)
            self.ui.FilterTabs.setDisabled(False)
            self.ui.comboBox_filterType.setDisabled(False)
            self.ui.comboBox_filterType.setDisabled(False)
            self.ui.Window0.canvas.axes.clear()
            self.ui.Window0.canvas.draw()

            try:  # Create video object
                update_loading_progress(2, "Opening video file...")
                self.video = Video(self.path)
                self.videometa = VideoMeta(
                    self.video,
                    os.path.join(self.folder, self.filename + ".meta"),
                    progress_callback=update_loading_progress,
                )

                if hasattr(self, "pixels_per_mm"):  # Ensure it's defined
                    self.videometa["PIXELS_PER_MM"] = self.pixels_per_mm
                    print(
                        f"DEBUG: Stored pixels_per_mm in videometa: {self.videometa['PIXELS_PER_MM']}"
                    )  # Debugging
                else:
                    print("WARNING: pixels_per_mm not set before loading video.")
                print("Number of Frames: ", self.video.nframes)
                print(
                    f"First-Last Good Frame: {self.videometa['FIRST_GOOD_FRAME']}-{self.videometa['LAST_GOOD_FRAME']}"
                )
                print(
                    f"Min-Max Brightness: {min(self.videometa['BRIGHTNESS'])}-{max(self.videometa['BRIGHTNESS'])}"
                )

                # Initialize UI
                self.ui.Window3.canvas.axes.clear()
                self.ax1 = self.ui.Window3.canvas.axes
                self.ui.spinBox_FrameIndex.setRange(0, self.video.nframes - 1)
                self.ui.spinBox_FrameIndex.setValue(self.videometa["FIRST_GOOD_FRAME"])
                self.ui.spinBox_FirstGoodFrame.setValue(
                    self.videometa["FIRST_GOOD_FRAME"]
                )
                self.ui.spinBox_LastGoodFrame.setValue(
                    self.videometa["LAST_GOOD_FRAME"]
                )
                self.ui.comboBox_flowDirection.setCurrentText(
                    self.videometa["FLOW_DIRECTION"]
                )
                self.ui.lineEdit_filename.setText(self.video.name)
                self.ui.spinBox_LastGoodFrame.setRange(0, self.video.nframes - 1)
                self.ui.spinBox_FirstGoodFrame.setRange(0, self.video.nframes - 1)

                # Plot the brightness signal
                self.ax1.plot(self.videometa["BRIGHTNESS"])
                self.ax1.autoscale_view()
                self.ax1.set_xlim([0, self.video.nframes - 1])

                # Plot the red point for the specific n value and start and stop
                self.plot_location(reset=True)
                self.plot_start_stop()

                self.ui.Window3.canvas.draw()

                self.ui.Window3.canvas.clicked.connect(
                    self.brightness_click_slot, Qt.UniqueConnection
                )
                self.ui.spinBox_FirstGoodFrame.valueChanged.connect(
                    self.plot_start_stop, Qt.UniqueConnection
                )
                self.ui.spinBox_LastGoodFrame.valueChanged.connect(
                    self.plot_start_stop, Qt.UniqueConnection
                )
                if self.processor is None:
                    update_loading_progress(99, "Initializing processor...")
                    self.processor = ArcjetProcessor(
                        self.videometa, progress_bar=self.ui.progressBar
                    )
                else:  # avoid reimporting CNN model each time
                    self.processor.update_video_meta(self.videometa)

                # Connect UI only on first video load
                if self.VIDEO_LOADED is False:
                    self.connect_elements()
                    self.VIDEO_LOADED = True

                    self.update_frame_index()
                    self.NEW_VIDEO = False
                update_loading_progress(100, "Video loaded.")

            except Exception as e:
                if self.testing:
                    print("Could not load video:\n" + str(e))
                else:
                    self.arcjetcv_message_box(
                        "Warning", "Could not load video:\n" + str(e)
                    )
            finally:
                if loading_dialog is not None:
                    loading_dialog.close()

    def load_calibration(self):
        """
        Load a JSON calibration file and store the available calibration parameters.

        Only existing keys in the JSON file will be loaded. Missing keys are skipped.
        The calibration data is stored in self.calibration_data and self.pixels_per_mm.
        """
        file_path, _ = QFileDialog.getOpenFileName(
            None, "Select Calibration File", "", "JSON Files (*.json)"
        )

        if not file_path:
            return

        try:
            with open(file_path, "r") as file:
                raw_data = json.load(file)

            # Initialize calibration data dictionary
            self.calibration_data = {}

            # Load each expected key if it exists
            key_types = {
                "camera_matrix": (np.float32, (3, 3)),
                "dist_coeffs": (np.float32, None),
                "rvec": (np.float32, None),
                "tvec": (np.float32, None),
                "pixels_per_mm": (float, None),
                "centers": (np.float32, None),
                "square_layout": (np.float32, None),
                "homography": (np.float32, None),
                "homography_inverse": (np.float32, None),
                "pattern_size": (tuple, None),
                "affine_matrix": (np.float32, None),
            }

            for key, (dtype, shape) in key_types.items():
                if key in raw_data:
                    value = raw_data[key]
                    if dtype == np.float32:
                        arr = np.array(value, dtype=dtype)
                        self.calibration_data[key] = arr
                    elif dtype == float:
                        self.calibration_data[key] = float(value)
                    elif dtype == tuple:
                        self.calibration_data[key] = tuple(value)

            # Set ppm and status flag
            self.pixels_per_mm = self.calibration_data.get("pixels_per_mm", None)
            self.calibrated = True

            # Propagate scaling immediately when a video/processor is already active.
            if self.pixels_per_mm is not None:
                if self.videometa is not None:
                    self.videometa["PIXELS_PER_MM"] = self.pixels_per_mm
                    self.videometa.write()
                if self.processor is not None:
                    self.processor.pixels_per_mm = self.pixels_per_mm

            # Update label
            shortpath = self.shorten_path(file_path, 60)
            self.ui.label_calibrationPath.setText(f"Calibration Path: {shortpath}")
            self._set_calibration_resolution_label()

            print("✅ Calibration data loaded:", list(self.calibration_data.keys()))

        except Exception as e:
            self.arcjetcv_message_box(
                "Error", f"Failed to load calibration file: {str(e)}"
            )
            self.calibration_data = None
            self.calibrated = False
            self.pixels_per_mm = None
            self._set_calibration_resolution_label()
        if self.VIDEO_LOADED:
            self.update_frame_index()

    def plot_location(self, reset=False):
        if self.video is None:
            return
        n = self.ui.spinBox_FrameIndex.value()
        current_axes = self.ui.Window3.canvas.axes
        if (
            self._tplot_ref is None
            or reset
            or getattr(self._tplot_ref, "axes", None) is not current_axes
        ):
            self._tplot_ref = self.ui.Window3.canvas.axes.axvline(
                x=n, color="red", linestyle="-"
            )
        else:
            self._tplot_ref.set_xdata([n, n])
        self.ui.Window3.canvas.draw()

    def plot_start_stop(self):
        if self.video is None:
            return

        start = self.ui.spinBox_FirstGoodFrame.value()
        stop = self.ui.spinBox_LastGoodFrame.value()
        current_axes = self.ui.Window3.canvas.axes

        if 0 <= start < self.video.nframes:
            if (
                self.start_line is None
                or getattr(self.start_line, "axes", None) is not current_axes
            ):
                self.start_line = current_axes.axvline(
                    x=start, color="green", linestyle="-"
                )
            else:
                self.start_line.set_xdata([start, start])
        else:
            if self.start_line is not None:
                try:
                    self.start_line.remove()
                except (NotImplementedError, ValueError):
                    pass
            self.start_line = None

        if 0 <= stop < self.video.nframes:
            if (
                self.stop_line is None
                or getattr(self.stop_line, "axes", None) is not current_axes
            ):
                self.stop_line = current_axes.axvline(
                    x=stop, color="green", linestyle="-"
                )
            else:
                self.stop_line.set_xdata([stop, stop])
        else:
            if self.stop_line is not None:
                try:
                    self.stop_line.remove()
                except (NotImplementedError, ValueError):
                    pass
            self.stop_line = None

        self.ui.Window3.canvas.draw()

    def process_all(self):
        """Starts video processing in a separate QThread."""
        if not self.VIDEO_LOADED:
            return

        # Get UI values
        ilow, ihigh = (
            self.ui.spinBox_FirstGoodFrame.value(),
            self.ui.spinBox_LastGoodFrame.value(),
        )
        frame_stride = self.ui.spinBox_frame_skips.value()

        # Update metadata
        self.videometa["FIRST_GOOD_FRAME"] = ilow
        self.videometa["LAST_GOOD_FRAME"] = ihigh
        self.videometa.write()

        # ✅ Stop existing thread before starting a new one
        if hasattr(self, "thread") and isinstance(self.thread, QThread):
            if self.thread.isRunning():
                print("⚠️ Stopping existing worker thread...")
                self.worker.stop()  # Ensure worker has a `stop()` method
                self.thread.quit()
                self.thread.wait(5000)  # Wait for thread to exit
                print("✅ Previous worker thread stopped successfully")

        # ✅ Ensure proper cleanup of old worker
        if hasattr(self, "worker") and self.worker is not None:
            self.worker.deleteLater()
            self.worker = None

        # ✅ Create a new Worker Thread safely
        self.worker = ProcessorWorker(
            self.processor,
            self.video,
            self.grab_ui_values(),
            ilow,
            ihigh,
            frame_stride,
            self.ui.lineEdit_filename.text(),
            write_json=True,
            write_video=self.ui.checkBox_writeVideo.isChecked(),
            display_shock=self.ui.checkBox_display_shock.isChecked(),
        )

        self.thread = QThread()
        self.worker.moveToThread(self.thread)
    
        # ✅ Connect signals (only once)
        self.thread.started.connect(self.worker.run)
        self.worker.progress_updated.connect(self.ui.progressBar.setValue)
        self.worker.finished.connect(self.on_processing_complete)
        self.worker.finished.connect(self.thread.quit)
        self.thread.finished.connect(lambda: print("✅ Worker thread fully terminated"))

        # ✅ Start the thread
        self.thread.start()

    def on_processing_complete(self):
        """Handles UI updates when processing is completed."""
        if self.thread:
            print("Stopping thread safely...")
            try:
                self.worker.finished.disconnect()
            except:
                pass
            self.thread.quit()
            self.thread.wait(5000)
            print("✅ Worker thread stopped successfully")

            self.worker.deleteLater()
            self.worker = None
            self.thread.deleteLater()
            self.thread = None

        self.processing_done = True
        self.ui.progressBar.setValue(0)
        self.arcjetcv_message_box("Video Processed", "The video has been processed.")


    def grab_ui_values(self):
        inputdict = {"SEGMENT_METHOD": str(self.ui.comboBox_filterType.currentText())}
        inputdict["HSV_MODEL_RANGE"] = [
            (
                self.ui.minHue.value(),
                self.ui.minSaturation.value(),
                self.ui.minIntensity.value(),
            ),
            (
                self.ui.maxHue.value(),
                self.ui.maxSaturation.value(),
                self.ui.maxIntensity.value(),
            ),
        ]
        inputdict["HSV_SHOCK_RANGE"] = [
            (
                self.ui.minHue_2.value(),
                self.ui.minSaturation_2.value(),
                self.ui.minIntensity_2.value(),
            ),
            (
                self.ui.maxHue_2.value(),
                self.ui.maxSaturation_2.value(),
                self.ui.maxIntensity_2.value(),
            ),
        ]
        inputdict["THRESHOLD"] = self.ui.minIntensity.value()
        return inputdict

    # PLOTTING DATA TAB
    def load_outputs(self):
        try:
            # create fileDialog to select file
            options = QtWidgets.QFileDialog.Options()
            files, _ = QtWidgets.QFileDialog.getOpenFileNames(
                self,
                "Load ouput files",
                "",
                "Output Files (*.json);;All Files (*)",
                options=options,
            )
            self.ui.basebar.setText("Loading %i files" % len(files))

            # Load all files & concatenate
            self.raw_outputs = []
            for fname in files:
                with open(fname, "rb") as file:
                    opl = json.load(file)
                    self.raw_outputs.extend(opl)

            if len(files) > 0:
                fpath, name, ext = splitfn(files[0])
                shortpath = self.shorten_path(fpath, 30)

                # Show summary of loaded data
                summary = "Loaded %i files\n" % len(files)
                summary += "Folder: %s\n" % shortpath
                for fname in files:
                    fpath, name, ext = splitfn(fname)
                    summary += "File: %s\n" % name
                summary += "Total frames: %i\n" % len(self.raw_outputs)

                self.ui.label_data_summary.setText(summary)
                self.ui.basebar.setText("Finished loading files")

        except Exception as e:
            if self.testing:
                print("File loading failed:\n" + str(e))
            else:
                self.arcjetcv_message_box("Warning", "File loading failed:\n" + str(e))

        self.plot_outputs()

    @staticmethod
    def shorten_path(path, max_length):
        if len(path) <= max_length:
            return path  # Return the original path if it is already short enough

        max_visible_length = max_length - 3  # Reserve 3 characters for '...'
        first_part = path[: max_visible_length // 2]
        last_part = path[-(max_visible_length - len(first_part)) :]
        return first_part + "..." + last_part

    def _apply_kalman_filter(self, series):
        if series is None:
            return series

        is_masked = isinstance(series, np.ma.MaskedArray)
        if is_masked:
            data = np.array(series.filled(np.nan), dtype=float)
            mask = np.ma.getmaskarray(series).copy()
        else:
            data = np.array(series, dtype=float)
            mask = np.zeros_like(data, dtype=bool)

        non_finite_mask = ~np.isfinite(data)
        if non_finite_mask.any():
            mask = mask | non_finite_mask
            data[non_finite_mask] = np.nan

        valid_mask = ~mask
        if np.count_nonzero(valid_mask) <= 1:
            return series

        valid_values = data[valid_mask]
        measurement_var = np.nanvar(valid_values)
        if not np.isfinite(measurement_var) or measurement_var <= 1e-12:
            measurement_var = 1.0

        process_var = measurement_var * 0.05
        if not np.isfinite(process_var) or process_var <= 0:
            process_var = 1e-3

        error_cov = measurement_var
        filtered = data.copy()
        valid_indices = np.flatnonzero(valid_mask)
        estimate = data[valid_indices[0]]
        filtered[valid_indices[0]] = estimate

        for idx in valid_indices[1:]:
            error_cov += process_var
            kalman_gain = error_cov / (error_cov + measurement_var)
            estimate = estimate + kalman_gain * (data[idx] - estimate)
            error_cov = (1.0 - kalman_gain) * error_cov
            filtered[idx] = estimate

        if is_masked:
            return np.ma.masked_array(filtered, mask=mask)

        filtered[~valid_mask] = data[~valid_mask]
        return filtered

    def plot_outputs(self):
        self.ui.basebar.setText("Setting up plots...")

        # Reset plotting windows
        self.ui.Window1.canvas.axes.clear()
        self.ui.Window2.canvas.axes.clear()

        self.ax1 = self.ui.Window1.canvas.axes
        self.ax2 = self.ui.Window2.canvas.axes

        # Plotting params
        index, m95, m50, mc, p50, p95, radius = [], [], [], [], [], [], []
        time, sarea, marea, sc, sm, ypos = [], [], [], [], [], []

        units_selected = self.ui.comboBox_units.currentText()
        units_lower = units_selected.lower()
        use_mm_requested = "mm" in units_lower

        pixels_per_mm_default = None
        if self.raw_outputs and "PIXELS_PER_MM" in self.raw_outputs[0]:
            pixels_per_mm_default = self.raw_outputs[0]["PIXELS_PER_MM"]

        default_pixel_length = (
            1 / pixels_per_mm_default
            if pixels_per_mm_default not in (None, 0, float("inf"))
            else None
        )

        use_mm = use_mm_requested and default_pixel_length is not None
        units = "[mm]" if use_mm else ("[px]" if use_mm_requested else units_selected)
        xy_scale_default = default_pixel_length if use_mm else 1.0
        length_unit_label = units.strip()
        if length_unit_label.startswith("[") and length_unit_label.endswith("]"):
            length_unit_label = length_unit_label[1:-1]
        if not length_unit_label:
            length_unit_label = units

        fps = self.ui.doubleSpinBox_fps.value()
        maskn = self.ui.spinBox_mask_frames.value()
        flow_sign = self._flow_axis_sign(self.ui.comboBox_flowDirection.currentText())

        self.ui.basebar.setText("Plotting data...")
        try:
            if len(self.raw_outputs) > 10:
                numframes = len(self.raw_outputs)
                max_traces = 10
                first_color = "red"
                swell_color = "#2e8b57"
                last_color = "#1f77b4"
                mid_color = "#999999"

                model_traces = []
                shock_traces = {}
                swell_trace_index = None

                for loop_idx, i in enumerate(range(0, len(self.raw_outputs), maskn)):
                    # Save frame index, time
                    index.append(self.raw_outputs[i]["INDEX"])
                    time.append(self.raw_outputs[i]["INDEX"] / fps)

                    frame_scale = xy_scale_default
                    frame_pixels_per_mm = self.raw_outputs[i].get(
                        "PIXELS_PER_MM", pixels_per_mm_default
                    )
                    if use_mm and frame_pixels_per_mm not in (None, 0, float("inf")):
                        frame_scale = 1 / frame_pixels_per_mm

                    # Model positions (-95%, -50%, center, 50%, 95% radius)
                    if self.raw_outputs[i]["MODEL"] is not None:
                        xpos = np.asarray(
                            self.raw_outputs[i]["MODEL_INTERP_XPOS"], dtype=float
                        )
                        signed_xpos = flow_sign * xpos
                        center = self.raw_outputs[i]["MODEL_YCENTER"]
                        model_array = np.array(self.raw_outputs[i]["MODEL"])
                        self.raw_outputs[i]["MODEL"] = model_array
                        model_coords = np.array(model_array[:, 0, :], dtype=float)
                        if use_mm:
                            model_coords *= frame_scale
                        model_traces.append(
                            {
                                "loop_idx": loop_idx,
                                "coords": model_coords,
                                "centroid": self.raw_outputs[i].get(
                                    "MODEL_CENTROID_X"
                                ),
                            }
                        )
                        m95.append(signed_xpos[0])
                        m50.append(signed_xpos[1])
                        mc.append(signed_xpos[2])
                        p50.append(signed_xpos[3])
                        p95.append(signed_xpos[4])
                        ypos.append(center)
                        radius.append(self.raw_outputs[i]["MODEL_RADIUS"])
                    else:
                        m95.append(np.nan)
                        m50.append(np.nan)
                        mc.append(np.nan)
                        p50.append(np.nan)
                        p95.append(np.nan)
                        ypos.append(np.nan)
                        radius.append(np.nan)

                    # Shock center x-position
                    if self.raw_outputs[i]["SHOCK"] is not None:
                        sc.append(
                            flow_sign
                            * float(self.raw_outputs[i]["SHOCK_INTERP_XPOS"][0])
                        )
                        shock_array = np.array(self.raw_outputs[i]["SHOCK"])
                        self.raw_outputs[i]["SHOCK"] = shock_array
                        shock_coords = np.array(shock_array[:, 0, :], dtype=float)
                        if use_mm:
                            shock_coords *= frame_scale
                        shock_traces[loop_idx] = shock_coords
                    else:
                        sc.append(np.nan)

                    # Shock and model area
                    if self.raw_outputs[i]["SHOCK"] is not None:
                        sarea.append(self.raw_outputs[i]["SHOCK_AREA"])
                    else:
                        sarea.append(np.nan)

                    if self.raw_outputs[i]["MODEL"] is not None:
                        marea.append(self.raw_outputs[i]["MODEL_AREA"])
                    else:
                        marea.append(np.nan)

                    # Shock-model separation, center
                    if (
                        self.raw_outputs[i]["MODEL"] is not None
                        and self.raw_outputs[i]["SHOCK"] is not None
                    ):
                        sm.append(abs(sc[-1] - mc[-1]))
                    else:
                        sm.append(np.nan)

                if model_traces:
                    # Determine swell index based on minimum centroid value
                    valid_centroids = [
                        (idx, trace["centroid"])
                        for idx, trace in enumerate(model_traces)
                        if trace["centroid"] is not None
                    ]
                    if valid_centroids:
                        swell_trace_index = min(valid_centroids, key=lambda x: x[1])[0]

                    num_model_traces = len(model_traces)
                    stride = max(1, num_model_traces // max_traces)
                    selected_indices = set(range(0, num_model_traces, stride))
                    selected_indices.add(0)
                    selected_indices.add(num_model_traces - 1)
                    if swell_trace_index is not None:
                        selected_indices.add(swell_trace_index)
                    selected_indices = sorted(selected_indices)

                    used_labels = set()

                    def style_for_trace(trace_idx):
                        if trace_idx == 0:
                            return first_color, 2.5, "Initial Profile"
                        if (
                            swell_trace_index is not None
                            and trace_idx == swell_trace_index
                        ):
                            return swell_color, 1.0, "Growth Profile"
                        if trace_idx == num_model_traces - 1:
                            return last_color, 1.0, "Final Profile"
                        return mid_color, 1.0, "Model Profile"

                    for trace_idx in selected_indices:
                        trace = model_traces[trace_idx]
                        coords = trace["coords"]
                        if self.ui.checkBox_kalman_filter.isChecked():
                            model_x = self._apply_kalman_filter(coords[:, 0])
                            model_y = self._apply_kalman_filter(coords[:, 1])
                            coords = np.column_stack(
                                (np.asarray(model_x, dtype=float), np.asarray(model_y, dtype=float))
                            )
                        color, linewidth, label = style_for_trace(trace_idx)
                        display_label = (
                            label if label and label not in used_labels else None
                        )
                        if display_label:
                            used_labels.add(display_label)

                        self.ax1.plot(
                            coords[:, 0],
                            coords[:, 1],
                            color=color,
                            linewidth=linewidth,
                            alpha=0.9 if display_label else 0.5,
                            label=display_label,
                        )

                        if (
                            self.ui.checkBox_display_shock2.isChecked()
                            and trace["loop_idx"] in shock_traces
                        ):
                            shock_coords = shock_traces[trace["loop_idx"]]
                            if self.ui.checkBox_kalman_filter.isChecked():
                                shock_x = self._apply_kalman_filter(shock_coords[:, 0])
                                shock_y = self._apply_kalman_filter(shock_coords[:, 1])
                                shock_coords = np.column_stack(
                                    (
                                        np.asarray(shock_x, dtype=float),
                                        np.asarray(shock_y, dtype=float),
                                    )
                                )
                            shock_label = (
                                "Shock Profile"
                                if "Shock Profile" not in used_labels
                                else None
                            )
                            if shock_label:
                                used_labels.add(shock_label)
                            self.ax1.plot(
                                shock_coords[:, 0],
                                shock_coords[:, 1],
                                color="black",
                                linewidth=0.8,
                                linestyle="--",
                                alpha=0.5,
                                label=shock_label,
                            )

                    self.swell_loop_idx = (
                        model_traces[swell_trace_index]["loop_idx"]
                        if swell_trace_index is not None
                        else None
                    )
                    self.swell_time = (
                        time[self.swell_loop_idx]
                        if self.swell_loop_idx is not None
                        and self.swell_loop_idx < len(time)
                        else None
                    )
                else:
                    self.swell_loop_idx = None
                    self.swell_time = None

                handles, labels = self.ax1.get_legend_handles_labels()
                if labels:
                    legend_map = {}
                    for handle, label in zip(handles, labels):
                        if not label or label in legend_map:
                            continue
                        legend_map[label] = handle

                    legend = self.ax1.legend(
                        list(legend_map.values()),
                        list(legend_map.keys()),
                        loc="best",
                    )
                    legend.get_title().set_fontweight("bold")

                self.ax1.set_xlabel(f"X ({length_unit_label})")
                self.ax1.set_ylabel(f"Y ({length_unit_label})")
                self.ax1.invert_yaxis()
                self.ax1.figure.canvas.draw()

                # Mask outliers
                metrics = [marea, ypos, radius, mc]
                mask = getOutlierMask(metrics)
                self.XT_MASK = mask

                # Infer pixel length from raw outputs
                radius_masked = np.ma.masked_where(mask < 0, radius)
                if default_pixel_length is not None:
                    pixel_length = default_pixel_length
                else:
                    pixel_length = 1  # Default to pixel units if calibration missing
                    if use_mm_requested:
                        self.arcjetcv_message_box(
                            "Warning",
                            "Pixels per mm not available in raw outputs. "
                            "Spatial data will be shown in pixels.",
                        )
                self.pixel_length = pixel_length
                series_scale = pixel_length if use_mm else 1.0
                radius_series = radius_masked * series_scale

                # Plot XT series
                ym95 = np.ma.masked_where(mask < 0, m95) * series_scale
                ym50 = np.ma.masked_where(mask < 0, m50) * series_scale
                ymc = np.ma.masked_where(mask < 0, mc) * series_scale
                yp50 = np.ma.masked_where(mask < 0, p50) * series_scale
                yp95 = np.ma.masked_where(mask < 0, p95) * series_scale
                ysarea = np.ma.masked_where(mask < 0, sarea)
                ymarea = np.ma.masked_where(mask < 0, marea)
                ysc = np.ma.masked_where(mask < 0, sc) * series_scale
                ysm = np.ma.masked_where(mask < 0, sm) * series_scale
                yypos = np.ma.masked_where(mask < 0, ypos) * series_scale

                if self.ui.checkBox_kalman_filter.isChecked():
                    ym95 = self._apply_kalman_filter(ym95)
                    ym50 = self._apply_kalman_filter(ym50)
                    ymc = self._apply_kalman_filter(ymc)
                    yp50 = self._apply_kalman_filter(yp50)
                    yp95 = self._apply_kalman_filter(yp95)
                    ysarea = self._apply_kalman_filter(ysarea)
                    ymarea = self._apply_kalman_filter(ymarea)
                    ysc = self._apply_kalman_filter(ysc)
                    ysm = self._apply_kalman_filter(ysm)
                    yypos = self._apply_kalman_filter(yypos)
                    radius_series = self._apply_kalman_filter(radius_series)

                self.PLOTKEYS = []
                self.max_expansion_info = None
                self.max_recession_info = None

                if self.ui.checkBox_m95_radius.isChecked():
                    self.ax2.plot(
                        time,
                        ym95,
                        color="magenta",
                        linewidth=0,
                        marker="o",
                        markersize=4,
                        label="Model -95%R",
                    )
                    self.PLOTKEYS.append("MODEL_-0.95R " + units)

                if self.ui.checkBox_m50_radius.isChecked():
                    self.ax2.plot(
                        time,
                        ym50,
                        color="blue",
                        linewidth=0,
                        marker="o",
                        markersize=4,
                        label="Model -50%R",
                    )
                    self.PLOTKEYS.append("MODEL_-0.50R " + units)

                if self.ui.checkBox_model_center.isChecked():
                    time_arr = np.asarray(time, dtype=float)
                    center_mask = ~np.ma.getmaskarray(ymc)
                    time_valid = time_arr[center_mask]
                    center_valid = np.asarray(ymc)[center_mask]

                    if center_valid.size > 0:
                        center_zeroed = center_valid - center_valid[0]
                        exp_idx = int(np.argmin(center_zeroed))
                        rec_idx = int(np.argmax(center_zeroed))
                        self.max_expansion_info = (
                            float(center_zeroed[exp_idx]),
                            float(time_valid[exp_idx]),
                        )
                        self.max_recession_info = (
                            float(center_zeroed[rec_idx]),
                            float(time_valid[rec_idx]),
                        )
                        swell_idx = int(np.argmin(center_zeroed))
                        swell_time_value = time_valid[swell_idx]

                        self.ax2.axvline(
                            swell_time_value,
                            ymin=0.0,
                            ymax=1.0,
                            linestyle="--",
                            color="black",
                            linewidth=1,
                            label=f"Recess Start: {swell_time_value:.2f} [s]",
                        )

                        if swell_idx >= 0:
                            self.ax2.plot(
                                time_valid[: swell_idx + 1],
                                center_zeroed[: swell_idx + 1],
                                color=swell_color,
                                linewidth=0,
                                marker="o",
                                markersize=5,
                                label="Growth",
                            )

                        if swell_idx + 1 < center_zeroed.size:
                            recession_time = time_valid[swell_idx + 1 :]
                            recession_values = center_zeroed[swell_idx + 1 :]
                            self.ax2.plot(
                                recession_time,
                                recession_values,
                                color=last_color,
                                linewidth=0,
                                marker="o",
                                markersize=5,
                                label="Recess",
                            )
                            if recession_time.size >= 2:
                                coeffs = np.polyfit(
                                    recession_time, recession_values, 1
                                )
                                fit_line = np.polyval(coeffs, recession_time)
                                self.ax2.plot(
                                    recession_time,
                                    fit_line,
                                    color="red",
                                    linewidth=1,
                                    label=f"Fit: y = {coeffs[0]:.2f}x+{coeffs[1]:.2f}",
                                )

                        valid_indices = np.flatnonzero(center_mask)
                        if swell_idx < valid_indices.size:
                            self.swell_loop_idx = int(valid_indices[swell_idx])
                            self.swell_time = time[self.swell_loop_idx]

                        # Keep dynamic plot limits; fixed ticks can force misleading ranges.
                        y_min = float(np.min(center_zeroed))
                        y_max = float(np.max(center_zeroed))
                        if np.isfinite(y_min) and np.isfinite(y_max):
                            lower_data = min(y_min, 0.0)
                            upper_data = max(y_max, 0.0)
                            span = upper_data - lower_data
                            pad = max(0.25, 0.1 * span) if span > 0 else 0.5
                            self.ax2.set_ylim(lower_data - pad, upper_data + pad)
                        self.ax2.locator_params(axis="x", nbins=6)
                        self.ax2.locator_params(axis="y", nbins=6)

                    self.PLOTKEYS.append("MODEL_CENTER " + units)

                if self.ui.checkBox_50_radius.isChecked():
                    self.ax2.plot(
                        time,
                        yp50,
                        color="cyan",
                        linewidth=0,
                        marker="o",
                        markersize=4,
                        label="Model +50%R",
                    )
                    self.PLOTKEYS.append("MODEL_0.50R " + units)

                if self.ui.checkBox_95_radius.isChecked():
                    self.ax2.plot(
                        time,
                        yp95,
                        color="red",
                        linewidth=0,
                        marker="o",
                        markersize=4,
                        label="Model +95%R",
                    )
                    self.PLOTKEYS.append("MODEL_0.95R " + units)

                if self.ui.checkBox_shock_area.isChecked():
                    self.ax2.plot(
                        time,
                        ysarea,
                        color="gold",
                        linewidth=0,
                        marker="o",
                        markersize=4,
                        label="Shock area (px)",
                    )
                    self.PLOTKEYS.append("SHOCK_AREA [px]")

                if self.ui.checkBox_model_rad.isChecked():
                    self.ax2.plot(
                        time,
                        radius_series,
                        color="navy",
                        linewidth=0,
                        marker="o",
                        markersize=4,
                        label=f"Model radius ({length_unit_label})",
                    )
                    self.PLOTKEYS.append("MODEL_RADIUS " + units)

                if self.ui.checkBox_shock_center.isChecked():
                    self.ax2.plot(
                        time,
                        ysc,
                        color="darkcyan",
                        linewidth=0,
                        marker="o",
                        markersize=4,
                        label="Shock center",
                    )
                    self.PLOTKEYS.append("SHOCK_CENTER " + units)

                if self.ui.checkBox_shockmodel.isChecked():
                    self.ax2.plot(
                        time,
                        ysm,
                        color="darkred",
                        linewidth=0,
                        marker="o",
                        markersize=4,
                        linestyle="--",
                        label="Shock-model distance",
                    )
                    self.PLOTKEYS.append("SHOCK_TO_MODEL " + units)

                if self.ui.checkBox_ypos.isChecked():
                    self.ax2.plot(
                        time,
                        yypos,
                        color="black",
                        linewidth=0,
                        marker="o",
                        markersize=4,
                        label="Axis position",
                    )
                    self.PLOTKEYS.append("MODEL_YPOS " + units)

                self.ax2.set_xlabel("Time (s)")
                self.ax2.set_ylabel(units)

                legend = self.ax2.legend()
                legend.set_draggable(True)
                self.ax2.figure.canvas.draw()

                # Save to dictionary data structure
                output_dict = {"TIME [s]": time}
                self.length_units = [ym95, ym50, ymc, yp50, yp95, ysc, ysm, yypos]
                self.length_labels = [
                    "MODEL_-0.95R",
                    "MODEL_-0.50R",
                    "MODEL_CENTER",
                    "MODEL_0.50R",
                    "MODEL_0.95R",
                    "SHOCK_CENTER",
                    "SHOCK_TO_MODEL",
                    "MODEL_YPOS",
                ]
                for k in range(0, len(self.length_units)):
                    output_dict[self.length_labels[k] + " " + units] = (
                        self.length_units[k]
                    )

                self.px_units = [ymarea, ysarea, radius_series]
                self.px_labels = [
                    "MODEL_AREA [px]",
                    "SHOCK_AREA [px]",
                    "MODEL_RADIUS " + units,
                ]
                for k in range(0, len(self.px_units)):
                    output_dict[self.px_labels[k]] = self.px_units[k]

                output_dict["CONFIG"] = [
                    "UNITS: %s" % units,
                    "FPS: %.2f" % fps,
                    "MASK_NFRAMES: %i" % maskn,
                ]
                self.time_series = output_dict.copy()

                # Update ui metrics
                self.ui.doubleSpinBox_fit_start_time.setMinimum(min(time))
                self.ui.doubleSpinBox_fit_start_time.setMaximum(max(time))
                self.ui.doubleSpinBox_fit_last_time.setMinimum(min(time))
                self.ui.doubleSpinBox_fit_last_time.setMaximum(max(time))
                self.ui.doubleSpinBox_fit_start_time.setValue(min(time))
                self.ui.doubleSpinBox_fit_last_time.setValue(max(time))

                # Update data summary with pixel length and displacement statistics
                summary = self.ui.label_data_summary.text().strip()
                lines = summary.split("\n") if summary else []
                px_per_mm = (
                    1 / pixel_length if pixel_length != 0 else float("inf")
                )  # Avoid division by zero

                def upsert_line(prefix, text):
                    for idx, line in enumerate(lines):
                        if line.startswith(prefix):
                            lines[idx] = text
                            return
                    lines.append(text)

                def remove_line(prefix):
                    lines[:] = [line for line in lines if not line.startswith(prefix)]

                upsert_line(
                    "Pixel length",
                    "Pixel length %s: %.4f mm/px | %.4f px/mm"
                    % (
                        self.ui.comboBox_units.currentText(),
                        pixel_length,
                        px_per_mm,
                    ),
                )

                if self.max_expansion_info is not None:
                    exp_val, exp_time = self.max_expansion_info
                    upsert_line(
                        "Max expansion",
                        "Max expansion: %+0.3f %s at %.2f s"
                        % (exp_val, length_unit_label, exp_time),
                    )
                else:
                    remove_line("Max expansion")

                if self.max_recession_info is not None:
                    rec_val, rec_time = self.max_recession_info
                    upsert_line(
                        "Max recession",
                        "Max recession: %+0.3f %s at %.2f s"
                        % (rec_val, length_unit_label, rec_time),
                    )
                else:
                    remove_line("Max recession")

                self.ui.label_data_summary.setText("\n".join(lines))
                self.ui.basebar.setText("Finished plotting data")
            else:
                self.ui.basebar.setText("Not enough data to plot")
                if self.testing:
                    print("Not enough data to plot")
                else:
                    self.arcjetcv_message_box(
                        "Warning",
                        "Not enough data to plot: only %i points"
                        % len(self.raw_outputs),
                    )

        except Exception as e:
            self.ui.basebar.setText("Plotting failed !")
            if self.testing:
                print("Warning", "Plotting failed :\n" + str(e))
            else:
                self.arcjetcv_message_box("Warning", "Plotting failed:\n" + str(e))

        self.ui.Window1.repaint()
        self.ui.Window2.repaint()
        self.fit_data()

    def save_plots(self, name):
        try:
            # Extract necessary data from self.time_series
            time = self.time_series["TIME [s]"]
            units = self.ui.comboBox_units.currentText()

            # Plot each dataset in a separate graph
            for data, label in zip(self.length_units[5:], self.length_labels[5:]):
                fig, ax = plt.subplots()
                ax.scatter(time, data, label=label)
                ax.set_xlabel("Time (s)")
                ax.set_ylabel(f"{label} ({units})")
                ax.legend()
                filename = f"_{label.replace(' ', '_').lower()}.png"
                fig.savefig(str(name + filename))
                plt.close(fig)

                # Plot each dataset in a separate graph
            for data, label in zip(self.px_units, self.px_labels):
                fig, ax = plt.subplots()
                ax.scatter(time, data, label=label)
                ax.set_xlabel("Time (s)")
                ax.set_ylabel(f"{label} ({units})")
                ax.legend()
                filename = f"_{label.replace(' ', '_').lower()}.png"
                fig.savefig(str(name + filename))
                plt.close(fig)

            # Plot "MODEL_-0.95R", "MODEL_-0.50R", "MODEL_CENTER", "MODEL_0.50R", "MODEL_0.95R" in one graph
            fig, ax_combined = plt.subplots()
            for i, data in enumerate(self.length_units[:5]):
                label = f"Model {['-0.95R', '-0.50R', 'Center', '0.50R', '0.95'][i]}"
                ax_combined.scatter(time, data, label=label)
            ax_combined.set_xlabel("Time (s)")
            ax_combined.set_ylabel(f"Model Positions ({units})")
            ax_combined.legend()
            fig.savefig(name + "_radius.png")
            plt.close(fig)

        except Exception as e:
            self.ui.basebar.setText("Plot saving failed !")
            if self.testing:
                print("Warning", "Plot saving failed:\n" + str(e))
            else:
                self.arcjetcv_message_box(
                    "Warning", "Plot saving failed:\n" + str(e)
                )

    def export_to_csv(self):
        if self.time_series is not None:
            try:
                dialog = QtWidgets.QFileDialog()
                pathmask = dialog.getExistingDirectory(None, "Export CSV/plots", "")
                os.makedirs(pathmask, exist_ok=True)
                folder_path = os.path.abspath(pathmask)
                folder_name = os.path.basename(pathmask)

                self.ax1.figure.savefig(
                    os.path.join(folder_path, folder_name + "_xy.png")
                )
                self.ax2.figure.savefig(
                    os.path.join(folder_path, folder_name + "_xt.png")
                )
                self.save_plots(os.path.join(folder_path, folder_name))

                # Convert time series into dataframe
                df = pd.DataFrame(
                    dict([(k, pd.Series(v)) for k, v in self.time_series.items()])
                )

                # Convert fits into dataframe
                if self.fit_dict is not None:
                    param_dict = {"Pixel length": self.pixel_length}
                    # Create pandas Series from fitparam dictionary
                    keys, vals, errs = [], [], []
                    for key, val in self.fit_dict.items():
                        if "LINEAR_FIT" in key:
                            keys.append(key + "_slope")
                            keys.append(key + "_intercept")
                            vals.append(val[0][0])
                            vals.append(val[0][1])
                            errs.append(val[1][0])
                            errs.append(val[1][1])

                    # Create pandas Series from fitparam dictionary
                    keys_series = pd.Series(keys)
                    values_series = pd.Series(vals)
                    err_series = pd.Series(errs)
                    param_dict.update(self.fit_dict)

                    # Assign the Series as new columns to the DataFrame
                    df = df.assign(
                        Keys=keys_series, Values=values_series, Error=err_series
                    )

                    csv_file_path = os.path.join(folder_path, f"{folder_name}.csv")
                    df.to_csv(csv_file_path, index=False)

            except Exception as e:
                self.ui.basebar.setText("CSV export failed !")
                if self.testing:
                    print("Warning", "CSV export failed:\n" + str(e))
                else:
                    self.arcjetcv_message_box(
                        "Warning", "CSV export failed:\n" + str(e)
                    )

    def fit_data(self):
        if self.time_series is not None:

            try:
                # Retrieve user inputs for time/units
                units = self.ui.comboBox_units.currentText()
                time = np.array(self.time_series["TIME [s]"])
                mask = np.array(self.XT_MASK)
                t0 = self.ui.doubleSpinBox_fit_start_time.value()
                t1 = self.ui.doubleSpinBox_fit_last_time.value()

                # identify relevant indicies
                # Include interval endpoints to avoid dropping valid edge samples.
                inds = (time >= t0) * (time <= t1) * (mask > 0)

                # Initialize data structure
                fit_dict = {}

                # Get list of keys to loop through
                keys = list(self.time_series.keys())
                keys.remove("TIME [s]")
                keys.remove("CONFIG")

                # Linear fits
                if self.ui.comboBox_fit_type.currentText() == "linear":
                    longstring = ""
                    for key in keys:
                        t = time[inds]
                        y = self.time_series[key][inds]
                        if len(t) > 2 and len(y) > 2:
                            fitp, cov = np.ma.polyfit(t, y, 1, cov=True)
                            err = np.sqrt(np.diag(cov))
                            m, b = fitp[0], fitp[1]
                            fit_dict[key + "_LINEAR_FIT"] = (fitp, err)
                            if key in self.PLOTKEYS:
                                longstring += (
                                    key
                                    + "\tMin = %f\t Max = %f\t Delta = %f\n"
                                    % (y.min(), y.max(), y.max() - y.min())
                                )
                                longstring += "\tLINEAR FIT: y = mx+b \t"
                                longstring += (
                                    "m = %f+-%f %s" % (m, err[0], units + "/s") + "\t"
                                )
                                longstring += (
                                    "b = %f+-%f %s" % (b, err[1], units) + "\n\n"
                                )
                        else:
                            longstring = "not enough points to fit data"

                    self.ui.textBrowser.setText(longstring)
                    self.ui.textBrowser.repaint()

                self.fit_dict = fit_dict

            except Exception as e:
                self.ui.basebar.setText("Fitting failed !")
                if self.testing:
                    print("Warning", "Fitting failed :\n" + str(e))
                else:
                    self.arcjetcv_message_box(
                        "Warning", "Fitting failed :\n" + str(e)
                    )

    def arcjetcv_message_box(self, title, message):
        msg_box = QMessageBox(self)
        msg_box.setIconPixmap(QPixmap(self.logo_path))
        msg_box.setWindowTitle(title)
        msg_box.setText(message)
        msg_box.setStandardButtons(QMessageBox.Ok)
        msg_box.exec()
