from PyQt6.QtWidgets import (
    QTabWidget,
    QWidget,
    QHBoxLayout,
    QMainWindow,
    QMessageBox,
    QVBoxLayout,
    QPushButton,
    QDialog,
    QGroupBox,
    QLabel,
    QDateEdit,
    QComboBox,
    QLineEdit,
    QFormLayout,
    QDateTimeEdit,
    QProgressBar,
    QScrollArea,
)
from functions import create_setting_field
from PyQt6.QtGui import QAction, QPixmap
from PyQt6.QtCore import Qt, pyqtSignal, QDate, QDateTime, QSize
from PyQt6 import QtCore
import matplotlib
from matplotlib import colormaps, colors
import matplotlib.cm as cm
import numpy as np

matplotlib.use("QtAgg")
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
from matplotlib.backends.backend_qt5agg import NavigationToolbar2QT as NavigationToolbar
from matplotlib.figure import Figure

from matplotlib.dates import DateFormatter
from matplotlib.colorbar import ColorbarBase
from matplotlib.colors import Normalize

from datetime import datetime, timedelta


class MainWindow(QMainWindow):
    showSettingsRequested = pyqtSignal()
    showHelpRequested = pyqtSignal()
    showSaveRequested = pyqtSignal()
    showLoadRequested = pyqtSignal()

    def __init__(self):
        super(MainWindow, self).__init__()
        self._setupUI()

    def _setupUI(self):
        self.setWindowTitle("High Altitude Balloon Calculation")
        # self.setStyleSheet(
        #     """
        #     QMainWindow {
        #         background-color: #F0F0F0;
        #     }
        #     QMenuBar {
        #         background-color: #E1E1E1;
        #     }
        #     QMenuBar::item {
        #         spacing: 3px;
        #         padding: 2px 10px;
        #         background-color: #E1E1E1;
        #         border-radius: 4px;
        #         color: #333333;
        #     }
        #     QMenuBar::item:selected {
        #         background-color: #DDEEFF;
        #     }
        #     QPushButton {
        #         background-color: #DDEEFF;
        #         border: 1px solid #BBDDEE;
        #         border-radius: 5px;
        #         padding: 5px;
        #         color: #333333;
        #     }
        #     QPushButton:hover {
        #         background-color: #BBDDEE;
        #     }
        #     QPushButton:pressed {
        #         background-color: #AACCEE;
        #     }
        #     QGroupBox {
        #         border: 2px solid #CCCCCC;
        #         border-radius: 5px;
        #         margin-top: 20px;
        #     }
        #     QGroupBox::title {
        #         subcontrol-origin: margin;
        #         subcontrol-position: top center;
        #         padding: 2px 15px 2px 15px;
        #         background-color: #E6E6E6;
        #         color: #333333;
        #         border-radius: 10px;
        #     }
        # """
        # )
        self._initialize_menu_bar()
        self._initialize_tabs()

    def _initialize_menu_bar(self):
        menu_bar = self.menuBar()

        file_menu = menu_bar.addMenu("&File")

        # Save
        save_action = QAction("Save", self)
        save_action.setObjectName("saveAction")
        save_action.triggered.connect(self.showSaveRequested.emit)
        file_menu.addAction(save_action)

        # Load
        load_action = QAction("Load", self)
        load_action.setObjectName("loadAction")
        load_action.triggered.connect(self.showLoadRequested.emit)
        file_menu.addAction(load_action)

        # Settings
        settings_action = QAction("Settings", self)
        settings_action.triggered.connect(self.showSettingsRequested.emit)
        menu_bar.addAction(settings_action)

        # Help
        help_action = QAction("Help", self)
        help_action.setObjectName("helpAction")
        help_action.triggered.connect(self.showHelpRequested.emit)
        menu_bar.addAction(help_action)

    def _initialize_tabs(self):
        self.tab_widget = QTabWidget()
        self.setCentralWidget(self.tab_widget)

    def add_tab(self, widget, title):
        self.tab_widget.addTab(widget, title)

    def show_maximized(self):
        self.showMaximized()


