import threading
import time
import tkinter as tk
from tkinter import ttk

import cv2
import numpy as np
import screeninfo  # Import the screeninfo library
import queue  # Import queue
import os
import logging

# Configure logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# Constants
MFD_WIDTH = 484
MFD_HEIGHT = 484
SAR_WIDTH = 2048
SAR_HEIGHT = 2048
SAR_DATA_TYPE = np.uint16
INITIAL_SAR_WIDTH = 1024  # Initial SAR display width
INITIAL_SAR_HEIGHT = 1024  # Initial SAR display height
MFD_FPS = 25  # Target MFD FPS
TKINTER_MIN_WIDTH = 484
TKINTER_MIN_HEIGHT = 484

# --- SAR Georeferencing Information (Example) ---
SAR_CENTER_LAT = 40.7128  # Example latitude of the SAR image center
SAR_CENTER_LON = -74.0060  # Example longitude of the SAR image center
SAR_IMAGE_SIZE_KM = 50.0  # Example size of the SAR image (50km x 50km)

# --- Available Color Palettes ---
COLOR_PALETTES = ["GRAY", "AUTUMN", "BONE", "JET", "WINTER", "RAINBOW", "OCEAN", "SUMMER", "SPRING", "COOL", "HSV",
                  "HOT"]

# --- Image File Paths ---
MFD_IMAGE_PATH = "MFD_img8_000000.bmp"  # Replace with your MFD image path
SAR_IMAGE_PATH = "SAR_geo_img16_000000.tif"  # Replace with your SAR image path


