"""
Filename: exportcsv.py
Description: Export data as csv analysis window for FLARES

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

import os

import numpy as np
import pandas as pd

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

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


class ExportDataAsCSVViewerWidget(FlaresBaseWidget):
    def __init__(self, haemo_dict, cha_dict, df_ind, design_matrix, group, contrast_results_dict):
        super().__init__("ExportDataAsCSVViewer")
        self.setWindowTitle(f"Export Data As CSV Viewer - {APP_NAME.upper()}")
        self.haemo_dict = haemo_dict
        self.cha_dict = cha_dict
        self.df_ind = df_ind
        self.design_matrix = design_matrix
        self.group = group
        self.contrast_results_dict = contrast_results_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)
              
        self.index_texts = [
            "0 (Export Data to CSV)",
            "1 (CSV for SPARKS)",
            # "2 (third image)",
            # "3 (fourth image)",
            ]

        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.generate_and_save_csv)

        self.top_bar.addWidget(QLabel("Participants:"))
        self.top_bar.addWidget(self.participant_dropdown)
        self.top_bar.addWidget(QLabel("Export Type:"))
        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 generate_and_save_csv(self):
        
        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)
        ]

        if not selected_file_paths or not selected_indexes:
            QMessageBox.warning(self, "Selection Missing", "Please select at least one participant and one export type.")
            return

        # 2. ASK ONCE: Select Output Directory
        output_dir = QFileDialog.getExistingDirectory(self, "Select Output Folder for CSV Exports")
        
        if not output_dir:
            print("Export cancelled: No folder selected.")
            return

        success_count = 0

        # Pass the necessary arguments to each method
        for file_path in selected_file_paths:
            base_filename = os.path.splitext(os.path.basename(file_path))[0]
            haemo_obj = self.haemo_dict.get(file_path)
            if haemo_obj is None:
                continue

            cha = self.cha_dict.get(file_path)

            for idx in selected_indexes:
                try:
                    if idx == 0:
                        save_path = os.path.join(output_dir, f"{base_filename}_exported.csv")
                        if cha is not None:
                            cha.to_csv(save_path)
                            success_count += 1


                    elif idx == 1:
                            # SPARKS Export
                            save_path = os.path.join(output_dir, f"{base_filename}_sparks.csv")
                            if haemo_obj is not None:
                                raw = haemo_obj
                                data, times = raw.get_data(return_times=True)
                                ann_col = np.full(times.shape, "", dtype=object)

                                if raw.annotations is not None and len(raw.annotations) > 0:
                                    for onset, duration, desc in zip(
                                        raw.annotations.onset,
                                        raw.annotations.duration,
                                        raw.annotations.description
                                    ):
                                        mask = (times >= onset) & (times < onset + duration)
                                        ann_col[mask] = desc

                                df = pd.DataFrame(data.T, columns=raw.ch_names)
                                df.insert(0, "annotation", ann_col)
                                df.insert(0, "time", times)
                                df.to_csv(save_path, index=False)
                                success_count += 1
                        
                    else:
                        print(f"No method defined for index {idx}")
                
                except Exception as e:
                    print(f"Failed to export {file_path} (Type {idx}): {e}")

        # 4. Final Notification
        if success_count > 0:
            QMessageBox.information(self, "Export Complete", f"Successfully saved {success_count} CSV files to:\n{output_dir}")
            
            # # If SPARKS export was included, show the Event Window once at the end
            # if 1 in selected_indexes:
            #     win = UpdateEventsWindow(
            #         parent=self,
            #         mode=EventUpdateMode.WRITE_JSON,
            #         caller="Video Alignment Tool"
            #     )
            #     win.show()