class OptimizeTabView(QWidget):
    calculateOptimizeRequested = pyqtSignal()
    calculationCompleted = pyqtSignal()
    calculationFailed = pyqtSignal(str)
    dropdownSelectionChanged = pyqtSignal(str)
    saveBalloonRequested = pyqtSignal()

    # View for each calculation tab
    def __init__(self):
        super().__init__()

        self.input_widgets = {}

        self.main_layout = QHBoxLayout(self)
        self.input_layout = QVBoxLayout()
        self.output_layout = QVBoxLayout()

        self.input_container = QWidget()  # Container for input_layout
        self.output_container = QWidget()  # Container for output_layout
        self.output_container.setEnabled(False)
        self.input_container.setLayout(self.input_layout)
        self.output_container.setLayout(self.output_layout)

        ## Input Layout Groups
        self.mission_group = QGroupBox("Mission")
        self.balloon_group = QGroupBox("Balloon")
        self.payload_group = QGroupBox("Payload")
        self.solar_group = QGroupBox("Solar Panels")
        self.battery_group = QGroupBox("Battery")
        self.other_group = QGroupBox("Others")

        ## Output Layout Groups
        self.general_output_layout = QFormLayout()
        self.balloon1_output_layout = QFormLayout()
        self.balloon2_output_layout = QFormLayout()
        self.solar_panel_output_layout = QFormLayout()
        self.battery_output_layout = QFormLayout()
        self.other_output_layout = QFormLayout()
        self.message_output_layout = QLabel()

        # Plot and dropdown
        self.dropdown_plot = QComboBox()
        self._plot_dropdown_setup()
        self.canvas = PlotCanvasOptimize(self, width=100, height=100, dpi=100)

        self.other_outputs_layout = QFormLayout()

        ## Calculate Button
        self.calculate_button = QPushButton("Calculate")
        self.calculate_button.clicked.connect(self.calculateOptimizeRequested.emit)

        ## Save Balloon Button
        self.save_balloon_button = QPushButton("Save Balloon")
        self.save_balloon_button.clicked.connect(self.saveBalloonRequested.emit)

        # Image label
        self.image_label = ImageLabel()

        # Adding input UI to input_layout and output UI to output_layout
        self._create_input_UI()
        self._create_output_UI()

        self.main_layout.addWidget(self.input_container)
        self.main_layout.addWidget(self.output_container)

        self.main_layout.setStretch(0, 1)  # Input layout takes up 1/3 of space
        self.main_layout.setStretch(1, 2)  # Output layout takes up 2/3 of space
        # Setup tab UI (inputs, button, output display)

    def _plot_dropdown_setup(self):
        self.dropdown_plot.addItems(
            [
                "Balloon 1",
                "Balloon 2",
                "Solar Panel Power Curve",
                "Energy Balance over 3 Days",
                "Volumes over Altitude",
                "Gas Weights over Altitude",
                "Balloon 1 Stress",
                "Balloon 2 Stress",
                "Balloon 1 Optimization Shapes",
                "Balloon 2 Optimization Shapes",
                "Optimization Balance",
            ]
        )
        self.dropdown_plot.currentTextChanged.connect(
            self.dropdownSelectionChanged.emit
        )

    def populate_input_group(self, group, definitions):
        """
        Creates and adds widgets for a whole group of inputs.

        Parameters:
        - group: An input group parent layout
        - definitions: A dictionary of input group definitions.
        """
        layout = QVBoxLayout()
        widgets = {}
        for key, definition in definitions.items():
            label_text, description = definition[0], definition[1]
            options = definition[2] if len(definition) > 2 else None
            input_type = definition[3] if len(definition) > 3 else None

            label = QLabel(label_text)
            layout.addWidget(label)

            if input_type == "date":
                # Handle creation of a QDateEdit for date inputs
                date_input = QDateEdit()
                date_input.setCalendarPopup(True)
                date_input.setDate(QDate.currentDate())
                layout.addWidget(date_input)
                widgets[key] = date_input
            elif options:
                # Handle creation of a QComboBox for settings with options
                combo_box = QComboBox()
                for option in options:
                    combo_box.addItem(option)
                layout.addWidget(combo_box)
                widgets[key] = combo_box
            else:
                # Default to QLineEdit for standard text inputs
                input_field = QLineEdit()
                input_field.setPlaceholderText(description)
                layout.addWidget(input_field)
                widgets[key] = input_field

        group.setLayout(layout)
        self.input_widgets.update(widgets)

    def _create_input_UI(self):

        # Horizontal grouping mission and balloon groups
        mission_payload_layout = QHBoxLayout()
        mission_payload_layout.addWidget(self.mission_group)
        mission_payload_layout.addWidget(self.balloon_group)

        # Layout equal stretch factor
        mission_payload_layout.setStretch(0, 1)
        mission_payload_layout.setStretch(1, 1)

        # Horizontal grouping Solar Panels, Battery
        solar_battery_layout = QHBoxLayout()
        solar_battery_layout.addWidget(self.solar_group)
        solar_battery_layout.addWidget(self.battery_group)

        # Layout equal stretch factor
        solar_battery_layout.setStretch(0, 1)
        solar_battery_layout.setStretch(1, 1)

        # Horizontal grouping Payload, Others

        payload_others_layout = QHBoxLayout()
        payload_others_layout.addWidget(self.payload_group)
        payload_others_layout.addWidget(self.other_group)

        # Layout equal stretch factor
        payload_others_layout.setStretch(0, 1)
        payload_others_layout.setStretch(1, 1)

        self.input_layout.addLayout(mission_payload_layout)
        self.input_layout.addLayout(solar_battery_layout)
        self.input_layout.addLayout(payload_others_layout)
        self.input_layout.addWidget(self.calculate_button)

    def reset_inputs(self):
        for key, widget in self.input_widgets.items():
            if isinstance(widget, QLineEdit):
                widget.clear()
            elif isinstance(widget, QComboBox):
                widget.setCurrentIndex(0)
            elif isinstance(widget, QDateEdit):
                widget.setDate(QDate.currentDate())

    def update_input_widget(self, key, value):
        if key in self.input_widgets:
            widget = self.input_widgets[key]
            if isinstance(widget, QLineEdit):
                widget.setText(str(value))
            elif isinstance(widget, QComboBox):
                index = widget.findText(str(value), QtCore.Qt.MatchFixedString)
                if index >= 0:
                    widget.setCurrentIndex(index)
            elif isinstance(widget, QDateEdit):
                if value:
                    date = QDate.fromString(value, "yyyy-MM-dd")
                    widget.setDate(date)

    def _create_output_UI(self):

        # Group boxes for various system components
        system_layout_group = QGroupBox("General Output:")
        balloon_1_layout_group = QGroupBox("Balloon 1 Output:")
        balloon_2_layout_group = QGroupBox("Balloon 2 Output:")
        solar_panel_layout_group = QGroupBox("Solar Panel Output:")
        battery_layout_group = QGroupBox("Battery Output:")
        plot_layout_group = QGroupBox("Plots:")
        other_parameters_group = QGroupBox("Other Parameters:")
        summary_group = QGroupBox("Summary:")
        # Add simple layout to summary group to show message
        summary_group_layout = QVBoxLayout()
        summary_group_layout.addWidget(self.message_output_layout)
        summary_group.setLayout(summary_group_layout)

        # Layouts for balloon plots
        plot_layout = QVBoxLayout()
        toolbar = NavigationToolbar(self.canvas, self)

        # Setting layouts for balloon output groups
        system_layout_group.setLayout(self.general_output_layout)
        balloon_1_layout_group.setLayout(self.balloon1_output_layout)
        balloon_2_layout_group.setLayout(self.balloon2_output_layout)
        solar_panel_layout_group.setLayout(self.solar_panel_output_layout)
        battery_layout_group.setLayout(self.battery_output_layout)
        other_parameters_group.setLayout(self.other_output_layout)

        plot_layout_group.setLayout(plot_layout)
        plot_layout.addWidget(toolbar)
        plot_layout.addWidget(self.canvas)
        plot_layout.addWidget(self.dropdown_plot)

        # Left vertical layout containing various group boxes
        left_vertical_layout = QVBoxLayout()
        left_vertical_layout.addWidget(system_layout_group)
        left_vertical_layout.addWidget(balloon_1_layout_group)
        left_vertical_layout.addWidget(balloon_2_layout_group)
        left_vertical_layout.addWidget(solar_panel_layout_group)
        left_vertical_layout.addWidget(battery_layout_group)
        left_vertical_layout.addWidget(other_parameters_group)
        left_vertical_layout.addWidget(self.save_balloon_button)

        # Bottom layout for summary and image
        bottom_layout = QHBoxLayout()
        bottom_layout.addWidget(summary_group)
        bottom_layout.addWidget(self.image_label)
        bottom_layout.setStretch(0, 3)
        bottom_layout.setStretch(1, 1)

        # Main horizontal layout to include the left layout and plots
        horizontal_layout = QHBoxLayout()
        horizontal_layout.addLayout(left_vertical_layout)
        horizontal_layout.addWidget(plot_layout_group)
        horizontal_layout.setStretch(0, 1)
        horizontal_layout.setStretch(1, 3)

        self.output_layout.addLayout(horizontal_layout)
        self.output_layout.addLayout(bottom_layout)

        self.output_layout.setStretch(0, 5)
        self.output_layout.setStretch(1, 1)

    def fill_message_output(self, message):
        self.message_output_layout.setText(message)

    def clear_form_layout(self, form_layout):
        """Removes all rows from a QFormLayout."""
        while form_layout.rowCount():
            form_layout.removeRow(0)

    def set_balloon_image(self, image_path):
        pixmap = QPixmap(image_path)
        self.image_label.setPixmap(pixmap)

    def populate_output_form(self, group: QFormLayout, parameters: dict, names: dict):
        """
        Populates a QFormLayout with labels and values based on parameter names and values.

        Args:
            group (QFormLayout): The form layout to populate.
            parameters (dict): Dictionary containing parameter values.
            names (dict): Dictionary containing human-readable names for the parameters.
        """

        # Clear the layout first
        self.clear_form_layout(group)
        # Iterate through the parameter names dictionary
        for parameter, name in names.items():
            # Check if the parameter exists in the parameters dictionary and get its value
            value = parameters.get(
                parameter, "N/A"
            )  # Use 'N/A' or similar if value is None or missing

            # Format the value if it's a float, allowing up to 5 decimal places
            if isinstance(value, (float, int)):  # Check if value is a number
                value_str = f"{value:.5f}"  # Format with five decimal places
            else:
                value_str = str(
                    value
                )  # Use the string representation for non-numeric values

            # Add a row to the form layout with the friendly name and the corresponding value
            group.addRow(QLabel(name), QLabel(value_str))


