"""
Filename: participantfunctionalconnectivity.py
Description: Participant functional connectivity analysis window for FLARES

Author: Tyler de Zeeuw
License: GPL-3.0
"""

import os

from PySide6.QtWidgets import QComboBox, QDialog, QGridLayout, QHBoxLayout, QPushButton, QScrollArea, QWidget, QVBoxLayout, QLabel, QMessageBox
from PySide6.QtCore import QSize

from src.shared.flaresbasewidget import FlaresBaseWidget, ParameterInputDialog
from src.shared.shareddata import APP_NAME


class ParticipantFunctionalConnectivityWidget(FlaresBaseWidget):
    def __init__(self, haemo_dict, epochs_dict):
        super().__init__("FunctionalConnectivityWidget")
        self.setWindowTitle(f"Functional Connectivity Viewer [BETA] - {APP_NAME.upper()}")
        self.haemo_dict = haemo_dict
        self.epochs_dict = epochs_dict

        QMessageBox.warning(self, f"Warning - {APP_NAME.upper()}", f"Functional Connectivity is still in development and the results should currently be taken with a grain of salt. "
                            "By clicking OK, you accept that the images generated may not be factual.")

        # Create mappings: file_path -> participant label and dropdown display text
        self.participant_map = {}  # file_path -> "Participant 1"
        self.participant_dropdown_items = []  # "Participant 1 (filename)"

        for i, file_path in enumerate(self.haemo_dict.keys(), start=1):
            short_label = f"Participant {i}"
            display_label = f"{short_label} ({os.path.basename(file_path)})"
            self.participant_map[file_path] = short_label
            self.participant_dropdown_items.append(display_label)

        self.layout = QVBoxLayout(self)
        self.top_bar = QHBoxLayout()
        self.layout.addLayout(self.top_bar)

        self.participant_dropdown = self._create_multiselect_dropdown(self.participant_dropdown_items)
        self.participant_dropdown.currentIndexChanged.connect(self.update_participant_dropdown_label)

        self.event_dropdown = QComboBox()
        self.event_dropdown.addItem("<None Selected>")

                
        self.index_texts = [
            "0 (Spectral Connectivity Epochs)",
            "1 (Envelope Correlation)",
            "2 (Betas)",
            "3 (Spectral Connectivity Epochs)",
            ]

        self.image_index_dropdown = self._create_multiselect_dropdown(self.index_texts)
        self.image_index_dropdown.currentIndexChanged.connect(self.update_image_index_dropdown_label)

        self.submit_button = QPushButton("Submit")
        self.submit_button.clicked.connect(self.show_brain_images)

        self.top_bar.addWidget(QLabel("Participants:"))
        self.top_bar.addWidget(self.participant_dropdown)
        self.top_bar.addWidget(QLabel("Event:"))
        self.top_bar.addWidget(self.event_dropdown)
        self.top_bar.addWidget(QLabel("Image Indexes:"))
        self.top_bar.addWidget(self.image_index_dropdown)
        self.top_bar.addWidget(self.submit_button)

        self.scroll = QScrollArea()
        self.scroll.setWidgetResizable(True)
        self.scroll_content = QWidget()
        self.grid_layout = QGridLayout(self.scroll_content)
        self.scroll.setWidget(self.scroll_content)
        self.layout.addWidget(self.scroll)

        self.thumb_size = QSize(280, 180)
        self.showMaximized()


    def show_brain_images(self):
        import flares as flares

        selected_event = self.event_dropdown.currentText()
        if selected_event == "<None Selected>":
            selected_event = None

        selected_display_names = self._get_checked_items(self.participant_dropdown)
        selected_file_paths = []
        for display_name in selected_display_names:
            for fp, short_label in self.participant_map.items():
                expected_display = f"{short_label} ({os.path.basename(fp)})"
                if display_name == expected_display:
                    selected_file_paths.append(fp)
                    break

        selected_indexes = [
            int(s.split(" ")[0]) for s in self._get_checked_items(self.image_index_dropdown)
        ]


        parameterized_indexes = {
            0: [
                {
                    "key": "n_lines",
                    "label": "<Description>",
                    "default": "20",
                    "type": int,
                },
                {
                    "key": "vmin",
                    "label": "<Description>",
                    "default": "0.9",
                    "type": float,
                },
            ],
            1: [
                {
                    "key": "n_lines",
                    "label": "<Description>",
                    "default": "20",
                    "type": int,
                },
                {
                    "key": "vmin",
                    "label": "<Description>",
                    "default": "0.9",
                    "type": float,
                },
             
            ],
            2: [
                {
                    "key": "n_lines",
                    "label": "<Description>",
                    "default": "20",
                    "type": int,
                },
                {
                    "key": "vmin",
                    "label": "<Description>",
                    "default": "0.9",
                    "type": float,
                },
             
            ],
            3: [
                {
                    "key": "n_lines",
                    "label": "<Description>",
                    "default": "20",
                    "type": int,
                },
                {
                    "key": "vmin",
                    "label": "<Description>",
                    "default": "0.9",
                    "type": float,
                },
             
            ],
        }

        # Inject full_text from index_texts
        for idx, params_list in parameterized_indexes.items():
            full_text = self.index_texts[idx] if idx < len(self.index_texts) else f"{idx} (No label found)"
            for param_info in params_list:
                param_info["full_text"] = full_text

        indexes_needing_params = {idx: parameterized_indexes[idx] for idx in selected_indexes if idx in parameterized_indexes}

        param_values = {}
        if indexes_needing_params:
            dialog = ParameterInputDialog(indexes_needing_params, parent=self)
            if dialog.exec_() == QDialog.Accepted:
                param_values = dialog.get_values()
                if param_values is None:
                    return
            else:
                return

        # Pass the necessary arguments to each method
        for file_path in selected_file_paths:
            haemo_obj = self.haemo_dict.get(file_path)
            epochs_obj = self.epochs_dict.get(file_path)

            if selected_event:
                participant_events = set(haemo_obj.annotations.description)
                if selected_event not in participant_events:
                    print(f"Skipping {self.participant_map[file_path]}: Event '{selected_event}' not found.")
                    continue

            if haemo_obj is None:
                raise Exception("How did we get here?")


            for idx in selected_indexes:
                if idx == 0:

                    params = param_values.get(idx, {})
                    n_lines = params.get("n_lines", None)
                    vmin = params.get("vmin", None)

                    if n_lines is None or vmin is None:
                        print(f"Missing parameters for index {idx}, skipping.")
                        continue
                    flares.functional_connectivity_spectral_epochs(epochs_obj, n_lines, vmin)

                elif idx == 1:
                    params = param_values.get(idx, {})
                    n_lines = params.get("n_lines", None)
                    vmin = params.get("vmin", None)

                    if n_lines is None or vmin is None:
                        print(f"Missing parameters for index {idx}, skipping.")
                        continue
                    flares.functional_connectivity_envelope(epochs_obj, n_lines, vmin)
                    
                elif idx == 2:
                    params = param_values.get(idx, {})
                    n_lines = params.get("n_lines", None)
                    vmin = params.get("vmin", None)

                    if n_lines is None or vmin is None:
                        print(f"Missing parameters for index {idx}, skipping.")
                        continue
                    flares.functional_connectivity_betas(haemo_obj, n_lines, vmin, selected_event)

                elif idx == 3:
                    params = param_values.get(idx, {})
                    n_lines = params.get("n_lines", None)
                    vmin = params.get("vmin", None)

                    if n_lines is None or vmin is None:
                        print(f"Missing parameters for index {idx}, skipping.")
                        continue
                    flares.functional_connectivity_spectral_time(epochs_obj, n_lines, vmin)

                else:
                    print(f"No method defined for index {idx}")