"""
Filename: participant.py
Description: Participant analysis window for FLARES

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

import os
from pathlib import Path
from datetime import datetime

from PySide6.QtWidgets import QGridLayout, QHBoxLayout, QMessageBox, QPushButton, QScrollArea, QWidget, QVBoxLayout, QLabel
from PySide6.QtCore import Qt, QSize
from PySide6.QtGui import QPixmap

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


class ParticipantImageViewerWidget(FlaresBaseWidget):
    def __init__(self, haemo_dict, fig_bytes_dict):
        super().__init__("ParticipantImage")
        self.setAttribute(Qt.WidgetAttribute.WA_DeleteOnClose)
        self.setWindowTitle(f"Participant Viewer - {APP_NAME.upper()}")
        self.haemo_dict = haemo_dict
        self.fig_bytes_dict = fig_bytes_dict

        # 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)
                
        first_fig_dict = next(iter(self.fig_bytes_dict.values()))
        image_label_items = list(first_fig_dict.keys())

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

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

        self.top_bar.addWidget(QLabel("Participants:"))
        self.top_bar.addWidget(self.participant_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.save_button = QPushButton("Save Displayed Images")
        self.save_button.clicked.connect(self.save_displayed_images)
        self.top_bar.addWidget(self.save_button)

        self.showMaximized()



    def show_selected_images(self):
        # Clear previous images
        while self.grid_layout.count():
            item = self.grid_layout.takeAt(0)
            widget = item.widget()
            if widget:
                widget.deleteLater()

        selected_display_names = self._get_checked_items(self.participant_dropdown)
        # Map from display names back to file paths
        selected_file_paths = []
        for display_name in selected_display_names:
            # Find file_path by matching display name
            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_labels = self._get_checked_items(self.image_index_dropdown)

        row, col = 0, 0
        for file_path in selected_file_paths:
            fig_list = self.fig_bytes_dict.get(file_path, [])
            participant_label = self.participant_map[file_path]
            for label in selected_labels:
                fig_bytes = fig_list.get(label)
                if not fig_bytes:
                    continue

                full_pixmap = QPixmap()
                full_pixmap.loadFromData(fig_bytes)

                thumbnail_pixmap = full_pixmap.scaled(
                    self.thumb_size,
                    Qt.AspectRatioMode.KeepAspectRatio,
                    Qt.TransformationMode.SmoothTransformation
                )

                container = QWidget()
                hlayout = QHBoxLayout(container)
                hlayout.setContentsMargins(0, 0, 0, 0)
                hlayout.setSpacing(0)
                hlayout.setAlignment(Qt.AlignmentFlag.AlignCenter)

                image_label = ClickableLabel(full_pixmap, thumbnail_pixmap)
                image_label.setToolTip(f"{participant_label}\n{label}")
                hlayout.addWidget(image_label)

                self.grid_layout.addWidget(container, row, col)

                col += 1
                if col >= 6:
                    col = 0
                    row += 1

        # Update dropdown labels after display
        self.update_participant_dropdown_label()
        self.update_image_index_dropdown_label()


    def save_displayed_images(self):
        # Ensure the folder exists
        save_dir = Path("individual_images")
        save_dir.mkdir(exist_ok=True)

        selected_display_names = self._get_checked_items(self.participant_dropdown)
        selected_image_labels = self._get_checked_items(self.image_index_dropdown)

        for display_name in selected_display_names:
            # Match display name to file path
            for file_path, short_label in self.participant_map.items():
                expected_display = f"{short_label} ({os.path.basename(file_path)})"
                if display_name == expected_display:
                    fig_dict = self.fig_bytes_dict.get(file_path, {})
                    for label in selected_image_labels:
                        if label not in fig_dict:
                            continue
                        fig_bytes = fig_dict[label]
                        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                        filename = f"{os.path.basename(file_path)}_{label}_{timestamp}.png"
                        output_path = save_dir / filename
                        with open(output_path, "wb") as f:
                            f.write(fig_bytes)
                    break  # file_path matched; stop loop

        QMessageBox.information(self, "Save Complete", f"Images saved to {save_dir.resolve()}")