class FlyTabView(QWidget):
    loadBalloonRequested = pyqtSignal()
    loadControllerRequested = pyqtSignal()
    calculateFlyRequested = pyqtSignal()
    calculationCompleted = pyqtSignal()
    calculationFailed = pyqtSignal(str)
    dropdownSelectionChanged = pyqtSignal(str)

    # View for each calculation tab
    def __init__(self):
        super().__init__()

        self.input_widgets = {}

        self.main_layout = QHBoxLayout(self)
        self.input_layout = QVBoxLayout()
        self.output_layout = QVBoxLayout()

        self.input_container = QWidget()  # Container for input_layout
        self.output_container = QWidget()  # Container for output_layout
        self.output_container.setEnabled(False)
        self.input_container.setLayout(self.input_layout)
        self.output_container.setLayout(self.output_layout)

        ## Input Layout Groups
        self.flight_group = QGroupBox("Flight Specification")
        self.balloon_group = QGroupBox("Balloon Specification")

        self.balloon_layout = QFormLayout()

        ## Output Layout Groups
        self.general_output_layout = QFormLayout()

        # Plot and dropdown
        self.dropdown_plot = QComboBox()
        self._plot_dropdown_setup()
        self.canvas = PlotCanvasFly(self, width=100, height=100, dpi=100)

        ## Load Button
        self.load_button = QPushButton("Load Balloon")
        self.load_button.clicked.connect(self.loadBalloonRequested.emit)

        ## Add Control Button
        self.controller_button = QPushButton("Load Controller")
        self.controller_button.clicked.connect(self.loadControllerRequested.emit)

        ## Calculate Button
        self.calculate_button = QPushButton("Calculate")
        self.calculate_button.clicked.connect(self.calculateFlyRequested.emit)

        # Adding input UI to input_layout and output UI to output_layout
        self._create_input_UI()
        self._create_output_UI()

        self.main_layout.addWidget(self.input_container)
        self.main_layout.addWidget(self.output_container)

        self.main_layout.setStretch(0, 1)  # Input layout takes up 1/3 of space
        self.main_layout.setStretch(1, 3)  # Output layout takes up 2/3 of space
        # Setup tab UI (inputs, button, output display)

    def _plot_dropdown_setup(self):
        self.dropdown_plot.addItems(
            [
                "3d Trajectory",
                "Altitude vs. Time",
                "Vertical Speed vs. Time",
                "Position vs. Time",
            ]
        )
        self.dropdown_plot.currentTextChanged.connect(
            self.dropdownSelectionChanged.emit
        )

    def populate_input_group(self, group, definitions):
        """
        Creates and adds widgets for a whole group of inputs.

        Parameters:
        - group: An input group parent layout
        - definitions: A dictionary of input group definitions.
        """
        layout = QVBoxLayout()
        widgets = {}
        for key, definition in definitions.items():
            label_text, description = definition[0], definition[1]
            options = definition[2] if len(definition) > 2 else None
            input_type = definition[3] if len(definition) > 3 else None

            label = QLabel(label_text)
            layout.addWidget(label)

            if input_type == "date":
                # Handle creation of a QDateEdit for date inputs
                date_input = QDateTimeEdit()
                date_input.setCalendarPopup(True)
                date_input.setDateTime(QDateTime.currentDateTime())
                date_input.setMinimumDateTime(QDateTime.currentDateTime())
                date_input.setMaximumDateTime(QDateTime.currentDateTime().addDays(7))
                layout.addWidget(date_input)
                widgets[key] = date_input
            elif options:
                # Handle creation of a QComboBox for settings with options
                combo_box = QComboBox()
                for option in options:
                    combo_box.addItem(option)
                layout.addWidget(combo_box)
                widgets[key] = combo_box
            else:
                # Default to QLineEdit for standard text inputs
                input_field = QLineEdit()
                input_field.setPlaceholderText(description)
                layout.addWidget(input_field)
                widgets[key] = input_field

        group.setLayout(layout)
        self.input_widgets.update(widgets)

    def _create_input_UI(self):
        self.balloon_group.setLayout(self.balloon_layout)

        self.input_layout.addWidget(self.flight_group)
        self.input_layout.addWidget(self.load_button)
        self.input_layout.addWidget(self.controller_button)
        self.input_layout.addWidget(self.balloon_group)
        self.input_layout.addWidget(self.calculate_button)

    def reset_inputs(self):
        for key, widget in self.input_widgets.items():
            if isinstance(widget, QLineEdit):
                widget.clear()
            elif isinstance(widget, QComboBox):
                widget.setCurrentIndex(0)
            elif isinstance(widget, QDateEdit):
                widget.setDate(QDate.currentDate())

    def update_input_widget(self, key, value):
        if key in self.input_widgets:
            widget = self.input_widgets[key]
            if isinstance(widget, QLineEdit):
                widget.setText(str(value))
            elif isinstance(widget, QComboBox):
                index = widget.findText(str(value), QtCore.Qt.MatchFixedString)
                if index >= 0:
                    widget.setCurrentIndex(index)
            elif isinstance(widget, QDateEdit):
                if value:
                    date = QDate.fromString(value, "yyyy-MM-dd")
                    widget.setDate(date)

    def _create_output_UI(self):

        plot_layout = QVBoxLayout()
        toolbar = NavigationToolbar(self.canvas, self)

        # Setting layouts for fly output groups
        plot_layout.addWidget(toolbar)
        plot_layout.addWidget(self.canvas)
        plot_layout.addWidget(self.dropdown_plot)

        self.output_layout.addLayout(plot_layout)

    def clear_form_layout(self, form_layout):
        """Removes all rows from a QFormLayout."""
        while form_layout.rowCount():
            form_layout.removeRow(0)

    def populate_output_form(self, group: QFormLayout, parameters: dict, names: dict):
        """
        Populates a QFormLayout with labels and values based on parameter names and values.

        Args:
            group (QFormLayout): The form layout to populate.
            parameters (dict): Dictionary containing parameter values.
            names (dict): Dictionary containing human-readable names for the parameters.
        """

        # Clear the layout first
        self.clear_form_layout(group)
        # Iterate through the parameter names dictionary
        for parameter, name in names.items():
            # Check if the parameter exists in the parameters dictionary and get its value
            value = parameters.get(
                parameter, "N/A"
            )  # Use 'N/A' or similar if value is None or missing

            # Format the value if it's a float, allowing up to 5 decimal places
            if isinstance(value, (float, int)):  # Check if value is a number
                value_str = f"{value:.5f}"  # Format with five decimal places
            else:
                value_str = str(
                    value
                )  # Use the string representation for non-numeric values

            # Add a row to the form layout with the friendly name and the corresponding value
            group.addRow(QLabel(name), QLabel(value_str))