class App:
    def __init__(self, root):
        self.root = root
        self.root.title("Grifo E Control Panel")  # Set application title
        self.root.minsize(TKINTER_MIN_WIDTH, TKINTER_MIN_HEIGHT)  # Set minimum size

        # --- Queues for SAR Images and Mouse Coordinates ---
        self.sar_queue = queue.Queue(maxsize=5)
        self.mouse_queue = queue.Queue(maxsize=20)

        # --- Get Screen Information ---
        screen = screeninfo.get_monitors()[0]  # Use the primary screen
        screen_width = screen.width
        screen_height = screen.height

        # --- Window Placement (Initial values) ---
        self.tkinter_x = 10
        self.tkinter_y = MFD_HEIGHT + 20  # Place Tkinter window below MFD
        self.root.geometry(f"+{self.tkinter_x}+{self.tkinter_y}")

        self.mfd_x = self.tkinter_x
        self.mfd_y = 10

        self.sar_x = self.tkinter_x + TKINTER_MIN_WIDTH + 20  # Place SAR window to the right of Tkinter
        self.sar_y = 10

        # --- Initial SAR Contrast and Brightness ---
        self.sar_contrast = 1.0
        self.sar_brightness = 0
        self.sar_palette = "GRAY"

        # --- Status Bar ---
        self.statusbar = ttk.Label(self.root, text="Ready | MFD FPS: 0 | SAR FPS: N/A", relief=tk.SUNKEN,
                                   anchor=tk.W)
        self.statusbar.pack(side=tk.BOTTOM, fill=tk.X)

        # --- Control Frame ---
        self.control_frame = ttk.Frame(self.root, padding=10)
        self.control_frame.pack(side=tk.TOP, fill=tk.X)

        # --- SAR Parameters Frame ---
        self.sar_params_frame = ttk.Labelframe(self.control_frame, text="SAR Parameters", padding=10)
        self.sar_params_frame.pack(side=tk.LEFT, padx=10, fill=tk.X)

        # --- SAR Info Frame ---
        self.sar_info_frame = ttk.Labelframe(self.control_frame, text="SAR Info", padding=10)
        self.sar_info_frame.pack(side=tk.TOP, padx=10, fill=tk.X)

        # --- Layout all the parameter widgets in a grid ---
        row = 0

        # --- Test Image Checkbox ---
        self.test_image_var = tk.IntVar(value=0)  # Initial value: Static image
        self.test_image_check = ttk.Checkbutton(self.sar_params_frame, text="Test Image",
                                                 variable=self.test_image_var, command=self.update_image_mode)
        self.test_image_check.grid(row=row, column=0, columnspan=2, sticky=tk.W, padx=5, pady=2)
        row += 1

        # --- SAR Size Control ---
        self.sar_size_label = ttk.Label(self.sar_params_frame, text="SAR Size:")
        self.sar_size_label.grid(row=row, column=0, sticky=tk.W, padx=5, pady=2)

        self.sar_size_factors = ["1:1", "1:2", "1:3", "1:5", "1:10"]
        self.sar_size_combo = ttk.Combobox(self.sar_params_frame, values=self.sar_size_factors, state="readonly",
                                            width=5)
        self.sar_size_combo.set("1:2")  # Initial value (1/2 of the original size)
        self.sar_size_combo.grid(row=row, column=1, sticky=tk.E, padx=5, pady=2)
        self.sar_size_combo.bind("<<ComboboxSelected>>", self.update_sar_size)
        row += 1

        # --- Contrast Control ---
        self.contrast_label = ttk.Label(self.sar_params_frame, text="Contrast:")
        self.contrast_label.grid(row=row, column=0, sticky=tk.W, padx=5, pady=2)
        self.contrast_scale = ttk.Scale(self.sar_params_frame, orient=tk.HORIZONTAL, length=200,
                                         from_=0.1, to=3.0, command=self.update_contrast, value=1.0)
        self.contrast_scale.grid(row=row, column=1, sticky=tk.E, padx=5, pady=2)
        row += 1

        # --- Brightness Control ---
        self.brightness_label = ttk.Label(self.sar_params_frame, text="Brightness:")
        self.brightness_label.grid(row=row, column=0, sticky=tk.W, padx=5, pady=2)
        self.brightness_scale = ttk.Scale(self.sar_params_frame, orient=tk.HORIZONTAL, length=200,
                                           from_=-100, to=100, command=self.update_brightness, value=0)
        self.brightness_scale.grid(row=row, column=1, sticky=tk.E, padx=5, pady=2)
        row += 1

        # --- Color Palette Control ---
        self.palette_label = ttk.Label(self.sar_params_frame, text="Palette:")
        self.palette_label.grid(row=row, column=0, sticky=tk.W, padx=5, pady=2)
        self.palette_combo = ttk.Combobox(self.sar_params_frame, values=COLOR_PALETTES, state="readonly", width=8)
        self.palette_combo.set("GRAY")  # Initial value
        self.palette_combo.grid(row=row, column=1, sticky=tk.E, padx=5, pady=2)
        self.palette_combo.bind("<<ComboboxSelected>>", self.update_sar_palette)
        row += 1

        # --- SAR Center Lat/Lon Label ---
        self.sar_center_label = ttk.Label(self.sar_info_frame,
                                           text=f"Center: Lat={SAR_CENTER_LAT:.4f}, Lon={SAR_CENTER_LON:.4f}")
        self.sar_center_label.pack(side=tk.TOP, anchor=tk.W)

        # --- Mouse Lat/Lon Label ---
        self.mouse_latlon_label = ttk.Label(self.sar_info_frame, text="Mouse: Lat=N/A, Lon=N/A")
        self.mouse_latlon_label.pack(side=tk.TOP, anchor=tk.W)

        # --- Initial Image Data ---
        self.mfd_image_data = self.load_image(MFD_IMAGE_PATH, expected_dtype=np.uint8)
        self.sar_image_data = self.load_image(SAR_IMAGE_PATH, expected_dtype=SAR_DATA_TYPE)
        self.current_sar = cv2.normalize(self.sar_image_data, None, 0, 255, cv2.NORM_MINMAX,
                                          cv2.CV_8U)  # Current SAR for display
        self.adjusted_sar = self.current_sar.copy()  # Create a copy for contrast/brightness adjustments

        # --- Translation Offset (for simulation) ---
        self.mfd_x_offset = 0
        self.sar_x_offset = 0

        # --- FPS Counter ---
        self.mfd_frame_count = 0
        self.mfd_start_time = time.time()
        self.sar_update_time = time.time()  # Time of last SAR update
        self.sar_frame_count = 0
        self.mfd_fps = 0.0  # Initial MFD FPS
        self.sar_fps = 0.0  # Initial SAR FPS
        self.last_mouse_update = time.time()

        # --- Initial SAR Size ---
        self.sar_display_width = INITIAL_SAR_WIDTH
        self.sar_display_height = INITIAL_SAR_HEIGHT
        self.resized_sar = cv2.resize(self.adjusted_sar, (self.sar_display_width, self.sar_display_height),
                                       interpolation=cv2.INTER_AREA)

        # --- Window Initialization Flags ---
        self.mfd_window_initialized = False
        self.sar_window_initialized = False
        self.sar_mouse_callback_set = False  # Flag to track if the callback has been set

        # --- LUT Initialization ---
        self.brightness_contrast_lut = None  # Initialize LUT
        self.update_brightness_contrast_lut()  # Initial LUT calculation

        # --- Test Image Buffers ---
        self.test_mfd_image = np.zeros((MFD_HEIGHT, MFD_WIDTH), dtype=np.uint8)
        self.test_sar_image = np.zeros((SAR_HEIGHT, SAR_WIDTH), dtype=SAR_DATA_TYPE)

        # --- Image Loading Thread ---
        self.image_loading_thread = threading.Thread(target=self.load_images_and_generate_test_data, daemon=True)
        self.image_loading_thread.start()

        # --- Start Image Display ---
        self.update_mfd()
        self.update_sar()
        self.process_sar_queue()
        self.process_mouse_queue()

    def load_images_and_generate_test_data(self):
        """Loads images and generates test data in a separate thread."""
        self.set_status("Loading images...")
        self.mfd_image_data = self.load_image(MFD_IMAGE_PATH, expected_dtype=np.uint8)
        self.sar_image_data = self.load_image(SAR_IMAGE_PATH, expected_dtype=SAR_DATA_TYPE)
        self.generate_test_images()  # Generate initial test images
        self.set_initial_sar_image() # Set the initial SAR image
        self.set_status("Ready")  # Back to ready status

    def set_initial_sar_image(self):
        """Sets the initial SAR image for display."""
        self.current_sar = cv2.normalize(self.sar_image_data, None, 0, 255, cv2.NORM_MINMAX,
                                          cv2.CV_8U)  # Current SAR for display
        self.adjusted_sar = self.current_sar.copy()  # Create a copy for contrast/brightness adjustments
        self.resized_sar = cv2.resize(self.adjusted_sar, (self.sar_display_width, self.sar_display_height),
                                       interpolation=cv2.INTER_AREA)

    def generate_test_images(self):
        """Generates random test images for MFD and SAR."""
        self.test_mfd_image = np.random.randint(0, 256, size=(MFD_HEIGHT, MFD_WIDTH), dtype=np.uint8)
        self.test_sar_image = np.random.randint(0, 65536, size=(SAR_HEIGHT, SAR_WIDTH), dtype=SAR_DATA_TYPE).astype(
            SAR_DATA_TYPE)
        self.resized_sar = cv2.resize(cv2.normalize(self.test_sar_image, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U), (self.sar_display_width, self.sar_display_height),
                                       interpolation=cv2.INTER_AREA)

    def load_image(self, path, expected_dtype):
        """Loads an image from the given path and ensures the correct data type."""
        if not os.path.exists(path):
            print(f"Error: Image file not found at {path}")
            logging.error(f"Image file not found at {path}")
            # Create a placeholder image instead
            if expected_dtype == np.uint8:
                return np.zeros((MFD_HEIGHT, MFD_WIDTH), dtype=np.uint8)
            else:
                return np.zeros((SAR_HEIGHT, SAR_WIDTH), dtype=SAR_DATA_TYPE)

        img = cv2.imread(path, cv2.IMREAD_ANYDEPTH)  # Load as is

        if img is None:
            print(f"Error: Could not load image at {path}")
            logging.error(f"Could not load image at {path}")
            # Create a placeholder image instead
            if expected_dtype == np.uint8:
                return np.zeros((MFD_HEIGHT, MFD_WIDTH), dtype=np.uint8)
            else:
                return np.zeros((SAR_HEIGHT, SAR_WIDTH), dtype=SAR_DATA_TYPE)

        if img.dtype != expected_dtype:
            print(f"Warning: Converting image at {path} from {img.dtype} to {expected_dtype}")
            logging.warning(f"Converting image at {path} from {img.dtype} to {expected_dtype}")
            img = img.astype(expected_dtype)
        return img

    def update_mfd(self):
        """Updates the MFD image with translation and calculates FPS, limiting to MFD_FPS."""
        try:
            start_time = time.time()
            if self.test_image_var.get() == 1:
                # Use test image and simulate translation
                self.mfd_x_offset = (self.mfd_x_offset + 1) % MFD_WIDTH
                translated_image = np.roll(self.test_mfd_image, -self.mfd_x_offset, axis=1)
                delay = max(1, int(1000 / MFD_FPS))  # Target delay for MFD_FPS
            else:
                # Use loaded image
                translated_image = self.mfd_image_data
                delay = 1  # Update as fast as possible when not in test mode

            # Create and move window only once
            if not self.mfd_window_initialized:
                cv2.imshow("MFD", translated_image)
                try:
                    cv2.moveWindow("MFD", self.mfd_x, self.mfd_y)
                    self.mfd_window_initialized = True
                except cv2.error as e:
                    print(f"Error moving MFD window on initialization: {e}")
                    logging.error(f"Error moving MFD window on initialization: {e}")
            else:
                cv2.imshow("MFD", translated_image)  # Just update the image

            # Calculate FPS
            self.mfd_frame_count += 1
            elapsed_time = time.time() - self.mfd_start_time
            if elapsed_time >= 1.0:  # Update FPS every second
                self.mfd_fps = self.mfd_frame_count / elapsed_time
                self.mfd_start_time = time.time()
                self.mfd_frame_count = 0
                self.update_status()  # Update status bar

            self.root.after(delay, self.update_mfd)  # Schedule the next update

        except Exception as e:
            print(f"Error updating MFD: {e}")
            logging.exception(f"Error updating MFD: {e}")

    def update_sar(self):
        """Updates the SAR image with translation."""
        try:
            if self.test_image_var.get() == 1:
                # Use test image and simulate translation
                self.sar_x_offset = (self.sar_x_offset + 1) % self.sar_display_width
                translated_image = np.roll(self.resized_sar, -self.sar_x_offset, axis=1)  # shift columns
            else:
                # Use loaded image
                translated_image = self.resized_sar

            # Create and move window only once
            if not self.sar_window_initialized:
                cv2.imshow("SAR", translated_image)
                try:
                    cv2.moveWindow("SAR", self.sar_x, self.sar_y)
                    self.sar_window_initialized = True
                except cv2.error as e:
                    print(f"Error moving SAR window on initialization: {e}")
                    logging.error(f"Error moving SAR window on initialization: {e}")

                # Set mouse callback *after* the window is created
                cv2.setMouseCallback("SAR", self.sar_mouse_callback)
                self.sar_mouse_callback_set = True

            else:
                cv2.imshow("SAR", translated_image)  # Just update the image

            # Calculate SAR FPS
            self.sar_frame_count += 1
            elapsed_time = time.time() - self.sar_update_time
            if elapsed_time >= 1.0:
                self.sar_fps = self.sar_frame_count / elapsed_time
                self.sar_update_time = time.time()
                self.sar_frame_count = 0
                self.update_status()  # Update status bar

            self.root.after(5000, self.update_sar)

        except Exception as e:
            print(f"Error updating SAR: {e}")
            logging.exception(f"Error updating SAR: {e}")

    def update_sar_size(self, event=None):
        """Updates the SAR image size based on the combobox selection."""
        selected_size = self.sar_size_combo.get()
        if selected_size == "1:1":
            self.sar_display_width = SAR_WIDTH
            self.sar_display_height = SAR_HEIGHT
        elif selected_size == "1:2":
            self.sar_display_width = INITIAL_SAR_WIDTH
            self.sar_display_height = INITIAL_SAR_HEIGHT
        elif selected_size == "1:3":
            self.sar_display_width = SAR_WIDTH // 3
            self.sar_display_height = SAR_HEIGHT // 3
        elif selected_size == "1:5":
            self.sar_display_width = SAR_WIDTH // 5
            self.sar_display_height = SAR_HEIGHT // 5
        elif selected_size == "1:10":
            self.sar_display_width = SAR_WIDTH // 10
            self.sar_display_height = SAR_HEIGHT // 10

        if self.test_image_var.get() == 1:
            self.generate_test_images() #Regenerate the test images
        else:
            self.set_initial_sar_image()

        # Resize the adjusted image
        self.resized_sar = cv2.resize(self.adjusted_sar, (self.sar_display_width, self.sar_display_height),
                                       interpolation=cv2.INTER_AREA)
        print(f"New SAR size: {self.sar_display_width}x{self.sar_display_height}")
        self.update_sar()  # Update display

        # Update SAR window position in case the size changed
        self.sar_x = self.tkinter_x + TKINTER_MIN_WIDTH + 20  # Place SAR window to the right of Tkinter
        self.sar_y = 10

    def update_contrast(self, value):
        """Updates the SAR image contrast based on the slider value."""
        try:
            self.sar_contrast = float(value)
            self.update_brightness_contrast_lut()
            self.update_sar_display()
        except ValueError:
            pass

    def update_brightness(self, value):
        """Updates the SAR image brightness based on the slider value."""
        try:
            self.sar_brightness = int(float(value))
            self.update_brightness_contrast_lut()
            self.update_sar_display()
        except ValueError:
            pass

    def update_sar_palette(self, event=None):
        """Updates the SAR image palette based on the combobox selection."""
        self.sar_palette = self.palette_combo.get()
        self.update_sar_display()

    def update_brightness_contrast_lut(self):
        """Calculates the Look-Up Table (LUT) for brightness and contrast adjustment."""
        self.brightness_contrast_lut = np.array([
            np.clip(i * self.sar_contrast + self.sar_brightness, 0, 255).astype(np.uint8) for i in range(256)
        ])

    def update_sar_display(self):
        """Applies contrast, brightness, and palette adjustments to the SAR image and triggers SAR update."""

        # Apply contrast and brightness using LUT
        if self.test_image_var.get() == 1:
            adjusted_image = cv2.LUT(cv2.normalize(self.test_sar_image, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U),
                                      self.brightness_contrast_lut)
        else:
            adjusted_image = cv2.LUT(self.current_sar, self.brightness_contrast_lut)

        # Apply color Palette
        if self.sar_palette != "GRAY":
            adjusted_image = self.apply_color_palette(adjusted_image, self.sar_palette)

        self.adjusted_sar = adjusted_image.copy()  # Store the adjusted image
        # Resize the adjusted image
        self.resized_sar = cv2.resize(self.adjusted_sar, (self.sar_display_width, self.sar_display_height),
                                       interpolation=cv2.INTER_AREA)
        # Force image display update
        self.put_sar_queue(self.resized_sar)

    def apply_color_palette(self, image, palette):
        """Applies the given color palette to the image"""

        try:
            colormap = getattr(cv2, f"COLORMAP_{palette.upper()}")
            colorized_image = cv2.applyColorMap(image, colormap)
            return colorized_image
        except AttributeError:
            print(f"Error: Colormap '{palette}' not found in OpenCV.")
            logging.error(f"Colormap '{palette}' not found in OpenCV.")
            return image  # Return the original image if the colormap is invalid
        except Exception as e:
            print(f"Error applying colormap: {e}")
            logging.exception(f"Error applying colormap: {e}")
            return image

    def set_sar_info(self, message):
        """Updates the SAR info label in Tkinter."""
        self.sar_info_label.config(text=message)

    def update_status(self, ):
        """Updates the status bar with MFD and SAR FPS (and seconds per image for SAR)."""
        status_text = f"Ready | MFD FPS: {self.mfd_fps:.2f} | SAR FPS: {self.sar_fps:.2f}"

        # Add seconds per image for SAR
        if self.sar_fps > 0:
            seconds_per_image = 1 / self.sar_fps
            status_text += f" ({seconds_per_image:.2f} s/img)"
        else:
            status_text += " (N/A s/img)"

        self.statusbar.config(text=status_text)

    def set_status(self, message):
        """Updates the status bar."""
        self.statusbar.config(text=message)

    def set_new_sar_image(self, image):
        """Function to set a new Sar image"""
        self.sar_image_data = image
        self.current_sar = cv2.normalize(self.sar_image_data, None, 0, 255, cv2.NORM_MINMAX, cv2.CV_8U)
        self.adjusted_sar = self.current_sar.copy()  # Update the adjusted image too
        self.resized_sar = cv2.resize(self.adjusted_sar, (self.sar_display_width, self.sar_display_height),
                                       interpolation=cv2.INTER_AREA)
        # Force update to image
        self.put_sar_queue(self.resized_sar)

    def sar_mouse_callback(self, event, x, y, flags, param):
        """Callback function for mouse events on the SAR window."""
        if event == cv2.EVENT_MOUSEMOVE:
            # Convert pixel coordinates to normalized coordinates (0.0 - 1.0)
            normalized_x = x / self.sar_display_width
            normalized_y = y / self.sar_display_height

            # Convert normalized coordinates to kilometers from the center
            km_x = (normalized_x - 0.5) * SAR_IMAGE_SIZE_KM
            km_y = (0.5 - normalized_y) * SAR_IMAGE_SIZE_KM  # Y axis is inverted

            # Convert kilometers to latitude and longitude offset
            lat_offset = km_y / 111.0  # Approximate km to latitude conversion
            lon_offset = km_x / (111.0 * np.cos(np.radians(SAR_CENTER_LAT)))  # Approximate km to longitude

            # Calculate latitude and longitude
            latitude = SAR_CENTER_LAT + lat_offset
            longitude = SAR_CENTER_LON + lon_offset

            # Put the coordinates in the queue
            try:
                self.mouse_queue.put((latitude, longitude), block=False)
            except queue.Full:
                pass  # Drop value

    def update_mouse_latlon_label(self, latitude, longitude):
        """Updates the Tkinter label with the mouse latitude and longitude."""
        self.mouse_latlon_label.config(text=f"Mouse: Lat={latitude:.4f}, Lon={longitude:.4f}")

    def put_sar_queue(self, image):
        """Function to put the image on the queue"""
        try:
            self.sar_queue.put(image, block=False)
        except queue.Full:
            pass

    def process_sar_queue(self):
        """This function get the new Sar image from queue and displays"""
        try:
            image = self.sar_queue.get(block=False)
            cv2.imshow("SAR", image)
        except queue.Empty:
            pass
        self.root.after(1, self.process_sar_queue)  # Loop in the main Thread

    def process_mouse_queue(self):
        """Processes the mouse coordinate queue and updates the Tkinter label."""
        try:
            latitude, longitude = self.mouse_queue.get(block=False)
            # Limit the update rate to 10 times per second (100 ms interval)
            current_time = time.time()
            if current_time - self.last_mouse_update >= 0.1:
                self.update_mouse_latlon_label(latitude, longitude)
                self.last_mouse_update = current_time
        except queue.Empty:
            pass
        self.root.after(1, self.process_mouse_queue)

    def update_image_mode(self):
        """Updates the image mode based on the checkbox value and reload images if needed."""
        if self.test_image_var.get() == 1:
            self.generate_test_images()  # Generate test images when test mode is enabled
        else:
            self.set_initial_sar_image() # Reset the SAR image
        self.update_mfd()
        self.update_sar()

if __name__ == "__main__":
    root = tk.Tk()
    app = App(root)
    root.mainloop()
    cv2.destroyAllWindows()