"""
Filename: crossgroupbrainimage.py
Description: Logic for the Cross-Group Brain & Image analysis window

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

# External library imports
import pandas as pd

from flares import aggregate_fnirs_group_geometry, plot_2d_3d_contrasts_between_groups
from src.shared.flaresbasewidget import CrossGroupUIMixin, FlaresBaseWidget
from src.shared.shareddata import APP_NAME


PARAMETERIZED_INDEXES = {
    0: [
        {
            "key": "show_optodes",
            "label": "Determine what is rendered above the brain. Valid values are 'sensors', 'labels', 'none', 'all'.",
            "default": "all",
            "type": str,
        },
        {
            "key": "t_or_theta",
            "label": "Specify if t values or theta values should be plotted. Valid values are 't', 'theta'",
            "default": "theta",
            "type": str,
        },
        {
            "key": "show_text",
            "label": "Display informative text on the top left corner about the contrast.",
            "default": "True",
            "type": bool,
        },
        {
            "key": "brain_bounds",
            "label": "Graph Upper/Lower Limit",
            "default": "1.0",
            "type": float,
        },
        {
            "key": "is_3d",
            "label": "Should we display the results in a 3D interactive window?",
            "default": "True",
            "type": bool,
        }
    ],
}



class CrossGroupBrainImageWidget(CrossGroupUIMixin, FlaresBaseWidget):
    def __init__(self, haemo_dict, df_ind_dict, design_matrix_dict, contrast_results_dict, group_dict):
        super().__init__("CrossGroupBrainImage")
        self.setWindowTitle(f"Cross-Group Brain & Image Viewer - {APP_NAME.upper()}")
        self.haemo_dict = haemo_dict
        self.df_ind_dict = df_ind_dict
        self.design_matrix_dict = design_matrix_dict
        self.contrast_results_dict = contrast_results_dict
        self.group_dict = group_dict

        self.setup_cross_group_ui(["0 (Contrast Image)"])


    def proccess_request(self):
              
        request = self.get_common_request_data(PARAMETERIZED_INDEXES)
        if request is None:
            return

        (selected_event, file_paths_a, file_paths_b, all_selected_paths, selected_indexes, param_values,) = request


        # Build group-level contrast DataFrames
        def concat_group_contrasts(file_paths: list[str], event: str | None) -> pd.DataFrame:
            group_df = pd.DataFrame()
            for fp in file_paths:
                print(f"Looking up contrast for: {fp}")
                event_con_dict = self.contrast_results_dict.get(fp, {})
                print("Available events for this file:", list(event_con_dict.keys()))
                if event and event in event_con_dict:
                    df = event_con_dict[event]
                    print(f"Appending contrast df for event: {event}")
                    group_df = pd.concat([group_df, df], ignore_index=True)
                else:
                    print(f"Event '{event}' not found for {fp}")
            return group_df

        print("Selected event:", selected_event)
        print("File paths A:", file_paths_a)
        print("File paths B:", file_paths_b)

        contrast_df_a = concat_group_contrasts(file_paths_a, selected_event)
        contrast_df_b = concat_group_contrasts(file_paths_b, selected_event)

        print("contrast_df_a empty?", contrast_df_a.empty)
        print("contrast_df_b empty?", contrast_df_b.empty)

        all_raw_objs = [self.haemo_dict.get(fp) for fp in all_selected_paths if self.haemo_dict.get(fp)]
        
        if len(all_raw_objs) > 1:
            processed_raw = aggregate_fnirs_group_geometry(all_raw_objs)
        else:
            processed_raw = all_raw_objs[0].copy().pick(picks="hbo")

        # Visualizations
        for idx in selected_indexes:
            if idx == 0:
                params = param_values.get(idx, {})
                show_optodes = params.get("show_optodes", None)
                t_or_theta = params.get("t_or_theta", None)
                show_text = params.get("show_text", None)
                brain_bounds = params.get("brain_bounds", None)
                is_3d = params.get("is_3d", None)

                if show_optodes is None or t_or_theta is None or show_text is None or brain_bounds is None or is_3d is None:
                    print(f"Missing parameters for index {idx}, skipping.")
                    continue

                if not contrast_df_a.empty and not contrast_df_b.empty and processed_raw:

                    plot_2d_3d_contrasts_between_groups(
                        contrast_df_a,
                        contrast_df_b,
                        raw_haemo=processed_raw,
                        group_a_name=self.group_a_dropdown.currentText(),
                        group_b_name=self.group_b_dropdown.currentText(),
                        is_3d=is_3d,
                        t_or_theta=t_or_theta,
                        show_optodes=show_optodes,
                        show_text=show_text,
                        brain_bounds=brain_bounds
                    )
            else:
                print(f"No method defined for index {idx}")