class SettingsWindow(QDialog):
    saveSettingsRequested = pyqtSignal()

    def __init__(self):
        super().__init__()
        self.setWindowTitle("Settings")
        self.setGeometry(100, 100, 500, 400)
        self.layout = QVBoxLayout(self)

        ## Settings Groups
        self.optimization_ZP_group = QGroupBox("Zero Pressure Balloon")
        self.optimization_SP_group = QGroupBox("Super Pressure Balloon")
        self.optimization_others_group = QGroupBox("Other Parameters")

        # Create layouts for each group
        self.optimization_ZP_layout = QVBoxLayout()
        self.optimization_SP_layout = QVBoxLayout()
        self.optimization_others_layout = QVBoxLayout()

        self.optimization_ZP_group.setLayout(self.optimization_ZP_layout)
        self.optimization_SP_group.setLayout(self.optimization_SP_layout)
        self.optimization_others_group.setLayout(self.optimization_others_layout)

        # Add group boxes to the main layout
        self.layout.addWidget(self.optimization_ZP_group)
        self.layout.addWidget(self.optimization_SP_group)
        self.layout.addWidget(self.optimization_others_group)

        self.input_widgets = {}
        # The save button
        self.save_button = QPushButton("Save")
        self.layout.addWidget(self.save_button)
        self.save_button.clicked.connect(self.saveSettingsRequested.emit)

    def add_setting(self, group_layout, label_text, key, value):
        label, input_field = create_setting_field(self, label_text, value)
        self.input_widgets[key] = input_field
        group_layout.addWidget(label)
        group_layout.addWidget(input_field)

    def populate_settings_group(self, group_layout, settings):
        for key, (label_text, value) in settings.items():
            self.add_setting(group_layout, label_text, key, value)

    def display_validation_error(self, key, message):
        if key in self.input_widgets:
            self.input_widgets[key].setStyleSheet(
                "border: 2px solid red;"
            )  # Highlight field in red
            self.input_widgets[key].setToolTip(message)  # Show error message as tooltip
        else:
            QMessageBox.warning(self, "Invalid Input", message)  # General error message


