"""
Dynamic Spectrum Viewer & RFI Cleaning Tool
=============================================

Description:
    This graphical user interface (GUI) application is designed for viewing,
    cleaning, and processing dynamic spectra generated by the pipeline or
    'make_dynamic_spectra.py'. Specially tested with LOFAR data, the tool
    provides users with the ability to:
        1. View the dynamic spectrum.
        2. Flag and mask regions within the dynamic spectrum.
        3. Perform bandpass normalization.
        4. Extract active solar radio emissions from the dynamic spectrum.
        5. Save the cleaned dynamic spectrum as a FITS file.

Features:
    - Interactive visualization using Matplotlib integrated with PyQt5.
    - ROI (Region of Interest) selection for RFI (Radio Frequency Interference) flagging.
    - Cross-section mode for extracting and analyzing time and frequency slices.
    - Auto-scaling and adjustable visualization parameters.
    - Undo/Redo support for modification steps.
    - Option to view FITS file metadata via a menu command.

Dependencies:
    - Python 3.x
    - NumPy, SciPy, OpenCV (cv2)
    - Astropy
    - Matplotlib
    - PyQt5

Tested Environment:
    - This tool has been specifically tested with LOFAR data.

Usage:
    - Launch the application to view dynamic spectra.
    - Use the provided controls to mask unwanted regions, normalize the bandpass,
      and extract specific source features.
    - Use Undo (Ctrl+Z) and Redo (Ctrl+Y) to revert modifications.
    - Select "View Metadata" from the menu to see the FITS file header.
    - Save the processed dynamic spectrum as a FITS file for further analysis.

Authors:
    Soham Dey, Deepan Patra

Version:
    1.0

Date:
    6th February 2025

Notes:
    - Ensure that all required dependencies are installed.
    - For any queries or issues, please refer to the documentation.
"""

import sys
import os
import numpy as np
import numpy.ma as ma
import cv2

from astropy.io import fits
from astropy.time import Time

from PyQt5.QtWidgets import (
    QMainWindow, QApplication, QWidget, QVBoxLayout, QHBoxLayout, QSplitter,
    QAction, QFileDialog, QMessageBox, QStatusBar, QMenuBar,
    QLabel, QSlider, QComboBox, QFormLayout, QPushButton,
    QDockWidget, QDoubleSpinBox, QSpinBox, QCheckBox, QInputDialog, QGroupBox,
    QDialog, QTextEdit, QVBoxLayout as QVBoxLayoutDialog, QScrollArea
)
from PyQt5 import QtWidgets
from PyQt5.QtCore import Qt, pyqtSlot

import matplotlib
import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm, Normalize, PowerNorm
from matplotlib.backends.backend_qt5agg import (
    FigureCanvasQTAgg as FigureCanvas,
    NavigationToolbar2QT as NavigationToolbar
)
from matplotlib.widgets import RectangleSelector
from matplotlib.dates import date2num, DateFormatter


###############################################################################
#                           RFI CLEANING UTILITIES                            #
###############################################################################

def create_binary(data, thresh):
    """Create a binary DS: 0 where data < thresh, 1 where data >= thresh."""
    return np.where(data < thresh, 0.0, 1.0)