class HelpWindow(QDialog):
    def __init__(self, html_content):
        super(HelpWindow, self).__init__()

        self.setWindowTitle("Help")
        self.setGeometry(100, 100, 500, 400)

        # Set the dialog layout
        layout = QVBoxLayout(self)

        # Create a QScrollArea
        scroll_area = QScrollArea()
        scroll_area.setWidgetResizable(True)

        # Create a QLabel
        label = QLabel()

        # Ensure the QLabel interprets HTML content
        label.setTextFormat(Qt.TextFormat.RichText)
        label.setText(html_content)
        label.setWordWrap(True)  # Ensure text wraps within the QLabel

        # Set the QLabel as the widget of the QScrollArea
        scroll_area.setWidget(label)

        # Add the QScrollArea to the layout
        layout.addWidget(scroll_area)

        # Add a QPushButton to close the dialog
        close_button = QPushButton("Close")
        close_button.clicked.connect(self.close)
        layout.addWidget(close_button)

        self.setLayout(layout)


class PlotCanvasOptimize(FigureCanvas):
    def __init__(self, parent=None, width=5, height=4, dpi=100):
        self.fig = Figure(figsize=(width, height), dpi=dpi, constrained_layout=True)
        self.axes = None
        super(PlotCanvasOptimize, self).__init__(self.fig)

    def clear_plot(self):
        """Clear the plot and the color bar if it exists."""
        self.fig.clear()
        self.axes = None
        self.draw()

    def plot_balloon(self, radius, height, title_suffix):
        self.clear_plot()  # Clear existing data
        self.axes = self.fig.add_subplot(111)
        self.axes.axis("equal")
        self.axes.grid(True)
        self.axes.set_xlabel("Radius (m)")
        self.axes.set_ylabel("Height (m)")
        self.axes.set_title(f"{title_suffix} Shape")
        self.axes.plot(radius, height, "b")
        self.axes.plot([-x for x in radius], height, "b")
        self.draw()

    def plot_balloon_iterations(self, balloon_radii, balloon_heights, title_suffix):
        self.clear_plot()  # Clear existing data for new plot
        self.axes = self.fig.add_subplot(111)
        self.axes.axis("equal")
        self.axes.grid(True)
        self.axes.set_xlabel("Radius (m)")
        self.axes.set_ylabel("Height (m)")
        self.axes.set_title(f"{title_suffix} Shape Over Iterations")

        # Color map
        # color_map = cm.get_cmap("viridis", len(balloon_radii))
        # norm = Normalize(vmin=0, vmax=len(balloon_radii) - 1)
        color_map = colormaps.get_cmap("viridis")
        norm = colors.Normalize(vmin=0, vmax=len(balloon_radii) - 1)

        # Plotting balloons for each iteration
        for i in range(len(balloon_radii)):
            color = color_map(norm(i))  # Get color for current iteration
            radii = balloon_radii[i]
            heights = balloon_heights[i]
            self.axes.plot(radii, heights, color=color)
            self.axes.plot(
                [-r for r in radii], heights, color=color
            )  # Reflecting over y-axis

        # Create a colorbar
        sm = cm.ScalarMappable(cmap=color_map, norm=norm)
        # sm = colors.ScalarMappable(cmap=color_map, norm=norm)
        sm.set_array([])
        colorbar = self.figure.colorbar(sm, ax=self.axes, orientation="vertical")
        colorbar.set_label("Iteration")

        self.draw()

    def plot_weight_balance(self, weight_balances):
        self.clear_plot()
        self.axes = self.fig.add_subplot(111)
        self.axes.grid(True)
        self.axes.set_xlabel("Iteration")
        self.axes.set_ylabel("Weight Balance (kg)")
        self.axes.set_title("Weight Balance Over Iterations")

        iterations = np.arange(
            len(weight_balances)
        )  # Create an array of iteration indices
        self.axes.plot(
            iterations,
            weight_balances,
            marker="o",
            linestyle="-",
            color="b",
            label="Weight Balance",
        )
        self.axes.legend()
        self.draw()

    def plot_solar_power_curve(self, power_curve):
        self.clear_plot()
        self.axes = self.fig.add_subplot(111)
        self.axes.grid(True)

        num_points = len(power_curve)
        start_time = datetime.now().replace(hour=0, minute=0, second=0, microsecond=0)
        times = [start_time + timedelta(minutes=i) for i in range(num_points)]
        self.axes.plot(times, power_curve, color="orange")
        self.axes.set_xlabel("Time")
        self.axes.set_ylabel("Power (W)")
        self.axes.set_title("Solar Panel Power Output on the Least Sunny Day")

        # Customize the x-axis tick labels to display hours and minutes
        time_format = DateFormatter("%H:%M")
        self.axes.xaxis.set_major_formatter(time_format)

        self.draw()

    def plot_balloon_volumes(
        self, altitude_array, volume_array_balloon_1, volume_array_balloon_2
    ):
        self.clear_plot()
        self.axes = self.fig.add_subplot(111)
        self.axes.plot(altitude_array, volume_array_balloon_1, label="Balloon 1")
        self.axes.plot(
            altitude_array,
            volume_array_balloon_2,
            label="Balloon 2",
        )
        self.axes.set_xlabel("Altitude (m)")
        self.axes.set_ylabel("Volume (m3)")
        self.axes.legend()
        self.axes.grid(True)
        self.draw()

    def plot_balloon_gas_weights(
        self, altitude_array, gas_weight_array_balloon_1, gas_weight_array_balloon_2
    ):
        self.clear_plot()
        self.axes = self.fig.add_subplot(111)
        self.axes.plot(altitude_array, gas_weight_array_balloon_1, label="Balloon 1")
        self.axes.plot(
            altitude_array,
            gas_weight_array_balloon_2,
            label="Balloon 2",
        )
        self.axes.set_xlabel("Altitude (m)")
        self.axes.set_ylabel("Gas Weight (kg)")
        self.axes.legend()
        self.axes.grid(True)
        self.draw()

    def plot_energy_balance(
        self,
        charge_state,
        battery_capacity,
        min_discharge_level,
        power_balance,
        solar_power_array,
        payload_power_array,
        altitude_power_array,
        heating_auxiliary_power_array,
    ):
        self.clear_plot()

        # Creating two subplots, one on top of the other, sharing the same x-axis
        ax2 = self.fig.add_subplot(212)
        ax1 = self.fig.add_subplot(211, sharex=ax2)

        # Hide x-tick labels on ax1
        ax1.tick_params(axis="x", which="both", labelbottom=False)
        # Determine the active range of the solar power array
        active_start = 0
        active_end = len(solar_power_array) - 1
        # Set the x-axis limits
        ax1.set_xlim(active_start, active_end)
        ax2.set_xlim(active_start, active_end)
        # Plotting

        ax1.plot(charge_state, label="Battery State of Charge")
        ax1.axhline(y=battery_capacity, color="r", linestyle="--", label="Max Charge")
        ax1.axhline(
            y=min_discharge_level,
            color="g",
            linestyle="--",
            label="Min Discharge Level",
        )
        ax1.set_title("State of Charge")
        ax1.set_ylabel("Battery Charge (Wh)")
        ax1.legend()
        ax1.grid(True)

        # Power balance plot
        ax2.plot(power_balance, label="Net Power Balance", color="black")
        ax2.plot(solar_power_array, label="Solar Power", color="orange")
        ax2.plot(-payload_power_array, label="Payload Power Consumption", color="blue")
        ax2.plot(
            -altitude_power_array,
            label="Altitude Control Power Consumption",
            color="purple",
        )
        ax2.plot(
            -heating_auxiliary_power_array,
            label="Heating and Auxiliaries Power Consumption",
            color="red",
        )
        ax2.set_title("Power Balance and Consumption")
        ax2.set_xlabel("Time (HH:MM)")
        ax2.set_ylabel("Power (W)")
        ax2.legend()
        ax2.grid(True)

        # Create time labels for 24-hour format for the x-axis
        time_labels = [
            f"{(minute // 60) % 24:02d}:{minute % 60:02d}"
            for minute in range(len(solar_power_array))
        ]
        ticks_position = np.arange(active_start, active_end + 1, 120)

        ax2.set_xticks(ticks_position)
        ax2.set_xticklabels([time_labels[i] for i in ticks_position], rotation=45)

        self.draw()

    def plot_balloon_stress(
        self,
        skin_length,
        stress_meridional,
        stress_film,
        stress_tape,
        stress_total,
        title,
    ):
        self.clear_plot()
        axes = self.fig.add_subplot(111)

        axes.plot(skin_length, stress_meridional, label="Stress Meridional")

        axes.plot(
            skin_length,
            stress_film,
            linestyle="dotted",
            label="Max. Stress Film",
        )
        axes.plot(
            skin_length,
            stress_tape,
            linestyle="dotted",
            label="Max. Stress Tape",
        )
        axes.plot(skin_length, stress_total, label="Max. Stress Total")
        axes.set_xlabel("Skin length (m)")
        axes.set_ylabel("Meridional Stress (N/m)")
        axes.set_title(f"{title} Meridional Stress vs. Skin Length")
        if np.median(stress_total) > np.median(stress_meridional):
            axes.set_ylim(0, np.median(stress_total) * 1.2)
        else:
            axes.set_ylim(0, np.median(stress_meridional) * 1.2)

        axes.legend()

        self.draw()


class PlotCanvasFly(FigureCanvas):
    def __init__(self, parent=None, width=5, height=4, dpi=100):
        self.fig = Figure(figsize=(width, height), dpi=dpi, constrained_layout=True)
        self.axes = None
        super(PlotCanvasFly, self).__init__(self.fig)

    def clear_plot(self):
        """Clear the plot and the color bar if it exists."""
        self.fig.clear()
        self.axes = None
        self.draw()


class ProgressDialog(QDialog):
    # terminationRequested = pyqtSignal()

    def __init__(self, parent=None):
        super(ProgressDialog, self).__init__(parent)
        self.setWindowTitle("Optimization Progress")
        self.initUI()
        self.setFixedSize(300, 120)
        self.setWindowFlags(
            self.windowFlags() & ~Qt.WindowType.WindowCloseButtonHint
        )  # Disable the close button

    def initUI(self):
        layout = QVBoxLayout()
        self.label = QLabel("Optimization is running...")
        layout.addWidget(self.label)

        # self.terminateButton = QPushButton("Terminate")
        # self.terminateButton.clicked.connect(self.terminationRequested.emit)
        # layout.addWidget(self.terminateButton)

        self.setLayout(layout)

    def set_message(self, message):
        self.label.setText(message)


class ProgressFlyDialog(QDialog):
    canceled = pyqtSignal()

    def __init__(self, parent=None):
        super(ProgressFlyDialog, self).__init__(parent)
        self.setWindowTitle("Progress")

        self.layout = QVBoxLayout()
        self.setLayout(self.layout)

        self.label = QLabel("Calculating...")
        self.layout.addWidget(self.label)

        self.progress_bar = QProgressBar(self)
        self.progress_bar.setRange(0, 100)
        self.layout.addWidget(self.progress_bar)

        self.message_label = QLabel("")
        self.layout.addWidget(self.message_label)

        # self.cancel_button = QPushButton("Cancel")
        # self.cancel_button.clicked.connect(self.on_cancel)
        # self.layout.addWidget(self.cancel_button)

    def update_progress(self, value):
        self.progress_bar.setValue(value)

    def update_message(self, message):
        self.message_label.setText(message)

    def on_cancel(self):
        self.canceled.emit()

    def show_completion_message(self):
        self.label.setText("Fly simulation completed.")
        # self.cancel_button.setText("OK")
        # self.cancel_button.clicked.disconnect(self.on_cancel)
        # self.cancel_button.clicked.connect(self.accept)


# class FileDialog(QDialog):
#     # File dialog for save/load operations
#     def __init__(self, mode):
#         super().__init__()
#         # Setup file dialog UI


class ImageLabel(QLabel):
    def __init__(self, parent=None, max_height=200):
        super().__init__(parent)
        self.original_pixmap = None
        self.max_height = max_height  # Maximum height for the image

    def setPixmap(self, pixmap):
        super().setPixmap(pixmap)
        self.original_pixmap = pixmap  # Store the original pixmap for resizing

    def resizeEvent(self, event):
        if self.original_pixmap:
            # Calculate the maximum size with restricted height
            max_size = QSize(self.width(), self.max_height)

            # Resize the pixmap proportionally within the max size, limiting height
            scaled_pixmap = self.original_pixmap.scaled(
                max_size,
                Qt.AspectRatioMode.KeepAspectRatio,
                Qt.TransformationMode.SmoothTransformation,
            )

            # Update the pixmap on the QLabel
            super().setPixmap(scaled_pixmap)
        super().resizeEvent(event)