def region_detection(original_image, binary_image, min_width=1, min_height=5):
    """
    Finds closed contours in the binary image.
    Returns (closed_image, valid_contours, overlay_image).
    """
    binary_image_uint8 = np.uint8(binary_image * 255)
    kernel = np.ones((2, 2), np.uint8)
    closed_image = cv2.morphologyEx(binary_image_uint8, cv2.MORPH_CLOSE, kernel)
    contours, _ = cv2.findContours(closed_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    if len(original_image.shape) == 2:
        overlay_image = cv2.cvtColor(original_image.astype(np.float32), cv2.COLOR_GRAY2BGR)
    else:
        overlay_image = original_image.copy()

    valid_contours = []
    for contour in contours:
        x, y, w, h = cv2.boundingRect(contour)
        if w >= min_width and h >= min_height:
            valid_contours.append(contour)
            cv2.rectangle(overlay_image, (x, y), (x + w, y + h), (255, 0, 0), 3)
    return closed_image, valid_contours, overlay_image

def create_mask(original_binary, contours):
    """
    Create a mask with detected contours filled.
    """
    mask = np.zeros(original_binary.shape, dtype=np.uint8)
    for contour in contours:
        cv2.drawContours(mask, [contour], -1, 255, thickness=cv2.FILLED)
    return mask

def subtract_contours(original_image, mask):
    """
    Zero out the flagged (RFI) region in original_image.
    """
    result_image = original_image.copy()
    result_image[mask > 0] = 0
    return result_image


###############################################################################
#              MATPLOTLIB CANVAS FOR DYNAMIC SPECTRUM DISPLAY                 #
###############################################################################

class DynamicSpectrumCanvas(FigureCanvas):
    """
    A Matplotlib canvas for displaying a dynamic spectrum (time vs freq).
    """
    def __init__(self, parent=None):
        # Setup logging
        import logging
        self.logger = logging.getLogger("DynamicSpectrumCanvas")
        self.logger.setLevel(logging.DEBUG)
        
        # Apply initial theme settings to matplotlib
        self._apply_initial_theme()
        
        # Create figure with better DPI and constrained layout for better auto-resizing
        self.fig = plt.figure(figsize=(9, 6), dpi=100, constrained_layout=True)
        self.ax = self.fig.add_subplot(111)
        super().__init__(self.fig)
        self.setParent(parent)
        
        # Set minimum size to avoid layout issues
        self.setMinimumSize(400, 300)

        self.ax.set_title("Dynamic Spectrum Viewer", fontsize=14)
        self.ax.set_xlabel("Time (UTC)", fontsize=12)
        self.ax.set_ylabel("Frequency (MHz)", fontsize=12)
        
        # Apply theme to the newly created canvas
        self._apply_canvas_theme_colors()

        # Internal references
        self._data = None               # Float array with NaNs
        self._time_axis = None
        self._freq_axis = None
        self._extent = None

        self._colorbar = None
        self._im = None

        # Visualization parameters
        self._scale_mode = "Linear"     # Options: "Linear", "Log", "Sqrt", "Gamma"
        self._gamma = 1.0
        self._vmin = None
        self._vmax = None
        self._cmap = 'inferno'
        self._smart_scale = "0.5-99.5%"     # Options: "0-100%", "0.1-99.9%", "0.5-99.5%", "1-99%"

        # Interactive ROI
        self.rect_selector = None
        self.roi_active = False
        self.roi_callback = None

        # Cross-section mode
        self.cross_section_active = False
        
        # Mouse hover info control
        self.hover_info_enabled = False  # Default disabled
        self._last_hover_text = ""  # Cache last hover text to avoid unnecessary redraws
        self._hover_throttle_counter = 0  # Throttle hover updates
        self._hover_throttle_limit = 3  # Update every N mouse events
        
        # Store event connection IDs for reconnection
        self._event_connections = {}
        self._connect_events()

        # Mouse readout: text annotation on the axes for amplitude values (default to dark theme)
        self._hover_text = self.ax.text(
            1.0, 1.05, "",
            transform=self.ax.transAxes, ha="right", va="top",
            fontsize=10, color="white",
            bbox=dict(facecolor="black", alpha=0.8)
        )
        
        self.fig.tight_layout()

    def _apply_initial_theme(self):
        """Apply dark theme as default for matplotlib."""
        import matplotlib.pyplot as plt
        # Set dark theme as default
        plt.rcParams.update({
            'figure.facecolor': '#2b2b2b',
            'axes.facecolor': '#2b2b2b',
            'axes.edgecolor': 'white',
            'axes.labelcolor': 'white',
            'xtick.color': 'white',
            'ytick.color': 'white',
            'text.color': 'white',
            'axes.grid': True,
            'grid.color': '#505050',
            'grid.alpha': 0.3
        })

    def _apply_canvas_theme_colors(self):
        """Apply theme colors to the canvas elements."""
        # Get parent window to access theme information
        parent = self.parent()
        while parent is not None and not hasattr(parent, 'current_theme'):
            parent = parent.parent()
            
        # Default to dark theme if no parent found
        if parent is None or not hasattr(parent, 'themes'):
            theme = {
                'plot_bg': '#2b2b2b',
                'plot_text': 'white',
                'plot_grid': '#505050'
            }
        else:
            theme = parent.themes[parent.current_theme]
        
        # Set figure and axes background
        self.fig.patch.set_facecolor(theme['plot_bg'])
        self.ax.set_facecolor(theme['plot_bg'])
        
        # Set text colors
        self.ax.tick_params(colors=theme['plot_text'])
        self.ax.xaxis.label.set_color(theme['plot_text'])
        self.ax.yaxis.label.set_color(theme['plot_text'])
        self.ax.title.set_color(theme['plot_text'])
        
        # Set border colors
        for spine in self.ax.spines.values():
            spine.set_edgecolor(theme['plot_text'])
        
        # Enable grid with theme colors
        self.ax.grid(True, color=theme['plot_grid'], alpha=0.3)
        
        # Update hover text colors
        if hasattr(self, '_hover_text'):
            if parent and parent.current_theme == "dark":
                self._hover_text.set_color("white")
                self._hover_text.set_bbox(dict(facecolor="black", alpha=0.8))
            else:
                self._hover_text.set_color("black")
                self._hover_text.set_bbox(dict(facecolor="white", alpha=0.8))

    def _connect_events(self):
        """Connect all necessary event handlers"""
        # Disconnect existing connections if they exist
        self._disconnect_events()
        
        # Connect mouse events
        self._event_connections['button_press'] = self.mpl_connect("button_press_event", self.on_mouse_click)
        self._event_connections['motion_notify'] = self.mpl_connect("motion_notify_event", self.on_mouse_move)
        
        # If ROI selector is active, reconnect it
        if self.roi_active and self.roi_callback:
            self.enable_roi_selector(True, self.roi_callback)

    def _disconnect_events(self):
        """Disconnect all matplotlib event connections"""
        for event_id in self._event_connections.values():
            try:
                self.mpl_disconnect(event_id)
            except:
                pass
        self._event_connections = {}

    def clear_plot(self):
        """Clear the axes and remove the colorbar."""
        # Store current states before clearing
        roi_was_active = self.roi_active
        roi_callback = self.roi_callback
        cross_was_active = self.cross_section_active
        
        # Turn off rectangle selector before clearing
        if self.rect_selector is not None:
            try:
                self.rect_selector.set_active(False)
                self.rect_selector.set_visible(False)
                # For Matplotlib 3.x, disconnect events
                if hasattr(self.rect_selector, 'disconnect_events'):
                    self.rect_selector.disconnect_events()
                # Invalidate the selector
                self.rect_selector = None
            except Exception as e:
                print(f"Error disabling selector: {e}")
            
        # Disconnect existing events before clearing
        self._disconnect_events()

        # Clear the axis
        self.ax.clear()
        
        # Reset basic axis properties
        self.ax.set_title(" ", fontsize=14)
        self.ax.set_xlabel("Time (UTC)", fontsize=12)
        self.ax.set_ylabel("Frequency (MHz)", fontsize=12)
        
        # Handle colorbar removal
        if self._colorbar is not None:
            try:
                self._colorbar.remove()
            except (AttributeError, ValueError):
                # Handle case where colorbar removal fails - recreate the figure
                self.fig.clf()
                self.ax = self.fig.add_subplot(111)
                
                # Reset basic axis properties
                self.ax.set_title(" ", fontsize=14)
                self.ax.set_xlabel("Time (UTC)", fontsize=12)
                self.ax.set_ylabel("Frequency (MHz)", fontsize=12)
            self._colorbar = None
            
        self._im = None
        self._extent = None
        
        # Recreate hover text after clearing with theme-aware colors
        parent = self.parent()
        while parent is not None and not hasattr(parent, 'current_theme'):
            parent = parent.parent()
        
        # Set colors based on current theme
        if parent and hasattr(parent, 'current_theme') and parent.current_theme == "light":
            text_color = "black"
            bg_color = "white"
        else:
            text_color = "white"
            bg_color = "black"
            
        self._hover_text = self.ax.text(
            1.0, 1.05, "",
            transform=self.ax.transAxes, ha="right", va="top",
            fontsize=10, color=text_color,
            bbox=dict(facecolor=bg_color, alpha=0.8)
        )
        
        # Reconnect event handlers
        self._connect_events()
        
        # Restore interactive features if they were active
        if roi_was_active and roi_callback:
            self.enable_roi_selector(True, roi_callback)
        
        self.cross_section_active = cross_was_active

    def set_data(self, data, time_axis=None, freq_axis=None):
        """Set data and optional time and frequency axes."""
        self._data = data
        self._time_axis = time_axis
        self._freq_axis = freq_axis

    def draw_spectrum(self):
        """Plot the dynamic spectrum using the current visualization parameters."""
        self.clear_plot()
        if self._data is None:
            self.draw()
            return

        # Reset axis limits before plotting
        self.ax.set_xlim(auto=True)
        self.ax.set_ylim(auto=True)

        data_ma = ma.masked_invalid(self._data)
        data_flat = data_ma.compressed()
        if len(data_flat) == 0:
            self._vmin, self._vmax = 0, 1
        else:
            if self._smart_scale == "0-100%":
                self._vmin, self._vmax = data_flat.min(), data_flat.max()
            elif self._smart_scale == "0.1-99.9%":
                p01 = np.percentile(data_flat, 0.1)
                p999 = np.percentile(data_flat, 99.9)
                self._vmin, self._vmax = p01, p999
            elif self._smart_scale == "0.5-99.5%":
                p05 = np.percentile(data_flat, 0.5)
                p995 = np.percentile(data_flat, 99.5)
                self._vmin, self._vmax = p05, p995
            elif self._smart_scale == "1-99%":
                p1 = np.percentile(data_flat, 1)
                p99 = np.percentile(data_flat, 99)
                self._vmin, self._vmax = p1, p99
            

        if self._vmin == self._vmax:
            self._vmin -= 1e-9
            self._vmax += 1e-9

        if self._scale_mode == "Log":
            norm = LogNorm(vmin=max(self._vmin, 1e-12), vmax=self._vmax)
        elif self._scale_mode == "Sqrt":
            norm = PowerNorm(gamma=0.5, vmin=self._vmin, vmax=self._vmax)
        elif self._scale_mode == "Gamma":
            norm = PowerNorm(gamma=self._gamma, vmin=self._vmin, vmax=self._vmax)
        else:
            norm = Normalize(vmin=self._vmin, vmax=self._vmax)

        if (self._time_axis is None) or (self._freq_axis is None):
            self._im = self.ax.imshow(data_ma.T, aspect='auto', origin='lower',
                                     cmap=self._cmap, norm=norm)
            self.ax.set_xlabel("Time index")
            self.ax.set_ylabel("Frequency index")
        else:
            time_mjd = self._time_axis / 86400.0
            utc_dt = Time(time_mjd, format='mjd', scale='utc').to_datetime()
            utc_num = np.array([date2num(dt) for dt in utc_dt])
            self._extent = [utc_num[0], utc_num[-1], self._freq_axis[0], self._freq_axis[-1]]
            self._im = self.ax.imshow(data_ma.T, aspect='auto', origin='lower',
                                     extent=self._extent, cmap=self._cmap, norm=norm)
            date_formatter = DateFormatter('%Y-%m-%d\n%H:%M:%S')
            self.ax.xaxis.set_major_formatter(date_formatter)
            self.fig.autofmt_xdate()
            self.ax.set_xlabel("Time (UTC)")
            self.ax.set_ylabel("Frequency (MHz)")

        self._colorbar = self.fig.colorbar(self._im, ax=self.ax, label="Amplitude")
        
        # Apply theme colors to the newly created colorbar
        self._apply_colorbar_theme()
        
        # Ensure the figure layout is updated
        self.fig.tight_layout()
        self.draw()

    def _apply_colorbar_theme(self):
        """Apply theme colors to the colorbar."""
        if self._colorbar is None:
            return
            
        # Get parent window to access theme information
        parent = self.parent()
        while parent is not None and not hasattr(parent, 'current_theme'):
            parent = parent.parent()
            
        if parent is None or not hasattr(parent, 'themes'):
            return
            
        theme = parent.themes[parent.current_theme]
        
        # Update colorbar text colors
        self._colorbar.ax.tick_params(colors=theme['plot_text'])
        self._colorbar.ax.yaxis.label.set_color(theme['plot_text'])
        
        # Update colorbar outline
        for spine in self._colorbar.ax.spines.values():
            spine.set_edgecolor(theme['plot_text'])

    def set_scale_mode(self, mode):
        self._scale_mode = mode
        self.draw_spectrum()

    def set_gamma(self, gamma):
        self._gamma = gamma
        if self._scale_mode == "Gamma":
            self.draw_spectrum()

    def set_cmap(self, cmap_name):
        self._cmap = cmap_name
        self.draw_spectrum()

    def set_smart_scale(self, scale_option):
        """Set the percentile scale option for auto scaling."""
        self._smart_scale = scale_option
        self.draw_spectrum()

    def get_normalization(self, vmin, vmax):
        if self._scale_mode == "Log":
            return LogNorm(vmin=max(vmin, 1e-12), vmax=vmax)
        elif self._scale_mode == "Sqrt":
            return PowerNorm(gamma=0.5, vmin=vmin, vmax=vmax)
        elif self._scale_mode == "Gamma":
            return PowerNorm(gamma=self._gamma, vmin=vmin, vmax=vmax)
        else:
            return Normalize(vmin=vmin, vmax=vmax)

    # -------------------------- ROI Selector ------------------------------------
    def enable_roi_selector(self, enable, callback=None):
        """
        Enable or disable the ROI selection mode.
        When enabled, allows the user to draw a rectangle to select a region of interest.
        """
        self.roi_active = enable
        self.roi_callback = callback

        if enable:
            # Disconnect any existing rectangle selector
            if self.rect_selector is not None:
                try:
                    if hasattr(self.rect_selector, 'disconnect_events'):
                        self.rect_selector.disconnect_events()
                    self.rect_selector = None
                except:
                    pass
                    
            # Clear any existing rectangles
            self._clear_all_rects()
            
            # Use direct event handling rather than RectangleSelector
            self.logger.info("Setting up direct ROI event handlers")
            
            # Set up our own event handlers for more direct control
            self._roi_start_point = None
            self._roi_current_rect = None
            
            # Connect direct mouse events for ROI selection
            if 'roi_press' in self._event_connections:
                self.mpl_disconnect(self._event_connections['roi_press'])
            if 'roi_release' in self._event_connections:
                self.mpl_disconnect(self._event_connections['roi_release'])
            if 'roi_motion' in self._event_connections:
                self.mpl_disconnect(self._event_connections['roi_motion'])
                
            self._event_connections['roi_press'] = self.mpl_connect('button_press_event', self._roi_on_press)
            self._event_connections['roi_release'] = self.mpl_connect('button_release_event', self._roi_on_release)
            self._event_connections['roi_motion'] = self.mpl_connect('motion_notify_event', self._roi_on_motion)
            
            # Add a highly visible message to indicate ROI selection is active
            self._hover_text.set_text("ROI selection active: Click and drag to select a region")
            self._hover_text.set_bbox(dict(facecolor='orange', alpha=0.9, boxstyle='round'))
            self._hover_text.set_color('black')
            self._hover_text.set_fontsize(12)
            
            self.draw_idle()
        else:
            # Disable all ROI events
            for event_name in ['roi_press', 'roi_release', 'roi_motion']:
                if event_name in self._event_connections:
                    try:
                        self.mpl_disconnect(self._event_connections[event_name])
                        del self._event_connections[event_name]
                    except:
                        pass
            
            # Clean up any remaining rectangles
            self._clear_all_rects()
            
            # Reset hover text with theme-aware colors
            parent = self.parent()
            while parent is not None and not hasattr(parent, 'current_theme'):
                parent = parent.parent()
            
            if parent and hasattr(parent, 'current_theme') and parent.current_theme == "light":
                text_color = "black"
                bg_color = "white"
            else:
                text_color = "white"
                bg_color = "black"
                
            self._hover_text.set_text("")
            self._hover_text.set_bbox(dict(facecolor=bg_color, alpha=0.8))
            self._hover_text.set_color(text_color)
            self._hover_text.set_fontsize(10)
            
            self._roi_start_point = None
            self._roi_current_rect = None
            
            self.draw_idle()
    
    def _clear_all_rects(self):
        """Clear all rectangle patches from the axes."""
        try:
            # Remove all rectangle patches that are not the axes background
            for patch in self.ax.patches[:]:
                if isinstance(patch, matplotlib.patches.Rectangle):
                    if patch != self.ax.patch:  # Don't remove the axes background
                        patch.remove()
            self.draw_idle()
        except Exception as e:
            self.logger.error(f"Error clearing rectangles: {e}")
    
    def _roi_on_press(self, event):
        """Handle mouse button press for ROI selection."""
        if not self.roi_active or event.inaxes != self.ax or event.button != 1:
            return
            
        # Clear any existing rectangles
        self._clear_all_rects()
        
        # Store the starting point
        self._roi_start_point = (event.xdata, event.ydata)
        
        # Update status message
        self._hover_text.set_text("Dragging: release to confirm selection")
        self.draw_idle()
    
    def _roi_on_motion(self, event):
        """Handle mouse motion for ROI selection."""
        if not self.roi_active or self._roi_start_point is None or event.inaxes != self.ax:
            return
            
        # Remove the previous rectangle if it exists
        if self._roi_current_rect is not None:
            self._roi_current_rect.remove()
            self._roi_current_rect = None
        
        # Create a new rectangle from start point to current position
        start_x, start_y = self._roi_start_point
        current_x, current_y = event.xdata, event.ydata
        
        # Don't draw if we're outside the axes
        if None in (start_x, start_y, current_x, current_y):
            return
            
        # Calculate rectangle parameters
        x = min(start_x, current_x)
        y = min(start_y, current_y)
        width = abs(current_x - start_x)
        height = abs(current_y - start_y)
        
        # Create a very visible rectangle
        from matplotlib.patches import Rectangle
        self._roi_current_rect = Rectangle(
            (x, y), width, height,
            facecolor='red', edgecolor='yellow',
            alpha=0.4, fill=True, linewidth=3.0,
            zorder=1000  # Ensure it's on top of everything
        )
        
        # Add the rectangle to the plot
        self.ax.add_patch(self._roi_current_rect)
        
        # Draw immediately
        self.draw_idle()
    
    def _roi_on_release(self, event):
        """Handle mouse button release for ROI selection."""
        if not self.roi_active or self._roi_start_point is None or event.inaxes != self.ax or event.button != 1:
            return
            
        # Get the final coordinates
        start_x, start_y = self._roi_start_point
        end_x, end_y = event.xdata, event.ydata
        
        # Don't process if we're outside the axes
        if None in (start_x, start_y, end_x, end_y):
            self._roi_start_point = None
            if self._roi_current_rect is not None:
                self._roi_current_rect.remove()
                self._roi_current_rect = None
            self.draw_idle()
            return
            
        # Process the selection
        self.logger.info(f"ROI selection: ({start_x}, {start_y}) to ({end_x}, {end_y})")
        
        # Keep the rectangle visible for feedback
        if self._roi_current_rect is not None:
            # Make it slightly transparent to show it's being processed
            self._roi_current_rect.set_alpha(0.6)
            self._roi_current_rect.set_edgecolor('lime')  # Change color to indicate processing
            self.draw_idle()
        
        # Calculate indices from data coordinates
        try:
            nt, nf = self._data.shape
            if self._extent is not None:
                x0, x1_, y0, y1_ = self._extent
                
                # Get selection bounds
                xmin, xmax = sorted([start_x, end_x])
                ymin, ymax = sorted([start_y, end_y])
                
                # Check for division by zero
                if x1_ == x0 or y1_ == y0:
                    self.logger.warning("Invalid extent: division by zero")
                    return
                
                frac_xmin = (xmin - x0) / (x1_ - x0)
                frac_xmax = (xmax - x0) / (x1_ - x0)
                frac_ymin = (ymin - y0) / (y1_ - y0)
                frac_ymax = (ymax - y0) / (y1_ - y0)
                
                # Check for values outside valid range
                if (frac_xmin < 0 and frac_xmax < 0) or (frac_xmin > 1 and frac_xmax > 1) or \
                   (frac_ymin < 0 and frac_ymax < 0) or (frac_ymin > 1 and frac_ymax > 1):
                    self.logger.warning("Selection outside valid range")
                    return
                
                ixmin = int(max(0, min(1, frac_xmin)) * (nt - 1))
                ixmax = int(max(0, min(1, frac_xmax)) * (nt - 1))
                iymin = int(max(0, min(1, frac_ymin)) * (nf - 1))
                iymax = int(max(0, min(1, frac_ymax)) * (nf - 1))
            else:
                # Direct pixel coordinates
                xmin, xmax = sorted([int(start_x), int(end_x)])
                ymin, ymax = sorted([int(start_y), int(end_y)])
                ixmin, ixmax = xmin, xmax
                iymin, iymax = ymin, ymax
            
            # Ensure indices are within valid range
            ixmin = max(0, min(ixmin, nt - 1))
            ixmax = max(0, min(ixmax, nt - 1))
            iymin = max(0, min(iymin, nf - 1))
            iymax = max(0, min(iymax, nf - 1))
            
            # Call the callback function if defined
            if self.roi_callback and ixmin <= ixmax and iymin <= iymax:
                self.roi_callback(ixmin, ixmax, iymin, iymax)
        except Exception as e:
            self.logger.error(f"Error processing ROI selection: {e}")
        
        # Clean up for next selection
        self._roi_start_point = None
        if self._roi_current_rect is not None:
            self._roi_current_rect.remove()
            self._roi_current_rect = None
        
        # Reset hover text
        self._hover_text.set_text("ROI selection active: Click and drag to select a region")
        self.draw_idle()

    # --------------------- Cross-Section Mode -----------------------------------
    def enable_cross_section(self, enable):
        self.cross_section_active = enable

    def on_mouse_click(self, event):
        if not self.cross_section_active or self._data is None:
            return
        if event.inaxes != self.ax or event.button != 1:
            return

        try:
            nt, nf = self._data.shape
            if self._extent is not None:
                x0, x1, y0, y1 = self._extent
                # Ensure valid division
                if x1 == x0 or y1 == y0:
                    return
                    
                frac_x = (event.xdata - x0) / (x1 - x0)
                frac_y = (event.ydata - y0) / (y1 - y0)
                
                # Check if click is within valid data range
                if not (0 <= frac_x <= 1) or not (0 <= frac_y <= 1):
                    return
                    
                time_idx = int(frac_x * (nt - 1))
                freq_idx = int(frac_y * (nf - 1))
            else:
                # Handle integer indices case
                if event.xdata is None or event.ydata is None:
                    return
                time_idx = int(event.xdata)
                freq_idx = int(event.ydata)

            # Ensure indices are within valid range
            time_idx = max(0, min(time_idx, nt - 1))
            freq_idx = max(0, min(freq_idx, nf - 1))

            options = [
                f"Time slice at freq_idx = {freq_idx}",
                f"Freq slice at time_idx = {time_idx}"
            ]
            choice, ok = QInputDialog.getItem(
                None, "Cross Section", "Which slice to plot?", options, 0, False
            )
            if not ok:
                return

            if choice.startswith("Time slice"):
                data_slice = self._data[:, freq_idx]
                self._plot_1d_time(data_slice, freq_idx)
            else:
                data_slice = self._data[time_idx, :]
                self._plot_1d_freq(data_slice, time_idx)
        except (ValueError, TypeError, IndexError) as e:
            # Handle any errors in coordinate conversion
            QMessageBox.warning(None, "Cross Section Error", 
                                f"Could not extract cross-section: {str(e)}")

    def _plot_1d_time(self, data_slice, freq_idx):
        nt = data_slice.shape[0]
        
        # Get theme colors from parent window
        parent = self.parent()
        while parent is not None and not hasattr(parent, 'current_theme'):
            parent = parent.parent()
        
        # Determine colors based on theme
        if parent and hasattr(parent, 'themes') and hasattr(parent, 'current_theme'):
            theme = parent.themes[parent.current_theme]
            plot_bg = theme['plot_bg']
            plot_text = theme['plot_text']
            line_color = plot_text  # Use theme text color for line
        else:
            # Default to dark theme
            plot_bg = '#2b2b2b'
            plot_text = 'white'
            line_color = 'white'
        
        fig, ax = plt.subplots(facecolor=plot_bg)
        ax.set_facecolor(plot_bg)
        yvals = np.ma.masked_invalid(data_slice)

        if self._time_axis is not None:
            time_mjd = self._time_axis / 86400.0
            utc_dt = Time(time_mjd, format='mjd', scale='utc').to_datetime()
            xvals = np.array([date2num(dt) for dt in utc_dt])
            if np.any(np.diff(xvals) <= 0):
                xvals = np.linspace(xvals[0], xvals[-1], nt)
            ax.plot_date(xvals, yvals, '-', lw=1.5, color=line_color)
            date_formatter = DateFormatter('%H:%M:%S')
            ax.xaxis.set_major_formatter(date_formatter)
            fig.autofmt_xdate()
            xlabel = "Time (UTC)"
        else:
            xvals = np.linspace(0, nt - 1, nt)
            ax.plot(xvals, yvals, '-', lw=1.5, color=line_color)
            xlabel = "Time index"

        # Apply theme colors to all text elements
        ax.set_title(f"Time Slice @ freq_idx={freq_idx}", color=plot_text)
        ax.set_xlabel(xlabel, color=plot_text)
        ax.set_ylabel("Amplitude", color=plot_text)
        ax.tick_params(colors=plot_text)
        
        # Set spine colors
        for spine in ax.spines.values():
            spine.set_edgecolor(plot_text)
        
        fig.tight_layout()
        fig.show()

    def _plot_1d_freq(self, data_slice, time_idx):
        nf = data_slice.shape[0]
        
        # Get theme colors from parent window
        parent = self.parent()
        while parent is not None and not hasattr(parent, 'current_theme'):
            parent = parent.parent()
        
        # Determine colors based on theme
        if parent and hasattr(parent, 'themes') and hasattr(parent, 'current_theme'):
            theme = parent.themes[parent.current_theme]
            plot_bg = theme['plot_bg']
            plot_text = theme['plot_text']
            line_color = plot_text  # Use theme text color for line
        else:
            # Default to dark theme
            plot_bg = '#2b2b2b'
            plot_text = 'white'
            line_color = 'white'
        
        fig, ax = plt.subplots(facecolor=plot_bg)
        ax.set_facecolor(plot_bg)
        yvals = np.ma.masked_invalid(data_slice)

        if self._freq_axis is not None:
            xvals = self._freq_axis
            ax.plot(xvals, yvals, '-', lw=1.5, color=line_color)
            xlabel = "Frequency (MHz)"
        else:
            xvals = np.linspace(0, nf - 1, nf)
            ax.plot(xvals, yvals, '-', lw=1.5, color=line_color)
            xlabel = "Frequency index"

        # Apply theme colors to all text elements
        ax.set_title(f"Freq Slice @ time_idx={time_idx}", color=plot_text)
        ax.set_xlabel(xlabel, color=plot_text)
        ax.set_ylabel("Amplitude", color=plot_text)
        ax.tick_params(colors=plot_text)
        
        # Set spine colors
        for spine in ax.spines.values():
            spine.set_edgecolor(plot_text)
        
        fig.tight_layout()
        fig.show()

    # ------------------ Mouse Hover: Display Amplitude --------------------------
    def on_mouse_move(self, event):
        # Check if hover info is enabled
        if not self.hover_info_enabled:
            # Only clear text and redraw if text was previously showing
            if self._last_hover_text:
                self._hover_text.set_text("")
                self._last_hover_text = ""
                self.draw_idle()
            return
            
        if event.inaxes != self.ax or self._data is None:
            # Only clear text and redraw if text was previously showing
            if self._last_hover_text:
                self._hover_text.set_text("")
                self._last_hover_text = ""
                self.draw_idle()
            return

        # Throttle hover updates to reduce CPU usage
        self._hover_throttle_counter += 1
        if self._hover_throttle_counter < self._hover_throttle_limit:
            return
        self._hover_throttle_counter = 0

        try:
            nt, nf = self._data.shape
            if self._extent is not None:
                x0, x1, y0, y1 = self._extent
                frac_x = (event.xdata - x0) / (x1 - x0)
                time_idx = int(frac_x * (nt - 1))
                frac_y = (event.ydata - y0) / (y1 - y0)
                freq_idx = int(frac_y * (nf - 1))
            else:
                time_idx = int(event.xdata)
                freq_idx = int(event.ydata)

            time_idx = max(0, min(time_idx, nt - 1))
            freq_idx = max(0, min(freq_idx, nf - 1))

            val = self._data[time_idx, freq_idx]
            if np.isnan(val):
                msg = f"Time={time_idx}, Freq={freq_idx}, Amp=NaN"
            else:
                msg = f"Time={time_idx}, Freq={freq_idx}, Amp={val:.3f}"
            
            # Only update and redraw if the text has changed
            if msg != self._last_hover_text:
                self._hover_text.set_text(msg)
                self._last_hover_text = msg
                self.draw_idle()
                
        except (TypeError, ValueError, IndexError) as e:
            # Handle any errors in coordinate conversion - only clear if needed
            if self._last_hover_text:
                self._hover_text.set_text("")
                self._last_hover_text = ""
                self.draw_idle()


###############################################################################
#                       MAIN WINDOW - PyQt5 APPLICATION                       #
###############################################################################

class MainWindow(QMainWindow):
    def __init__(self, theme="dark"):
        super(MainWindow, self).__init__()
        self.setWindowTitle("Dynamic Spectra Viewer")
        self.resize(1500, 850)

        # Store theme from parameter (allows external theme override)
        self._external_theme = theme

        # Data placeholders
        self._original_unmodified = None  # For revert
        self._original_data = None        # Working data (float array with NaNs)
        self._time_axis = None
        self._freq_axis = None

        # Metadata string (FITS header)
        self._metadata = ""

        # Undo/Redo stacks (store copies of working data)
        self.undo_stack = []
        self.redo_stack = []

        # Theme management - use theme from constructor parameter
        self.current_theme = self._external_theme if self._external_theme in ["dark", "light"] else "dark"
        
        # Theme palettes matching solarviewer's styles.py for consistency
        self.themes = {
            "light": {
                # Solarviewer LIGHT_PALETTE colors
                "main_bg": "#A59D84",        # window
                "panel_bg": "#ECEBDE",       # base/surface
                "panel_border": "#b0b0b0",   # border
                "text_color": "#1a1a1a",     # text
                "button_bg": "#f0f0f0",      # button
                "button_hover": "#e0e0e0",   # button_hover
                "button_pressed": "#d0d0d0", # button_pressed
                "menubar_bg": "#D7D3BF",     # toolbar_bg
                "statusbar_bg": "#ECEBDE",   # surface
                "groupbox_bg": "#ECEBDE",    # surface
                "input_bg": "#ECEBDE",       # base
                "input_border": "#b0b0b0",   # border
                "highlight": "#0066cc",      # highlight
                "plot_bg": "#ECEBDE",        # plot_bg
                "plot_text": "#1a1a1a",      # plot_text
                "plot_grid": "#d0d0d0"       # plot_grid
            },
            "dark": {
                # Solarviewer DARK_PALETTE colors
                "main_bg": "#1a1a2e",        # window
                "panel_bg": "#1f2940",       # surface
                "panel_border": "#2a3f5f",   # border
                "text_color": "#eeeeee",     # text
                "button_bg": "#0f3460",      # button
                "button_hover": "#1a4a7a",   # button_hover
                "button_pressed": "#0a2540", # button_pressed
                "menubar_bg": "#1a1a2e",     # window
                "statusbar_bg": "#1f2940",   # surface
                "groupbox_bg": "#1f2940",    # surface
                "input_bg": "#16213e",       # base
                "input_border": "#2a3f5f",   # border
                "highlight": "#e94560",      # highlight
                "plot_bg": "#16213e",        # base
                "plot_text": "#eeeeee",      # text
                "plot_grid": "#2a3f5f"       # border
            }
        }

        self._setupLogger()
        self._createActions()
        self._createMenuBar()
        self._createMainWidgets()
        self._createStatusBar()
        self._applyStyle()
        
        # Connect resize event 
        self.resizeTimer = None

    def _setupLogger(self):
        import logging
        self.logger = logging.getLogger("DynamicSpectrumViewer")
        self.logger.setLevel(logging.DEBUG)
        if not self.logger.handlers:
            ch = logging.StreamHandler()
            ch.setLevel(logging.DEBUG)
            formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
            ch.setFormatter(formatter)
            self.logger.addHandler(ch)

    # ------------------------- Actions and Menu -------------------------------
    def _createActions(self):
        self.openAct = QAction("&Open FITS...", self)
        self.openAct.setShortcut("Ctrl+O")
        self.openAct.triggered.connect(self.openFile)

        self.saveAct = QAction("&Save FITS...", self)
        self.saveAct.setShortcut("Ctrl+S")
        self.saveAct.triggered.connect(self.saveCleanedData)

        self.exitAct = QAction("E&xit", self)
        self.exitAct.setShortcut("Ctrl+Q")
        self.exitAct.triggered.connect(self.close)

        self.compareAct = QAction("Compare &Original vs Cleaned", self)
        self.compareAct.setShortcut("Ctrl+C")
        self.compareAct.triggered.connect(self.compareOriginalCleaned)

        self.revertAct = QAction("&Revert to Original", self)
        self.revertAct.setShortcut("Ctrl+R")
        self.revertAct.triggered.connect(self.revertToOriginal)

        # New Undo/Redo actions
        self.undoAct = QAction("&Undo", self)
        self.undoAct.setShortcut("Ctrl+Z")
        self.undoAct.triggered.connect(self.undo)

        self.redoAct = QAction("&Redo", self)
        self.redoAct.setShortcut("Ctrl+Y")
        self.redoAct.triggered.connect(self.redo)

        # New action for viewing metadata
        self.viewMetaAct = QAction("View &Metadata", self)
        self.viewMetaAct.triggered.connect(self.viewMetadata)

        # Theme toggle action
        self.toggleThemeAct = QAction("Toggle &Dark/Light Theme", self)
        self.toggleThemeAct.setShortcut("Ctrl+T")
        self.toggleThemeAct.triggered.connect(self.toggleTheme)

    def _createMenuBar(self):
        menubar = QMenuBar(self)
        self.setMenuBar(menubar)

        fileMenu = menubar.addMenu("&File")
        fileMenu.addAction(self.openAct)
        fileMenu.addAction(self.saveAct)
        fileMenu.addSeparator()
        fileMenu.addAction(self.exitAct)

        editMenu = menubar.addMenu("&Edit")
        editMenu.addAction(self.undoAct)
        editMenu.addAction(self.redoAct)

        viewMenu = menubar.addMenu("&View")
        viewMenu.addAction(self.viewMetaAct)
        viewMenu.addSeparator()
        viewMenu.addAction(self.toggleThemeAct)

        toolsMenu = menubar.addMenu("&Tools")
        toolsMenu.addAction(self.compareAct)
        toolsMenu.addAction(self.revertAct)

    # -------------------------- Central Widgets ---------------------------------
    def _createMainWidgets(self):
        centralWidget = QWidget()
        self.setCentralWidget(centralWidget)
        layout = QHBoxLayout(centralWidget)

        # Left side: Matplotlib canvas and toolbar
        leftWidget = QWidget()
        leftLayout = QVBoxLayout(leftWidget)
        
        self.canvas = DynamicSpectrumCanvas(self)
        self.navbar = NavigationToolbar(self.canvas, self)

        leftLayout.addWidget(self.canvas)
        leftLayout.addWidget(self.navbar)
        
        # Connect canvas resize events to ensure proper redrawing
        self.canvas.setSizePolicy(
            QtWidgets.QSizePolicy.Expanding,
            QtWidgets.QSizePolicy.Expanding
        )
        self.canvas.updateGeometry()
        
        # Right side: Control Panel
        rightWidget = QWidget()
        rightLayout = QVBoxLayout(rightWidget)

        # Set size policies for the left and right panels
        leftWidget.setSizePolicy(
            QtWidgets.QSizePolicy.Expanding,
            QtWidgets.QSizePolicy.Expanding
        )
        # Right panel: allow horizontal expansion to fill splitter space
        rightWidget.setSizePolicy(
            QtWidgets.QSizePolicy.Preferred,
            QtWidgets.QSizePolicy.Preferred
        )
        rightWidget.setMinimumWidth(250)

        # (A) Scale mode
        scaleLabel = QLabel("Intensity Scale:")
        self.scaleCombo = QComboBox()
        self.scaleCombo.addItems(["Linear", "Log", "Sqrt", "Gamma"])
        self.scaleCombo.setCurrentText("Linear")
        self.scaleCombo.currentTextChanged.connect(self.onScaleModeChanged)
        rightLayout.addWidget(scaleLabel)
        rightLayout.addWidget(self.scaleCombo)

        # (B) Gamma slider
        self.gammaLabel = QLabel("Gamma:")
        self.gammaSlider = QSlider(Qt.Horizontal)
        self.gammaSlider.setRange(1, 300)
        self.gammaSlider.setValue(100)
        self.gammaSlider.valueChanged.connect(self.onGammaChanged)
        self.gammaLabel.setEnabled(False)
        self.gammaSlider.setEnabled(False)
        rightLayout.addWidget(self.gammaLabel)
        rightLayout.addWidget(self.gammaSlider)

        # (C) Colormap
        cmapLabel = QLabel("Colormap:")
        self.cmapCombo = QComboBox()
        cmaps = ["viridis", "plasma", "inferno", "magma", "cividis",
                 "jet", "gray", "bone", "turbo", "afmhot", "afmhot_r", "cubehelix", "Greens", "gist_heat", "gist_heat_r"]
        self.cmapCombo.addItems(cmaps)
        self.cmapCombo.setCurrentText("inferno")
        self.cmapCombo.currentTextChanged.connect(self.onCmapChanged)
        rightLayout.addWidget(cmapLabel)
        rightLayout.addWidget(self.cmapCombo)

        # (D) Smart Auto-Scale - replace checkbox with dropdown
        scaleRangeLabel = QLabel("Intensity Range:")
        self.scaleRangeCombo = QComboBox()
        scale_options = ["0-100%", "0.1-99.9%", "0.5-99.5%", "1-99%"]
        self.scaleRangeCombo.addItems(scale_options)
        self.scaleRangeCombo.setCurrentText("0.5-99.5%")
        self.scaleRangeCombo.currentTextChanged.connect(self.onScaleRangeChanged)
        rightLayout.addWidget(scaleRangeLabel)
        rightLayout.addWidget(self.scaleRangeCombo)

        # (E) ROI selection (mask region)
        self.roiButton = QPushButton("Mask Region (ROI)")
        self.roiButton.setCheckable(True)
        self.roiButton.toggled.connect(self.onRoiToggled)
        rightLayout.addWidget(self.roiButton)

        # (F) Mouse Hover Info
        hoverGroup = QGroupBox("Mouse Hover Settings")
        hoverLayout = QVBoxLayout(hoverGroup)
        
        self.hoverInfoBox = QCheckBox("Show Mouse Hover Info")
        self.hoverInfoBox.setChecked(False)  # Default disabled
        self.hoverInfoBox.toggled.connect(self.onHoverInfoToggled)
        hoverLayout.addWidget(self.hoverInfoBox)
        
        # Hover update frequency (throttling)
        hoverFreqLayout = QHBoxLayout()
        hoverFreqLabel = QLabel("Update Rate:")
        self.hoverFreqSpin = QSpinBox()
        self.hoverFreqSpin.setRange(1, 20)
        self.hoverFreqSpin.setValue(3)  # Default: update every 3rd mouse event
        self.hoverFreqSpin.setSuffix(" events")
        self.hoverFreqSpin.setToolTip("Higher values = less frequent updates = better performance")
        self.hoverFreqSpin.valueChanged.connect(self.onHoverFrequencyChanged)
        hoverFreqLayout.addWidget(hoverFreqLabel)
        hoverFreqLayout.addWidget(self.hoverFreqSpin)
        hoverLayout.addLayout(hoverFreqLayout)
        
        rightLayout.addWidget(hoverGroup)

        # (G) Theme Toggle
        # self.themeToggleBtn = QPushButton("☀️ Light Theme")  # Start with dark theme, so button shows "switch to light"
        # self.themeToggleBtn.setToolTip("Toggle between Dark and Light themes (Ctrl+T)")
        # self.themeToggleBtn.clicked.connect(self.toggleTheme)
        # rightLayout.addWidget(self.themeToggleBtn)

        # (H) Cross Section
        self.crossSectionBtn = QPushButton("Cross Section Mode")
        self.crossSectionBtn.setCheckable(True)
        self.crossSectionBtn.toggled.connect(self.onCrossSectionToggled)
        rightLayout.addWidget(self.crossSectionBtn)

        # (H) Region Detection box
        detectGroup = QGroupBox("Region Detection")
        detectLayout = QVBoxLayout(detectGroup)
        form = QFormLayout()
        self.threshSpin = QDoubleSpinBox()
        self.threshSpin.setRange(0.0, 1e9)
        self.threshSpin.setValue(5.0)
        form.addRow("Threshold:", self.threshSpin)
        self.minWidthSpin = QSpinBox()
        self.minWidthSpin.setRange(1, 9999)
        self.minWidthSpin.setValue(1)
        form.addRow("Min Width:", self.minWidthSpin)
        self.minHeightSpin = QSpinBox()
        self.minHeightSpin.setRange(1, 9999)
        self.minHeightSpin.setValue(5)
        form.addRow("Min Height:", self.minHeightSpin)
        regionDetectBtn = QPushButton("Region Detection")
        regionDetectBtn.clicked.connect(self.onCleanRFIRegionDetect)
        form.addRow(regionDetectBtn)
        detectLayout.addLayout(form)
        rightLayout.addWidget(detectGroup)

        # (I) Bandpass Normalization
        self.bandpassBtn = QPushButton("Bandpass Norm")
        self.bandpassBtn.clicked.connect(self.onBandpassNorm)
        rightLayout.addWidget(self.bandpassBtn)

        rightLayout.addStretch()
        
        # Wrap right widget in a scroll area for long control panels
        scrollArea = QScrollArea()
        scrollArea.setWidget(rightWidget)
        scrollArea.setWidgetResizable(True)
        scrollArea.setMinimumWidth(250)
        
        # Use QSplitter for draggable resizable panels
        splitter = QSplitter(Qt.Horizontal)
        splitter.addWidget(leftWidget)
        splitter.addWidget(scrollArea)
        splitter.setStretchFactor(0, 1)  # Canvas gets more space
        splitter.setStretchFactor(1, 0)  # Control panel stays compact
        splitter.setSizes([900, 300])    # Initial sizes
        
        layout.addWidget(splitter)

    # ---------------------------- Status Bar ------------------------------------
    def _createStatusBar(self):
        self.statusBar = QStatusBar()
        self.setStatusBar(self.statusBar)
        self.amplitudeLabel = QLabel("")
        self.statusBar.addPermanentWidget(self.amplitudeLabel)

    # --------------------- View Metadata Dialog ---------------------------------
    def viewMetadata(self):
        if not self._metadata:
            QMessageBox.information(self, "Metadata", "No metadata available.")
            return
        dlg = QDialog(self)
        dlg.setWindowTitle("FITS File Metadata")
        dlg.resize(600, 400)
        layout = QVBoxLayout(dlg)
        textEdit = QTextEdit(dlg)
        textEdit.setReadOnly(True)
        textEdit.setPlainText(self._metadata)
        layout.addWidget(textEdit)
        dlg.exec_()

    def _applyStyle(self):
        """Apply the current theme's styling to all UI elements."""
        theme = self.themes[self.current_theme]
        
        style = f"""
            QMainWindow {{
                background-color: {theme['main_bg']};
                color: {theme['text_color']};
            }}
            
            QMenuBar {{
                background-color: {theme['menubar_bg']};
                color: {theme['text_color']};
                font-size: 12pt;
                border-bottom: 1px solid {theme['panel_border']};
            }}
            
            QMenuBar::item {{
                background-color: transparent;
                padding: 4px 8px;
            }}
            
            QMenuBar::item:selected {{
                background-color: {theme['button_hover']};
            }}
            
            QMenu {{
                background-color: {theme['panel_bg']};
                color: {theme['text_color']};
                border: 1px solid {theme['panel_border']};
            }}
            
            QMenu::item:selected {{
                background-color: {theme['button_hover']};
            }}
            
            QStatusBar {{
                background-color: {theme['statusbar_bg']};
                color: {theme['text_color']};
                font-size: 10pt;
                border-top: 1px solid {theme['panel_border']};
            }}
            
            QWidget {{
                background-color: {theme['panel_bg']};
                color: {theme['text_color']};
            }}
            
            QPushButton {{
                background-color: {theme['button_bg']};
                color: {theme['text_color']};
                border: 1px solid {theme['panel_border']};
                padding: 6px 12px;
                border-radius: 4px;
                font-weight: bold;
            }}
            
            QPushButton:hover {{
                background-color: {theme['button_hover']};
            }}
            
            QPushButton:pressed {{
                background-color: {theme['button_pressed']};
            }}
            
            QPushButton:checked {{
                background-color: {theme['button_pressed']};
                border: 2px solid {theme['text_color']};
            }}
            
            QGroupBox {{
                background-color: {theme['groupbox_bg']};
                color: {theme['text_color']};
                border: 2px solid {theme['panel_border']};
                border-radius: 6px;
                margin-top: 6px;
                font-weight: bold;
            }}
            
            QGroupBox::title {{
                subcontrol-origin: margin;
                left: 10px;
                padding: 0 5px 0 5px;
            }}
            
            QComboBox {{
                background-color: {theme['input_bg']};
                color: {theme['text_color']};
                border: 1px solid {theme['input_border']};
                padding: 4px;
                border-radius: 4px;
            }}
            
            QComboBox:drop-down {{
                border: none;
            }}
            
            QComboBox::drop-down:hover {{
                background-color: {theme['button_hover']};
            }}
            
            QSpinBox, QDoubleSpinBox {{
                background-color: {theme['input_bg']};
                color: {theme['text_color']};
                border: 1px solid {theme['input_border']};
                padding: 4px;
                border-radius: 4px;
            }}
            
            QSlider::groove:horizontal {{
                border: 1px solid {theme['input_border']};
                height: 6px;
                background: {theme['input_bg']};
                border-radius: 3px;
            }}
            
            QSlider::sub-page:horizontal {{
                background: {theme['highlight']};
                border-radius: 3px;
            }}
            
            QSlider::handle:horizontal {{
                background: {theme['highlight']};
                border: 2px solid {theme['text_color']};
                width: 16px;
                height: 16px;
                margin: -6px 0;
                border-radius: 9px;
            }}
            
            QSlider::handle:horizontal:hover {{
                background: {theme['button_hover']};
                border: 2px solid {theme['text_color']};
            }}
            
            QCheckBox {{
                color: {theme['text_color']};
            }}
            
            QCheckBox::indicator {{
                width: 16px;
                height: 16px;
                background-color: {theme['input_bg']};
                border: 1px solid {theme['input_border']};
                border-radius: 3px;
            }}
            
            QCheckBox::indicator:checked {{
                background-color: {theme['button_pressed']};
                border: 2px solid {theme['text_color']};
            }}
            
            QLabel {{
                color: {theme['text_color']};
            }}
            
            QTextEdit {{
                background-color: {theme['input_bg']};
                color: {theme['text_color']};
                border: 1px solid {theme['input_border']};
                border-radius: 4px;
            }}
        """
        
        self.setStyleSheet(style)
        
        # Apply theme to matplotlib canvas
        self._applyCanvasTheme()

    def _applyCanvasTheme(self):
        """Apply the current theme to the matplotlib canvas."""
        if not hasattr(self, 'canvas'):
            return
            
        theme = self.themes[self.current_theme]
        
        # Set matplotlib style parameters
        import matplotlib.pyplot as plt
        plt.rcParams.update({
            'figure.facecolor': theme['plot_bg'],
            'axes.facecolor': theme['plot_bg'],
            'axes.edgecolor': theme['plot_text'],
            'axes.labelcolor': theme['plot_text'],
            'xtick.color': theme['plot_text'],
            'ytick.color': theme['plot_text'],
            'text.color': theme['plot_text'],
            'axes.grid': True,
            'grid.color': theme['plot_grid'],
            'grid.alpha': 0.3
        })
        
        # Always update the canvas theme, regardless of whether data is loaded
        self.canvas.fig.patch.set_facecolor(theme['plot_bg'])
        self.canvas.ax.set_facecolor(theme['plot_bg'])
        
        # Update axis colors and spines (borders)
        self.canvas.ax.tick_params(colors=theme['plot_text'])
        self.canvas.ax.xaxis.label.set_color(theme['plot_text'])
        self.canvas.ax.yaxis.label.set_color(theme['plot_text'])
        self.canvas.ax.title.set_color(theme['plot_text'])
        
        # Update spine (border) colors
        for spine in self.canvas.ax.spines.values():
            spine.set_edgecolor(theme['plot_text'])
        
        # Update grid
        self.canvas.ax.grid(True, color=theme['plot_grid'], alpha=0.3)
        
        # Update colorbar if it exists
        if hasattr(self.canvas, '_colorbar') and self.canvas._colorbar is not None:
            # Update colorbar text colors
            self.canvas._colorbar.ax.tick_params(colors=theme['plot_text'])
            self.canvas._colorbar.ax.yaxis.label.set_color(theme['plot_text'])
            
            # Update colorbar outline
            for spine in self.canvas._colorbar.ax.spines.values():
                spine.set_edgecolor(theme['plot_text'])
        
        # Update hover text colors if it exists
        if hasattr(self.canvas, '_hover_text'):
            if self.current_theme == "dark":
                self.canvas._hover_text.set_color("white")
                self.canvas._hover_text.set_bbox(dict(facecolor="black", alpha=0.8))
            else:
                self.canvas._hover_text.set_color("black")
                self.canvas._hover_text.set_bbox(dict(facecolor="white", alpha=0.8))
        
        # Always redraw the canvas to show theme changes
        self.canvas.draw()

    def toggleTheme(self):
        """Toggle between dark and light themes."""
        if self.current_theme == "light":
            self.current_theme = "dark"
            self.themeToggleBtn.setText("☀️ Light Theme")
        else:
            self.current_theme = "light"
            self.themeToggleBtn.setText("🌙 Dark Theme")
            
        # Apply the new theme
        self._applyStyle()
        
        # Update status message
        self.statusBar.showMessage(f"Switched to {self.current_theme} theme", 3000)

    # ---------------------------- File Operations -------------------------------
    def openFile(self):
        fileName, _ = QFileDialog.getOpenFileName(
            self, "Open FITS File", "", "FITS Files (*.fits *.fts);;All Files (*)"
        )
        if not fileName:
            return
        try:
            self.logger.info(f"Opening FITS file: {fileName}")
            hdul = fits.open(fileName)
            data = hdul[0].data
            header = hdul[0].header  # Get primary HDU header as metadata
            self._metadata = str(header)
            time_axis = None
            freq_axis = None
            
            # First try to get axes from dedicated extensions (LOFAR/Learmonth format)
            for hdu in hdul:
                if hdu.name.upper() == "TIME_AXIS":
                    time_mjd = hdu.data["TIME_MJD"]
                    time_axis = time_mjd * 86400.0
                if hdu.name.upper() == "FREQ_AXIS":
                    freq_axis = hdu.data["FREQ_MHz"]
            
            # Fallback: try to extract from WCS headers if extensions not found
            if time_axis is None or freq_axis is None:
                self.logger.info("No axis extensions found, trying WCS headers...")
                try:
                    # Get data shape - shape is (freq, time) or (time, freq) depending on file
                    naxis1 = header.get('NAXIS1', data.shape[1] if len(data.shape) > 1 else data.shape[0])
                    naxis2 = header.get('NAXIS2', data.shape[0])
                    
                    # Try to get frequency axis from CTYPE1/CRVAL1/CDELT1
                    ctype1 = header.get('CTYPE1', '')
                    if freq_axis is None and 'FREQ' in ctype1.upper():
                        crval1 = header.get('CRVAL1', 0)
                        cdelt1 = header.get('CDELT1', 1)
                        crpix1 = header.get('CRPIX1', 1)
                        freq_axis = crval1 + (np.arange(naxis1) - (crpix1 - 1)) * cdelt1
                        self.logger.info(f"Extracted freq_axis from WCS: {freq_axis[0]:.2f} to {freq_axis[-1]:.2f} MHz")
                    
                    # Try to get time axis from CTYPE2/CRVAL2/CDELT2
                    ctype2 = header.get('CTYPE2', '')
                    if time_axis is None and 'TIME' in ctype2.upper():
                        # Get start time from DATE-OBS
                        date_obs = header.get('DATE-OBS', None)
                        cdelt2 = header.get('CDELT2', 3.0)  # seconds between samples
                        if date_obs:
                            from astropy.time import Time as AstroTime
                            t_start = AstroTime(date_obs)
                            # Create time array in MJD seconds
                            time_axis = (t_start.mjd * 86400.0) + np.arange(naxis2) * cdelt2
                            self.logger.info(f"Extracted time_axis from WCS: {naxis2} samples, dt={cdelt2}s")
                    
                    # Alternative: check if axes are swapped (CTYPE1=TIME, CTYPE2=FREQ)
                    if freq_axis is None and 'FREQ' in ctype2.upper():
                        crval2 = header.get('CRVAL2', 0)
                        cdelt2 = header.get('CDELT2', 1)
                        crpix2 = header.get('CRPIX2', 1)
                        freq_axis = crval2 + (np.arange(naxis2) - (crpix2 - 1)) * cdelt2
                        self.logger.info(f"Extracted freq_axis from WCS (CTYPE2): {freq_axis[0]:.2f} to {freq_axis[-1]:.2f} MHz")
                    
                    if time_axis is None and 'TIME' in ctype1.upper():
                        date_obs = header.get('DATE-OBS', None)
                        cdelt1 = header.get('CDELT1', 3.0)
                        if date_obs:
                            from astropy.time import Time as AstroTime
                            t_start = AstroTime(date_obs)
                            time_axis = (t_start.mjd * 86400.0) + np.arange(naxis1) * cdelt1
                            self.logger.info(f"Extracted time_axis from WCS (CTYPE1): {naxis1} samples")
                except Exception as wcs_err:
                    self.logger.warning(f"Could not extract axes from WCS: {wcs_err}")
            
            hdul.close()

            data = np.array(data, dtype=np.float32)
            data[np.isnan(data)] = np.nan
            data[np.isinf(data)] = np.nan

            self.undo_stack.clear()
            self.redo_stack.clear()

            self._original_unmodified = data.copy()
            self._original_data = data.copy()
            self._time_axis = time_axis
            self._freq_axis = freq_axis

            self.canvas.set_data(self._original_data, self._time_axis, self._freq_axis)
            self.canvas.draw_spectrum()
            self.statusBar.showMessage(f"Loaded {os.path.basename(fileName)}", 5000)
        except Exception as e:
            self.logger.exception(f"Failed to open or plot FITS file: {e}")
            QMessageBox.critical(self, "Error", f"Failed to open FITS file:\n{e}")

    def saveCleanedData(self):
        if self._original_data is None:
            QMessageBox.information(self, "Info", "No data to save.")
            return
        fileName, _ = QFileDialog.getSaveFileName(
            self, "Save Cleaned FITS", "", "FITS Files (*.fits)"
        )
        if not fileName:
            return
        try:
            hdu = fits.PrimaryHDU(self._original_data)
            hdul = fits.HDUList([hdu])
            if self._time_axis is not None:
                from astropy.table import Table
                t = Table()
                t["TIME_MJD"] = self._time_axis / 86400.0
                time_hdu = fits.BinTableHDU(t, name="TIME_AXIS")
                hdul.append(time_hdu)
            if self._freq_axis is not None:
                from astropy.table import Table
                t = Table()
                t["FREQ_MHz"] = self._freq_axis
                freq_hdu = fits.BinTableHDU(t, name="FREQ_AXIS")
                hdul.append(freq_hdu)
            hdul.writeto(fileName, overwrite=True)
            self.statusBar.showMessage(f"Saved cleaned data to {fileName}", 5000)
        except Exception as e:
            QMessageBox.critical(self, "Error", f"Failed to save FITS:\n{e}")

    def compareOriginalCleaned(self):
        if self._original_unmodified is None or self._original_data is None:
            QMessageBox.information(self, "Info", "No data loaded.")
            return
        fig, axes = plt.subplots(1, 2, figsize=(12, 5), sharey=True)
        axes[0].set_title("Original Unmodified")
        axes[1].set_title("Current Cleaned")
        orig_ma = ma.masked_invalid(self._original_unmodified)
        cle_ma = ma.masked_invalid(self._original_data)
        
        # Apply the same scaling as the current view
        if self.canvas._smart_scale == "0-100%":
            orig_vals = orig_ma.compressed()
            cle_vals = cle_ma.compressed()
            orig_vmin, orig_vmax = orig_vals.min(), orig_vals.max()
            cle_vmin, cle_vmax = cle_vals.min(), cle_vals.max()
        elif self.canvas._smart_scale == "0.1-99.9%":
            orig_vals = orig_ma.compressed()
            cle_vals = cle_ma.compressed()
            orig_vmin = np.percentile(orig_vals, 0.1)
            orig_vmax = np.percentile(orig_vals, 99.9)
            cle_vmin = np.percentile(cle_vals, 0.1)
            cle_vmax = np.percentile(cle_vals, 99.9)
        elif self.canvas._smart_scale == "0.5-99.5%":
            orig_vals = orig_ma.compressed()
            cle_vals = cle_ma.compressed()
            orig_vmin = np.percentile(orig_vals, 0.5)
            orig_vmax = np.percentile(orig_vals, 99.5)
            cle_vmin = np.percentile(cle_vals, 0.5)
            cle_vmax = np.percentile(cle_vals, 99.5)
        else:  # "1-99%"
            orig_vals = orig_ma.compressed()
            cle_vals = cle_ma.compressed()
            orig_vmin = np.percentile(orig_vals, 1)
            orig_vmax = np.percentile(orig_vals, 99)
            cle_vmin = np.percentile(cle_vals, 1)
            cle_vmax = np.percentile(cle_vals, 99)
        
        norm1 = self.canvas.get_normalization(orig_vmin, orig_vmax)
        norm2 = self.canvas.get_normalization(cle_vmin, cle_vmax)
        im1 = axes[0].imshow(orig_ma.T, origin='lower', aspect='auto',
                             norm=norm1, cmap=self.canvas._cmap)
        fig.colorbar(im1, ax=axes[0], label="Amplitude")
        im2 = axes[1].imshow(cle_ma.T, origin='lower', aspect='auto',
                             norm=norm2, cmap=self.canvas._cmap)
        fig.colorbar(im2, ax=axes[1], label="Amplitude")
        axes[0].set_xlabel("Time Index")
        axes[1].set_xlabel("Time Index")
        axes[0].set_ylabel("Frequency Index")
        axes[1].set_ylabel("Frequency Index")
        fig.suptitle("Original vs Cleaned")
        fig.tight_layout()
        fig.show()

    def revertToOriginal(self):
        if self._original_unmodified is None:
            QMessageBox.warning(self, "Warning", "No original data to revert to.")
            return
        self._original_data = self._original_unmodified.copy()
        self.canvas.set_data(self._original_data, self._time_axis, self._freq_axis)
        self.canvas.draw_spectrum()
        self.statusBar.showMessage("Reverted to original data.", 5000)

    # ----------------------- Undo/Redo Functionality ----------------------------
    def undo(self):
        if self.undo_stack:
            self.redo_stack.append(self._original_data.copy())
            self._original_data = self.undo_stack.pop()
            self.canvas.set_data(self._original_data, self._time_axis, self._freq_axis)
            self.canvas.draw_spectrum()
        else:
            QMessageBox.information(self, "Undo", "No more actions to undo.")

    def redo(self):
        if self.redo_stack:
            self.undo_stack.append(self._original_data.copy())
            self._original_data = self.redo_stack.pop()
            self.canvas.set_data(self._original_data, self._time_axis, self._freq_axis)
            self.canvas.draw_spectrum()
        else:
            QMessageBox.information(self, "Redo", "No more actions to redo.")

    # In each modifying operation, push current state to undo stack and clear redo stack.
    def onCleanRFIRegionDetect(self):
        if self._original_data is None:
            QMessageBox.information(self, "Info", "No data loaded.")
            return
        self.undo_stack.append(self._original_data.copy())
        self.redo_stack.clear()
        thresh = self.threshSpin.value()
        min_w = self.minWidthSpin.value()
        min_h = self.minHeightSpin.value()
        data_2d = self._original_data.copy()
        data_2d[np.isnan(data_2d)] = 0
        data_T = data_2d.T
        med_band = np.nanmedian(data_T, axis=1, keepdims=True)
        normed_data = data_T / (med_band + 1e-20)
        binary_image = create_binary(normed_data, thresh=thresh)
        _, valid_contours, _ = region_detection(normed_data, binary_image,
                                                 min_width=min_w, min_height=min_h)
        mask = create_mask(binary_image, valid_contours)
        rfi_map = subtract_contours(normed_data, mask=~mask)
        cleaned_data = rfi_map.T
        cleaned_data[np.isnan(cleaned_data)] = np.nan
        self._original_data = cleaned_data
        self.canvas.set_data(self._original_data, self._time_axis, self._freq_axis)
        self.canvas.draw_spectrum()
        self.statusBar.showMessage("Region detection done.", 5000)

    def onBandpassNorm(self):
        if self._original_data is None:
            QMessageBox.information(self, "Info", "No data loaded.")
            return
        self.undo_stack.append(self._original_data.copy())
        self.redo_stack.clear()
        data_2d = self._original_data
        freq_profile = np.nanmedian(data_2d, axis=0)
        freq_profile[freq_profile == 0] = 1e-20
        for i in range(len(freq_profile)):
            if np.isnan(freq_profile[i]):
                freq_profile[i] = 1e-20
        normed = data_2d / freq_profile
        normed[np.isnan(normed)] = np.nan
        self._original_data = normed
        self.canvas.set_data(self._original_data, self._time_axis, self._freq_axis)
        self.canvas.draw_spectrum()
        self.statusBar.showMessage("Bandpass normalization applied.", 5000)

    def _mask_roi_values(self, ixmin, ixmax, iymin, iymax):
        """Mask a region of interest with NaN values."""
        try:
            self.logger.info(f"Masking ROI: time=[{ixmin},{ixmax}], freq=[{iymin},{iymax}]")
            
            # Validate indices
            if ixmax < ixmin or iymax < iymin:
                self.logger.warning("Invalid ROI coordinates: Empty region")
                return
                
            # Validate against data dimensions
            if self._original_data is None:
                self.logger.error("No data loaded")
                return
                
            # Save current state for undo
            self.undo_stack.append(self._original_data.copy())
            self.redo_stack.clear()
            
            # Apply the mask
            self._original_data[ixmin:ixmax+1, iymin:iymax+1] = np.nan
            
            # Update the display
            self.canvas.set_data(self._original_data, self._time_axis, self._freq_axis)
            self.canvas.draw_spectrum()
            
            # Update status
            self.statusBar.showMessage(f"ROI masked: time=[{ixmin},{ixmax}], freq=[{iymin},{iymax}] => NaN", 5000)
            self.logger.info("ROI masked successfully")
        except Exception as e:
            self.logger.exception(f"Error masking ROI: {e}")
            QMessageBox.warning(self, "Error", f"Failed to mask region: {str(e)}")
            # Disable ROI selector on error
            self.roiButton.setChecked(False)
            self.canvas.enable_roi_selector(False)
            self.statusBar.showMessage("ROI masking failed", 5000)

    # ----------------------- Scale / Colormap Controls --------------------------
    @pyqtSlot(str)
    def onScaleModeChanged(self, mode):
        self.canvas.set_scale_mode(mode)
        if mode == "Gamma":
            self.gammaLabel.setEnabled(True)
            self.gammaSlider.setEnabled(True)
        else:
            self.gammaLabel.setEnabled(False)
            self.gammaSlider.setEnabled(False)

    @pyqtSlot(int)
    def onGammaChanged(self, sliderValue):
        gamma = sliderValue / 100.0
        self.canvas.set_gamma(gamma)

    @pyqtSlot(str)
    def onCmapChanged(self, cmap_name):
        self.canvas.set_cmap(cmap_name)

    @pyqtSlot(str)
    def onScaleRangeChanged(self, scale_option):
        """Handle changes to the scale range dropdown."""
        self.canvas.set_smart_scale(scale_option)

    # ------------------------- ROI / Cross-Section -----------------------------
    def onRoiToggled(self, checked):
        if self._original_data is None:
            QMessageBox.information(self, "Info", "No data loaded.")
            self.roiButton.setChecked(False)
            return
        
        try:
            if checked:
                self.logger.info("Enabling ROI selector...")
                self.canvas.enable_roi_selector(True, self._mask_roi_values)
                self.statusBar.showMessage("Mask Region ON: Click and drag to select an area to mask.")
                self.logger.info("ROI selector enabled successfully")
            else:
                self.logger.info("Disabling ROI selector...")
                self.canvas.enable_roi_selector(False)
                self.statusBar.clearMessage()
                self.logger.info("ROI selector disabled")
        except Exception as e:
            # Handle any exceptions and inform the user
            error_msg = f"There was an error with the ROI selector: {str(e)}"
            self.logger.exception(error_msg)
            QMessageBox.warning(
                self, 
                "ROI Selection Error", 
                f"{error_msg}\n\n"
                "This may be due to compatibility issues with your Matplotlib version.\n"
                "Please try updating Matplotlib or using a different selection method."
            )
            
            # Reset button state
            self.roiButton.setChecked(False)
            
            # Disable the selector and clean up
            try:
                self.canvas.roi_active = False
                if hasattr(self.canvas, 'rect_selector') and self.canvas.rect_selector is not None:
                    self.canvas.rect_selector.set_active(False)
                    if hasattr(self.canvas.rect_selector, 'disconnect_events'):
                        self.canvas.rect_selector.disconnect_events()
                    self.canvas.rect_selector = None
            except Exception as cleanup_error:
                self.logger.error(f"Error during cleanup: {cleanup_error}")
            
            # Update the status bar
            self.statusBar.showMessage("ROI selection disabled due to an error", 5000)

    def onHoverInfoToggled(self, checked):
        """Toggle the mouse hover info display on/off."""
        if hasattr(self.canvas, 'hover_info_enabled'):
            self.canvas.hover_info_enabled = checked
        else:
            # For backwards compatibility, set the attribute
            self.canvas.hover_info_enabled = checked
        
        # Clear hover text and reset cache if disabled
        if not checked and hasattr(self.canvas, '_hover_text'):
            self.canvas._hover_text.set_text("")
            if hasattr(self.canvas, '_last_hover_text'):
                self.canvas._last_hover_text = ""
            self.canvas.draw_idle()

    def onHoverFrequencyChanged(self, value):
        """Update the hover throttle limit for performance tuning."""
        if hasattr(self.canvas, '_hover_throttle_limit'):
            self.canvas._hover_throttle_limit = value
            # Reset counter to apply change immediately
            if hasattr(self.canvas, '_hover_throttle_counter'):
                self.canvas._hover_throttle_counter = 0

    def onCrossSectionToggled(self, checked):
        if self._original_data is None:
            QMessageBox.information(self, "Info", "No data loaded.")
            self.crossSectionBtn.setChecked(False)
            return
        self.canvas.enable_cross_section(checked)
        if checked:
            self.statusBar.showMessage("Cross-section mode ON. Click on data.")
        else:
            self.statusBar.clearMessage()

    def resizeEvent(self, event):
        """Handle window resize events to ensure canvas redraws properly"""
        super().resizeEvent(event)
        # Use a timer to avoid excessive redrawing during resize
        if self.resizeTimer is not None:
            self.resizeTimer.stop()
            
        # Wait for resize to complete before redrawing
        from PyQt5.QtCore import QTimer
        self.resizeTimer = QTimer()
        self.resizeTimer.setSingleShot(True)
        self.resizeTimer.timeout.connect(self._delayed_redraw)
        self.resizeTimer.start(200)  # 200ms delay
    
    def _delayed_redraw(self):
        """Redraw the canvas after resize is complete"""
        if hasattr(self, 'canvas') and self.canvas._data is not None:
            self.canvas.draw_idle()

###############################################################################
#                                  MAIN                                       #
###############################################################################

def main():
    """Entry point for viewds command."""
    app = QApplication(sys.argv)
    
    # Apply dark theme from solarviewer
    from solar_radio_image_viewer.from_simpl.simpl_theme import apply_theme
    apply_theme(app, "dark")
    
    w = MainWindow(theme="dark")
    w.show()
    sys.exit(app.exec_())

if __name__ == '__main__':
    main()
