flares/flares.py

"""
Filename: flares.py
Description: Core functionality for FLARES

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

# Built-in imports
import os
import gc
import re
import sys
import json
import time
import logging
import warnings
import threading
import traceback
import itertools
import os.path as op
from io import BytesIO
from pathlib import Path
from copy import deepcopy
import multiprocessing as mp
from itertools import compress
from queue import Empty, Queue
from typing import Any, Optional, Sequence, cast, Literal, Union, List

# External library imports
import matplotlib.pyplot as plt
import matplotlib.colors as mcolors
from matplotlib.axes import Axes
from matplotlib.lines import Line2D
from matplotlib.figure import Figure
from matplotlib.colors import LinearSegmentedColormap

import numpy as np
from numpy.typing import NDArray
from numpy import float64, floating

import pandas as pd
from pandas import DataFrame

import h5py # type: ignore
import seaborn as sns

from nilearn.plotting import plot_design_matrix  # type: ignore
from nilearn.glm.regression import OLSModel # type: ignore

import statsmodels.formula.api as smf  # type: ignore
from statsmodels.stats.multitest import multipletests # type: ignore
from statsmodels.tools.sm_exceptions import ConvergenceWarning # type: ignore

from scipy.spatial.distance import cdist
from scipy.signal import welch, butter, filtfilt, periodogram # type: ignore
from scipy.stats import pearsonr, zscore, ttest_1samp, ttest_ind, sem, t as t_dist # type: ignore

from openpyxl import Workbook
from openpyxl.styles import Font, PatternFill, Alignment
from openpyxl.formatting.rule import ColorScaleRule
from openpyxl.utils import get_column_letter

import pywt # type: ignore
import neurokit2 as nk # type: ignore

# Backend visualization needed to be defined for pyinstaller
import pyvistaqt # type: ignore
import vtkmodules.util.data_model
import vtkmodules.util.execution_model
import xlrd

# External library imports for mne
from mne import (
    EvokedArray, SourceEstimate, Info, Epochs, Label, Annotations,
    events_from_annotations, read_source_spaces, create_info, # type: ignore
    stc_near_sensors, pick_types, grand_average, get_config, set_config, read_labels_from_annot # type: ignore
)
from mne.source_space import SourceSpaces
from mne.transforms import Transform  # type: ignore
from mne.io import BaseRaw, RawArray, read_raw_snirf  # type: ignore
from mne.preprocessing.nirs import (
    beer_lambert_law, optical_density, # type: ignore
    temporal_derivative_distribution_repair, # type: ignore
    source_detector_distances, short_channels # type: ignore
)
from mne.viz import Brain, plot_events, plot_evoked_topo, plot_compare_evokeds # type: ignore
from mne.filter import filter_data  # type: ignore
from mne.utils import _check_fname, _validate_type, warn # type: ignore
from mne.channels import make_standard_montage # type: ignore
from mne.datasets.sample import data_path # type: ignore

from mne_nirs.visualisation import plot_glm_group_topo  # type: ignore
from mne_nirs.channels import get_long_channels, get_short_channels  # type: ignore
from mne_nirs.experimental_design import make_first_level_design_matrix  # type: ignore
from mne_nirs.statistics import run_glm, statsmodels_to_results  # type: ignore
from mne_nirs.signal_enhancement import ( # type: ignore
    enhance_negative_correlation, short_channel_regression # type: ignore
)
from mne_nirs.io.fold import fold_channel_specificity  # type: ignore
from mne_nirs.preprocessing import peak_power  # type: ignore
from mne.preprocessing.nirs.nirs import _validate_nirs_info # type: ignore
from mne_nirs.statistics._glm_level_first import RegressionResults  # type: ignore

from mne_connectivity.viz import plot_connectivity_circle # type: ignore
from mne_connectivity import envelope_correlation, spectral_connectivity_epochs, spectral_connectivity_time # type: ignore

from src.shared.shareddata import PLATFORM_NAME, resource_path, get_app_dir



PRIMARY_COLORS = {
    "SCI only": "skyblue",      # Scalp Coupling Index (Standard MNE)
    "SNR only": "lightgreen",   # Signal-to-Noise Ratio (Original)
    "PSP only": "salmon",       # Power Spectral Peak (Original Noise check)
    "Coeff_var only": "yellow", # Relative Noise (The coeff_var-only check)
    "Range only": "coral",      # Z-Swing (The Range Outlier check)
    "Noise only": "plum",       # High-Freq PSD (The Noise check)
    "Disp. only": "palegreen",  # Sensor Displacement (Variance Drop)
    "Multiple": "orangered"     # Failed 2+ categories
}

COMBINATION_COLOR = "gray"
NUISANCE_EXCLUDE = ("drift", "constant", "short")

def get_category_color(label):
    """Returns the primary color if it's a single failure, otherwise gray."""
    return PRIMARY_COLORS.get(label, COMBINATION_COLOR)


# direction: True = lower value is better (green), False = higher is better (green)
QC_METRIC_DIRECTIONS = {
    "n_bad_sci": True,
    "n_bad_snr": True,
    "n_bad_psp": True,
    "n_bad_coeff_var": True,
    "n_bad_mad": True,
    "n_bad_psd_noise": True,
    "n_bad_dropout": True,
    "n_bad_channels_total": True,
    "pct_bad_channels": True,
    "total_processing_seconds": True,
    "n_epochs_final": False,
}

# metrics with no inherent "good direction" - colored by deviation from the
# group median instead (outliers flagged, not high/low values per se)
QC_METRIC_DEVIATION_BASED = {"final_hr_bpm"}

QC_METRIC_LABELS = {
    "n_bad_sci": "Bad Channels - SCI",
    "n_bad_snr": "Bad Channels - SNR",
    "n_bad_psp": "Bad Channels - PSP",
    "n_bad_coeff_var": "Bad Channels - Coeff. Var",
    "n_bad_mad": "Bad Channels - MAD",
    "n_bad_psd_noise": "Bad Channels - PSD Noise",
    "n_bad_dropout": "Bad Channels - Dropout",
    "n_bad_channels_total": "Bad Channels - Total (union)",
    "pct_bad_channels": "% Channels Bad",
    "final_hr_bpm": "Final Heart Rate (BPM)",
    "n_epochs_final": "Epochs Retained",
    "total_processing_seconds": "Processing Time (s)",
}

DOWNSAMPLE: bool
DOWNSAMPLE_FREQUENCY: int

TRIM: bool
SECONDS_TO_KEEP: float

OPTODE_PLACEMENT: bool
SHOW_OPTODE_NAMES: bool

SHORT_CHANNELS: bool
LONG_CHANNELS: bool
SHORT_CHANNELS_THRESHOLD: float
LONG_CHANNELS_THRESHOLD: float

HEART_RATE: bool
SECONDS_TO_STRIP_HR: int
HR_LOW_FREQ: int
HR_HIGH_FREQ: int
HR_SEARCH_MIN: int
HR_SEARCH_MAX: int
MAX_LOW_HR: int
MAX_HIGH_HR: int
SMOOTHING_WINDOW_HR: int
HEART_RATE_WINDOW: int

SCI: bool
SCI_USE_HEART_RATE_BAND: bool
SCI_LOW_FREQ: float
SCI_HIGH_FREQ: float
SCI_TIME_WINDOW: int
SCI_THRESHOLD: float

SNR: bool
SNR_THRESHOLD: float
SNR_SIGNAL_LOW_FREQ: float
SNR_SIGNAL_HIGH_FREQ: float
SNR_NOISE_LOW_FREQ: float
SNR_NOISE_HIGH_FREQ: float

PSP: bool
PSP_USE_HEART_RATE_BAND: bool
PSP_LOW_FREQ: float
PSP_HIGH_FREQ: float
PSP_TIME_WINDOW: int
PSP_THRESHOLD: float

COEFF_VAR: bool
COEFF_VAR_THRESHOLD: int

MAD: bool
MAD_THRESHOLD: int

PSD_NOISE: bool
TARGET_FREQ_DIV: int
DB_LIMIT: int
PSD_MIN_FREQ: float
PSD_TARGET_BANDWIDTH: float

SENSOR_DROPOUT: bool
SENSOR_DROPOUT_VARIANCE_THRESHOLD: float

BAD_CHANNELS_HANDLING: str
MAX_DIST: float
MIN_NEIGHBORS: int
MAX_BAD_CHANNELS: int

TDDR: bool

WAVELET: bool
IQR: float
WAVELET_TYPE: str
WAVELET_LEVEL: int

OVERRIDE_PPF: bool
PPF_LOWER_WAVELENGTH: float
PPF_UPPER_WAVELENGTH: float

ENHANCE_NEGATIVE_CORRELATION: bool

FILTER: bool
FILTER_ALGORITHM: list
L_FREQ: float
H_FREQ: float
L_TRANS_BANDWIDTH: float
H_TRANS_BANDWIDTH: float
IIR_TYPE: list
IIR_ORDER: int
FILTER_LENGTH: str
FILTER_PHASE: list
FIR_WINDOW: list
FIR_DESIGN: list
IIR_OUTPUT: list
PASSBAND_RIPPLE: float
STOPBAND_ATTENUATION: float
FILTER_PAD: list
SKIP_BY_ANNOTATION: list
FILTER_N_JOBS: int

EVENTS: bool
EVENT_ID: str
EVENT_REGEX: str
EVENT_CHUNK_DURATION: float

EPOCHS: bool
EPOCH_HANDLING: str
MAX_SHIFT: int
T_MIN: int
T_MAX: int
BASELINE: list
REJECT_EPOCHS: bool
REJECT_HBO_THRESHOLD: float

RESAMPLE: bool
RESAMPLE_FREQ: int
HRF_MODEL: str
STIM_DUR: float
FIR_DELAYS: range
DRIFT_MODEL: str
HIGH_PASS: float
DRIFT_ORDER: int
MIN_ONSET: int
OVERSAMPLING: int
SHORT_CHANNEL_REGRESSION: bool

NOISE_MODEL: str
BINS: int
N_JOBS: int

JSON_LOCATION: str

MAX_WORKERS: int
VERBOSITY: bool

AGE: int = 25  # Assume 25 if not set from the GUI. This will result in a reasonable PPF if calculated dynamically
GENDER: str = ""
GROUP: str = "Default"

FOLDING_BYP: bool = False

FEATURE_1: bool = False
FEATURE_2: bool = False

# Ensure that we are working in the directory of this file
script_dir = os.path.dirname(os.path.abspath(__file__))
os.chdir(script_dir)


# Configure logging to file with timestamps and realtime flush
if PLATFORM_NAME == 'darwin':
    log_path = os.path.abspath(os.path.join(os.path.dirname(sys.executable), "../../../fnirs_analysis.log"))
else:
    log_path = os.path.join(get_app_dir(), "fnirs_analysis.log")

logging.basicConfig(
    filename=log_path,
    level=logging.INFO,
    format='%(asctime)s - %(processName)s - %(levelname)s - %(message)s',
    datefmt='%Y-%m-%d %H:%M:%S',
    filemode='a'
)

logger = logging.getLogger()



def set_config_me(config: dict[str, Any]) -> None:
    """
    Validates and applies the given configuration dictionary.

    Parameters
    ----------
    config : dict[str, Any]
        Dictionary containing configuration keys and their values.
    """
    logger.info(f"[DEBUG] set_config called")

    globals().update(config)


    
def set_metadata(file_path, metadata: dict[str, Any]) -> None:
    """
    Validates and applies the given configuration dictionary.

    Parameters
    ----------
    config : dict[str, Any]
        Dictionary containing configuration keys and their values.
    """
    logger.info(f"[DEBUG] set_metadata called")

    globals()['AGE'] = 25
    globals()['GENDER'] = ""
    globals()['GROUP'] = "Default"
    
    if metadata.get(file_path) is not None:
                
        file_metadata = metadata.get(file_path, {})

        for key in ("AGE", "GENDER", "GROUP"):
            val = file_metadata.get(key, None)
            if val not in (None, '', [], {}, ()):  # check for "empty" values
                globals()[key] = val



def gui_entry(config: dict[str, Any], gui_queue: mp.Queue, progress_queue: mp.Queue, ack_queue: mp.Queue) -> None:
    start_time = time.time()
    try:
        file_paths = config['SNIRF_FILES']
        file_params = config['PARAMS']
        file_metadata = config['METADATA']
        max_workers = file_params.get("MAX_WORKERS", int(os.cpu_count()/4))

        results = process_multiple_participants(
            file_paths, file_params, file_metadata, progress_queue, gui_queue, max_workers
        )

        elapsed = time.time() - start_time
        success_count = getattr(process_multiple_participants, "_success_count", 0)
        total_duration = getattr(process_multiple_participants, "_duration_total", 0.0)
        failed_stages = getattr(process_multiple_participants, "_failed_stages", [])

        speedup = None
        if success_count > 0 and elapsed > 0:
            avg_success_duration = total_duration / success_count
            failed_credit = sum(stage * avg_success_duration for stage in failed_stages)
            naive_serial_estimate = total_duration + failed_credit
            speedup = min(naive_serial_estimate / elapsed, max_workers)

        gui_queue.put({
            "type": "FINISHED_SUCCESSFULLY",
            "success": True,
            "elapsed": elapsed,
            "speedup": speedup,
        })
        
        try:
            print("CHILD: Waiting for GUI acknowledgment...")
            ack_queue.get(timeout=10)
        except:
            print("CHILD: Ack timeout, exiting anyway.")

    except Exception as e:
        gui_queue.put({
            "type": "FINISHED_SUCCESSFULLY",
            "success": False,
            "error": str(e),
            "traceback": traceback.format_exc(),
            "elapsed": time.time() - start_time
        })

    finally:
        pass



def process_participant_worker(file_path, file_params, file_metadata, result_queue, progress_queue):
    file_start = time.time()
    stage_tracker = {"value": 0.0}
    try:
        set_config_me(file_params)
        set_metadata(file_path, file_metadata)

        def progress_callback(step_idx):
            stage_tracker["value"] = min(step_idx / 28, 1.0)
            if progress_queue:
                try:
                    progress_queue.put_nowait(('progress', file_path, step_idx))
                except Exception:
                    pass

        result = process_participant(file_path, file_start, progress_callback=progress_callback)
        duration = time.time() - file_start
        result_queue.put((file_path, result, None, duration, 1.0))

    except Exception as e:
        duration = time.time() - file_start
        try:
            result_queue.put((file_path, None, f"{e}\n{traceback.format_exc()}", duration, stage_tracker["value"]))
        except Exception:
            pass

    finally:
        try:
            plt.close('all')
            gc.collect()
        except Exception:
            pass



def process_multiple_participants(file_paths, file_params, file_metadata,
                                   progress_queue=None, gui_queue=None, max_workers=6, qc_summary_path: str | None = None):
    ctx = mp.get_context("spawn")
    result_queue = ctx.Queue()

    pending_files = list(file_paths)
    pending_lock = threading.Lock()
    results_by_file = {}
    results_lock = threading.Lock()
    stop_event = threading.Event()

    active_lock = threading.Lock()
    active_processes = []   # tracked only for emergency cleanup on error

    start_time = time.time()
    duration_total = {"value": 0.0}
    success_count = {"value": 0}
    failed_stages = {"value": []}
    qc_rows: list[dict[str, Any]] = []
    if qc_summary_path is None:
        qc_summary_path = os.path.join(get_app_dir(), "qc_summary.xlsx")

    def elapsed_heartbeat():
        # Ticks once a second so the GUI can show a live-updating timer,
        # independent of when files actually finish.
        while not stop_event.is_set():
            if gui_queue:
                try:
                    gui_queue.put({"type": "elapsed", "seconds": time.time() - start_time})
                except Exception:
                    pass
            stop_event.wait(1.0)

    def relay_progress():
        while not stop_event.is_set():
            try:
                msg = progress_queue.get(timeout=0.05)
            except Empty:
                continue
            except (EOFError, OSError):
                break
            if gui_queue:
                try:
                    gui_queue.put(msg)
                except Exception:
                    pass

    def result_collector():
        # Runs continuously in its own thread, fully decoupled from spawning.
        while not stop_event.is_set():
            try:
                res_path, result, error, duration, stage = result_queue.get(timeout=0.05)
            except Empty:
                continue
            except (EOFError, OSError):
                break
            if error is None:
                duration_total["value"] += duration
                success_count["value"] += 1
                qc = result[-2] if isinstance(result, tuple) else None # accessing by index of return
                if isinstance(qc, dict):
                    qc["status"] = "success"
                    qc["duration_seconds"] = round(duration, 2)
                    qc_rows.append(qc)
            else:
                failed_stages["value"].append(stage)
                qc_rows.append({
                    "file_path": res_path,
                    "status": "FAILED",
                    "error": error.splitlines()[0] if error else "unknown error",
                    "duration_seconds": round(duration, 2),
                    "total_processing_seconds": None,
                })
            if gui_queue:
                try:
                    gui_queue.put({
                        "type": "file_done",
                        "file": res_path,
                        "success": error is None,
                        "result": result if error is None else None,
                        "error": error,
                    })
                except Exception:
                    pass
            else:
                with results_lock:
                    results_by_file[res_path] = result

    def slot_worker():
        # Each slot independently: grab a file, spawn+join, repeat.
        # Replacement happens immediately on exit -- no shared polling loop.
        while not stop_event.is_set():
            with pending_lock:
                if not pending_files:
                    return
                file_path = pending_files.pop(0)

            p = ctx.Process(
                target=process_participant_worker,
                args=(file_path, file_params, file_metadata, result_queue, progress_queue)
            )
            p.daemon = False
            p.start()

            with active_lock:
                active_processes.append(p)

            p.join()

            with active_lock:
                if p in active_processes:
                    active_processes.remove(p)

    relay_thread = None
    if progress_queue is not None:
        relay_thread = threading.Thread(target=relay_progress, daemon=True)
        relay_thread.start()

    collector_thread = threading.Thread(target=result_collector, daemon=True)
    collector_thread.start()

    heartbeat_thread = threading.Thread(target=elapsed_heartbeat, daemon=True)
    heartbeat_thread.start()

    n_slots = max(1, min(max_workers, len(pending_files)))
    slot_threads = [threading.Thread(target=slot_worker, daemon=True) for _ in range(n_slots)]

    try:
        for t in slot_threads:
            t.start()          # all slots fire their first spawn essentially at once
        for t in slot_threads:
            t.join()
    except Exception as e:
        print(f"MAIN LOOP ERROR: {e}")
    finally:
        stop_event.set()
        with active_lock:
            for p in active_processes:
                try:
                    if p.is_alive():
                        p.terminate()
                        p.join(timeout=1)
                except Exception:
                    pass
        if relay_thread:
            relay_thread.join(timeout=2)
        collector_thread.join(timeout=2)
        heartbeat_thread.join(timeout=2)

    process_multiple_participants._duration_total  = duration_total ["value"]
    process_multiple_participants._success_count = success_count["value"]
    process_multiple_participants._failed_stages = failed_stages["value"]

    successes = [r for r in qc_rows if r.get("status") == "success"]
    population_flags = flag_population_outliers(successes)

    for path, flags in population_flags.items():
        metric_summary = ", ".join(f"{f['metric']}={f['value']} (median={f['median']:.1f}, z={f['z']})" for f in flags)
        n_metrics_flagged = len(flags)
        severity = "STRONG" if n_metrics_flagged >= 3 else "possible"
        logger.warning(f"{severity} outlier: {path} - flagged on {n_metrics_flagged} metric(s): {metric_summary}")

    if qc_summary_path and qc_rows:
        if FEATURE_1:
            try:
                write_qc_excel_summary(qc_rows, qc_summary_path, population_flags=population_flags)
                logger.info(f"QC summary written to {qc_summary_path} ({len(qc_rows)} participant(s))")
            except Exception as e:
                logger.error(f"Failed to write QC summary: {e}")

    return results_by_file



def flag_population_outliers(
    qc_rows: list[dict],
    metrics: list[str] | None = None,
    z_threshold: float = 2.5,
) -> dict[str, list[dict]]:
    """
    Given QC metrics collected across an entire batch (each participant
    processed with identical, fixed thresholds - nothing adaptive changes
    per-participant), flags participants whose value on a given metric is
    a statistical outlier relative to the rest of the batch.

    This does NOT change what gets marked as a bad channel, and does NOT
    imply anything is wrong with the metric thresholds themselves - it
    answers a different question: "given everyone went through the same
    pipeline with the same settings, does this participant's result look
    unusual compared to everyone else who did." A real, uniformly bad
    dataset (e.g. wrong population for the config) would show up as
    outliers across MULTIPLE metrics; a participant with one borderline
    metric and nothing else is much weaker evidence of a real problem.

    Uses a modified z-score (median + MAD-based, not mean + std) since
    QC metric distributions across a real population are often skewed by
    a few genuinely bad files - MAD-based scoring is robust to those
    outliers dominating the very estimate used to detect them, unlike a
    standard mean/std z-score.

    Parameters
    ----------
    qc_rows : list of per-participant QC dicts (successful runs only -
        filter out FAILED entries before calling this).
    metrics : list of QC dict keys to check. Defaults to the standard
        bad-channel-count metrics if not specified.
    z_threshold : float, default 2.5
        Modified z-score magnitude above which a participant is flagged
        for that metric. 2.5 is a commonly used, moderately conservative
        starting point for outlier flagging - not validated against your
        specific data, worth adjusting based on what you see in practice.

    Returns
    -------
    dict[str, list[dict]]
        Keyed by file_path, listing which metric(s) that participant was
        flagged as an outlier on and by how much, e.g.
        {"sub-07.snirf": [{"metric": "n_bad_sci", "z": 3.1, "value": 22, "median": 3}]}
    """
    if metrics is None:
        metrics = [
            "n_bad_sci", "n_bad_snr", "n_bad_psp", "n_bad_coeff_var",
            "n_bad_mad", "n_bad_psd_noise", "n_bad_dropout",
            "n_bad_channels_total", "pct_bad_channels",
        ]

    flagged: dict[str, list[dict]] = {}

    for metric in metrics:
        values = np.array([row.get(metric) for row in qc_rows if row.get(metric) is not None], dtype=float)
        if len(values) < 4:
            logger.info(f"Skipping outlier check for '{metric}' - too few participants ({len(values)}) for a meaningful comparison.")
            continue

        median = np.median(values)
        mad = np.median(np.abs(values - median))
        if mad == 0:
            continue  # every participant identical on this metric - nothing to flag

        for row in qc_rows:
            val = row.get(metric)
            if val is None:
                continue
            modified_z = 0.6745 * (val - median) / mad
            if abs(modified_z) >= z_threshold:
                path = row.get("file_path", "unknown")
                flagged.setdefault(path, []).append({
                    "metric": metric, "z": round(float(modified_z), 2),
                    "value": val, "median": float(median),
                })

    return flagged


def markbad(data, ax, ch_names: list[str]) -> None:
    """
    Add a strikethrough to a plot for channels marked as bad.

    Parameters
    ----------
    data : BaseRaw
        The loaded data object to process.
    ax : Axes
        Matplotlib Axes object where the strikethrough lines will be drawn.
    ch_names : list[str]
        List of channel names corresponding to the y-axis of the plot.
    """

    # Iterate over all the channels
    for i, ch in enumerate(ch_names):

        # If it is marked as bad, place a strikethrough on the channel
        if ch in data.info["bads"]:
            ax.axhline(i + 0.5, ls="solid", lw=4, color="black", zorder=10) # type: ignore



def plot_timechannel_quality_metrics(data, scores, times: list[tuple[float]], color_stops: tuple[list[float], list[float]], threshold: float, title: Optional[str] = None):

    """
    Generate two heatmaps visualizing channel quality metrics over time.

    Parameters
    ----------
    data : BaseRaw
        The loaded data object to process.
    scores : NDArray[float64]
        A 2D array of quality scores for each channel over time.
    times : list[tuple[float]]
        List of time boundaries used to label each score column.
    color_stops : tuple[list[float], list[float]]
        Two lists of color values for custom colormaps.
    threshold : float, 
        Threshold value for the color bar.
    title : Optional[str], optional
        Base title for the figures, (default is None).

    Returns
    -------
    tuple[Figure, Figure]
        - Figure: Heatmap of all scores across channels and time.
        - Figure: Binary heatmap showing only scores above the threshold.
    """

    # Get only the hbo / hbr channels once as we dont need to see the same results twice
    half_ch = len(getattr(data, "ch_names")) // 2
    ch_names = getattr(data, "ch_names")[:half_ch]
    scores = scores[:half_ch, :]
    
    # Extract rounded time points to use as column headers
    cols = [np.round(t[0]) for t in times]
    n_chans = len(ch_names)
    vsize = 0.2 * n_chans

    # Create the first figure
    fig1, ax1 = plt.subplots(figsize=(10, vsize), layout="constrained")  # type: ignore
    fig1.suptitle(title + " - All Scores", fontsize=16, fontweight="bold")  # type: ignore

    # Create a DataFrame to structure data for the heatmap
    data_to_plot = DataFrame(
        data=scores,
        columns=pd.Index(cols, name="Time (s)"),
        index=pd.Index(ch_names, name="Channel"),
    )

    # Define a custom colormap using provided color stops and base colors
    base_colors = ['red', 'red', 'yellow', 'green', 'green']
    colors = list(zip(color_stops[0], base_colors[:len(color_stops[0])]))
    cmap = mcolors.LinearSegmentedColormap.from_list('gyr', colors)

    # Plot heatmap of scores
    sns.heatmap(  # type: ignore
        data=data_to_plot,
        cmap=cmap,
        vmin=0,
        vmax=1,
        cbar_kws=dict(label="Score"),
        ax=ax1,
    )

    # Add vertical dashed lines at each time boundary, sit the title, and place a black strikethrough through a bad channel
    for x in range(1, len(times)):
        ax1.axvline(x, ls="dashed", lw=0.25, dashes=(25, 15), color="gray") # type: ignore
    ax1.set_title("All Scores", fontweight="bold") # type: ignore
    markbad(data, ax1, ch_names)

    # Calculate average score per channel and annotate to the right of the heatmap
    avg_sci_subset: pd.Series[float] = data_to_plot.mean(axis=1) # type: ignore
    norm = mcolors.Normalize(vmin=0, vmax=1)
    text_x = data_to_plot.shape[1] + 0.5
    for i, val in enumerate(avg_sci_subset):
        color = cmap(norm(val))
        ax1.text( # type: ignore
            text_x,
            i + 0.5,
            f"{val:.3f}",
            va='center',
            ha='left',
            fontsize=9,
            color=color
        )
    ax1.set_xlim(right=text_x + 1.5)
    
    plt.close(fig1)

    # Create the second figure
    fig2, ax2 = plt.subplots(figsize=(10, vsize), layout="constrained") # type: ignore
    fig2.suptitle(title + " - Scores Above Threshold", fontsize=16, fontweight="bold") # type: ignore
   
    # Create a DataFrame to structure data for the heatmap
    data_to_plot = DataFrame(
        data=scores > threshold,
        columns=pd.Index(cols, name="Time (s)"),
        index=pd.Index(ch_names, name="Channel"),
    )

    # Define a custom colormap using provided color stops and base colors
    base_colors = ['red', 'red', 'white', 'white']
    colors = list(zip(color_stops[1], base_colors[:len(color_stops[1])]))
    cmap = mcolors.LinearSegmentedColormap.from_list('gyr', colors)

    # Plot heatmap of scores
    sns.heatmap( # type: ignore
        data=data_to_plot,
        vmin=0,
        vmax=1,
        cmap=cmap,
        cbar_kws=dict(label="Score"),
        ax=ax2,
    )

    # Add vertical dashed lines at each time boundary, sit the title, and place a black strikethrough through a bad channel
    for x in range(1, len(times)):
        ax2.axvline(x, ls="dashed", lw=0.25, dashes=(25, 15), color="gray") # type: ignore
    ax2.set_title("Scores > Threshold", fontweight="bold") # type: ignore
    markbad(data, ax2, ch_names)
    
    plt.close(fig2)

    return fig1, fig2



def scalp_coupling_index_windowed_raw(data, time_window: float = 3.0, l_freq: float = 0.7, h_freq: float = 1.5, l_trans_bandwidth: float = 0.3, h_trans_bandwidth: float = 0.3):
    """
    Compute windowed scalp coupling index (SCI) across fNIRS channels.

    Parameters
    ----------
    data : BaseRaw
        The loaded data object to process.
    time_window : float, optional
        Length of each time window in seconds (default is 3.0).
    l_freq : float, optional
        Low cutoff frequency for filtering in Hz (default is 0.7).
    h_freq : float, optional
        High cutoff frequency for filtering in Hz (default is 1.5).
    l_trans_bandwidth : float, optional
        Transition bandwidth for the low cutoff in Hz (default is 0.3).
    h_trans_bandwidth : float, optional
        Transition bandwidth for the high cutoff in Hz (default is 0.3).

    Returns
    -------
    tuple[BaseRaw, NDArray[float64], list[tuple[float, float]]]
        - BaseRaw: The original data object (unchanged). Ensures compatibility with peak_power().
        - NDArray[float64]: Correlation scores for each channel and time window.
        - list[tuple[float, float]]: Time intervals for each window in seconds.
    """

    # Pick only fNIRS channels and sort them by channel name
    sfreq: float = data.info["sfreq"]
    ch_names = data.ch_names
    times_arr = data.times

    # 1. Pick and sort fNIRS channel indices
    picks: NDArray[np.intp] = pick_types(cast(Info, data.info), fnirs=True)
    sort_idx = np.argsort([ch_names[p] for p in picks])
    picks = picks[sort_idx]

    # FIXME: This may happen if the heart rate calculation tries to set a value way too low
    if l_freq < 0.3:
        l_freq = 0.3

    # Band-pass filter the selected fNIRS channels
    filtered_data = filter_data(
        getattr(data, "_data"), 
        getattr(data, "info")["sfreq"],
        l_freq,
        h_freq,
        picks=picks,
        verbose=False,
        l_trans_bandwidth=l_trans_bandwidth, # type: ignore
        h_trans_bandwidth=h_trans_bandwidth, # type: ignore
    )

    # Calculate number of samples per time window, the total number of windows, and prepare output variables
    window_samples = int(np.ceil(time_window * sfreq))
    n_windows = int(np.floor(filtered_data.shape[1] / window_samples))
    total_samples = n_windows * window_samples

    starts = np.arange(n_windows) * window_samples
    stops = np.minimum(starts + window_samples, len(times_arr) - 1)
    times = [(times_arr[s], times_arr[e]) for s, e in zip(starts, stops)]

    # 5. Vectorized Correlation Calculation
    # Truncate to exact window boundary: shape (n_channels, n_windows, window_samples)
    truncated = filtered_data[picks, :total_samples].reshape(len(picks), n_windows, window_samples)

    # Pair channels: c1 (even rows), c2 (odd rows) -> shape (n_pairs, n_windows, window_samples)
    c1 = truncated[0::2]
    c2 = truncated[1::2]

    # Zero-mean center along the window sample axis
    c1_zero = c1 - np.mean(c1, axis=-1, keepdims=True)
    c2_zero = c2 - np.mean(c2, axis=-1, keepdims=True)

    # Standard deviations along window sample axis
    std1 = np.std(c1, axis=-1)
    std2 = np.std(c2, axis=-1)

    # Covariance along window sample axis
    cov = np.mean(c1_zero * c2_zero, axis=-1)
    denom = std1 * std2

    # Vectorized Pearson r calculation with NaN/Zero-std handling
    with np.errstate(divide="ignore", invalid="ignore"):
        corrs = np.where((denom == 0) | np.isnan(denom), 0.0, cov / denom)

    # Assign pair correlations back to output score matrix
    scores = np.zeros((len(picks), n_windows))
    scores[0::2, :] = corrs
    scores[1::2, :] = corrs

    # Revert scores to the original pick ordering if needed
    inv_sort = np.argsort(sort_idx)
    scores = scores[inv_sort]

    return data, scores, times



def calculate_scalp_coupling(data, l_freq: float = 0.7, h_freq: float = 1.5, time_window: int = 3, threshold: float = 0.6):
    """
    Calculate the scalp coupling index (SCI) and identify bad channels based on a threshold.

    Parameters
    ----------
    data : BaseRaw
        The loaded data object to process.
    l_freq : float, optional
        Low cutoff frequency for bandpass filtering in Hz (default is 0.7).
    h_freq : float, optional
        High cutoff frequency for bandpass filtering in Hz (default is 1.5)

    Returns
    -------
    tuple[list[str], Figure, Figure]
        - list[str]: Channel names identified as bad based on SCI threshold.
        - Figure: Heatmap of all SCI scores across time and channels.
        - Figure: Binary heatmap of SCI scores exceeding the threshold.
    """

    print("Calculating scalp coupling index...")

    # Compute the SCI
    _, scores, times = scalp_coupling_index_windowed_raw(data, time_window=time_window, l_freq=l_freq, h_freq=h_freq)

    # Identify channels that don't meet the provided threshold
    print("Identifying channels that do not meet the threshold...")
    sci = scores.mean(axis=1)
    data.info["bads"] = list(compress(cast(list[str], getattr(data, "ch_names")), sci < threshold))

    # Determine the colors based on the threshold, and create the figures
    print("Creating the figures...")
    color_stops = ([0.0, threshold, threshold+0.1, 0.8, 1.0], [0.0, threshold, threshold, 1.0])
    fig1, fig2 = plot_timechannel_quality_metrics(data, scores, times, color_stops, threshold, "Scalp Coupling Index")

    print("Successfully calculated scalp coupling index.")

    return list(compress(cast(list[str], getattr(data, "ch_names")), sci < threshold)), fig1, fig2



def build_fnirs_adjacency(raw, threshold_meters=0.03):
    """Build an adjacency dictionary for fNIRS channels using 3D distance."""
    # Extract channel positions
    ch_locs = []
    ch_names = []
    
    for ch in raw.info['chs']:
        loc = ch['loc'][:3]  # Get x, y, z coordinates
        if not np.isnan(loc).any():
            ch_locs.append(loc)
            ch_names.append(ch['ch_name'])
    
    ch_locs = np.array(ch_locs)
    
    # Compute pairwise distances
    dists = cdist(ch_locs, ch_locs)
    
    # Build adjacency dictionary
    adjacency = {}
    for i, ch_name in enumerate(ch_names):
        neighbors = [ch_names[j] for j in range(len(ch_names))
                     if 0 < dists[i, j] < threshold_meters]
        adjacency[ch_name] = neighbors
    
    return adjacency



def get_hbo_hbr_picks(raw):
    # Pick all fNIRS channels
    fnirs_picks = pick_types(raw.info, fnirs=True, exclude=[])

    # Extract wavelengths from channel names (expecting something like 'S6_D4 763' or 'S6_D4 841')
    wavelengths = []
    for idx in fnirs_picks:
        ch_name = raw.ch_names[idx]
        # Extract last 3 digits from channel name using regex
        match = re.search(r'(\d{3})$', ch_name)
        if match:
            wavelengths.append(int(match.group(1)))
        else:
            raise ValueError(f"Channel name '{ch_name}' does not end with 3 digits.")

    wavelengths = np.array(wavelengths)
    unique_wavelengths = np.unique(wavelengths)
    if len(unique_wavelengths) != 2:
        raise RuntimeError(f"Expected exactly 2 distinct wavelengths, found {unique_wavelengths}")

    # Determine which is HbO (larger) and which is HbR (smaller)
    hbr_wl = unique_wavelengths.min()
    hbo_wl = unique_wavelengths.max()

    print(f"HbR wavelength: {hbr_wl}, HbO wavelength: {hbo_wl}")

    # Find picks corresponding to each wavelength
    hbr_picks = [fnirs_picks[i] for i, wl in enumerate(wavelengths) if wl == hbr_wl]
    hbo_picks = [fnirs_picks[i] for i, wl in enumerate(wavelengths) if wl == hbo_wl]

    print(f"Found {len(hbr_picks)} HbR channels and {len(hbo_picks)} HbO channels.")

    return hbo_picks, hbr_picks, hbo_wl, hbr_wl



def interpolate_fNIRS_bads_weighted_average(raw, max_dist=0.03, min_neighbors=2, short_channels_threshold=0.015):
    """
    Interpolate bad fNIRS channels using a distance-weighted average of nearby good channels.

    Parameters
    ----------
    raw : mne.io.Raw
        The raw fNIRS data with bads marked in raw.info['bads'].
    max_dist : float
        Maximum distance (in meters) to consider for neighboring good channels.
    min_neighbors : int
        Minimum number of neighbors required to interpolate a bad channel.

    Returns
    -------
    raw : mne.io.Raw
        Modified raw object with bads interpolated (in-place).
    """

    print("Finding fNIRS channels...")
    hbo_picks, hbr_picks, hbo_wl, hbr_wl = get_hbo_hbr_picks(raw)


    if len(hbo_picks) != len(hbr_picks):
        raise RuntimeError("Number of HbO and HbR channels must be the same.")

    # Base names without wavelength for pairing
    def base_name(ch_name):
        # Strip last 4 chars assuming format ' <wavelength>'
        # e.g. "S6_D6 841" -> "S6_D6"
        return ch_name[:-4]

    hbo_names = [base_name(raw.ch_names[i]) for i in hbo_picks]
    hbr_names = [base_name(raw.ch_names[i]) for i in hbr_picks]

    # Sanity check: pairs must match
    for i in range(len(hbo_names)):
        if hbo_names[i] != hbr_names[i]:
            raise RuntimeError(f"Channel pairs do not match: {hbo_names[i]} vs {hbr_names[i]}")

    all_distances = source_detector_distances(raw.info)
    pair_distances = all_distances[hbo_picks]

    # Identify bad pairs if either channel in pair is bad
    bad_pairs = []
    good_pairs = []
    n_short_excluded = 0
    for i, base in enumerate(hbo_names):
        hbo_ch = raw.ch_names[hbo_picks[i]]
        hbr_ch = raw.ch_names[hbr_picks[i]]
        is_bad = (hbo_ch in raw.info['bads']) or (hbr_ch in raw.info['bads'])
        is_short = pair_distances[i] < short_channels_threshold

        if is_short:
            n_short_excluded += 1
        elif is_bad:
            bad_pairs.append(i)
        else:
            good_pairs.append(i)
 
    print(f"Total pairs: {len(hbo_names)}")
    print(f"Good LONG pairs (eligible donors): {len(good_pairs)}")
    print(f"Good SHORT pairs (excluded from donor pool): {n_short_excluded}")
    print(f"Bad pairs to interpolate: {len(bad_pairs)}")

    if len(bad_pairs) == 0:
        print("No bad pairs found. Skipping interpolation.")
        return raw, None, None

    raw_before_data = raw.get_data().copy()

    # Extract locations (use HbO channel loc as pair location)
    locs = np.array([raw.info['chs'][hbo_picks[i]]['loc'][:3] for i in range(len(hbo_names))])
    good_locs = locs[good_pairs]
    bad_locs = locs[bad_pairs]

    # Compute distance matrix between bad and good pairs
    dist_matrix = cdist(bad_locs, good_locs)

    interpolated_pairs = []

    for i, bad_idx in enumerate(bad_pairs):
        bad_base = hbo_names[bad_idx]
        distances = dist_matrix[i]
        close_idxs = np.where(distances < max_dist)[0]

        print(f"\nInterpolating pair {bad_base} (index {bad_idx})")
        print(f"  Nearby good pairs found: {len(close_idxs)}")

        if len(close_idxs) < min_neighbors:
            print(f"  Skipping {bad_base}: not enough neighbors (found {len(close_idxs)} < {min_neighbors})")
            continue

        weights = 1 / (distances[close_idxs] + 1e-6)
        weights /= weights.sum()

        neighbor_hbo_indices = [hbo_picks[good_pairs[idx]] for idx in close_idxs]
        neighbor_hbr_indices = [hbr_picks[good_pairs[idx]] for idx in close_idxs]

        neighbor_hbo_data = raw._data[neighbor_hbo_indices, :]
        neighbor_hbr_data = raw._data[neighbor_hbr_indices, :]

        interpolated_hbo = np.average(neighbor_hbo_data, axis=0, weights=weights)
        interpolated_hbr = np.average(neighbor_hbr_data, axis=0, weights=weights)

        raw._data[hbo_picks[bad_idx]] = interpolated_hbo
        raw._data[hbr_picks[bad_idx]] = interpolated_hbr

        interpolated_pairs.append(bad_base)

    n_bad = len(bad_pairs)
    n_cols = 4  # Fixed width for horizontal scaling
    n_rows = int(np.ceil(n_bad / n_cols))
    
    # Calculate height: 2.5 inches per row is usually enough for readability
    fig_height = max(4, n_rows * 2.5)
    fig_compare, axes = plt.subplots(n_rows, n_cols, figsize=(15, fig_height), 
                                     constrained_layout=True)
    if n_bad == 1: axes = [axes] # Handle single subplot case
    
    axes_flat = np.asarray(axes).ravel()
    for j in range(n_bad, len(axes_flat)):
        if j >= n_bad:
            axes_flat[j].axis('off')

    times = raw.times
    for i, bad_idx in enumerate(bad_pairs):
        ax = axes_flat[i]
        base = hbo_names[bad_idx]
        
        # Plot "Before" (Dirty data) in light gray/red
        ax.plot(times, raw_before_data[hbo_picks[bad_idx]], color='red', alpha=0.3, label='Original HbO')
        
        # Plot "After" (Interpolated data) in solid blue/green
        if base in interpolated_pairs:
            ax.plot(times, raw._data[hbo_picks[bad_idx]], color='blue', label='Interpolated HbO')
            status = "SUCCESS"
            color = "green"
        else:
            status = "FAILED (Isolated)"
            color = "red"
            
        ax.set_title(f"Channel {base} | Status: {status}", color=color, fontweight='bold')
        ax.set_ylabel("Amplitude")
        if i == 0: ax.legend(loc='upper right')

    plt.close(fig_compare)

    if interpolated_pairs:
        bad_ch_to_remove = []
        for base_ in interpolated_pairs:
            bad_ch_to_remove.append(base_ + f" {hbr_wl}")  # HbR
            bad_ch_to_remove.append(base_ + f" {hbo_wl}")  # HbO

        raw.info['bads'] = [ch for ch in raw.info['bads'] if ch not in bad_ch_to_remove]

    print("\nInterpolation complete.\n")
    print("Bads cleared:", raw.info['bads'])
    raw.info['bads'] = []

    fig_raw_after = raw.plot(duration=raw.times[-1], n_channels=raw.info['nchan'], title="After interpolation", show=False)

    return raw, fig_raw_after, fig_compare



def calculate_signal_noise_ratio(data):
    """
    Calculates the signal-to-noise ratio (SNR) for each channel and identifies those below a defined threshold.

    Parameters
    ----------
    data : BaseRaw
        The loaded data object to process.

    Returns
    -------
    tuple[list[str], Figure]
        - list[str]: A list of channel names that fall below the SNR threshold and are considered bad.
        - Figure: A matplotlib Figure showing the channels' SNR values.
    """
    
    print("Calculating signal to noise ratio...")

    # Compute the signal-to-noise ratio values
    print("Computing the signal to noise ratio...")
    signal_band=(0.01, 0.5)
    noise_band=(1.0, 10.0)
    data_signal = data.copy().filter(*signal_band, verbose=False) #type: ignore
    data_noise = data.copy().filter(*noise_band, verbose=False) #type: ignore
    signal_power = np.mean(data_signal.get_data()**2, axis=1) #type: ignore
    noise_power = np.mean(data_noise.get_data()**2, axis=1) #type: ignore

    # Calculate the snr using the standard formula for dB
    snr = 10 * np.log10(signal_power / (noise_power + np.finfo(float).eps))

    # TODO: Understand what this does
    groups: dict[str, list[str]] = {}
    for ch in getattr(data, "ch_names"):
        # Look for the space in the channel names and remove the characters after
        # This is so we can get both oxy and deoxy to remove, as they will have the same source and detector
        base = ch.rsplit(' ', 1)[0]
        groups.setdefault(base, []).append(ch) # type: ignore

    # If any of the channels do not meet our threshold, they will get inserted into the bad_channels set
    bad_channels: set[str] = set()
    for base, ch_list in groups.items():
        if any(s < SNR_THRESHOLD for s, ch in zip(snr, getattr(data, "ch_names")) if ch in ch_list):
            bad_channels.update(ch_list)

    # Design and create the figure
    print("Creating the figure...")
    snr_fig, ax = plt.subplots(figsize=(12, 4), layout="constrained") # type: ignore
    colors = [(0/20, 'red'), (SNR_THRESHOLD/20, 'red'), ((SNR_THRESHOLD+.5)/20, 'yellow'), ((SNR_THRESHOLD+1)/20, 'green'), (20/20, 'green')]
    cmap = LinearSegmentedColormap.from_list('custom_snr_cmap', colors)
    norm = mcolors.Normalize(vmin=0, vmax=20)
    scatter = ax.scatter(range(len(snr)), snr, c=snr, cmap=cmap, alpha=0.8, s=100, norm=norm) # type: ignore
    ax.set(xlabel="Channel Number",  ylabel="Signal-to-Noise Ratio (dB)", xlim=[0, len(snr)], ylim=[0, 20])
    ax.axhline(SNR_THRESHOLD, color='black', linestyle='--', alpha=0.3, linewidth=1) # type: ignore
    cbar = snr_fig.colorbar(scatter, ax=ax, label="SNR Thresholds (dB)") # type: ignore
    cbar.set_ticks([0, SNR_THRESHOLD, SNR_THRESHOLD+1, 20]) # type: ignore
    cbar.set_ticklabels(['0', str(SNR_THRESHOLD), str(SNR_THRESHOLD+1), '20']) # type: ignore

    plt.close()

    print("Successfully calculated signal to noise ratio.")

    return list(bad_channels), snr_fig



def calculate_peak_power(data: BaseRaw, time_window: int = 3, threshold: float = 0.1, l_freq: float = 0.7, h_freq: float = 1.5) -> tuple[list[str], Figure, Figure]:
    """
    Calculate peak spectral power (PSP) for fNIRS channels and identify bad channels.

    Parameters
    ----------
    data : BaseRaw
        The loaded data object to process.
    l_freq : float, optional
        Low cutoff frequency for filtering in Hz (default is 0.7)
    h_freq : float, optional
        High cutoff frequency for filtering in Hz (default is 1.5)

    Returns
    -------
    tuple[list[str], Figure, Figure]
        - list[str]: Names of channels below the PSP threshold.
        - Figure: Heatmap of all PSP scores.
        - Figure: Heatmap of scores above the PSP threshold.
    """

    
    # Compute the PSP
    _, scores, times = cast(tuple[NDArray[float64], NDArray[float64], list[tuple[float]]], peak_power_fast(data, time_window=time_window, threshold=threshold, l_freq=l_freq, h_freq=h_freq))

    # Identify channels that don't meet the provided threshold
    psp = scores.mean(axis=1)
    bad_channels = list(compress(cast(list[str], data.ch_names), psp < threshold))

    plot_data = data.copy()
    plot_data.info["bads"] = bad_channels

    existing_bads = set(data.info.get("bads", []))
    data.info["bads"] = list(existing_bads | set(bad_channels))

    # Determine the colors based on the threshold, and create the figures
    color_stops = ([0.0, threshold, threshold+0.1, threshold+0.2, 1.0], [0.0, threshold, threshold, 1.0])
    psp1, psp2 = plot_timechannel_quality_metrics(plot_data, scores, times, color_stops, threshold, "Peak Spectral Power")


    return list(compress(cast(list[str], getattr(data, "ch_names")), psp < threshold)), psp1, psp2



def mark_bads(raw, bad_sci, bad_snr, bad_psp,  bad_coeff_var, bad_range, bad_noise, bad_disp):
    print(bad_sci, bad_snr, bad_psp, bad_coeff_var, bad_range, bad_noise, bad_disp)
    bads_combined = list(set(bad_snr) | set(bad_sci) | set(bad_psp) | set(bad_coeff_var) | set(bad_range) | set(bad_noise) | set(bad_disp))
    print(f"Automatically marked bad channels based on SNR and SCI: {bads_combined}")

    raw.info['bads'].extend(bads_combined)

    # Organize channels into categories
    sets = [
        (bad_sci, "SCI"),
        (bad_psp, "PSP"),
        (bad_snr, "SNR"),
        (bad_coeff_var, "coeff_var"),
        (bad_range, "Range"),
        (bad_noise, "Noise"),
        (bad_disp, "Disp.")
    ]

    # Graph what channels were dropped and why they were dropped
    channel_categories: dict[str, str] = {}

    for ch in bads_combined:
        present_in = [name for s, name in sets if ch in s]
        # Create a label for the category
        if len(present_in) == 1:
            label = f"{present_in[0]} only"
        else:
            label = " + ".join(sorted(present_in))
        channel_categories[ch] = label

    # Sort channels alphabetically within categories for nicer visualization
    categories = sorted(set(channel_categories.values()))
    channel_names: list[str] = []
    category_labels: list[str] = []
    for cat in categories:
        chs_in_cat = sorted([ch for ch, c in channel_categories.items() if c == cat])
        channel_names.extend(chs_in_cat)
        category_labels.extend([cat] * len(chs_in_cat))

    colors = {cat: get_category_color(cat) for cat in categories}
    # Create the figure
    fig_dropped, ax = plt.subplots(figsize=(10, max(3, len(channel_names) * 0.3))) # type: ignore
    y_pos = range(len(channel_names))
    ax.barh(y_pos, [1]*len(channel_names), color=[colors[cat] for cat in category_labels]) # type: ignore
    ax.set_yticks(y_pos) # type: ignore
    ax.set_yticklabels(channel_names) # type: ignore
    ax.set_xlabel("Marked as Bad") # type: ignore
    ax.set_title(f"Bad Channels by Method for") # type: ignore
    ax.set_xlim(0, 1)
    ax.set_xticks([])  # type: ignore
    ax.grid(axis='x', linestyle='--', alpha=0.7) # type: ignore

    # Add a legend denoting why the channels were bad
    for label, color in colors.items():
        ax.bar(0, 0, color=color, label=label) # type: ignore
    ax.legend() # type: ignore

    fig_dropped.tight_layout()


    raw_before = raw.copy()
    bads_channels = [ch for ch in raw.ch_names if ch in raw.info['bads']]
    print(bads_channels)
    if bads_channels:
        fig_raw_before = raw_before.plot(duration=raw.times[-1], n_channels=raw.info['nchan'], picks=bads_channels, title="What they were BEFORE", show=False)
    else:
        fig_dropped = None
        fig_raw_before = None

    return raw, fig_dropped, fig_raw_before, bads_channels



def filter_the_data(
        raw_haemo,
        filter_algorithm,
        l_freq,
        h_freq,
        l_trans_bandwidth,
        h_trans_bandwidth,
        # iir_type,
        # iir_order,
        filter_length,
        filter_phase,
        fir_window,
        fir_design,
        # iir_output,
        # passband_ripple,
        # stopband_attenuation,
        filter_pad,
        # skip_by_annotation,
        filter_n_jobs,
        verbosity
    ):

    fig_filter = raw_haemo.compute_psd(fmax=3).plot(
        average=True, color="r", show=False, amplitude=True
    )
    
    if l_freq == 0 and h_freq != 0:
        if filter_algorithm == "fir":
            raw_haemo = raw_haemo.filter(l_freq=None, h_freq=h_freq, filter_length=filter_length, h_trans_bandwidth=h_trans_bandwidth, n_jobs=filter_n_jobs, method=filter_algorithm, phase=filter_phase, fir_window=fir_window, fir_design=fir_design, pad=filter_pad, verbose=verbosity)
        else:
            raw_haemo = raw_haemo.filter(l_freq=None, h_freq=h_freq, filter_length=filter_length, h_trans_bandwidth=h_trans_bandwidth, n_jobs=filter_n_jobs, method=filter_algorithm, pad=filter_pad, verbose=verbosity)

    elif l_freq != 0 and h_freq == 0:
        if filter_algorithm == "fir":
            raw_haemo = raw_haemo.filter(l_freq=l_freq, h_freq=None, filter_length=filter_length, l_trans_bandwidth=l_trans_bandwidth, n_jobs=filter_n_jobs, method=filter_algorithm, phase=filter_phase, fir_window=fir_window, fir_design=fir_design, pad=filter_pad, verbose=verbosity)
        else:
            raw_haemo = raw_haemo.filter(l_freq=l_freq, h_freq=None, filter_length=filter_length, l_trans_bandwidth=l_trans_bandwidth, n_jobs=filter_n_jobs, method=filter_algorithm, pad=filter_pad, verbose=verbosity)

    elif l_freq != 0 and h_freq != 0:
        if filter_algorithm == "fir":
            raw_haemo = raw_haemo.filter(l_freq=l_freq, h_freq=h_freq, filter_length=filter_length, l_trans_bandwidth=l_trans_bandwidth, h_trans_bandwidth=h_trans_bandwidth, n_jobs=filter_n_jobs, method=filter_algorithm, phase=filter_phase, fir_window=fir_window, fir_design=fir_design, pad=filter_pad, verbose=verbosity)
        else:
            raw_haemo = raw_haemo.filter(l_freq=l_freq, h_freq=h_freq, filter_length=filter_length, l_trans_bandwidth=l_trans_bandwidth, h_trans_bandwidth=h_trans_bandwidth, n_jobs=filter_n_jobs, method=filter_algorithm, pad=filter_pad, verbose=verbosity)
    else:
        print("No filter")

    raw_haemo.compute_psd(fmax=3).plot(
        average=True, axes=fig_filter.axes, color="g", amplitude=True, show=False
    )

    fig_raw_haemo_filter = raw_haemo.plot(duration=raw_haemo.times[-1], n_channels=raw_haemo.info['nchan'], title="Filtered HbO and HbR", show=False)

    return raw_haemo, fig_filter, fig_raw_haemo_filter



def safe_create_epochs(raw, events, event_dict, tmin, tmax, baseline, max_shift, reject_epochs, reject_hbo_threshold):
    """
    Attempts to create epochs, shifting event times slightly if 
    sample collisions are detected.
    """
    shift_increment = 1.0 / raw.info['sfreq']  # The duration of exactly one sample

    for attempt in range(max_shift):  # Limit attempts to avoid infinite loops
        try:
            if reject_epochs:
                epochs = Epochs(
                    raw, events, event_id=event_dict,
                    tmin=tmin, tmax=tmax, baseline=baseline,
                    reject=reject_hbo_threshold,
                    preload=True, verbose=False
                )
            else:
                epochs = Epochs(
                    raw, events, event_id=event_dict,
                    tmin=tmin, tmax=tmax, baseline=baseline,
                    preload=True, verbose=False
                )
            return epochs
        
        except RuntimeError as e:
            if "Event time samples were not unique" in str(e):
                # Find duplicates in the events array (column 0 is the sample index)
                vals, counts = np.unique(events[:, 0], return_counts=True)
                duplicates = vals[counts > 1]
                
                # Shift the second occurrence of every duplicate by 1 sample
                for dup in duplicates:
                    idx = np.where(events[:, 0] == dup)[0][1:] # Get all but the first
                    events[idx, 0] += 1 
                
                print(f"Collision detected. Nudging events by {shift_increment:.4f}s and retrying...")
                continue
            else:
                raise e # Raise if it's a different Runtime Error
                
    raise RuntimeError("Could not resolve event collisions after 10 attempts.")



def epochs_calculations(raw_haemo, events, event_dict, epoch_handling, max_shift, t_min, t_max, baseline, reject_epochs, reject_hbo_threshold, png_queue):

    if epoch_handling == 'shift':
        epochs = safe_create_epochs(raw=raw_haemo, events=events, event_dict=event_dict, tmin=t_min, tmax=t_max, baseline=baseline, max_shift=max_shift, reject_epochs=reject_epochs, reject_hbo_threshold=reject_hbo_threshold)
    else:
        if reject_epochs:
            epochs = Epochs(raw_haemo, events, event_id=event_dict, tmin=t_min, tmax=t_max, baseline=baseline, reject=reject_hbo_threshold)
        else:
            epochs = Epochs(raw_haemo, events, event_id=event_dict, tmin=t_min, tmax=t_max, baseline=baseline)

    # Make a copy of the epochs and drop bad ones
    epochs2 = epochs.copy()
    epochs2.drop_bad()

    # Plot drop log
    # TODO: Why show this if we never use epochs2?
    fig_epochs_dropped = epochs2.plot_drop_log(show=False)
    _enqueue("epochs_fig_epochs_dropped", fig_epochs_dropped, png_queue)

    conditions = list(epochs.event_id.keys())
    evoked_cache = {}

    # Plot for each condition
    for idx, condition in enumerate(epochs.event_id.keys()):
        logger.info(condition)
        logger.info(idx)

        epo_cond = epochs[condition]

        # Plot images for each condition
        fig_epochs_data = epochs[condition].plot_image(
            combine="mean",
            vmin=-1,
            vmax=1,
            ts_args=dict(ylim=dict(hbo=[-1, 1], hbr=[-1, 1])),
            show=False
        )
        for j, fig in enumerate(fig_epochs_data):
            logger.info("------------------------------------------")
            logger.info(j)
            logger.info(fig)

            ax = fig.axes[0]
            original_title = ax.get_title()
            ax.set_title(f"{condition}: {original_title}")
            _enqueue(f"epochs_fig_{condition}_data_{idx}_{j}", fig, png_queue)

        # Evoked average figure for each condition
        evoked_avg = epo_cond.average()
        evoked_cache[condition] = evoked_avg

        clims = dict(hbo=[-1, 1], hbr=[1, -1])
        condition_fig = evoked_avg.plot_image(clim=clims, show=False)
        
        for ax in condition_fig.axes:
            original_title = ax.get_title()
            ax.set_title(f"{original_title} - {condition}")
        _enqueue(f"epochs_evoked_avg_{condition}", condition_fig, png_queue)

    # Prepare evokeds and colors for topographic plot
    evokeds3 = []
    colors = []
    cmap = plt.get_cmap("tab10", len(conditions))

    for idx, cond in enumerate(conditions):
        evoked = epochs[cond].average(picks="hbo")
        evokeds3.append(evoked)
        colors.append(cmap(idx))

    # Create the topographic plot
    fig, axes = plt.subplots(nrows=1, ncols=1, figsize=(6, 4))
    help = plot_evoked_topo(evokeds3, color=colors, axes=axes, legend=False, show=False)

    # Build custom legend
    lines = []
    for color in colors:
        line = plt.Line2D([0], [0], color=color, lw=2)
        lines.append(line)

    fig.legend(lines, conditions, loc="lower right")
    _enqueue("epochs_evoked_topo", help, png_queue)

    unique_annotations = set(raw_haemo.annotations.description)

    for cond in unique_annotations:

        # Evoked response for specific condition ("Activity")
        evoked_stim1 = evoked_cache.get(cond)
        if evoked_stim1 is None:
            # Not one of the epoch conditions already averaged above - fall back
            # to computing it directly (matches original behavior in that case).
            evoked_stim1 = epochs[cond].average()

        fig_evoked_hbo = evoked_stim1.copy().pick(picks='hbo').plot(time_unit='s', show=False)
        fig_evoked_hbr = evoked_stim1.copy().pick(picks='hbr').plot(time_unit='s', show=False)
        
        _enqueue(f"epochs_fig_evoked_hbo_{cond}", fig_evoked_hbo, png_queue)
        _enqueue(f"epochs_fig_evoked_hbr_{cond}", fig_evoked_hbr, png_queue)

        print("Evoked HbO peak amplitude:", evoked_stim1.copy().pick(picks='hbo').data.max())

    evokeds = {}
    for condition in epochs2.event_id:
        evokeds[condition] = epochs2[condition].average()
        print(f"Condition '{condition}': {len(epochs2[condition])} epochs averaged.")

    all_evokeds = {}
    for condition in epochs.event_id:
        if condition not in all_evokeds:
            all_evokeds[condition] = []
            all_evokeds[condition].append(evoked_cache[condition])

    group_averages = {cond: evoked_cache[cond] for cond in conditions if cond in evoked_cache}
    group_aucs = {}

    for condition, evoked in group_averages.items():
        group_aucs[condition] = {}
        for pick in ["hbo", "hbr"]:
            picks_idx = [i for i, ch in enumerate(evoked.ch_names) if pick in ch]
            if not picks_idx:
                continue

            data = evoked.data[picks_idx, :].mean(axis=0)
            t_start, t_end = 0, 15 #TODO: Is this in seconds? or is it 1hz input that makes it 15s?
            times_mask = (evoked.times >= t_start) & (evoked.times <= t_end)
            data_segment = data[times_mask]
            times_segment = evoked.times[times_mask]

            auc = np.trapezoid(data_segment, times_segment)
            group_aucs[condition][pick] = auc

    # Final evoked comparison plot for each condition
    for condition in conditions:
        if condition not in evokeds:
            continue
        evoked = evokeds[condition]

        fig, ax = plt.subplots(figsize=(6, 5))
        legend_labels = ["Oxyhaemoglobin"]

        for pick, color in zip(["hbo", "hbr"], ["r", "b"]):
            plot_compare_evokeds(
                evoked,
                combine="mean",
                picks=pick,
                axes=ax,
                show=False,
                colors=[color],
                legend=False,
                title=f"Participant: nCondition: {condition}",
                ylim=dict(hbo=[-0.5, 1], hbr=[-0.5, 1]),
                show_sensors=False,
            )
            auc_value = group_aucs.get(condition, {}).get(pick, None)
            if auc_value is not None:
                label = f"{pick.upper()} AUC: {auc_value * 1e6:.4f} µM·s"
            else:
                label = f"{pick.upper()} AUC: N/A"
            legend_labels.append(label)
            if len(legend_labels) == 2:
                legend_labels.append("Deoxyhaemoglobin")

        ax.legend(legend_labels)

        _enqueue(f"epochs_fig_{condition}_compare_evokeds", fig, png_queue)

    return epochs



def make_design_matrix(
    raw_haemo,
    resample,
    resample_freq,
    stim_dur,
    hrf_model,
    drift_model,
    high_pass,
    drift_order,
    fir_delays,
    min_onset,
    oversampling,
    short_channel_regression,
    short_channels,
    long_channels,
    short_channels_threshold,
    long_channels_threshold,
    folding_bypass
):

    # events_to_remove = REMOVE_EVENTS
    events_to_remove = ""

    filtered_annotations = [ann for ann in raw_haemo.annotations if ann['description'] not in events_to_remove]
    
    new_annot = Annotations(
        onset=[ann['onset'] for ann in filtered_annotations],
        duration=[ann['duration'] for ann in filtered_annotations],
        description=[ann['description'] for ann in filtered_annotations]
    )
    
    if short_channels:
        short_chans = get_short_channels(raw_haemo, max_dist=short_channels_threshold)
        if long_channels:
            raw_haemo = get_long_channels(raw_haemo, min_dist=short_channels_threshold, max_dist=long_channels_threshold)

    else:
        short_chans = None

    raw_haemo.set_annotations(new_annot)
    raw_haemo_dm = raw_haemo.copy()

    if resample:
        raw_haemo_dm.resample(resample_freq, npad="auto")
        try:
            short_chans.resample(resample_freq)
        except:
            pass

    raw_haemo_dm._data = raw_haemo_dm._data * 1e6

    design_matrix = make_first_level_design_matrix(
        raw=raw_haemo_dm,
        stim_dur=stim_dur,
        hrf_model=hrf_model,
        drift_model=drift_model,
        high_pass=high_pass,
        drift_order=drift_order,
        fir_delays=fir_delays,
        min_onset=min_onset,
        oversampling=oversampling
    )

    # 3) Average and Append Short Channels
    if short_channel_regression and not folding_bypass:
        if short_chans is not None and len(short_chans.ch_names) > 0:
            ch_types = short_chans.get_channel_types()
            
            # Scenario A: Short channels are already converted to Hemoglobin (hbo/hbr)
            if "hbo" in ch_types or "hbr" in ch_types:
                hbo_chs = [ch for ch, t in zip(short_chans.ch_names, ch_types) if t == "hbo"]
                hbr_chs = [ch for ch, t in zip(short_chans.ch_names, ch_types) if t == "hbr"]
                
                if hbo_chs:
                    hbo_data = short_chans.copy().pick(hbo_chs).get_data()
                    design_matrix["ShortHbO"] = np.mean(hbo_data, axis=0)
                if hbr_chs:
                    hbr_data = short_chans.copy().pick(hbr_chs).get_data()
                    design_matrix["ShortHbR"] = np.mean(hbr_data, axis=0)
                print(f"Successfully added averaged ShortHbO ({len(hbo_chs)} chs) and ShortHbR ({len(hbr_chs)} chs) to the matrix.")
                
            # Scenario B: Short channels are raw wavelengths (760nm, 850nm, etc.)
            else:
                wavelength_groups = {}
                for ch_name in short_chans.ch_names:
                    # Look for the wavelength number (digits) at the end of the channel name
                    match = re.search(r'(\d+)$', ch_name)
                    if match:
                        wl = match.group(1)
                        wavelength_groups.setdefault(wl, []).append(ch_name)
                
                if wavelength_groups:
                    for wl, chs in wavelength_groups.items():
                        wl_data = short_chans.copy().pick(chs).get_data()
                        col_name = f"Short_{wl}"
                        design_matrix[col_name] = np.mean(wl_data, axis=0)
                    print(f"Successfully added averaged short channels by wavelength: {list(wavelength_groups.keys())}")
                else:
                    # Emergency fallback: if names have no digits, average all of them together
                    design_matrix["Short_Avg"] = np.mean(short_chans.get_data(), axis=0)
                    print("Could not detect wavelengths. Averaged all short channels into 'Short_Avg'.")
        else:
            print("Warning: SHORT_CHANNEL_REGRESSION is True, but no short channels were found.")

    print(design_matrix.head())
    print(design_matrix.columns)

    fig, ax1 = plt.subplots(figsize=(10, 6), constrained_layout=True)
    _ = plot_design_matrix(design_matrix, axes=ax1)

    return raw_haemo, raw_haemo_dm, design_matrix, fig



def generate_montage_locations():
    """Get standard MNI montage locations in dataframe.

    Data is returned in the same format as the eeg_positions library.
    """
    # standard_1020 and standard_1005 are in MNI (fsaverage) space already,
    # but we need to undo the scaling that head_scale will do
    montage = make_standard_montage(
        "standard_1005", head_size=0.09700884729534559
    )
    for d in montage.dig:
        d["coord_frame"] = 2003
    montage.dig[:] = montage.dig[3:]
    montage.add_mni_fiducials()  # now in fsaverage space
    coords = pd.DataFrame.from_dict(montage.get_positions()["ch_pos"]).T
    coords["label"] = coords.index
    coords = coords.rename(columns={0: "x", 1: "y", 2: "z"})

    return coords.reset_index(drop=True)



def _find_closest_standard_location(position, reference, *, out="label"):
    """Return closest montage label to coordinates.

    Parameters
    ----------
    position : array, shape (3,)
        Coordinates.
    reference : dataframe
        As generated by _generate_montage_locations.
    trans_pos : str
        Apply a transformation to positions to specified frame.
        Use None for no transformation.
    """

    p0 = np.array(position)
    p0.shape = (-1, 3)
    # head_mri_t, _ = _get_trans("fsaverage", "head", "mri")
    # p0 = apply_trans(head_mri_t, p0)
    dists = cdist(p0, np.asarray(reference[["x", "y", "z"]], float))

    if out == "label":
        min_idx = np.argmin(dists)
        return reference["label"][min_idx]
    else:
        assert out == "dists"
        return dists
    


def _source_detector_fold_table(raw, cidx, reference, fold_tbl, interpolate):
    src = raw.info["chs"][cidx]["loc"][3:6]
    det = raw.info["chs"][cidx]["loc"][6:9]

    ref_lab = list(reference["label"])
    dists = _find_closest_standard_location([src, det], reference, out="dists")
    src_min, det_min = np.argmin(dists, axis=1)
    src_name, det_name = ref_lab[src_min], ref_lab[det_min]

    tbl = fold_tbl.query("Source == @src_name and Detector == @det_name")
    dist = np.linalg.norm(dists[[0, 1], [src_min, det_min]])
    # Try reversing source and detector
    if len(tbl) == 0:
        tbl = fold_tbl.query("Source == @det_name and Detector == @src_name")
    if len(tbl) == 0 and interpolate:
        # Try something hopefully not too terrible: pick the one with the
        # smallest net distance
        good = np.isin(fold_tbl["Source"], reference["label"]) & np.isin(
            fold_tbl["Detector"], reference["label"]
        )
        assert good.any()
        tbl = fold_tbl[good]
        assert len(tbl)
        src_idx = [ref_lab.index(src) for src in tbl["Source"]]
        det_idx = [ref_lab.index(det) for det in tbl["Detector"]]
        # Original
        tot_dist = np.linalg.norm([dists[0, src_idx], dists[1, det_idx]], axis=0)
        assert tot_dist.shape == (len(tbl),)
        idx = np.argmin(tot_dist)
        dist_1 = tot_dist[idx]
        src_1, det_1 = ref_lab[src_idx[idx]], ref_lab[det_idx[idx]]
        # And the reverse
        tot_dist = np.linalg.norm([dists[0, det_idx], dists[1, src_idx]], axis=0)
        idx = np.argmin(tot_dist)
        dist_2 = tot_dist[idx]
        src_2, det_2 = ref_lab[det_idx[idx]], ref_lab[src_idx[idx]]
        if dist_1 < dist_2:
            new_dist, src_use, det_use = dist_1, src_1, det_1
        else:
            new_dist, src_use, det_use = dist_2, det_2, src_2


        tbl = fold_tbl.query("Source == @src_use and Detector == @det_use")
        tbl = tbl.copy()
        tbl["BestSource"] = src_name
        tbl["BestDetector"] = det_name
        tbl["BestMatchDistance"] = dist
        tbl["MatchDistance"] = new_dist
        assert len(tbl)
    else:
        tbl = tbl.copy()
        tbl["BestSource"] = src_name
        tbl["BestDetector"] = det_name
        tbl["BestMatchDistance"] = dist
        tbl["MatchDistance"] = dist

    tbl = tbl.copy()  # don't get warnings about setting values later
    return tbl



def _read_fold_xls(fname, atlas="Juelich"):
    """Read fOLD toolbox xls file.

    The values are then manipulated in to a tidy dataframe.

    Note the xls files are not included as no license is provided.

    Parameters
    ----------
    fname : str
        Path to xls file.
    atlas : str
        Requested atlas.
    """
    page_reference = {"AAL2": 2, "AICHA": 5, "Brodmann": 8, "Juelich": 11, "Loni": 14}

    tbl = pd.read_excel(fname, sheet_name=page_reference[atlas])

    # Remove the spacing between rows
    empty_rows = np.where(np.isnan(tbl["Specificity"]))[0]
    tbl = tbl.drop(empty_rows).reset_index(drop=True)

    # Empty values in the table mean its the same as above
    for row_idx in range(1, tbl.shape[0]):
        for col_idx, col in enumerate(tbl.columns):
            if not isinstance(tbl[col][row_idx], str):
                if np.isnan(tbl[col][row_idx]):
                    tbl.iloc[row_idx, col_idx] = tbl.iloc[row_idx - 1, col_idx]

    tbl["Specificity"] = tbl["Specificity"] * 100
    tbl["brainSens"] = tbl["brainSens"] * 100
    return tbl



def _check_load_fold(fold_files, atlas):
    # _validate_type(fold_files, (list, "path-like", None), "fold_files")
    if fold_files is None:
        fold_files = get_config("MNE_NIRS_FOLD_PATH")
        if fold_files is None:
            raise ValueError(
                "MNE_NIRS_FOLD_PATH not set, either set it using "
                "mne.set_config or pass fold_files as str or list"
            )
    if not isinstance(fold_files, list):  # path-like
        fold_files = _check_fname(
            fold_files,
            overwrite="read",
            must_exist=True,
            name="fold_files",
            need_dir=True,
        )
        fold_files = [op.join(fold_files, f"10-{x}.xls") for x in (5, 10)]

    fold_tbl = pd.DataFrame()
    for fi, fname in enumerate(fold_files):
        fname = _check_fname(
            fname, overwrite="read", must_exist=True, name=f"fold_files[{fi}]"
        )
        fold_tbl = pd.concat(
            [fold_tbl, _read_fold_xls(fname, atlas=atlas)], ignore_index=True
        )
    return fold_tbl

                

def fold_channels(raw: BaseRaw, p_name: str, atlas: str='Brodmann', progress_queue: Optional[Any]=None) -> dict[str, list[dict[str, Any]]]:
    """Runs in background process. 
    Does only heavy math/lookup. Returns data instead of a static image.
    """
    if getattr(sys, 'frozen', False):
        fold_dir = resource_path("./mne_data/fOLD/fOLD-public-master/Supplementary")
    else:
        path = os.path.expanduser("~") + "/mne_data/fOLD/fOLD-public-master/Supplementary"
        fold_dir = resource_path(path)

    set_config('MNE_NIRS_FOLD_PATH', fold_dir)

    hbo_channel_names = cast(list[str], getattr(raw.copy().pick(picks='hbo'), "ch_names"))
    
    _validate_type(raw, BaseRaw, "raw")

    reference_locations = generate_montage_locations()

    fold_tbl = _check_load_fold(fold_files=fold_dir, atlas=atlas)

    channel_results = {}
    step_idx = 0

    for channel_name in hbo_channel_names:
        cidx = raw.ch_names.index(channel_name)
        tbl = _source_detector_fold_table(
            raw, cidx, reference_locations, fold_tbl, interpolate=True
        )

        channel_results[channel_name] = []

        for _, row in tbl.iterrows():
            channel_results[channel_name].append({
                'Landmark': str(row['Landmark']),
                'Specificity': float(row['Specificity'])
            })

        step_idx += 1
        if progress_queue is not None:
            progress_queue.put((p_name, step_idx))

    # Return raw data dictionary to the result_queue
    return channel_results



def plot_glm_results(file_path, raw_haemo, glm_est, design_matrix):
    
    fig_glms = []  # List to store figures

    dm = design_matrix.copy()
    logger.info(design_matrix.shape)
    logger.info(design_matrix.columns)
    logger.info(design_matrix.head())

    rois = dict(AllChannels=range(len(raw_haemo.ch_names)))
    conditions = design_matrix.columns
    df_individual = glm_est.to_dataframe_region_of_interest(rois, conditions)
    
    df_individual["ID"] = file_path
    # df_individual["theta"] = [t * 1.0e6 for t in df_individual["theta"]]

    first_onset_for_cond = {}
    for onset, desc in zip(raw_haemo.annotations.onset, raw_haemo.annotations.description):
        if desc not in first_onset_for_cond:
            first_onset_for_cond[desc] = onset

    # Get unique condition names from annotations (descriptions)
    unique_annotations = set(raw_haemo.annotations.description)

    for cond in unique_annotations:
        logger.info(cond)
        df_individual_filtered = df_individual.copy()

        # Filter for the condition of interest and FIR delays
        df_individual_filtered["isCondition"] = [cond in n for n in df_individual_filtered["Condition"]]
        df_individual_filtered["isDelay"] = ["delay" in n for n in df_individual_filtered["Condition"]]
        df_individual_filtered = df_individual_filtered.query("isDelay and isCondition")
        
        # Remove other conditions from design matrix
        dm_condition_cols = [col for col in dm.columns if cond in col]
        dm_cond = dm[dm_condition_cols]
            
        # Add a numeric delay column
        def extract_delay_number(condition_str):
            # Extracts the number at the end of a string like 'Activity_delay_5'
            return int(condition_str.split("_")[-1])

        df_individual_filtered["DelayNum"] = df_individual_filtered["Condition"].apply(extract_delay_number)

        # Now separate and sort using numeric delay
        df_hbo = df_individual_filtered[df_individual_filtered["Chroma"] == "hbo"].sort_values("DelayNum")
        df_hbr = df_individual_filtered[df_individual_filtered["Chroma"] == "hbr"].sort_values("DelayNum")

        vals_hbo = df_hbo["theta"].values
        vals_hbr = df_hbr["theta"].values
        
        # Create the plot
        fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(19, 10))
        
        # Scale design matrix components using numpy arrays instead of pandas operations
        dm_cond_values = dm_cond.values
        dm_cond_scaled_hbo = dm_cond_values * vals_hbo.reshape(1, -1)
        dm_cond_scaled_hbr = dm_cond_values * vals_hbr.reshape(1, -1)
        
        # Create time axis relative to stimulus onset
        time = dm_cond.index - np.ceil(first_onset_for_cond.get(cond, 0))
        
        # Plot
        axes[0].plot(time, dm_cond_values)
        axes[1].plot(time, dm_cond_scaled_hbo)
        axes[2].plot(time, np.sum(dm_cond_scaled_hbo, axis=1), 'r')
        axes[2].plot(time, np.sum(dm_cond_scaled_hbr, axis=1), 'b')
        
        # Format plots
        for ax in range(3):
            axes[ax].set_xlim(-5, 25)
            axes[ax].set_xlabel("Time (s)")
        axes[0].set_ylim(-0.2, 1.2)
        axes[1].set_ylim(-0.5, 1)
        axes[2].set_ylim(-0.5, 1)
        axes[0].set_title(f"FIR Model (Unscaled)")
        axes[1].set_title(f"FIR Components (Scaled by {cond} GLM Estimates)")
        axes[2].set_title(f"Evoked Response ({cond})")
        axes[0].set_ylabel("FIR Model")
        axes[1].set_ylabel("Oxyhaemoglobin (ΔμMol)")
        axes[2].set_ylabel("Haemoglobin (ΔμMol)")
        axes[2].legend(["Oxyhaemoglobin", "Deoxyhaemoglobin"])

        
        print(f"Number of FIR bins: {len(vals_hbo)}")
        print(f"Mean theta (HbO): {np.mean(vals_hbo):.4f}")
        print(f"Sum of theta (HbO): {np.sum(vals_hbo):.4f}")
        print(f"Mean theta (HbR): {np.mean(vals_hbr):.4f}")
        print(f"Sum of theta (HbR): {np.sum(vals_hbr):.4f}")
        
        fig_glms.append((f"Condition {cond}", fig))

    return fig_glms



def plot_3d_evoked_array(
    inst: Union[BaseRaw, EvokedArray, Info],
    statsmodel_df: DataFrame,
    picks: Optional[Union[str, list[str]]] = "hbo",
    value: str = "Coef.",
    background: str = "w",
    figure: Optional[object] = None,
    clim: Union[str, dict[str, Union[str, list[float]]]] = "auto",
    mode: str = "weighted",
    colormap: str = "RdBu_r",
    surface: str = "pial",
    hemi: str = "both",
    size: int = 800,
    view: Optional[Union[str, dict[str, float]]] = None,
    colorbar: bool = True,
    distance: float = 0.03,
    subjects_dir: Optional[str] = None,
    src: Optional[SourceSpaces] = None,
    verbose: bool = False,
) -> Brain:
    '''Ported from MNE'''

    info: Info = cast(Info, deepcopy(inst if isinstance(inst, Info) else inst.info)) # type: ignore
    if not (getattr(info, "ch_names") == list(statsmodel_df["ch_name"].values)): # type: ignore
        raise RuntimeError(
            'MNE data structure does not match dataframe '
            f'results.\nMNE = {getattr(info, "ch_names")}.\n'
            f'GLM = {list(statsmodel_df["ch_name"].values)}' # type: ignore
        )

    ea = EvokedArray(np.tile(statsmodel_df[value].values.T, (1, 1)).T, info.copy()) # type: ignore

    # TODO: mimic behaviour of other MNE-NIRS glm plotting options
    if picks is not None:
        ea = ea.pick(picks=picks) # type: ignore

    if subjects_dir is None:
        subjects_dir = os.environ.get("SUBJECTS_DIR")
        if subjects_dir is None:
            subjects_dir = str(data_path()) + "/subjects" # type: ignore
            os.environ["SUBJECTS_DIR"] = subjects_dir
            
    if src is None:
        fname_src_fs = os.path.join(
            subjects_dir, "fsaverage", "bem", "fsaverage-ico-5-src.fif"
        )
        src = read_source_spaces(fname_src_fs, verbose=verbose)

    picks = getattr(ea, "info")["ch_names"]

    # Set coord frame
    for idx in range(len(getattr(ea, "ch_names"))):
        getattr(ea, "info")["chs"][idx]["coord_frame"] = 4

    # Generate source estimate
    kwargs = dict(
        evoked=ea,
        subject="fsaverage",
        trans=Transform('head', 'mri', np.eye(4)),
        distance=distance,
        mode=mode,
        surface=surface,
        subjects_dir=subjects_dir,
        src=src,
        project=True,
    )

    stc = stc_near_sensors(picks=picks, **kwargs, verbose=verbose) # type: ignore

    assert isinstance(stc, SourceEstimate)

    # Produce brain plot
    brain: Brain = stc.plot( # type: ignore
        src=src,
        subjects_dir=subjects_dir,
        hemi=hemi,
        surface=surface,
        initial_time=0,
        clim=clim, # type: ignore
        size=size,
        colormap=colormap,
        figure=figure,
        background=background,
        colorbar=colorbar,
        verbose=verbose,
    )
    if view is not None:
        brain.show_view(view) # type: ignore

    return brain



def aggregate_fnirs_group_geometry(raw_list: Sequence[BaseRaw | None]) -> BaseRaw:
    """
    Averages fNIRS geometry across participants in two tiers:
    1. Average by Channel Pairing (S_D).
    2. Average by Individual Optode (S, D) across all averaged pairings.
    Returns a unified MNE Raw object with exactly one dot per optode.
    """

    def _safe_nanmean_columns(arr: np.ndarray, label: str, relevant_slice: slice = slice(0, 9)) -> np.ndarray | None:
        """Column-wise nanmean, treating a column as invalid only if it's still
        NaN after averaging AND within the coordinate range this function
        actually uses (loc[0:9]: channel midpoint, source, detector). Slots
        9-11 are unused fNIRS loc fields (commonly NaN by design) and must not
        trigger exclusion."""
        if arr.size == 0:
            return None
        with np.errstate(invalid="ignore"):
            result = np.nanmean(arr, axis=0)
        if np.any(np.isnan(result[relevant_slice])):
            logger.warning(f"[aggregate_geometry] '{label}' has NaN in relevant coordinate "
                            f"slots (0:9) after averaging - excluding.")
            return None
        return result

    channel_locs = {}
    all_ch_names = []

    for raw in raw_list:
        if raw is None: continue
        raw_hbo = raw.copy().pick(picks="hbo")

        for i, ch_name in enumerate(raw_hbo.ch_names):
            if ch_name not in channel_locs:
                channel_locs[ch_name] = []
                all_ch_names.append(ch_name)

            channel_locs[ch_name].append(raw_hbo.info['chs'][i]['loc'])

    avg_pairings = {}
    skipped_channels = []
    for name, locs in channel_locs.items():
        locs_arr = np.array(locs)
        valid = locs_arr[~np.all(np.isnan(locs_arr), axis=1)]
        result = _safe_nanmean_columns(valid, name)
        if result is None:
            skipped_channels.append(name)
            continue
        avg_pairings[name] = result

    if skipped_channels:
        logger.warning(f"[aggregate_geometry] {len(skipped_channels)} channel(s) had no valid "
                        f"position data from any participant, excluded: {skipped_channels}")

    optode_collections = {'sources': {}, 'detectors': {}}

    for ch_name, loc in avg_pairings.items():
        parts = ch_name.split()[0].split('_')
        s_name, d_name = parts[0], parts[1]

        optode_collections['sources'].setdefault(s_name, []).append(loc[3:6])
        optode_collections['detectors'].setdefault(d_name, []).append(loc[6:9])

    final_sources = {}
    for s, coords in optode_collections['sources'].items():
        coords_arr = np.array(coords)
        valid = coords_arr[~np.all(np.isnan(coords_arr), axis=1)]
        result = _safe_nanmean_columns(valid, f"source '{s}'")
        if result is not None:
            final_sources[s] = result

    final_detectors = {}
    for d, coords in optode_collections['detectors'].items():
        coords_arr = np.array(coords)
        valid = coords_arr[~np.all(np.isnan(coords_arr), axis=1)]
        result = _safe_nanmean_columns(valid, f"detector '{d}'")
        if result is not None:
            final_detectors[d] = result

    ref_raw = raw_list[0].copy().pick(picks="hbo")
    template_lookup = {ch['ch_name']: ch for ch in ref_raw.info['chs']}
    final_chs = []

    for ch_name in all_ch_names:
        if ch_name not in avg_pairings:
            continue
        parts = ch_name.split()[0].split('_')
        s_name, d_name = parts[0], parts[1]

        if s_name not in final_sources or d_name not in final_detectors:
            logger.warning(f"[aggregate_geometry] channel '{ch_name}' references source/detector "
                            f"with no valid coordinates, excluding.")
            continue

        unified_loc = avg_pairings[ch_name].copy()
        unified_loc[3:6] = final_sources[s_name]
        unified_loc[6:9] = final_detectors[d_name]
        unified_loc[0:3] = (final_sources[s_name] + final_detectors[d_name]) / 2.0

        # Create the new channel object
        new_ch = template_lookup.get(ch_name, ref_raw.info['chs'][0]).copy()
        new_ch['ch_name'] = ch_name
        new_ch['loc'] = unified_loc
        final_chs.append(new_ch)

    # Create the final MNE Info
    final_ch_names = [ch['ch_name'] for ch in final_chs]
    fake_info = create_info(ch_names=final_ch_names, sfreq=ref_raw.info['sfreq'], ch_types='hbo')
    with fake_info._unlock():
        fake_info['chs'] = final_chs

    return RawArray(np.zeros((len(final_ch_names), 1)), fake_info)



def brain_3d_visualization(
    raw_haemo: BaseRaw | None,
    df_cha: DataFrame | None,
    selected_event: str | None,
    t_or_theta: Literal["t", "theta"] = "theta",
    show_optodes: Literal["sensors", "labels", "none", "all"] = "all",
    show_text: bool = True,
    brain_bounds: float | tuple[float, float] | Sequence[float] = 1.0,
) -> None:
    if raw_haemo is None:
        raise ValueError("No haemo data available for the selected participant(s) - cannot render brain visualization.")
    if df_cha is None or df_cha.empty:
        raise ValueError("No channel-level results (df_cha) available for the selected participant(s).")
    if selected_event is None:
        raise ValueError("No event selected - brain_3d_visualization requires a specific condition to plot.")

    if isinstance(brain_bounds, (tuple, list)):
        b_low, b_high = float(brain_bounds[0]), float(brain_bounds[-1])
    else:
        b_low, b_high = 0.0, float(brain_bounds)

    clim = dict(kind="value", pos_lims=(b_low, (b_low + b_high) / 2, b_high))

    cond = str(selected_event)

    ch_summary = df_cha.query("Condition == @cond and Chroma == 'hbo'", engine='python')
    if ch_summary.empty:
        raise ValueError(f"No hbo data found for condition '{cond}' in the selected participant(s).")


    n_participants = ch_summary["ID"].nunique() if "ID" in ch_summary.columns else 1
    formula = f"{t_or_theta} ~ -1 + ch_name"

    if n_participants > 1 and "ID" in ch_summary.columns:
        try:
            ch_model = smf.mixedlm(formula, ch_summary, groups=ch_summary["ID"]).fit()
            if ch_model.cov_re.values.flatten()[0] < 1e-6:
                print(f"WARNING: Random-effects variance near zero for condition '{cond}' - "
                      f"mixed-effects model may not be meaningfully different from pooled OLS here.")
            print(f"Mixed-effects model used across {n_participants} participants.")
        except Exception as e:
            print(f"Mixed-effects model failed ({e}), falling back to OLS - "
                  f"note: pooled OLS across {n_participants} participants does not "
                  f"account for repeated-measures structure and may understate uncertainty.")
            ch_model = smf.ols(formula, ch_summary).fit()
    else:
        ch_model = smf.ols(formula, ch_summary).fit()
        print("OLS model used (single participant).")

    model_df = cast(DataFrame, statsmodels_to_results(ch_model, order=ch_summary["ch_name"].unique()))

    
    valid_channels = ch_summary["ch_name"].unique().tolist() # type: ignore
    raw_for_plot = raw_haemo.copy().pick(picks=valid_channels) # type: ignore

    print(f"DEBUG: Model DF rows: {len(model_df)}")
    print(f"DEBUG: Raw channels: {len(raw_for_plot.ch_names)}")

    brain = plot_3d_evoked_array(raw_for_plot.pick(picks="hbo"), model_df, view="dorsal", distance=0.02, colorbar=True, clim=clim, mode="weighted", size=(800, 700)) # type: ignore
    
    if show_optodes == 'all' or show_optodes == 'sensors':
        brain.add_sensors(raw_for_plot.pick(picks="hbo").info, trans=Transform('head', 'mri', np.eye(4)), fnirs=["channels", "pairs", "sources", "detectors"], verbose=False) # type: ignore

    elif show_optodes == "labels":
        print("show_optodes='labels' is not currently implemented - no optode overlay shown.")

    if show_text:
        display_text = (
            f"Condition: {cond}\n"
            f"Model: {'Mixed-effects' if n_participants > 1 else 'OLS'} "
            f"(n={n_participants})\n"
            f"Looking at: {t_or_theta} values"
        )
        brain.add_text(0.12, 0.64, display_text, "title", font_size=11, color="k") # type: ignore

    return brain



def brain_landmarks_3d(raw_haemo: BaseRaw, show_optodes: Literal['sensors', 'labels', 'none', 'all'] = 'all', show_brodmann: bool = True, subjects_dir: Union[str, Path, None] = None) -> None:

    if subjects_dir is None:
        subjects_dir = os.environ.get("SUBJECTS_DIR")
        if subjects_dir is None:
            subjects_dir = str(data_path()) + "/subjects" # type: ignore
            os.environ["SUBJECTS_DIR"] = subjects_dir
            
    brain = Brain("fsaverage", background="white", size=(800, 700))  # type: ignore

    distances = source_detector_distances(raw_haemo.info)

    # Add optode text labels manually
    if show_optodes == 'all' or show_optodes == 'sensors':
        brain.add_sensors(getattr(raw_haemo, "info"), trans=Transform('head', 'mri', np.eye(4)), fnirs=["channels", "pairs", "sources", "detectors"], verbose=False) # type: ignore
    
    if show_optodes == 'all' or show_optodes == 'labels':
        labeled_srcs = set()
        labeled_dets = set()
        label_counts = {}

        for idx, ch in enumerate(raw_haemo.info['chs']):
            ch_name = ch['ch_name']

            if not ch_name.endswith('hbo'):
                continue

            loc = ch['loc']
            logger.info(f"Channel: {ch_name}")
            logger.info(f"loc length: {len(loc)}")
            logger.info("loc contents:")
            for i, val in enumerate(loc):
                logger.info(f"  loc[{i}]: {val}")
            logger.info("-" * 30)
            if not ch_name or not ch['loc'].any():
                continue
            
            parts = ch_name.split()[0]
            src_str, det_str = parts.split('_')
            
            src_num = int(src_str[1:])
            det_num = int(det_str[1:])
            
            if src_num not in labeled_srcs:
                src_xyz = ch['loc'][3:6] * 1000
                brain._renderer.text3d(src_xyz[0], src_xyz[1], src_xyz[2], src_str,
                                    color='red', scale=0.002)
                labeled_srcs.add(src_num)
            
            if det_num not in labeled_dets:
                det_xyz = ch['loc'][6:9] * 1000
                brain._renderer.text3d(det_xyz[0], det_xyz[1], det_xyz[2], det_str,
                                    color='blue', scale=0.002)
                labeled_dets.add(det_num)

            # Get the source-detector distance for this channel (in meters)
            dist_m = distances[idx]
            dist_mm = dist_m * 1000

            label_text = f"{dist_mm:.1f} mm"
            label_counts[label_text] = label_counts.get(label_text, 0) + 1
            if label_counts[label_text] > 1:
                label_text += f" ({label_counts[label_text]})"
            
            # Label at channel midpoint
            mid_xyz = loc[0:3] * 1000

            logger.info(f"Channel: {ch_name} | Midpoint (mm): x={mid_xyz[0]:.2f}, y={mid_xyz[1]:.2f}, z={mid_xyz[2]:.2f} | Distance: {dist_mm:.1f} mm")

            brain._renderer.text3d(
                mid_xyz[0], mid_xyz[1], mid_xyz[2],
                label_text,
                color='gray',
                scale=0.002
            )

    
    if show_brodmann:# Add Brodmann labels
        labels = cast(list[Label], read_labels_from_annot("fsaverage", "PALS_B12_Brodmann", "lh", verbose=False)) # type: ignore

        #TODO: This has been hardcoded here for the entire applications lifecycle. About time to user expose?
        label_colors = {
            "Brodmann.1-lh": "red",
            "Brodmann.2-lh": "red",
            "Brodmann.3-lh": "red",
            "Brodmann.4-lh": "orange",
            "Brodmann.5-lh": "green",
            "Brodmann.6-lh": "yellow",
            "Brodmann.7-lh": "green",
            "Brodmann.17-lh": "blue",
            "Brodmann.18-lh": "blue",
            "Brodmann.19-lh": "blue",
            "Brodmann.39-lh": "pink",
            "Brodmann.40-lh": "purple",
            "Brodmann.42-lh": "white",
            "Brodmann.44-lh": "white",
            "Brodmann.48-lh": "white",

        }

        for label in labels:
            name = getattr(label, "name", None)
            if not isinstance(name, str):
                continue
            if name in label_colors:
                brain.add_label(label, borders=False, color=label_colors[name]) # type: ignore


    return brain



def brain_3d_contrast(con_model_df: DataFrame, con_model_df_filtered: BaseRaw, common_channels: list[str], first_name: str, second_name: str, t_or_theta: Literal['t', 'theta'] = 'theta', show_optodes: Literal['sensors', 'labels', 'none', 'all'] = 'all', show_text: bool = True, brain_bounds: float = 1.0) -> None:
    # Filter DataFrame to only common channels, and sort by raw order
    con_model = con_model_df

    con_model["ch_name"] = pd.Categorical(
        con_model["ch_name"], categories=common_channels, ordered=True
    )
    con_model = con_model.sort_values("ch_name").reset_index(drop=True) # type: ignore


    clim=dict(kind="value", pos_lims=(0, brain_bounds/2, brain_bounds))

    # Plot brain figure
    brain = plot_3d_evoked_array(con_model_df_filtered.copy().pick(picks="hbo"), con_model, view="dorsal", distance=0.02, colorbar=True, mode="weighted", clim=clim, size=(800, 700), verbose=False) # type: ignore

    if show_optodes == 'all' or show_optodes == 'sensors':
        brain.add_sensors(getattr(con_model_df_filtered, "info"), trans=Transform('head', 'mri', np.eye(4)), fnirs=["channels", "pairs", "sources", "detectors"], verbose=False) # type: ignore
   
    display_text = ('Contrast: ' + first_name + ' - ' + second_name + '\nLooking at: ' + t_or_theta + ' values')

    # Apply the text onto the brain
    if show_text:
        brain.add_text(0.12, 0.70, display_text, "title", font_size=11, color="k") # type: ignore
 


def plot_2d_3d_contrasts_between_groups(
    contrast_df_a: pd.DataFrame,
    contrast_df_b: pd.DataFrame,
    raw_haemo: BaseRaw,
    group_a_name: str,
    group_b_name: str,
    is_3d: bool = True,
    t_or_theta: Literal['t', 'theta'] = 'theta',
    show_optodes: Literal['sensors', 'labels', 'none', 'all'] = 'all',
    show_text: bool = True,
    brain_bounds: float = 1.0,
    min_participants_per_group: int = 2,
) -> None:
    if raw_haemo is None:
        raise ValueError("raw_haemo is required for group contrast visualization.")
    if group_a_name == group_b_name:
        raise ValueError(f"group_a_name and group_b_name must differ (both were '{group_a_name}').")
    if contrast_df_a.empty or contrast_df_b.empty:
        raise ValueError("One or both contrast dataframes are empty - check group selection.")

    contrast_df_a = contrast_df_a.copy()
    contrast_df_a["group"] = group_a_name
    contrast_df_b = contrast_df_b.copy()
    contrast_df_b["group"] = group_b_name

    df_combined = pd.concat([contrast_df_a, contrast_df_b], ignore_index=True)
    con_summary = df_combined.query("Chroma == 'hbo'").copy()

    counts = pd.crosstab(con_summary["group"], con_summary["ch_name"])
    valid_mask = (counts >= min_participants_per_group).all()
    valid_channels = valid_mask[valid_mask].index.tolist()
    con_summary = con_summary[con_summary["ch_name"].isin(valid_channels)]

    if con_summary.empty:
        raise ValueError(
            f"No channels have >= {min_participants_per_group} participants in BOTH "
            f"'{group_a_name}' and '{group_b_name}' - cannot fit a group contrast model. "
            f"Select more participants per group, or lower min_participants_per_group."
        )

    model_formula = "effect ~ -1 + group:ch_name:Chroma"
    try:
        con_model = smf.mixedlm(model_formula, con_summary, groups=con_summary["ID"]).fit(method="nm")
        if not getattr(con_model, "converged", True):
            print(f"WARNING: Group contrast mixed-effects model did not converge cleanly "
                  f"({group_a_name} vs {group_b_name}) - results may be unreliable.")
    except Exception as e:
        print(f"Mixed-effects model failed ({e}), falling back to OLS - "
              f"note: pooled OLS does not account for repeated-measures structure "
              f"and may understate uncertainty.")
        con_model = smf.ols(model_formula, con_summary).fit()

    if t_or_theta == "t":
        group1_vals = con_model.tvalues.filter(like=f"group[{group_a_name}]")
        group2_vals = con_model.tvalues.filter(like=f"group[{group_b_name}]")
    else:
        group1_vals = con_model.params.filter(like=f"group[{group_a_name}]")
        group2_vals = con_model.params.filter(like=f"group[{group_b_name}]")

    def _extract_ch_names(vals):
        names = []
        for term in vals.index:
            parts = term.split(":")
            if len(parts) < 2 or "[" not in parts[1] or "]" not in parts[1]:
                raise ValueError(
                    f"Unexpected coefficient name format: '{term}'. Expected "
                    f"'group[...]:ch_name[...]:Chroma[...]' - patsy naming may "
                    f"have changed. Cannot safely extract channel names."
                )
            names.append(parts[1].split("[")[1].split("]")[0])
        return names

    group1_channels = _extract_ch_names(group1_vals)
    group2_channels = _extract_ch_names(group2_vals)

    known_channels = set(raw_haemo.copy().pick(picks="hbo").ch_names)
    unrecognized = (set(group1_channels) | set(group2_channels)) - known_channels
    if unrecognized:
        raise ValueError(
            f"Extracted channel name(s) not found in raw_haemo: {unrecognized}. "
            f"Coefficient-name parsing likely broke - verify model term format."
        )

    df_group1 = DataFrame({"Coef.": group1_vals.values}, index=group1_channels)
    df_group2 = DataFrame({"Coef.": group2_vals.values}, index=group2_channels)
    df_contrast = df_group1.join(df_group2, how="inner", lsuffix=f"_{group_a_name}", rsuffix=f"_{group_b_name}")

    if df_contrast.empty:
        raise ValueError(f"No channels in common between '{group_a_name}' and '{group_b_name}' model results.")

    mne_ch_names = raw_haemo.copy().pick(picks="hbo").ch_names

    def _plot_direction(coef_a_col, coef_b_col, name_first, name_second):
        df_contrast["Coef."] = df_contrast[coef_a_col] - df_contrast[coef_b_col]
        con_model_df = DataFrame({
            "ch_name": df_contrast.index,
            "Coef.": df_contrast["Coef."],
            "Chroma": "hbo"
        })
        glm_ch_names = con_model_df["ch_name"].tolist()
        common_channels = [ch for ch in mne_ch_names if ch in glm_ch_names]
        con_model_df_filtered = raw_haemo.copy().pick(picks=common_channels)
        con_model_df = con_model_df.set_index("ch_name").loc[common_channels].reset_index()

        if is_3d:
            brain_3d_contrast(
                con_model_df, con_model_df_filtered, common_channels,
                name_first, name_second, t_or_theta, show_optodes, show_text, brain_bounds
            )
        else:
            plot_glm_group_topo(
                con_model_df_filtered.copy().pick(picks="hbo"), con_model_df,
                names=True, res=128, vlim=(-brain_bounds, brain_bounds)
            )
            plt.title(f"Contrast: {name_first} vs {name_second}")
            plt.show()

    _plot_direction(f"Coef._{group_a_name}", f"Coef._{group_b_name}", group_a_name, group_b_name)
    _plot_direction(f"Coef._{group_b_name}", f"Coef._{group_a_name}", group_b_name, group_a_name)



def plot_fir_model_results(
    df: DataFrame,
    raw_haemo: BaseRaw | None,
    dm: DataFrame | None,
    selected_event: str | None,
    l_bound: float,
    u_bound: float,
) -> None:
    '''
    FIR Model Results requires per-delay Condition data, but the current
    df_ind_dict is pre-collapsed (delay information stripped) upstream in
    generate_roi_results. This method is not currently functional as it
    needs an uncollapsed per-delay ROI dataframe to be threaded through separately.
    '''

    df["isActivity"] = [f"{selected_event}" in n for n in df["Condition"]]
    df["isDelay"] = ["delay" in n for n in df["Condition"]]
    df = df.query("isDelay in [True]")
    df = df.query("isActivity in [True]")
    # Make a new column that stores the condition name for tidier model below
    df.loc[:, "TidyCond"] = ""
    df.loc[df["isActivity"] == True, "TidyCond"] = f"{selected_event}"  # noqa: E712
    # Finally, extract the FIR delay in to its own column in data frame
    df.loc[:, "delay"] = [n.split("_")[-1] for n in df.Condition]

    # To simplify this example we will only look at the activity
    # condition so we now remove the other conditions from the
    # design matrix and GLM results
    dm_cols_activity = np.where([f"{selected_event}" in c for c in dm.columns])[0]
    dm = dm[[dm.columns[i] for i in dm_cols_activity]]

    try:
        lme = smf.mixedlm("theta ~ -1 + delay:TidyCond:Chroma", df, groups=df["ID"]).fit()
    except:
        lme = smf.ols("theta ~ -1 + delay:TidyCond:Chroma", df, groups=df["ID"]).fit() # type: ignore

    df_sum = statsmodels_to_results(lme)
    df_sum["delay"] = [int(n) for n in df_sum["delay"]]
    df_sum = df_sum.sort_values("delay")

    # Print the result for the oxyhaemoglobin data in the target condition
    df_sum.query(f"TidyCond in ['{selected_event}']").query("Chroma in ['hbo']")


    fig, axes = plt.subplots(nrows=1, ncols=3, figsize=(19, 10))

    print("dm columns:", dm.columns.tolist())

    # Extract design matrix columns that correspond to the condition of interest
    dm_cond_idxs = np.where([f"{selected_event}" in n for n in dm.columns])[0]
    dm_cond = dm[[dm.columns[i] for i in dm_cond_idxs]]

    # Extract the corresponding estimates from the lme dataframe for hbo
    df_hbo = df_sum.query(f"TidyCond in ['{selected_event}']").query("Chroma in ['hbo']")
    vals_hbo = [float(v) for v in df_hbo["Coef."]]

    # print("--------------------------------------")
    # print(f"dm_cond shape: {dm_cond.shape}")
    # print(f"dm_cond columns: {dm_cond.columns.tolist()}")
    # print(f"vals_hbo length: {len(vals_hbo)}")
    # print(f"vals_hbo sample: {vals_hbo[:5]}")
    # print(f"vals_hbo type: {type(vals_hbo)}")
    # print(f"vals_hbo element type: {type(vals_hbo[0]) if len(vals_hbo) > 0 else 'N/A'}")

    dm_cond_scaled_hbo = dm_cond * vals_hbo

    # Extract the corresponding estimates from the lme dataframe for hbr
    df_hbr = df_sum.query(f"TidyCond in ['{selected_event}']").query("Chroma in ['hbr']")
    vals_hbr = [float(v) for v in df_hbr["Coef."]]
    dm_cond_scaled_hbr = dm_cond * vals_hbr

    first_onset = None
    for desc, onset in zip(raw_haemo.annotations.description, raw_haemo.annotations.onset):
        if selected_event in desc:
            first_onset = onset
            break

    if first_onset is None:
        raise ValueError(f"Selected event '{selected_event}' not found in annotations.")

    # Align index values (time axis) to the first occurrence of selected_event
    index_values = dm_cond_scaled_hbo.index - np.ceil(first_onset)
    index_values = np.asarray(index_values)

    # Plot the result
    axes[0].plot(index_values, np.asarray(dm_cond))
    axes[1].plot(index_values, np.asarray(dm_cond_scaled_hbo))
    axes[2].plot(index_values, np.sum(dm_cond_scaled_hbo, axis=1), "r")
    axes[2].plot(index_values, np.sum(dm_cond_scaled_hbr, axis=1), "b")

    valid_mask = (index_values >= 0) & (index_values <= 15)
    hbo_sum_window = np.sum(dm_cond_scaled_hbo.loc[valid_mask, :], axis=1)
    peak_idx_in_window = np.argmax(hbo_sum_window)
    peak_idx = np.where(valid_mask)[0][peak_idx_in_window]
    peak_time = float(round(index_values[peak_idx], 2)) # type: ignore

    axes[2].axvline(x=peak_time, color='k', linestyle='--', linewidth=1.5, label='Peak time') # type: ignore

    # Format the plot
    for ax in range(3):
        axes[ax].set_xlim(-5, 25)
        axes[ax].set_xlabel("Time (s)")
    axes[0].set_ylim(-0.1, 1.1)
    axes[1].set_ylim(l_bound, u_bound)
    axes[2].set_ylim(l_bound, u_bound)
    axes[0].set_title("FIR Model (Unscaled by GLM estimates)")
    axes[1].set_title(f"FIR Components (Scaled by {selected_event} GLM Estimates)")
    axes[2].set_title(f"Evoked Response {selected_event}")
    axes[0].set_ylabel("FIR Model")
    axes[1].set_ylabel("Oyxhaemoglobin (ΔμMol)")
    axes[2].set_ylabel("Haemoglobin (ΔμMol)")
    axes[2].legend(["Oyxhaemoglobin", "Deoyxhaemoglobin"])

    # We can also extract the 95% confidence intervals of the estimates too
    l95_hbo = [float(v) for v in df_hbo["[0.025"]] # type: ignore
    u95_hbo = [float(v) for v in df_hbo["0.975]"]] # type: ignore
    dm_cond_scaled_hbo_l95 = dm_cond * l95_hbo
    dm_cond_scaled_hbo_u95 = dm_cond * u95_hbo
    l95_hbr = [float(v) for v in df_hbr["[0.025"]] # type: ignore
    u95_hbr = [float(v) for v in df_hbr["0.975]"]] # type: ignore
    dm_cond_scaled_hbr_l95 = dm_cond * l95_hbr
    dm_cond_scaled_hbr_u95 = dm_cond * u95_hbr

    axes2: Axes
    fig2, axes2 = plt.subplots(nrows=1, ncols=1, figsize=(7, 7)) # type: ignore

    # Plot the result
    axes2.plot(index_values, np.sum(dm_cond_scaled_hbo, axis=1), "r") # type: ignore
    axes2.plot(index_values, np.sum(dm_cond_scaled_hbr, axis=1), "b") # type: ignore
    axes2.axvline(x=peak_time, color='k', linestyle='--', linewidth=1.5, label='Peak time') # type: ignore

    axes2.fill_between( # type: ignore
        index_values,
        np.asarray(np.sum(dm_cond_scaled_hbo_l95, axis=1)),
        np.asarray(np.sum(dm_cond_scaled_hbo_u95, axis=1)),
        facecolor="red",
        alpha=0.25,
    )
    axes2.fill_between( # type: ignore
        index_values,
        np.asarray(np.sum(dm_cond_scaled_hbr_l95, axis=1)),
        np.asarray(np.sum(dm_cond_scaled_hbr_u95, axis=1)),
        facecolor="blue",
        alpha=0.25,
    )

    # Format the plot
    axes2.set_xlim(-5, 20)
    axes2.set_ylim(l_bound, u_bound)
    axes2.set_title(f"Evoked Response with 95% confidence intervals for )") # type: ignore
    axes2.set_ylabel("Haemoglobin (ΔμMol)") # type: ignore
    axes2.legend(["Oyxhaemoglobin", "Deoyxhaemoglobin", f"Peak {peak_time}s"]) # type: ignore
    axes2.set_xlabel("Time (s)") # type: ignore

    fig2.tight_layout()

    fig.show()
    fig2.show()



def load_snirf(file_path: str, downsample_frequency: int, verbosity: bool) -> tuple[BaseRaw, Figure]:
    """
    Loads a snirf file, optionally drops channels, downsamples, and creates a figure showing the results.

    Parameters
    ----------
    file_path : str
        Path of the snirf file to load.
    ID : str
        File name of the the snirf file that was loaded.
    drop_prefixes : list[str]
        List of channel name prefixes to drop from the data.

    Returns
    -------
    tuple[BaseRaw, Figure]
        - BaseRaw: The processed data object.
        - Figure: The corresponding Matplotlib figure.
    """
    
    # Read the snirf file
    raw = read_raw_snirf(file_path, preload=True, verbose=verbosity) # type: ignore
    #raw.load_data(verbose=VERBOSITY) # type: ignore  redundant since preload is set to true

    # TODO: Why was this commented again?
    # Maybe this should be a bypass parameter?

    # If the user forcibly dropped channels, remove them now before any processing occurs
    # logger.info("Checking if there are channels to forcibly drop...")
    # if drop_prefixes:
    #     logger.info("Force dropped channels was specified.")
    #     channels_to_drop = [ch for ch in cast(list[str], getattr(raw, "ch_names")) if any(ch.startswith(prefix) for prefix in drop_prefixes)]
    #     raw.drop_channels(channels_to_drop, "raise") # type: ignore
    #     logger.info("Force dropped channels:", channels_to_drop)

    # If the user wants to downsample, do it right away
    logger.info("Checking if we should downsample...")
    if DOWNSAMPLE:
        logger.info("Downsample was specified.")
        sfreq_old = getattr(raw, "info")["sfreq"]
        raw.resample(downsample_frequency, verbose=verbosity) # type: ignore
        sfreq_new = getattr(raw, "info")["sfreq"]
        logger.info(f"Finished downsampling. Old frequency: {sfreq_old}. New frequency: {sfreq_new}.")

    logger.info("Successfully loaded the snirf file.")

    return raw



def run_roi_second_level_analysis(
    df_roi_all: DataFrame,
    condition: str,
    df_cha_all: DataFrame | None = None,
    raw_haemo: BaseRaw | None = None,
    p_threshold: float = 0.05,
    min_subjects: int = 5,
    correction_method: str | None = "fdr_bh",
    target_chroma: str = "hbo",
    graph_bounds: float | None = None,
    threshold_topo: bool = False,
) -> DataFrame:
    
    """
    Perform group-level ROI analysis, prints stats to console, plots the ROI bar chart,
    and dynamically plots isolated channel-level group topography maps based on a JSON config.
    """
    # 1. Validation checks
    required_cols = ['ROI', 'Condition', 'Chroma', 'theta', 'ID']
    if not all(col in df_roi_all.columns for col in required_cols):
        raise ValueError(f"Input ROI DataFrame must include: {required_cols}")

    # 2. Filter ROI data for the targeted chromophore
    df_chroma = df_roi_all[(df_roi_all['Chroma'] == target_chroma) & (df_roi_all['Condition'] == condition)].copy()
    df_chroma = df_chroma.dropna(subset=['theta'])

    # 3. Perform 1-sample t-test against zero for each ROI
    rois = df_chroma['ROI'].unique()
    group_results = []

    for roi in rois:
        roi_data = df_chroma[df_chroma['ROI'] == roi]
        sub_data = roi_data.groupby('ID', as_index=False)['theta'].mean()
        
        n_subs = sub_data['ID'].nunique()
        if n_subs < min_subjects:
            continue

        Y = sub_data['theta'].values
        t_val, p_val = ttest_1samp(Y, 0)
        mean_beta = np.mean(Y)
        std_err = sem(Y)

        group_results.append({
            'ROI': roi,
            't_val': t_val,
            'p_val': p_val,
            'mean_beta': mean_beta,
            'std_err': std_err,
            'n_subjects': n_subs
        })

    if not group_results:
        print("\n[ERROR] No ROIs met the minimum subject threshold.\n")
        return pd.DataFrame()

    df_group = pd.DataFrame(group_results)

    # 4. Multiple comparisons correction
    if correction_method is not None:
        reject, p_corrected, _, _ = multipletests(
            df_group['p_val'].values, method=correction_method
        )
        df_group['p_corrected'] = p_corrected
        df_group['significant'] = reject
    else:
        df_group['p_corrected'] = df_group['p_val']
        df_group['significant'] = df_group['p_val'] <= p_threshold

    # Print results table to terminal
    print("\n" + "="*65)
    print(f"       GROUP-LEVEL ROI STATISTICAL RESULTS ({target_chroma.upper()})")
    print("="*65)
    df_print = df_group.copy()
    df_print['mean_beta'] = df_print['mean_beta'].apply(lambda x: f"{x:.4f}")
    df_print['t_val'] = df_print['t_val'].apply(lambda x: f"{x:.3f}")
    df_print['p_val'] = df_print['p_val'].apply(lambda x: f"{x:.4f}")
    df_print['p_corrected'] = df_print['p_corrected'].apply(lambda x: f"{x:.4f}")
    print(df_print[['ROI', 'mean_beta', 't_val', 'p_val', 'p_corrected', 'significant']].to_string(index=False))
    print("="*65 + "\n")

    # 5. Plotting ROI Bar Chart
    sns.set_theme(style="whitegrid")
    fig, ax = plt.subplots(figsize=(8, 6))

    df_sub_avg = df_chroma.groupby(['ROI', 'ID'], as_index=False)['theta'].mean()

    sns.barplot(
        data=df_sub_avg, x='ROI', y='theta', 
        ax=ax, errorbar=('ci', 95), capsize=0.1, 
        color='lightgray', edgecolor='black', linewidth=1.5, zorder=1
    )
    sns.swarmplot(
        data=df_sub_avg, x='ROI', y='theta', 
        ax=ax, color='darkblue', size=8, alpha=0.7, zorder=2
    )

    global_max = df_sub_avg['theta'].max()
    global_min = df_sub_avg['theta'].min()
    y_top = global_max * 1.35 if global_max > 0 else 0.5e-6
    y_bottom = global_min * 1.1 if global_min < 0 else -0.1 * global_max
    ax.set_ylim(y_bottom, y_top)

    if graph_bounds is not None and graph_bounds > 0.0:
        if graph_bounds < 0.5:  
            ax.set_ylim(-graph_bounds, graph_bounds)

    for idx, row in df_group.iterrows():
        roi_name = row['ROI']
        p_val_corr = row['p_corrected']
        
        roi_points = df_sub_avg[df_sub_avg['ROI'] == roi_name]['theta']
        max_y = roi_points.max() if len(roi_points) > 0 else 0
        text_y = max_y + (global_max * 0.03) 
        
        if p_val_corr < 0.001:
            sig_symbol = "***"
        elif p_val_corr < 0.01:
            sig_symbol = "**"
        elif p_val_corr < p_threshold:
            sig_symbol = "*"
        else:
            sig_symbol = "n.s."

        sig_text = f"{sig_symbol}\np_corr = {p_val_corr:.3f}"
        x_pos = list(rois).index(roi_name)
        ax.text(
            x_pos, text_y, sig_text, 
            ha='center', va='bottom', fontsize=11, 
            fontweight='bold', color='red' if p_val_corr < p_threshold else 'gray'
        )

    ax.axhline(0, color='black', linewidth=1, linestyle='--')
    ax.set_ylabel(r'Hemodynamic Response ($\Delta$ HbO / $\mu$mol/L)' if global_max > 1e-3 else r'Hemodynamic Response ($\Delta$ HbO / mol/L)', fontsize=12)
    ax.set_xlabel('Region of Interest (ROI)', fontsize=12)
    
    correction_lbl = f"({correction_method} corrected)" if correction_method else "(uncorrected)"
    ax.set_title(
        f"Group-Level ROI Activation ({target_chroma.upper()})\nSignificance threshold: p < {p_threshold} {correction_lbl}", 
        fontsize=13, fontweight='bold', pad=15
    )
    plt.tight_layout()
    plt.show()

    # === 6. Segmented Topography Plotting (No Hardcoded Regions) ===
    if df_cha_all is not None and raw_haemo is not None:
        print(f"--> Fitting group-level channel LME for {target_chroma.upper()} topography...")
        try:
            val_col = 'effect' if 'effect' in df_cha_all.columns else 'theta'
            ch_col = 'ch_name' if 'ch_name' in df_cha_all.columns else 'channel'
            
            con_summary = df_cha_all[(df_cha_all['Chroma'] == target_chroma) & (df_cha_all['Condition'] == condition)].copy()
            raw_picked = raw_haemo.copy().pick(picks=target_chroma)
            
            # Fit channel LME (suppress ConvergenceWarning locally)
            model_formula = f"{val_col} ~ -1 + {ch_col}:Chroma"
            with warnings.catch_warnings():
                warnings.simplefilter("ignore", ConvergenceWarning)
                con_model = smf.mixedlm(model_formula, con_summary, groups=con_summary["ID"]).fit(method="nm")
            
            # Map statsmodels output to MNE result format
            con_model_df = statsmodels_to_results(con_model, order=raw_picked.ch_names)

            # --- DYNAMIC ROI PARSING ---
            if 'ROI' not in con_summary.columns or con_summary['ROI'].dropna().empty:
                if df_roi_all is not None and 'ROI' in df_roi_all.columns and ch_col in df_roi_all.columns:
                    # Create channel -> ROI mapping from df_roi_all
                    ch_to_roi = df_roi_all.dropna(subset=['ROI', ch_col]).set_index(ch_col)['ROI'].to_dict()
                    con_summary['ROI'] = con_summary[ch_col].apply(
                        lambda x: ch_to_roi.get(x, ch_to_roi.get(x.split()[0], None) if isinstance(x, str) else None)
                    )

            unique_rois = []
            if 'ROI' in con_summary.columns:
                unique_rois = [r for r in con_summary['ROI'].dropna().unique() if r != ""]

            if not unique_rois:
                print("--> Warning: No ROI mappings detected. Plotting as a unified grid.")
                unique_rois = ['All_Channels']
                con_summary['ROI'] = 'All_Channels'

            # Calculate shared symmetric color limits
            vlim = (None, None)
            if 'Coef.' in con_model_df.columns:
                clean_vals = con_model_df['Coef.'].dropna().values
                if len(clean_vals) > 0:
                    max_abs = np.max(np.abs(clean_vals))
                    if max_abs > 1e-9:
                        vlim = (-max_abs, max_abs)

            fig_topo, ax_topo = plt.subplots(figsize=(6, 6))
            
            # Dynamic loop: Plot each region independently
            for i, roi_name in enumerate(unique_rois):
                roi_ch_names = con_summary[con_summary['ROI'] == roi_name][ch_col].unique().tolist()
                roi_ch_names = [ch for ch in roi_ch_names if ch in raw_picked.ch_names]
                
                if not roi_ch_names:
                    continue
                
                raw_roi = raw_picked.copy().pick(picks=roi_ch_names)
                show_colorbar = (i == len(unique_rois) - 1)
                
                # === FIX 2: Filter stats dataframe first to prevent "Reducing GLM results..." warnings ===
                roi_con_model_df = con_model_df[con_model_df['ch_name'].isin(roi_ch_names)].copy()
                
                plot_glm_group_topo(
                    raw_roi, 
                    roi_con_model_df, 
                    colorbar=show_colorbar,
                    threshold=threshold_topo,  # Now uses the parameter!
                    axes=ax_topo,
                    cmap='RdBu_r',
                    vlim=vlim
                )

            threshold_text = "p < 0.05 Masked" if threshold_topo else "Unthresholded"
            ax_topo.set_title(
                f"Group-Level {target_chroma.upper()} Activation Map\n(Regions Isolated Dynamically, {threshold_text})", 
                fontsize=11, fontweight='bold', pad=10
            )
            plt.tight_layout()
            plt.show()
            
        except Exception as e:
            logger.error(f"Could not generate topography plot: {e}")

    return df_group



def clean_subject_id(path_or_id):
    """
    Cleans file paths and ID strings to get a standardized subject identifier.
    E.g., 'C:/path/Sub-01_haemo.snirf' -> 'Sub-01'
    """
    if not isinstance(path_or_id, str):
        return str(path_or_id)
    base = os.path.basename(path_or_id)
    for ext in ['.snirf', '.nirs', '.fif', '.csv', '.pkl', '_haemo']:
        if base.endswith(ext):
            base = base[:-len(ext)]
    if base.endswith('_haemo'):
        base = base[:-6]
    return base



def run_inter_group_second_level_analysis(
    df_roi_all: DataFrame,
    file_paths_a: list[str],
    file_paths_b: list[str],
    group_a_name: str = "Group A",
    group_b_name: str = "Group B",
    df_cha_all: DataFrame | None = None,
    raw_haemo: Any = None,
    p_threshold: float = 0.05,
    min_subjects: int = 3,
    correction_method: str | None = "fdr_bh",
    target_chroma: str = "hbo",
    selected_event: str | None = None,
    graph_bounds: tuple[float, float] | list[float] | None = None,
    roi_channel_maps: dict[str, dict[str, str]] | None = None,
    threshold_topo: bool = False,
) -> DataFrame:
    
    """
    Perform cross-group independent statistical analyses (Group A vs Group B),
    renders a grouped bar chart with significance brackets, and plots a group-contrast topography map.
    """
    # 1. Align IDs and filter dataset to selected Event & Chromophore
    clean_a = set(file_paths_a)
    clean_b = set(file_paths_b)
    
    df_roi_all = df_roi_all.copy()
    # df_roi_all['clean_ID'] = df_roi_all['ID'].apply(clean_subject_id)
    df_roi_all['clean_ID'] = df_roi_all['ID']

    # Filter for active experimental conditions
    df_filtered = df_roi_all[
        (df_roi_all['Chroma'] == target_chroma) & 
        (df_roi_all['Condition'] == selected_event)
    ].copy()
    

    df_a = df_filtered[df_filtered['clean_ID'].isin(clean_a)].copy()
    df_b = df_filtered[df_filtered['clean_ID'].isin(clean_b)].copy()

    print(f"DEBUG: Filtering for Event: {selected_event}")
    print(f"DEBUG: Unique IDs in df_filtered: {df_filtered['clean_ID'].unique()}")
    print(f"DEBUG: Clean IDs from Group A: {clean_a}")
    print(f"DEBUG: Clean IDs from Group B: {clean_b}")
    print(f"DEBUG: Rows in df_a: {len(df_a)}, Rows in df_b: {len(df_b)}")


    if df_a.empty or df_b.empty:
        print("[ERROR] Missing data for one or both cohorts. Check file selection/IDs.")
        return pd.DataFrame()

    # 2. ROI-Level Welch's T-Test (Independent Two-Sample)
    rois = df_filtered['ROI'].dropna().unique()
    group_results = []

    for roi in rois:
        vals_a = df_a[df_a['ROI'] == roi].groupby('clean_ID')['theta'].mean().values
        vals_b = df_b[df_b['ROI'] == roi].groupby('clean_ID')['theta'].mean().values
        
        n_a, n_b = len(vals_a), len(vals_b)
        if n_a < min_subjects or n_b < min_subjects:
            continue

        # Welch's t-test (assumes unequal variances)
        t_val, p_val = ttest_ind(vals_a, vals_b, equal_var=False)
        mean_a, mean_b = np.mean(vals_a), np.mean(vals_b)
        diff_val = mean_a - mean_b

        group_results.append({
            'ROI': roi,
            'mean_A': mean_a,
            'mean_B': mean_b,
            'mean_diff': diff_val,
            't_val': t_val,
            'p_val': p_val,
            'n_A': n_a,
            'n_B': n_b
        })

    if not group_results:
        print("\n[ERROR] No ROIs met the subject requirements for BOTH groups.\n")
        return pd.DataFrame()

    df_group = pd.DataFrame(group_results)

    # Apply FDR correction
    if correction_method is not None:
        reject, p_corrected, _, _ = multipletests(df_group['p_val'].values, method=correction_method)
        df_group['p_corrected'] = p_corrected
        df_group['significant'] = reject
    else:
        df_group['p_corrected'] = df_group['p_val']
        df_group['significant'] = df_group['p_val'] <= p_threshold

    # Print clean terminal report
    print("\n" + "="*85)
    print(f"       CROSS-GROUP ROI CONTRAST: {group_a_name.upper()} vs {group_b_name.upper()} ({target_chroma.upper()})")
    print(f"       Event Condition: {selected_event}")
    print("="*85)
    df_print = df_group.copy()
    df_print['mean_A'] = df_print['mean_A'].apply(lambda x: f"{x:.4f}")
    df_print['mean_B'] = df_print['mean_B'].apply(lambda x: f"{x:.4f}")
    df_print['mean_diff'] = df_print['mean_diff'].apply(lambda x: f"{x:.4f}")
    df_print['t_val'] = df_print['t_val'].apply(lambda x: f"{x:.3f}")
    df_print['p_corrected'] = df_print['p_corrected'].apply(lambda x: f"{x:.4f}")
    print(df_print[['ROI', 'mean_A', 'mean_B', 'mean_diff', 't_val', 'p_corrected', 'significant']].to_string(index=False))
    print("="*85 + "\n")

    # 3. Double Grouped Bar Plot (Side-by-Side)
    sns.set_theme(style="whitegrid")
    fig, ax = plt.subplots(figsize=(10, 6))

    # Construct unified dataframe for seaborn grouped layouts
    df_a_tidy = df_a.groupby(['ROI', 'clean_ID'], as_index=False)['theta'].mean()
    df_a_tidy['Group'] = group_a_name
    df_b_tidy = df_b.groupby(['ROI', 'clean_ID'], as_index=False)['theta'].mean()
    df_b_tidy['Group'] = group_b_name
    combined_df = pd.concat([df_a_tidy, df_b_tidy], ignore_index=True)

    # Plot Bars
    sns.barplot(
        data=combined_df, x='ROI', y='theta', hue='Group',
        hue_order=[group_a_name, group_b_name], order=rois,
        ax=ax, errorbar=('ci', 95), capsize=0.08,
        palette=['#2b5c8f', '#d95f02'], edgecolor='black', linewidth=1.5, zorder=1
    )
    # Plot Individual Dots (Dodged over the specific bar widths)
    sns.swarmplot(
        data=combined_df, x='ROI', y='theta', hue='Group',
        hue_order=[group_a_name, group_b_name], order=rois,
        ax=ax, size=6, color='black', alpha=0.5, dodge=True, zorder=2,
        legend=False
    )

    global_max = combined_df['theta'].max()
    global_min = combined_df['theta'].min()
    y_top = global_max * 1.45 if global_max > 0 else 0.5e-6
    y_bottom = global_min * 1.15 if global_min < 0 else -0.15 * global_max
    ax.set_ylim(y_bottom, y_top)

    if graph_bounds is not None and graph_bounds > 0.0 and graph_bounds < 0.5:
        ax.set_ylim(-graph_bounds, graph_bounds)

    # Draw professional brackets over paired bars
    for idx, row in df_group.iterrows():
        roi_name = row['ROI']
        p_val_corr = row['p_corrected']
        
        roi_points = combined_df[combined_df['ROI'] == roi_name]['theta']
        max_y = roi_points.max() if len(roi_points) > 0 else 0
        
        x_a = idx - 0.2  # Approximate left bar X offset
        x_b = idx + 0.2  # Approximate right bar X offset
        y_bracket = max_y + (global_max * 0.08)
        h_tick = global_max * 0.02

        if p_val_corr < p_threshold:
            sig_symbol = "***" if p_val_corr < 0.001 else "**" if p_val_corr < 0.01 else "*"
            # Draw standard bracket line
            ax.plot([x_a, x_a, x_b, x_b], [y_bracket - h_tick, y_bracket, y_bracket, y_bracket - h_tick], color='black', lw=1.2)
            ax.text(
                idx, y_bracket + (global_max * 0.02), f"{sig_symbol}\np_corr = {p_val_corr:.3f}",
                ha='center', va='bottom', fontsize=9, fontweight='bold', color='red'
            )
        else:
            ax.text(
                idx, y_bracket, "n.s.",
                ha='center', va='bottom', fontsize=9, color='gray'
            )

    ax.axhline(0, color='black', linewidth=1, linestyle='--')
    ax.set_ylabel(r'Hemodynamic Response ($\Delta$ HbO)', fontsize=12)
    ax.set_xlabel('Region of Interest (ROI)', fontsize=12)
    ax.set_title(f"Cross-Group Comparison: {group_a_name} vs {group_b_name}\n({target_chroma.upper()} - {selected_event})", fontsize=13, fontweight='bold', pad=15)
    plt.tight_layout()

    # 4. Channel-by-Channel Group-Contrast Topography Map (Zero Hardcoding)
    if df_cha_all is not None and raw_haemo is not None:
        print(f"--> Computing group-level channel contrasts for topography...")
        try:
            val_col = 'effect' if 'effect' in df_cha_all.columns else 'theta'
            ch_col = 'ch_name' if 'ch_name' in df_cha_all.columns else 'channel'
            
            # Match channel levels and clean IDs
            con_summary = df_cha_all[
                (df_cha_all['Chroma'] == target_chroma) & 
                (df_cha_all['Condition'] == selected_event)
            ].copy()
            con_summary['clean_ID'] = con_summary['ID']
            
            raw_picked = raw_haemo.copy().pick(picks=target_chroma)
            
            # --- Perform manual Channel-by-Channel Two-Sample t-tests ---
            contrast_data = []
            for ch in raw_picked.ch_names:
                ch_a = con_summary[(con_summary['clean_ID'].isin(clean_a)) & (con_summary[ch_col] == ch)]
                ch_b = con_summary[(con_summary['clean_ID'].isin(clean_b)) & (con_summary[ch_col] == ch)]
                
                vals_a = ch_a[val_col].dropna().values
                vals_b = ch_b[val_col].dropna().values
                
                if len(vals_a) >= min_subjects and len(vals_b) >= min_subjects:
                    t_stat, p_val = ttest_ind(vals_a, vals_b, equal_var=False)
                    mean_diff = np.mean(vals_a) - np.mean(vals_b)
                else:
                    t_stat, p_val, mean_diff = 0.0, 1.0, 0.0
                    
                contrast_data.append({
                    'ch_name': ch,
                    'Coef.': mean_diff,  # Represents Mean A - Mean B
                    't': t_stat,
                    'P>|t|': p_val,
                    'Chroma': target_chroma,        # For threshold masking
                })
                
            con_model_df = pd.DataFrame(contrast_data)

            # --- DYNAMIC ROI PARSING ---
            if roi_channel_maps:
                def _lookup_roi(row):
                    m = roi_channel_maps.get(row['clean_ID'], {})
                    ch = row[ch_col]
                    return m.get(ch, m.get(ch.split()[0]) if isinstance(ch, str) else None)

                con_summary['ROI'] = con_summary.apply(_lookup_roi, axis=1)

            unique_rois = [r for r in con_summary['ROI'].dropna().unique() if r != ""] if 'ROI' in con_summary.columns else ['All_Channels']

            # Shared symmetric limits for the color bar
            max_abs = np.max(np.abs(con_model_df['Coef.'].dropna().values)) if len(con_model_df['Coef.']) > 0 else 1.0
            vlim = (-max_abs, max_abs) if max_abs > 1e-9 else (None, None)

            fig_topo, ax_topo = plt.subplots(figsize=(6, 6))
            
            # Isolated dynamic plotting loop to prevent spatial bleeding
            for i, roi_name in enumerate(unique_rois):
                roi_ch_names = con_summary[con_summary['ROI'] == roi_name][ch_col].unique().tolist() if 'ROI' in con_summary.columns else raw_picked.ch_names
                roi_ch_names = [ch for ch in roi_ch_names if ch in raw_picked.ch_names]
                
                if not roi_ch_names:
                    continue
                
                raw_roi = raw_picked.copy().pick(picks=roi_ch_names)
                show_colorbar = (i == len(unique_rois) - 1)
                
                # Filter contrast DF to current ROI channels
                roi_con_df = con_model_df[con_model_df['ch_name'].isin(roi_ch_names)].copy()
                
                plot_glm_group_topo(
                    raw_roi, 
                    roi_con_df, 
                    colorbar=show_colorbar,
                    threshold=threshold_topo,
                    axes=ax_topo,
                    cmap='RdBu_r',
                    vlim=vlim
                )

            threshold_text = "p < 0.05 Masked" if threshold_topo else "Unthresholded Contrast"
            ax_topo.set_title(f"Group Contrast: {group_a_name} - {group_b_name}\n({target_chroma.upper()} - {threshold_text})", fontsize=11, fontweight='bold', pad=10)
            plt.tight_layout()
            plt.show()
            
        except Exception as e:
            logger.error(f"Could not generate group-contrast topography plot: {e}", exc_info=True)

    return df_group



def run_inter_group_laterality_analysis(
    df_roi_all_a: DataFrame,
    df_roi_all_b: DataFrame,
    roi_pairs: tuple[str, str] | None,
    condition: str | None,
    group_a_name: str = "Group A",
    group_b_name: str = "Group B",
    target_chroma: str = "hbo",
    min_subjects: int = 3,
    p_threshold: float = 0.05,
    correction_method: str | None = None,
    roi_contra_label: str | None = None,
    roi_ipsi_label: str | None = None,
) -> DataFrame:
    
    """
    Compare LATERALITY between two independent groups of subjects (e.g. a
    control group vs. a target group), using Welch's t-test on each
    subject's within-subject laterality index rather than on raw ROI values.
 
    --------------------------------------------------------------------
    HOW THIS DIFFERS FROM run_cross_group_second_level_analysis
    --------------------------------------------------------------------
    run_cross_group_second_level_analysis (existing):
      - Compares one ROI's raw theta value between two groups directly
        (Group A's Right_PFC vs Group B's Right_PFC, say).
      - CLAIM IF SIGNIFICANT: this ROI's response magnitude differs between
        the two populations, for this condition.
      - WHAT IT DOES NOT SAY: whether that difference reflects a real,
        localized neural difference or a generic between-population
        difference unrelated to the specific task — e.g. different overall
        vascular reactivity, arousal, skull/scalp optical properties, or
        anything else that would shift a group's numbers up or down
        everywhere, not just in this ROI. Two independently recruited
        groups (e.g. patients vs. healthy controls) are considerably more
        likely to differ in these generic ways than two subsets of one
        study population, which makes this ambiguity a real risk here, not
        a theoretical one.
 
    run_cross_group_laterality_analysis (this function):
      - First computes each subject's OWN laterality index
        (contralateral ROI theta - ipsilateral ROI theta, within that
        subject, for one condition) — the same computation as
        run_roi_paired_contrast_analysis, just not yet tested there.
      - Then compares those per-subject laterality indices between the two
        groups with Welch's t-test.
      - CLAIM IF SIGNIFICANT: the DEGREE OF SPATIAL SPECIFICITY (how much
        more one hemisphere responds than the other, within a person)
        differs between the two groups — a claim about lateralization
        itself, not raw magnitude. Subtracting within-subject first cancels
        out whatever's common to both hemispheres for that person (general
        reactivity, arousal, etc.) before ever comparing across groups, so
        a significant result here is harder to explain away as a generic
        between-population confound.
      - WHAT IT DOES NOT SAY: anything about whether overall response
        magnitude differs between groups (a group could have identical
        laterality but very different raw amplitude — that's what the
        existing cross-group function is for) — and it only uses subjects
        who have BOTH the contra and ipsi ROI valid, so it can lose
        subjects the raw-ROI comparison would have kept.
 
    Use both, for different questions: the existing function for "is the
    raw response different between groups," this one for "is the
    LATERALIZATION different between groups."
 
    Parameters
    ----------
    df_roi_all_a, df_roi_all_b : pd.DataFrame
        Individual-level ROI results (['ROI', 'Condition', 'Chroma',
        'theta', 'ID']) for Group A and Group B RESPECTIVELY. Keeping them
        as separate frames (rather than one combined frame + ID lists)
        avoids any risk of cross-dataset ID collisions when the two groups
        come from genuinely separate studies/exports.
    roi_pairs : tuple(str, str) or list of tuple(str, str)
        (roi_contra, roi_ipsi) pair(s). Each pair's laterality index is
        computed as theta(roi_contra) - theta(roi_ipsi), per subject.
        Pass a list to test multiple hand/condition combinations in one call.
    condition : str or list of str
        The 'Condition' value (e.g. contrast name or event code) to use for
        each pair. Single value applies to all pairs; otherwise must match
        len(roi_pairs).
    target_chroma : str, default 'hbo'
        Chromophore to test. Never mix hbo/hbr in one laterality index.
    min_subjects : int, default 3
        Minimum subjects required in EACH group (after requiring both ROIs
        be present) for a pair to be tested.
    p_threshold : float, default 0.05
        Significance threshold for the (optionally corrected) p-value.
    correction_method : str or None, default None
        Multiple comparisons correction across the pairs tested in this
        call. Off by default for a single pre-specified pair; turn on
        ('fdr_bh') if testing several pairs/conditions at once.
    roi_contra_label, roi_ipsi_label : str or list of str, optional
        Display labels for the contra/ipsi ROI in each pair.
 
    Returns
    -------
    pd.DataFrame, one row per tested pair:
        ['roi_contra', 'roi_ipsi', 'condition', 'mean_A', 'mean_B',
         'mean_diff', 't_val', 'p_val', 'p_corrected', 'significant',
         'n_A', 'n_B']
    """
 
    required_cols = ['ROI', 'Condition', 'Chroma', 'theta', 'ID']
    for name, df in [('df_roi_all_a', df_roi_all_a), ('df_roi_all_b', df_roi_all_b)]:
        if not all(col in df.columns for col in required_cols):
            raise ValueError(f"{name} must include: {required_cols}")
 
    if isinstance(roi_pairs, tuple):
        roi_pairs = [roi_pairs]
    n_pairs = len(roi_pairs)
 
    if isinstance(condition, str):
        conditions = [condition] * n_pairs
    else:
        if len(condition) != n_pairs:
            raise ValueError("If passing a list of conditions, it must match len(roi_pairs).")
        conditions = list(condition)
 
    def _expand(labels):
        if labels is None:
            return [None] * n_pairs
        if isinstance(labels, str):
            return [labels] * n_pairs
        if len(labels) != n_pairs:
            raise ValueError("Label list length must match len(roi_pairs).")
        return list(labels)
 
    contra_labels = _expand(roi_contra_label)
    ipsi_labels = _expand(roi_ipsi_label)
 
    def _laterality_per_subject(df_roi_all, roi_contra, roi_ipsi, cond):
        """Collapse to one laterality index per subject, for one group."""
        df_chroma = df_roi_all[df_roi_all['Chroma'] == target_chroma]
        df_cond = df_chroma[df_chroma['Condition'] == cond]
 
        contra_vals = df_cond[df_cond['ROI'] == roi_contra].groupby('ID', as_index=False)['theta'].mean()
        ipsi_vals = df_cond[df_cond['ROI'] == roi_ipsi].groupby('ID', as_index=False)['theta'].mean()
 
        merged = contra_vals.merge(ipsi_vals, on='ID', suffixes=('_contra', '_ipsi'))
        merged['laterality'] = merged['theta_contra'] - merged['theta_ipsi']
        return merged[['ID', 'laterality']]
 
    results = []
    plot_rows = []
 
    for (roi_contra, roi_ipsi), cond, lbl_c, lbl_i in zip(roi_pairs, conditions, contra_labels, ipsi_labels):
        lat_a = _laterality_per_subject(df_roi_all_a, roi_contra, roi_ipsi, cond)
        lat_b = _laterality_per_subject(df_roi_all_b, roi_contra, roi_ipsi, cond)
 
        n_a, n_b = lat_a['ID'].nunique(), lat_b['ID'].nunique()
        if n_a < min_subjects or n_b < min_subjects:
            logger.warning(
                f"Skipping pair ({roi_contra} - {roi_ipsi}) for condition '{cond}' — "
                f"{group_a_name} n={n_a}, {group_b_name} n={n_b}, need at least {min_subjects} in EACH."
            )
            continue
 
        vals_a = lat_a['laterality'].values
        vals_b = lat_b['laterality'].values
 
        t_val, p_val = ttest_ind(vals_a, vals_b, equal_var=False)
        mean_a, mean_b = np.mean(vals_a), np.mean(vals_b)
        mean_diff = mean_a - mean_b
 
        pair_label = f"{lbl_c or roi_contra} - {lbl_i or roi_ipsi}\n({cond})"
        results.append({
            'roi_contra': roi_contra,
            'roi_ipsi': roi_ipsi,
            'label': pair_label,
            'condition': cond,
            'mean_A': mean_a,
            'mean_B': mean_b,
            'mean_diff': mean_diff,
            't_val': t_val,
            'p_val': p_val,
            'n_A': n_a,
            'n_B': n_b,
        })
        plot_rows.append(lat_a.assign(pair=pair_label, Group=group_a_name))
        plot_rows.append(lat_b.assign(pair=pair_label, Group=group_b_name))
 
    if not results:
        print("\n[ERROR] No ROI pairs met the minimum subject threshold for BOTH groups.\n")
        return pd.DataFrame()
 
    df_group = pd.DataFrame(results)
 
    if correction_method is not None:
        reject, p_corrected, _, _ = multipletests(df_group['p_val'].values, method=correction_method)
        df_group['p_corrected'] = p_corrected
        df_group['significant'] = reject
    else:
        df_group['p_corrected'] = df_group['p_val']
        df_group['significant'] = df_group['p_val'] <= p_threshold
 
    # --- Print report ---
    print("\n" + "=" * 85)
    print(f"       CROSS-GROUP LATERALITY CONTRAST: {group_a_name.upper()} vs {group_b_name.upper()} ({target_chroma.upper()})")
    print("=" * 85)
    df_print = df_group.copy()
    for c in ['mean_A', 'mean_B', 'mean_diff']:
        df_print[c] = df_print[c].apply(lambda x: f"{x:.4f}")
    df_print['t_val'] = df_print['t_val'].apply(lambda x: f"{x:.3f}")
    df_print['p_corrected'] = df_print['p_corrected'].apply(lambda x: f"{x:.4f}")
    print(df_print[['label', 'condition', 'mean_A', 'mean_B', 'mean_diff',
                     't_val', 'p_corrected', 'significant', 'n_A', 'n_B']].to_string(index=False))
    print("=" * 85 + "\n")
 
    # --- Plot: grouped bar + swarm per pair, Group A vs Group B ---
    sns.set_theme(style="whitegrid")
    fig, ax = plt.subplots(figsize=(max(7, 3 * len(results)), 6))
 
    plot_df = pd.concat(plot_rows, ignore_index=True)
    pair_order = [r['label'] for r in results]
 
    sns.barplot(
        data=plot_df, x='pair', y='laterality', hue='Group',
        order=pair_order, hue_order=[group_a_name, group_b_name],
        ax=ax, errorbar=('ci', 95), capsize=0.08,
        palette=['#2b5c8f', '#d95f02'], edgecolor='black', linewidth=1.5, zorder=1
    )
    sns.swarmplot(
        data=plot_df, x='pair', y='laterality', hue='Group',
        order=pair_order, hue_order=[group_a_name, group_b_name],
        ax=ax, size=6, palette=['black', 'black'], alpha=0.5, dodge=True, zorder=2,
        legend=False
    )
 
    ax.axhline(0, color='black', linewidth=1, linestyle='--')
 
    global_max = plot_df['laterality'].max()
    for i, row in df_group.iterrows():
        pair_points = plot_df[plot_df['pair'] == row['label']]['laterality']
        max_y = pair_points.max() if len(pair_points) > 0 else 0
        text_y = max_y + (abs(global_max) * 0.1 if global_max else 0.1)
 
        p_val_corr = row['p_corrected']
        if p_val_corr < p_threshold:
            sig_symbol = "***" if p_val_corr < 0.001 else "**" if p_val_corr < 0.01 else "*"
            ax.text(i, text_y, f"{sig_symbol}\np = {p_val_corr:.3f}",
                    ha='center', va='bottom', fontsize=10, fontweight='bold', color='red')
        else:
            ax.text(i, text_y, f"n.s.\np = {p_val_corr:.3f}",
                    ha='center', va='bottom', fontsize=10, color='gray')
 
    ax.set_ylabel(r'Laterality Index ($\Delta$ HbO, Contra $-$ Ipsi)', fontsize=12)
    ax.set_xlabel('')
    correction_lbl = f"({correction_method} corrected)" if correction_method else "(uncorrected)"
    ax.set_title(
        f"Cross-Group Laterality Comparison ({target_chroma.upper()})\n"
        f"{group_a_name} vs {group_b_name}, p < {p_threshold} {correction_lbl}",
        fontsize=13, fontweight='bold', pad=15
    )
    plt.tight_layout()
    plt.show()
 
    return df_group



def run_inter_group_contrast_analysis(
    df_contrasts_a: DataFrame,
    df_contrasts_b: DataFrame,
    contrast_name: str,
    roi_channel_maps_a: dict[str, dict[str, str]],
    roi_channel_maps_b: dict[str, dict[str, str]],
    group_a_name: str = "Group A",
    group_b_name: str = "Group B",
    target_chroma: str = "hbo",
    min_subjects: int = 3,
    p_threshold: float = 0.05,
    correction_method: str = "fdr_bh",
    weighted: bool = True,
) -> DataFrame:
    
    """
    Compare a JOINT-FIT TASK CONTRAST (e.g. '2.0_vs_3.0'), aggregated to ROI
    level, between two independent groups. This is the cross-group analog
    of the inter-group joint-contrast method — where that method asks "does
    this contrast differ from zero within one group," this asks "does the
    SIZE of this contrast differ between two groups."

    --------------------------------------------------------------------
    HOW THIS DIFFERS FROM THE OTHER TWO CROSS-GROUP METHODS
    --------------------------------------------------------------------
    run_cross_group_second_level_analysis: compares raw single-condition
      ROI magnitude between groups — vulnerable to generic between-
      population differences (vascular reactivity, arousal, etc.) that
      have nothing to do with the task.
    run_cross_group_laterality_analysis: compares each subject's own
      contra-minus-ipsi laterality index between groups — asks whether
      spatial specificity differs, says nothing about overall magnitude.
    run_cross_group_contrast_analysis (this function): compares a
      jointly-fit task contrast (e.g. Task A minus Task B, estimated
      together within each subject's GLM) between groups — asks whether
      one group differentiates between the two tasks more/less than the
      other does, at this ROI. Systemic noise is cancelled at the
      model-fitting stage (same GLM, both conditions) rather than left in
      raw single-condition magnitude, or cancelled only by within-subject
      spatial subtraction as in the laterality method. This is generally
      the most statistically efficient of the three at detecting a real
      between-group difference in TASK-SPECIFIC response, but — like the
      inter-group version of this same idea — it does not by itself tell
      you WHERE that difference is spatially localized unless you compare
      the sign/pattern across multiple ROIs.

    Parameters
    ----------
    df_contrasts_a, df_contrasts_b : pd.DataFrame
        Combined CHANNEL-LEVEL contrast results (contrasts.csv format) for
        Group A and Group B respectively. Must include:
        ['ch_name', 'effect', 'stat', 'Chroma', 'contrast_name', 'ID']
        Kept as separate frames per group (not one combined frame + ID
        lists) for the same reason as run_cross_group_laterality_analysis —
        avoids any risk of ID-matching mismatches between two genuinely
        separate dataset exports.
    contrast_name : str
        Which contrast to test (e.g. '2.0_vs_3.0'). Must exist in both
        groups' df_contrasts for a fair comparison.
    roi_json_path : str
        Path to the regions.json used elsewhere in the pipeline.
    target_chroma : str, default 'hbo'
    min_subjects : int, default 3
        Minimum subjects required in EACH group, per ROI.
    p_threshold : float, default 0.05
    correction_method : str or None, default 'fdr_bh'
        Unlike the paired/laterality functions (which default to no
        correction, since they test one pre-specified pair), this defaults
        ON — this function screens across every ROI in regions.json, which
        is an open multiple-comparisons scan, not a single planned contrast.
    weighted : bool, default True
        Passed through to aggregate_channel_contrasts_to_roi (inverse-
        variance weighting vs. plain mean across channels within an ROI).

    Returns
    -------
    pd.DataFrame, one row per ROI:
        ['ROI', 'mean_A', 'mean_B', 'mean_diff', 't_val', 'p_val',
         'p_corrected', 'significant', 'n_A', 'n_B']
    """

    required_cols = ['ch_name', 'effect', 'stat', 'Chroma', 'contrast_name', 'ID']
    for name, df in [('df_contrasts_a', df_contrasts_a), ('df_contrasts_b', df_contrasts_b)]:
        if not all(col in df.columns for col in required_cols):
            raise ValueError(f"{name} must include: {required_cols}")

    # Filter to the requested contrast BEFORE aggregating, so a missing
    # contrast_name fails clearly here rather than silently downstream.
    df_a_filt = df_contrasts_a[df_contrasts_a['contrast_name'] == contrast_name]
    df_b_filt = df_contrasts_b[df_contrasts_b['contrast_name'] == contrast_name]

    if df_a_filt.empty:
        print(f"[ERROR] Contrast '{contrast_name}' not found anywhere in {group_a_name}'s data.")
        return DataFrame()
    if df_b_filt.empty:
        print(f"[ERROR] Contrast '{contrast_name}' not found anywhere in {group_b_name}'s data.")
        return DataFrame()

    roi_a = aggregate_channel_contrasts_to_roi(df_a_filt, roi_channel_maps_a, weighted=weighted)
    roi_b = aggregate_channel_contrasts_to_roi(df_b_filt, roi_channel_maps_b, weighted=weighted)

    roi_a = roi_a[roi_a['Chroma'] == target_chroma]
    roi_b = roi_b[roi_b['Chroma'] == target_chroma]

    if roi_a.empty or roi_b.empty:
        print(f"[ERROR] No ROI-aggregated values produced for one or both groups "
              f"(check regions.json channel names against this montage).")
        return DataFrame()

    all_rois = sorted(set(roi_a['ROI'].unique()) | set(roi_b['ROI'].unique()))
    results = []
    plot_rows = []

    for roi in all_rois:
        vals_a = roi_a[roi_a['ROI'] == roi]['theta'].values
        vals_b = roi_b[roi_b['ROI'] == roi]['theta'].values

        n_a, n_b = len(vals_a), len(vals_b)
        if n_a < min_subjects or n_b < min_subjects:
            logger.warning(
                f"Skipping ROI '{roi}' — {group_a_name} n={n_a}, {group_b_name} n={n_b}, "
                f"need at least {min_subjects} in EACH."
            )
            continue

        t_val, p_val = ttest_ind(vals_a, vals_b, equal_var=False)
        mean_a, mean_b = np.mean(vals_a), np.mean(vals_b)
        mean_diff = mean_a - mean_b

        results.append({
            'ROI': roi, 'mean_A': mean_a, 'mean_B': mean_b, 'mean_diff': mean_diff,
            't_val': t_val, 'p_val': p_val, 'n_A': n_a, 'n_B': n_b,
        })
        plot_rows.append(pd.DataFrame({'theta': vals_a, 'ROI': roi, 'Group': group_a_name}))
        plot_rows.append(pd.DataFrame({'theta': vals_b, 'ROI': roi, 'Group': group_b_name}))

    if not results:
        print("\n[ERROR] No ROIs met the subject requirements for BOTH groups.\n")
        return pd.DataFrame()

    df_group = pd.DataFrame(results)

    if correction_method is not None:
        reject, p_corrected, _, _ = multipletests(df_group['p_val'].values, method=correction_method)
        df_group['p_corrected'] = p_corrected
        df_group['significant'] = reject
    else:
        df_group['p_corrected'] = df_group['p_val']
        df_group['significant'] = df_group['p_val'] <= p_threshold

    # --- Print report ---
    print("\n" + "=" * 85)
    print(f"       CROSS-GROUP CONTRAST COMPARISON: {group_a_name.upper()} vs {group_b_name.upper()} ({target_chroma.upper()})")
    print(f"       Contrast: {contrast_name}")
    print("=" * 85)
    df_print = df_group.copy()
    for c in ['mean_A', 'mean_B', 'mean_diff']:
        df_print[c] = df_print[c].apply(lambda x: f"{x:.4f}")
    df_print['t_val'] = df_print['t_val'].apply(lambda x: f"{x:.3f}")
    df_print['p_corrected'] = df_print['p_corrected'].apply(lambda x: f"{x:.4f}")
    print(df_print[['ROI', 'mean_A', 'mean_B', 'mean_diff', 't_val',
                     'p_corrected', 'significant', 'n_A', 'n_B']].to_string(index=False))
    print("=" * 85 + "\n")

    # --- Plot: grouped bar + swarm per ROI, Group A vs Group B, with brackets ---
    sns.set_theme(style="whitegrid")
    fig, ax = plt.subplots(figsize=(max(7, 2.5 * len(results)), 6))

    plot_df = pd.concat(plot_rows, ignore_index=True)
    roi_order = [r['ROI'] for r in results]

    sns.barplot(
        data=plot_df, x='ROI', y='theta', hue='Group',
        order=roi_order, hue_order=[group_a_name, group_b_name],
        ax=ax, errorbar=('ci', 95), capsize=0.08,
        palette=['#2b5c8f', '#d95f02'], edgecolor='black', linewidth=1.5, zorder=1
    )
    sns.swarmplot(
        data=plot_df, x='ROI', y='theta', hue='Group',
        order=roi_order, hue_order=[group_a_name, group_b_name],
        ax=ax, size=6, palette=['black', 'black'], alpha=0.5, dodge=True, zorder=2,
        legend=False
    )

    global_max = plot_df['theta'].max()
    global_min = plot_df['theta'].min()
    y_top = global_max * 1.45 if global_max > 0 else 0.5e-6
    y_bottom = global_min * 1.15 if global_min < 0 else -0.15 * global_max
    ax.set_ylim(y_bottom, y_top)

    for i, row in df_group.iterrows():
        roi_points = plot_df[plot_df['ROI'] == row['ROI']]['theta']
        max_y = roi_points.max() if len(roi_points) > 0 else 0
        x_a, x_b = i - 0.2, i + 0.2
        y_bracket = max_y + (global_max * 0.08 if global_max else 0.05)
        h_tick = global_max * 0.02 if global_max else 0.01

        p_val_corr = row['p_corrected']
        if p_val_corr < p_threshold:
            sig_symbol = "***" if p_val_corr < 0.001 else "**" if p_val_corr < 0.01 else "*"
            ax.plot([x_a, x_a, x_b, x_b],
                     [y_bracket - h_tick, y_bracket, y_bracket, y_bracket - h_tick],
                     color='black', lw=1.2)
            ax.text(i, y_bracket + (global_max * 0.02 if global_max else 0.01),
                     f"{sig_symbol}\np_corr = {p_val_corr:.3f}",
                     ha='center', va='bottom', fontsize=9, fontweight='bold', color='red')
        else:
            ax.text(i, y_bracket, "n.s.", ha='center', va='bottom', fontsize=9, color='gray')

    ax.axhline(0, color='black', linewidth=1, linestyle='--')
    ax.set_ylabel(r'Contrast Effect ($\Delta$ HbO)', fontsize=12)
    ax.set_xlabel('Region of Interest (ROI)', fontsize=12)
    correction_lbl = f"({correction_method} corrected)" if correction_method else "(uncorrected)"
    ax.set_title(
        f"Cross-Group Contrast Comparison: {group_a_name} vs {group_b_name}\n"
        f"({target_chroma.upper()} - {contrast_name}) {correction_lbl}",
        fontsize=13, fontweight='bold', pad=15
    )
    plt.tight_layout()
    plt.show()

    return df_group



def run_roi_paired_contrast_analysis(
    df_roi_all: DataFrame,
    roi_pairs: Sequence[tuple[str, str]] | list[list[str]],
    condition: str,
    target_chroma: str = 'hbo',
    min_subjects: int = 5,
    p_threshold: float = 0.05,
    correction_method: str | None = None,
    roi_a_label: str | None = None,
    roi_b_label: str | None = None,
) -> DataFrame:
    """
    Paired within-subject ROI contrast (e.g. contralateral minus ipsilateral
    motor ROI), as a companion to run_roi_second_level_analysis rather than a
    replacement for it. Where run_roi_second_level_analysis tests each ROI's
    theta against zero independently (still contaminated by systemic/global
    physiology shared across the whole head), this function computes, per
    subject, (ROI_A theta - ROI_B theta) for a single condition and tests
    THAT difference against zero. Any systemic component that's roughly equal
    in both ROIs cancels out in the subtraction itself, rather than being
    inferred afterwards by comparing two separate p-values.

    This is the more powerful, more directly interpretable test whenever you
    already have a specific hypothesis about which two ROIs should differ
    (e.g. laterality) — use run_roi_second_level_analysis for open-ended
    per-ROI screening, and this function for a pre-specified paired
    comparison you want to report as a single confirmatory statistic.

    Parameters
    ----------
    df_roi_all : pd.DataFrame
        Combined individual-level ROI results across subjects.
        Must include: ['ROI', 'Condition', 'Chroma', 'theta', 'ID']
    roi_pairs : tuple(str, str) or list of tuple(str, str)
        One (roi_a, roi_b) pair, or several. Each pair is tested
        independently as (roi_a - roi_b). Passing several pairs lets you
        e.g. test left-hand-tap laterality and right-hand-tap laterality
        (different `condition` values) in one call/figure.
    condition : str or list of str
        The 'Condition' value to filter to for the paired test. If
        `roi_pairs` has multiple pairs and you want a different condition
        per pair, pass a list of the same length as `roi_pairs`; otherwise
        a single value is used for every pair.
    target_chroma : str, default 'hbo'
        Chromophore to test. HbO and HbR should never be tested together.
    min_subjects : int, default 5
        Minimum number of subjects with BOTH ROI_A and ROI_B present (after
        dropping NaNs) required to run the test. Below this, the pair is
        skipped with a warning rather than silently reported.
    p_threshold : float, default 0.05
        Significance threshold applied to the (optionally corrected) p-value.
    correction_method : str or None, default None
        Multiple comparisons correction across the pairs tested in this call
        (statsmodels.stats.multitest.multipletests method name, e.g.
        'fdr_bh'). Left off by default since a single pre-specified paired
        contrast typically doesn't need correction — turn it on if you're
        testing several pairs in the same call and want to control for that.
    roi_a_label, roi_b_label : str or list of str, optional
        Display labels for each pair's ROI_A/ROI_B (defaults to the raw ROI
        names). If testing multiple pairs, pass lists matching `roi_pairs`.

    Returns
    -------
    pd.DataFrame with one row per tested pair:
        ['roi_a', 'roi_b', 'condition', 't_val', 'p_val', 'p_corrected',
         'significant', 'mean_diff', 'n_subjects']
    """

    required_cols = ['ROI', 'Condition', 'Chroma', 'theta', 'ID']
    if not all(col in df_roi_all.columns for col in required_cols):
        raise ValueError(f"Input ROI DataFrame must include: {required_cols}")

    # Normalize inputs to lists so single-pair and multi-pair calls share code.
    if isinstance(roi_pairs, tuple):
        roi_pairs = [roi_pairs]
    n_pairs = len(roi_pairs)

    if isinstance(condition, str):
        conditions = [condition] * n_pairs
    else:
        if len(condition) != n_pairs:
            raise ValueError("If passing a list of conditions, it must match len(roi_pairs).")
        conditions = list(condition)

    def _expand_labels(labels, default_from):
        if labels is None:
            return [None] * n_pairs
        if isinstance(labels, str):
            return [labels] * n_pairs
        if len(labels) != n_pairs:
            raise ValueError("Label list length must match len(roi_pairs).")
        return list(labels)

    roi_a_labels = _expand_labels(roi_a_label, roi_pairs)
    roi_b_labels = _expand_labels(roi_b_label, roi_pairs)

    df_chroma = df_roi_all[df_roi_all['Chroma'] == target_chroma].copy()
    df_chroma = df_chroma.dropna(subset=['theta'])

    results = []
    diff_data_for_plot = []  # keep per-subject diffs around for plotting

    for (roi_a, roi_b), cond, lbl_a, lbl_b in zip(roi_pairs, conditions, roi_a_labels, roi_b_labels):
        df_cond = df_chroma[df_chroma['Condition'] == cond]

        a_vals = df_cond[df_cond['ROI'] == roi_a].groupby('ID', as_index=False)['theta'].mean()
        b_vals = df_cond[df_cond['ROI'] == roi_b].groupby('ID', as_index=False)['theta'].mean()

        # Inner join on ID: only subjects with BOTH ROIs present for this
        # condition contribute to the paired test.
        merged = a_vals.merge(b_vals, on='ID', suffixes=('_a', '_b'))
        merged['diff'] = merged['theta_a'] - merged['theta_b']

        n_subs = merged['ID'].nunique()
        if n_subs < min_subjects:
            logger.warning(
                f"Skipping pair ({roi_a} - {roi_b}) for condition '{cond}' — "
                f"only {n_subs} subject(s) have both ROIs, need at least {min_subjects}."
            )
            continue

        Y = merged['diff'].values
        t_val, p_val = ttest_1samp(Y, 0)
        mean_diff = np.mean(Y)

        results.append({
            'roi_a': roi_a,
            'roi_b': roi_b,
            'label_a': lbl_a or roi_a,
            'label_b': lbl_b or roi_b,
            'condition': cond,
            't_val': t_val,
            'p_val': p_val,
            'mean_diff': mean_diff,
            'n_subjects': n_subs,
        })
        diff_data_for_plot.append(merged.assign(pair=f"{lbl_a or roi_a} - {lbl_b or roi_b}\n({cond})"))

    if not results:
        print("\n[ERROR] No ROI pairs met the minimum subject threshold.\n")
        return pd.DataFrame()

    df_group = pd.DataFrame(results)

    if correction_method is not None:
        reject, p_corrected, _, _ = multipletests(df_group['p_val'].values, method=correction_method)
        df_group['p_corrected'] = p_corrected
        df_group['significant'] = reject
    else:
        df_group['p_corrected'] = df_group['p_val']
        df_group['significant'] = df_group['p_val'] <= p_threshold

    # --- Print report ---
    print("\n" + "=" * 70)
    print(f"       PAIRED ROI CONTRAST RESULTS ({target_chroma.upper()})")
    print("=" * 70)
    df_print = df_group.copy()
    df_print['mean_diff'] = df_print['mean_diff'].apply(lambda x: f"{x:.4f}")
    df_print['t_val'] = df_print['t_val'].apply(lambda x: f"{x:.3f}")
    df_print['p_val'] = df_print['p_val'].apply(lambda x: f"{x:.4f}")
    df_print['p_corrected'] = df_print['p_corrected'].apply(lambda x: f"{x:.4f}")
    print(df_print[['label_a', 'label_b', 'condition', 'mean_diff', 't_val',
                     'p_val', 'p_corrected', 'significant', 'n_subjects']].to_string(index=False))
    print("=" * 70 + "\n")

    # --- Plot: one bar per pair, individual subject differences overlaid ---
    sns.set_theme(style="whitegrid")
    fig, ax = plt.subplots(figsize=(max(6, 2.2 * len(results)), 6))

    plot_df = pd.concat(diff_data_for_plot, ignore_index=True)

    sns.barplot(
        data=plot_df, x='pair', y='diff', ax=ax,
        errorbar=('ci', 95), capsize=0.1,
        color='lightgray', edgecolor='black', linewidth=1.5, zorder=1
    )
    sns.swarmplot(
        data=plot_df, x='pair', y='diff', ax=ax,
        color='darkblue', size=8, alpha=0.7, zorder=2
    )

    ax.axhline(0, color='black', linewidth=1, linestyle='--')

    global_max = plot_df['diff'].max()
    for i, row in df_group.iterrows():
        pair_label = f"{row['label_a']} - {row['label_b']}\n({row['condition']})"
        pair_points = plot_df[plot_df['pair'] == pair_label]['diff']
        max_y = pair_points.max() if len(pair_points) > 0 else 0
        text_y = max_y + (abs(global_max) * 0.08 if global_max else 0.1)

        p_val_corr = row['p_corrected']
        if p_val_corr < 0.001:
            sig_symbol = "***"
        elif p_val_corr < 0.01:
            sig_symbol = "**"
        elif p_val_corr < p_threshold:
            sig_symbol = "*"
        else:
            sig_symbol = "n.s."

        ax.text(
            i, text_y, f"{sig_symbol}\np = {p_val_corr:.3f}",
            ha='center', va='bottom', fontsize=11, fontweight='bold',
            color='red' if p_val_corr < p_threshold else 'gray'
        )

    ax.set_ylabel(r'Paired ROI Difference ($\Delta$ HbO, A $-$ B)', fontsize=12)
    ax.set_xlabel('')
    correction_lbl = f"({correction_method} corrected)" if correction_method else "(uncorrected — single pre-specified contrast)"
    ax.set_title(
        f"Paired ROI Contrast ({target_chroma.upper()})\n"
        f"Significance threshold: p < {p_threshold} {correction_lbl}",
        fontsize=13, fontweight='bold', pad=15
    )
    plt.tight_layout()
    plt.show()

    return df_group



def aggregate_channel_contrasts_to_roi(
    df_contrasts: DataFrame,
    roi_channel_maps: dict[str, dict[str, str]],
    weighted: bool = True
) -> DataFrame:
    """
    Combine already-computed per-channel CONTRAST results (e.g. your
    '2.0_vs_3.0' rows from contrasts.csv / contrast_results) into
    per-subject, per-ROI values — so a joint-fit task contrast can be tested
    at the ROI level using the same one-sample machinery as
    run_roi_second_level_analysis / run_roi_paired_contrast_analysis.
 
    This exists because mne_nirs.statistics.RegressionResults has a built-in
    .to_dataframe_region_of_interest() that does inverse-variance-weighted
    channel combination, but the ContrastResults object returned by
    glm_est.compute_contrast() does NOT have that method. This function
    replicates the same weighting logic (weight each channel by the inverse
    of its GLM fit's variance) manually, on the already-exported contrast
    dataframe, rather than requiring you to go back and recompute anything
    from raw GLM objects.
 
    Parameters
    ----------
    df_contrasts : pd.DataFrame
        Combined per-channel contrast results across subjects/contrasts
        (i.e. your contrasts.csv format). Must include:
        ['ch_name', 'effect', 'stat', 'Chroma', 'contrast_name', 'ID']
        `stat` must be the t-statistic (ContrastType == 't'), since standard
        error is recovered as effect / stat.
    roi_channel_maps : dict[str, dict[str, str]]
        Per-subject channel-to-ROI mapping, keyed by subject ID (the same
        ID values used in df_contrasts['ID']), e.g.
        {"sub-01": {"S1_D1 hbo": "Left", "S1_D1 hbr": "Left", ...}, ...}.
        This is the actual mapping generate_roi_results used for that
        subject (whichever tier produced it — JSON, geometric split, or
        per-channel fallback) — not re-derived here, so ROI assignments
        stay consistent with df_ind_dict for the same subject.
    weighted : bool, default True
        If True, combine channels within an ROI using inverse-variance
        weighting (weight = 1 / se^2), matching MNE-NIRS's own default
        behavior for to_dataframe_region_of_interest. If False, channels are
        weighted equally (a plain mean).
 
    Returns
    -------
    pd.DataFrame with columns ['ROI', 'Condition', 'Chroma', 'theta', 'ID'],
    directly usable as `df_roi_all` in run_roi_second_level_analysis or
    run_roi_paired_contrast_analysis. 'Condition' holds the contrast name
    (e.g. '2.0_vs_3.0'), and 'theta' holds the ROI-combined contrast effect.
    """
 
    required_cols = ['ch_name', 'effect', 'stat', 'Chroma', 'contrast_name', 'ID']
    print(df_contrasts.columns)
    if not all(col in df_contrasts.columns for col in required_cols):
        raise ValueError(f"Input contrast DataFrame must include: {required_cols}")
 
    df = df_contrasts.copy()
    df['ch_base'] = df['ch_name'].str.split().str[0]

    def lookup(row):
        m = roi_channel_maps.get(row['ID'], {})
        return m.get(row['ch_name'], m.get(row['ch_base']))

    df['ROI'] = df.apply(lookup, axis=1)
    df = df.dropna(subset=['ROI'])
    
    if df.empty:
        raise ValueError("No channel contrasts matched any subject's ROI mapping.")
 
    # Recover standard error from the t-statistic: t = effect / se -> se = effect / t
    with np.errstate(divide='ignore', invalid='ignore'):
        df['se'] = df['effect'] / df['stat']
    # A zero or near-zero t-stat gives an undefined/huge se; drop those rows
    # from the weighting rather than let them explode the ROI average.
    bad_se = ~np.isfinite(df['se']) | (df['se'] == 0)
    if bad_se.any():
        logger.warning(f"Dropping {bad_se.sum()} channel-rows with non-finite "
                        f"standard error (t-stat ~ 0) from ROI aggregation.")
    df = df[~bad_se]
 
    if weighted:
        df['weight'] = 1.0 / (df['se'] ** 2)
    else:
        df['weight'] = 1.0
 
    group_cols = ['ROI', 'contrast_name', 'Chroma', 'ID']
 
    df['_effect_weight'] = df['effect'] * df['weight']
    roi_theta = df.groupby(group_cols, as_index=False).agg(
        _sum_ew=('_effect_weight', 'sum'),
        _sum_w=('weight', 'sum'),
    )
    roi_theta['theta'] = roi_theta['_sum_ew'] / roi_theta['_sum_w']
 
    roi_theta = roi_theta.rename(columns={'contrast_name': 'Condition'})
    return roi_theta[['ROI', 'Condition', 'Chroma', 'theta', 'ID']]

 
 
def collapse_fir_condition_column(df, value_col, condition_col='Condition',
                                   group_cols=None, delay_sep='_delay_'):
    """
    Collapse FIR delay-bin rows (Condition values like '2.0_delay_6') into a
    single row per base condition ('2.0'), using a PLAIN (equal-weighted)
    mean of `value_col` across whatever delay bins are actually present for
    that condition/subject/unit. SAFE NO-OP for non-FIR data: if no
    Condition value contains `delay_sep`, the input is returned unchanged —
    so this can always be called unconditionally, regardless of HRF_MODEL.
 
    No window (start/end) parameter, deliberately: the set of delay bins
    to average is read directly from whichever bins actually exist in the
    data for that condition — which is itself just a reflection of
    FIR_DELAYS/the design matrix that was actually built — rather than a
    second, separately-configured window that could drift out of sync with
    it. This also makes no assumption about response shape or timing,
    appropriate when response latency is unpredictable (e.g. infant fNIRS).
 
    This is what lets all six group-level statistics functions work
    identically whether the underlying GLM used HRF_MODEL='glover'/'spm'
    (one regressor per condition already) or 'fir' (many delay-bin
    regressors per condition) — the delay-bin collapsing happens once,
    upstream, rather than needing separate handling inside each function.
 
    Parameters
    ----------
    df : pd.DataFrame
        Must include `condition_col` and `value_col`, plus whatever
        `group_cols` you want preserved (e.g. ['ROI', 'Chroma', 'ID'] or
        ['ch_name', 'Chroma', 'ID']).
    value_col : str
        Column to average (e.g. 'theta' for ROI data, 'effect' for
        channel-level data).
    condition_col : str, default 'Condition'
        Column holding condition/delay-bin labels.
    group_cols : list of str, optional
        Columns that define one "unit" to collapse within (e.g. one
        ROI/subject, or one channel/subject). Required whenever FIR rows
        are present, to avoid accidentally collapsing across subjects/ROIs.
    delay_sep : str, default '_delay_'
        Separator used in FIR condition names, matching the naming already
        used elsewhere in the pipeline ("{condition}_delay_{n}").
 
    Returns
    -------
    pd.DataFrame with the same columns as the input, `condition_col`
    holding base condition names, and one row per (group_cols, base
    condition) combination.
    """
    df = df.copy()
    is_fir_row = df[condition_col].astype(str).str.contains(delay_sep, regex=False)
 
    if not is_fir_row.any():
        # Nothing FIR-shaped here — pass through unchanged.
        return df
 
    if group_cols is None:
        raise ValueError(
            "group_cols must be specified when FIR delay-bin rows are present "
            "(e.g. ['ROI', 'Chroma', 'ID'] or ['ch_name', 'Chroma', 'ID']) — "
            "otherwise rows from different subjects/ROIs/channels could be "
            "collapsed together incorrectly."
        )
 
    df_fir = df[is_fir_row].copy()
    df_other = df[~is_fir_row].copy()  # any non-FIR rows pass through untouched
 
    df_fir['_base_condition'] = df_fir[condition_col].astype(str).str.split(delay_sep, n=1).str[0]
 
    full_group = group_cols + ['_base_condition']
 
    collapsed = (
        df_fir.groupby(full_group, as_index=False)[value_col]
        .mean()
        .rename(columns={'_base_condition': condition_col})
    )
 
    if not df_other.empty:
        collapsed = pd.concat([collapsed, df_other], ignore_index=True)
 
    return collapsed



def _channel_midpoint(ch_info):
    """(x, y, z) midpoint between source and detector for one channel entry
    from raw_haemo.info['chs']. Returns None if location info is missing/
    degenerate (e.g. an aux/stim channel with no real optode geometry).
    MNE head-coordinate convention: x = left(-)/right(+),
    y = posterior(-)/anterior(+), z = inferior(-)/superior(+)."""
    loc = ch_info['loc']
    if loc is None or not np.asarray(loc).any():
        return None
    src = loc[3:6]
    det = loc[6:9]
    return tuple((s + d) / 2.0 for s, d in zip(src, det))



def _build_axis_split_rois(raw_haemo, axis, names, balance_threshold=0.5):
    """
    Split channels into two ROIs by the sign of one coordinate axis of each
    channel's source-detector midpoint.
 
    Parameters
    ----------
    axis : int
        0 = x (left/right), 1 = y (posterior/anterior). z (axis 2, superior/
        inferior) isn't offered as a fallback split — depth splits aren't a
        meaningful functional distinction the way left/right or front/back
        are for a 2D optode array.
    names : (str, str)
        Names for the (negative-side, positive-side) ROIs.
    balance_threshold : float, default 0.5
        Minimum acceptable ratio of (smaller side size / larger side size).
        0.5 means the smaller side must be at least half the size of the
        larger — rejects near-degenerate splits (e.g. 27 channels on one
        side, 1 on the other) where "two ROIs" isn't really giving you two
        usable regions, without demanding a perfect, unrealistic 50/50.
 
    Returns
    -------
    dict of two ROIs, or None if geometry is missing/degenerate, every
    channel falls on one side, or the split is too imbalanced to be useful
    — signaling the caller to try the next fallback tier.
    """
    neg_indices, pos_indices = [], []
 
    for idx, ch in enumerate(raw_haemo.info['chs']):
        mid = _channel_midpoint(ch)
        if mid is None:
            continue
        coord = mid[axis]
        if coord < 0:
            neg_indices.append(idx)
        elif coord > 0:
            pos_indices.append(idx)
        # coord == 0 (exact midline) excluded from both — genuinely
        # ambiguous, not worth guessing a side for.
 
    if not neg_indices or not pos_indices:
        return None
 
    balance = min(len(neg_indices), len(pos_indices)) / max(len(neg_indices), len(pos_indices))
    if balance < balance_threshold:
        logger.warning(
            f"Axis-{axis} split too imbalanced ({len(neg_indices)} vs "
            f"{len(pos_indices)}, ratio {balance:.2f} < {balance_threshold}) — rejecting."
        )
        return None
 
    return {names[0]: neg_indices, names[1]: pos_indices}



def _build_geometric_fallback_rois(raw_haemo):
    """
    Generalized, zero-configuration fallback: try a Left/Right split first
    (the most common and most interpretable axis for bilateral montages);
    if that's unavailable or too imbalanced (e.g. a montage covering only
    one cortical region, where every channel falls on the same side), try
    a Front/Back split instead, using the exact same geometry. Returns None
    if neither axis gives a usable split, signaling the caller to fall back
    further to one-ROI-per-channel.
 
    Like the hemisphere-only version this replaces, these are coarse
    geometric splits, not hand-drawn functional regions — labeled
    "_Auto" so they're never mistaken for real regions.json ROIs anywhere
    downstream (tables, plot titles, exported CSVs).
    """
    lr = _build_axis_split_rois(raw_haemo, axis=0, names=("Left_Auto", "Right_Auto"))
    if lr is not None:
        logger.info("Automatic fallback ROIs: Left/Right split (axis available and balanced).")
        return lr
 
    logger.warning("Left/Right fallback unavailable or too imbalanced — trying Front/Back split.")
    fb = _build_axis_split_rois(raw_haemo, axis=1, names=("Back_Auto", "Front_Auto"))
    if fb is not None:
        logger.info("Automatic fallback ROIs: Front/Back split (Left/Right was not usable).")
        return fb
 
    logger.warning("Neither Left/Right nor Front/Back split is usable for this montage.")
    return None



def _build_per_channel_rois(raw_haemo):
    """
    Last-resort failsafe: one ROI per physical channel (source-detector
    pair), each containing that channel's hbo AND hbr indices together —
    same grouping convention as regions.json (one name -> both
    chromophores), just derived automatically from whatever channels exist.
 
    NOT good for statistics — every "region" is a single channel, so this
    provides none of ROI aggregation's noise-reduction or multiple-
    comparisons benefit. Only reached if regions.json AND both automatic
    geometric splits are unavailable, so processing degrades gracefully to
    single-channel resolution rather than crashing, or silently averaging
    unrelated regions together into one meaningless number the way a
    single AllChannels ROI would.
    """
    rois_formatted = {}
    for ch_name in raw_haemo.ch_names:
        base_name = ch_name.split()[0]  # "S1_D1 hbo" -> "S1_D1"
        idx = raw_haemo.ch_names.index(ch_name)
        rois_formatted.setdefault(base_name, []).append(idx)
    return rois_formatted



def calculate_dpf(file_path):
    # order is hbo / hbr
    with h5py.File(file_path, 'r') as f:
        wavelengths = f['/nirs/probe/wavelengths'][:]
        logger.info(f"Wavelengths (nm): {wavelengths}")
    wavelengths = sorted(wavelengths, reverse=True)
    age = float(AGE)
    logger.info(f"Their age was {AGE}")
    # where the hell did I get these from again?
    a = 223.3
    b = 0.05624
    c = 0.8493
    d = -5.723e-7
    e = 0.001245
    f = -0.9025
    dpf = []
    for w in wavelengths:
        logger.info(w)
        dpf.append(a + b * (age**c) + d* (w**3) + e * (w**2) + f*w)
    logger.info(dpf)
    return dpf



def iqr_threshold(coeffs: NDArray[float64], k: float = 1.5) -> floating[Any]:

    """
    Calculate the interquartile range (IQR) threshold scaled by a factor, k.

    Parameters
    ----------
    coeffs : NDArray[float64]
        Array of coefficients to compute the IQR from.
    k : float, optional
        Scaling factor for the IQR (default is 1.5).

    Returns
    -------
    floating[Any]
        The scaled IQR threshold value.
    """

    # Calculate the IQR
    q1 = np.percentile(coeffs, 25)
    q3 = np.percentile(coeffs, 75)
    iqr = q3 - q1
 
    return k * iqr



def wavelet_iqr_denoise(signal: NDArray[float64], wavelet: str = 'db4', level: int = 3, iqr: float = 1.5) -> NDArray[float64]:
    """
    Denoises a signal using wavelet decomposition and IQR-based thresholding on detail coefficients.

    Parameters
    ----------
    signal : NDArray[float64]
        The input signal array to denoise.
    wavelet : str, optional
        The type of wavelet to use for decomposition (default is 'db4').
    level : int, optional
        Decomposition level for wavelet transform (default is 3).
 
    Returns
    -------
    NDArray[float64]
        The denoised signal array, with the same length as the input.
    """

    max_level = pywt.dwt_max_level(len(signal), pywt.Wavelet(wavelet).dec_len)
    if level > max_level:
        raise ValueError(
            f"wavelet_level={level} exceeds the maximum valid decomposition level "
            f"({max_level}) for a signal of length {len(signal)} with wavelet '{wavelet}'. "
            f"Reduce wavelet_level to at most {max_level}, or use a longer signal."
        )
    if level < 1:
        raise ValueError(f"wavelet_level must be >= 1, got {level}.")
    
    # Decompose the signal using wavelet transform and initialize a list with approximation coefficients
    coeffs: list[NDArray[float64]] = pywt.wavedec(signal, wavelet, level=level) # type: ignore
    cA = coeffs[0]
    denoised_coeffs = [cA]

    # Threshold detail coefficients to reduce noise
    for cD in coeffs[1:]:
        threshold = iqr_threshold(cD, iqr)
        cD_thresh = np.sign(cD) * np.maximum(np.abs(cD) - threshold, 0.0) # np.where((cD < lower) | (cD > upper), 0, cD)
        cD_thresh = cD_thresh.astype(float64)
        denoised_coeffs.append(cD_thresh)

    # Reconstruct the denoised signal
    denoised_signal = cast(NDArray[float64], pywt.waverec(denoised_coeffs, wavelet)) # type: ignore
    return denoised_signal[:len(signal)]
        


def calculate_and_apply_wavelet(data: BaseRaw, wavelet_type: str, wavelet_level: int, iqr: float, verbosity: bool) -> tuple[BaseRaw, Figure]:
    """
    Applies a wavelet IQR denoising filter to the data and generates a plot.

    Parameters
    ----------
    data : BaseRaw
        The loaded data object to process.
    ID : str
        File name of the the snirf file that was loaded.

    Returns
    -------
    tuple[BaseRaw, Figure]
        - BaseRaw: The processed data object.
        - Figure: The corresponding Matplotlib figure.
    """

    logger.info("Applying the wavelet filter...")

    # Denoise the data
    logger.info("Denoising the data...")
    loaded_data: NDArray[float64] = data.get_data(verbose=verbosity) # type: ignore
    denoised_data = np.zeros_like(loaded_data)

    logger.info("Calculating the IQR, decomposing the signal, and thresholding the coefficients...")
    for ch in range(loaded_data.shape[0]):
        denoised_data[ch, :] = wavelet_iqr_denoise(loaded_data[ch, :], wavelet=wavelet_type, level=wavelet_level, iqr=iqr)

    # Reconstruct the data with the annotations
    logger.info("Reconstructing the data with annotations...")
    raw_with_tddr_and_wavelet = RawArray(denoised_data, cast(Info, data.info), verbose=verbosity)
    raw_with_tddr_and_wavelet.set_annotations(data.annotations.copy(), verbose=verbosity) # type: ignore

    # Create a figure for the results
    logger.info("Creating the figure...")
    fig = cast(Figure, raw_with_tddr_and_wavelet.plot(show=False, n_channels=len(getattr(data, "ch_names")), duration=data.times[-1]).figure) # type: ignore
    fig.suptitle(f"Wavelet for ", fontsize=16) # type: ignore
    fig.subplots_adjust(top=0.92)
    plt.close(fig)

    logger.info("Successfully applied the wavelet filter.")

    return raw_with_tddr_and_wavelet, fig



def _select_hr_source_channels(
    channel_data: NDArray[float64],
    sfreq: float,
    cardiac_band: tuple[float, float] = (0.8, 3.0),
) -> NDArray[float64]:
    """
    Scores each channel by the fraction of its power spectral density
    falling within a plausible cardiac frequency band. Higher scores mean
    more of that channel's signal energy sits where a heartbeat would be -
    a cheap proxy for cardiac-signal quality that doesn't require peak
    detection to already have succeeded.
    """
    n_channels = channel_data.shape[0]
    nperseg = min(channel_data.shape[1], 2048)
    if nperseg < 8:
        return np.zeros(n_channels)

    # single vectorized call across all channels, instead of one welch()
    # call per channel - same math, same result, no Python-level loop
    freqs, psd = welch(channel_data, fs=sfreq, nperseg=nperseg, axis=1)  # psd shape: (n_channels, n_freqs)

    band_mask = (freqs >= cardiac_band[0]) & (freqs <= cardiac_band[1])
    total_power = np.sum(psd, axis=1)
    band_power = np.sum(psd[:, band_mask], axis=1)

    scores = np.divide(
        band_power, total_power,
        out=np.zeros(n_channels), where=total_power > 0
    )
    return scores



def short_channel_processing_for_hr(
    data: BaseRaw,
    short_chans: BaseRaw | None,
    seconds_to_strip_hr: int,
    verbosity: bool,
    cardiac_band: tuple[float, float] = (0.8, 3.0),
    max_channels_used: int = 5,
) -> tuple[float, NDArray[float64], NDArray[float64]]:
    """
    Builds a single combined signal for heart rate estimation from ALL
    available candidate channels, rather than an arbitrary single channel.

    Prefers short-separation channels (dominated by superficial/systemic
    signal, which includes the cardiac pulse, with little brain-hemodynamic
    contamination). Falls back to long channels only if no short channels
    are available - a weaker source, since long-channel signal is a mix of
    systemic AND real task/brain response, so the combined signal there is
    more at risk of being pulled around by task-locked activity unrelated
    to heart rate. Treat long-channel-derived HR estimates with more
    caution than short-channel ones.

    Every candidate channel is scored by how much of its power sits in a
    plausible cardiac band (cardiac_band), and the top max_channels_used
    channels are combined via a score-weighted average - channels with
    more cardiac-band power contribute more, channels with little to none
    contribute little to nothing.

    Parameters
    ----------
    data : BaseRaw
        The loaded data object (used for long-channel fallback and time axis).
    short_chans : BaseRaw | None
        Data object with only short separation channels, or None if unavailable.
    seconds_to_strip_hr : int
        Seconds to trim from each end of the signal to remove edge artifacts.
        0 disables trimming.
    cardiac_band : tuple[float, float], default (0.8, 3.0)
        Frequency range (Hz) used to score channel quality (48-180 BPM by default).
    max_channels_used : int, default 5
        Maximum number of top-scoring channels to combine.

    Returns
    -------
    tuple[float, NDArray[float64], NDArray[float64]]
        - float: Sampling frequency of the signal.
        - NDArray[float64]: Trimmed, channel-combined signal.
        - NDArray[float64]: Corresponding time values.
    """
    if short_chans is not None and len(short_chans.ch_names) > 0:
        source = short_chans
        source_label = "short"
    else:
        logger.warning(
            "No short channels available for heart rate estimation - falling back "
            "to long channels. Long-channel HR estimates are less reliable, since "
            "task/brain-hemodynamic signal can dominate over the cardiac pulse."
        )
        source = data
        source_label = "long"

    channel_data = cast(NDArray[float64], source.get_data(verbose=verbosity))
    sfreq = cast(float, source.info['sfreq'])
    n_available = channel_data.shape[0]

    if n_available == 0:
        raise ValueError(f"No channels available in the '{source_label}' source for heart rate estimation.")

    scores = _select_hr_source_channels(channel_data, sfreq, cardiac_band=cardiac_band)
    n_use = min(max_channels_used, n_available)
    top_idx = np.argsort(scores)[::-1][:n_use]

    used_names = [source.ch_names[i] for i in top_idx]
    logger.info(
        f"Heart rate: using {n_use}/{n_available} {source_label} channel(s), "
        f"selected by cardiac-band power: {list(zip(used_names, np.round(scores[top_idx], 3)))}"
    )

    weights = scores[top_idx]
    if weights.sum() > 0:
        weights = weights / weights.sum()
    else:
        logger.warning(
            f"All selected {source_label} channels have zero measurable power in "
            f"the {cardiac_band} Hz cardiac band - heart rate estimate is unlikely "
            f"to be meaningful. Falling back to a plain (unweighted) average."
        )
        weights = np.ones(n_use) / n_use

    signal = np.average(channel_data[top_idx, :], axis=0, weights=weights)

    if seconds_to_strip_hr > 0:
        strip_samples = int(sfreq * seconds_to_strip_hr)
        signal_trimmed = signal[strip_samples:-strip_samples]
        times_trimmed = data.times[strip_samples:-strip_samples]
    else:
        signal_trimmed = signal
        times_trimmed = data.times

    return sfreq, signal_trimmed, times_trimmed


def reconcile_heart_rate_estimates(
    mean_hr_scipy: float,
    psd_confidence: float,
    mean_hr_nk: float,
    mode_hr_nk: float,
    agreement_tolerance_bpm: float = 10.0,
    psd_confidence_threshold: float = 3.0,
) -> tuple[float, bool, str]:
    """
    Combines three independent heart rate estimates - PSD spectral peak
    (scipy), NeuroKit's cleaned/interpolated mean, and NeuroKit's mode -
    into one final value, using majority agreement rather than a fixed
    pairwise override rule.

    Logic, in order:
    1. If all three estimates agree within agreement_tolerance_bpm of each
       other, average them - strongest possible evidence, no single method
       is being trusted over the others.
    2. Otherwise, check if any TWO of the three agree with each other -
       if so, average that agreeing pair and discard the outlier. Two
       independent methods landing on the same value by coincidence is
       unlikely; the third is more likely the one that's wrong.
    3. If no two agree at all, fall back to whichever single estimate is
       most trustworthy: the PSD estimate if its confidence clears
       psd_confidence_threshold (a genuinely sharp, unambiguous spectral
       peak), otherwise the NeuroKit mode (more robust than its mean to
       a residual minority of bad samples, per _mode_hr's reasoning).

    Returns
    -------
    tuple[float, bool, str]
        - float: final reconciled heart rate (BPM).
        - bool: True if any disagreement/overruling occurred (for plotting/logging).
        - str: human-readable explanation of which path was taken, for logs.
    """
    estimates = {
        "psd": mean_hr_scipy,
        "nk_mean": mean_hr_nk,
        "nk_mode": mode_hr_nk,
    }

    def _fmt(d: dict[str, float]) -> str:
        return ", ".join(f"{k}={v:.1f}" for k, v in d.items())
    
    pairs = [("psd", "nk_mean"), ("psd", "nk_mode"), ("nk_mean", "nk_mode")]
    agreeing_pairs = [
        (a, b) for a, b in pairs
        if abs(estimates[a] - estimates[b]) <= agreement_tolerance_bpm
    ]

    if len(agreeing_pairs) == 3:
        final = float(np.mean(list(estimates.values())))
        return final, False, f"All three estimates agree (within {agreement_tolerance_bpm} BPM) - averaged: {_fmt(estimates)}"

    if len(agreeing_pairs) >= 1:
        a, b = agreeing_pairs[0]
        final = float((estimates[a] + estimates[b]) / 2.0)
        outlier = [k for k in estimates if k not in (a, b)][0]
        return final, True, (
            f"{a} and {b} agree ({estimates[a]:.1f}, {estimates[b]:.1f}); "
            f"{outlier} is an outlier ({estimates[outlier]:.1f}) - discarded."
        )

    # No two estimates agree at all - fall back to the single most trustworthy one
    if psd_confidence >= psd_confidence_threshold:
        return mean_hr_scipy, True, (
            f"No two estimates agree ({_fmt(estimates)}); PSD peak is clear "
            f"(confidence={psd_confidence:.2f}) - trusting PSD alone."
        )
    else:
        return mode_hr_nk, True, (
            f"No two estimates agree ({_fmt(estimates)}); PSD peak is ambiguous "
            f"(confidence={psd_confidence:.2f} < {psd_confidence_threshold}) - "
            f"trusting NeuroKit mode alone (more robust to residual dips than its mean)."
        )

    
def _mode_hr(hr_clean: NDArray[float64], bin_width_bpm: float = 2.0) -> float:
    """
    Histogram-based mode of the HR trace: the center of the most frequently
    occurring bin. More robust than a plain mean to a minority of corrupted
    (missed-beat) dips, since those dips only need to avoid being the
    single largest cluster - unlike a mean, which every dip pulls down
    proportionally regardless of how rare it is.

    bin_width_bpm : float, default 2.0
        Histogram bin width. Too narrow and there's no meaningful mode
        (every value nearly unique); too wide and you lose real precision
        in the estimate. 2 BPM is a reasonable starting point - worth
        checking against your actual HR distributions.
    """
    if len(hr_clean) == 0:
        return float('nan')
    bins = np.arange(hr_clean.min(), hr_clean.max() + bin_width_bpm, bin_width_bpm)
    counts, edges = np.histogram(hr_clean, bins=bins)
    mode_bin_idx = np.argmax(counts)
    return float((edges[mode_bin_idx] + edges[mode_bin_idx + 1]) / 2.0)


def calculate_heart_rate_neurokit(sfreq: float, signal_trimmed: NDArray[float64], hr_low_freq, hr_high_freq, max_low_hr, max_high_hr, smoothing_window_hr, short_channels) -> tuple[NDArray[float64], float]:
    """
    Calculate and smooths heart rate from a trimmed signal using NeuroKit.

    Parameters
    ----------
    sfreq : float
        Sampling frequency of the signal.
    signal_trimmed : NDArray[float64]
        Preprocessed and trimmed fNIRS signal.

    Returns
    -------
    tuple[NDArray[float64], float]
        - NDArray[float64]: Smoothed heart rate time series (BPM).
        - float: Mean heart rate.
    """
    logger.info("Calculating heart rate using NeuroKit...")

    logger.info("Filtering the signal and detecting pulsatile peaks...")
    signal_filtered = cast(NDArray[float64], nk.signal_filter(signal_trimmed, sampling_rate=sfreq, lowcut=hr_low_freq, highcut=hr_high_freq))

    # bishop works better with challenging datasets, but like it is super slow on good datasets?
    if short_channels:
        peaks_dict = cast(dict[str, Any], nk.ppg_findpeaks(signal_filtered, sampling_rate=sfreq, method="elgendi"))
    else:
        peaks_dict = cast(dict[str, Any], nk.ppg_findpeaks(signal_filtered, sampling_rate=sfreq, method="bishop"))
    peaks = peaks_dict['PPG_Peaks']
    logger.info(f"ppg_findpeaks found {len(peaks)} peaks over {len(signal_trimmed)/sfreq:.1f}s "
            f"(~{len(peaks) / (len(signal_trimmed)/sfreq) * 60:.1f} BPM implied by peak count alone)")

    if len(peaks) < 2:
        logger.warning("ppg_findpeaks found fewer than 2 peaks - heart rate estimate is unreliable.")
        return np.full(len(signal_trimmed), np.nan), float('nan')

    hr = cast(NDArray[float64], nk.signal_rate(peaks, sampling_rate=sfreq, desired_length=len(signal_trimmed)))
    logger.info(f"Pre-clip HR range: min={hr.min():.1f}, max={hr.max():.1f} BPM (max_low_hr={max_low_hr}, max_high_hr={max_high_hr})")

    hr_clean = np.clip(hr, max_low_hr, max_high_hr)

    logger.info("Smoothing the signal and calculating the mean...")
    hr_series = pd.Series(hr_clean)
    local_median = hr_series.rolling(window=smoothing_window_hr, center=True, min_periods=1).median()
    spikes = (hr_series > local_median + 10) | (hr_series < local_median - 10)  # was upward-only; catches drops too
    smoothed_values = hr_series.copy()
    smoothed_spikes = hr_series.rolling(window=smoothing_window_hr, center=True, min_periods=1).mean()
    smoothed_values[spikes] = smoothed_spikes[spikes]
    hr_smooth_nk = cast(NDArray[float64], smoothed_values.to_numpy())
    mean_hr_nk = hr_smooth_nk.mean()
    mode_hr_nk = _mode_hr(hr_clean, bin_width_bpm=2.0)
    logger.info("Original HR min/max: %f, %f", hr_clean.min(), hr_clean.max())
    logger.info("Smoothed HR min/max:%f, %f", hr_smooth_nk.min(), hr_smooth_nk.max())
    logger.info(f"Estimated mean HR nk: {mean_hr_nk:.1f} BPM, mode HR nk: {mode_hr_nk:.1f} BPM")


    return hr_smooth_nk, mean_hr_nk, mode_hr_nk



def calculate_heart_rate_scipy(
    sfreq: float, signal_trimmed: NDArray[float64], search_min, search_max,
    cluster_window_bpm: float = 6.0,
) -> tuple[NDArray[floating[Any]], NDArray[float64], np.ndarray[Any, np.dtype[np.bool_]], float, float]:
    """
    Estimate heart rate using spectral analysis on a high-pass filtered signal.

    Parameters
    ----------
    sfreq : float
        Sampling frequency of the input signal.
    signal_trimmed : NDArray[float64]
        Trimmed fNIRS signal to analyze.
    cluster_window_bpm : float, default 6.0
            Width (in BPM, +/- from the argmax) used to find nearby local peaks
            that likely belong to the same underlying cardiac frequency (spread
            by natural heart-rate variability/frequency modulation) rather than
            being genuinely separate candidates. The reported HR is the
            power-weighted centroid of all local peaks within this window of
            the strongest bin, not just the single tallest bin.

    Returns
    -------
    tuple[NDArray[floating[Any]], NDArray[float64], np.ndarray[Any, np.dtype[np.bool_]], float]
        - NDArray[floating[Any]]: Frequencies converted to beats per minute (BPM).
        - NDArray[float64]: Power spectral density (PSD) of the signal.
        - np.ndarray[Any, np.dtype[np.bool_]]: Boolean mask indicating frequencies within heart rate range.
        - float: Estimated mean heart rate in BPM corresponding to the PSD peak within the range.
    """
    logger.info("Calculating heart rate using SciPy...")

    b, a = cast(tuple[NDArray[float64], NDArray[float64]], butter(2, 0.5 / (sfreq / 2), btype='high'))
    signal_hp = cast(NDArray[float64], filtfilt(b, a, signal_trimmed))

    nperseg = min(len(signal_hp), 4096)
    frequencies_scipy, psd_scipy = cast(tuple[NDArray[float64], NDArray[float64]], welch(signal_hp, fs=sfreq, nperseg=nperseg, noverlap=nperseg // 2))

    freq_bpm_scipy = frequencies_scipy * 60
    freq_range_scipy = (freq_bpm_scipy > search_min) & (freq_bpm_scipy < search_max)

    band_bpm = freq_bpm_scipy[freq_range_scipy]
    band_psd = psd_scipy[freq_range_scipy]
    if len(band_psd) == 0:
        raise ValueError(f"No frequency bins fall within the search range ({search_min}-{search_max} BPM).")

    # Find ALL local peaks in the band (bins strictly greater than both neighbors)
    local_peak_mask = np.zeros(len(band_psd), dtype=bool)
    if len(band_psd) >= 3:
        local_peak_mask[1:-1] = (band_psd[1:-1] > band_psd[:-2]) & (band_psd[1:-1] > band_psd[2:])
    # Edge bins can't be evaluated as local peaks by this rule; if the true
    # peak sits at the very edge of the search range, argmax below still
    # catches it as a fallback.
    peak_indices = np.where(local_peak_mask)[0]
    if len(peak_indices) == 0:
        peak_indices = np.array([np.argmax(band_psd)])

    strongest_idx = peak_indices[np.argmax(band_psd[peak_indices])]
    strongest_bpm = band_bpm[strongest_idx]

    # Cluster: local peaks within cluster_window_bpm of the strongest one
    cluster_mask = np.abs(band_bpm[peak_indices] - strongest_bpm) <= cluster_window_bpm
    cluster_indices = peak_indices[cluster_mask]
    cluster_bpm = band_bpm[cluster_indices]
    cluster_power = band_psd[cluster_indices]

    # Power-weighted centroid across the cluster - this is the actual fix:
    # a tight group of near-equal peaks now contributes to ONE combined
    # estimate near their shared center, instead of a coin-flip winner-take-all.
    mean_hr_scipy = float(np.average(cluster_bpm, weights=cluster_power))

    # Confidence: strongest cluster's TOTAL power vs. the median power of
    # everything OUTSIDE the cluster - reflects how dominant the whole
    # cluster is, not just one bin within it. A 4-peak near-tie spread
    # across the band now scores lower confidence than a single sharp,
    # isolated peak of similar height, even though argmax alone couldn't
    # tell them apart.
    outside_cluster = np.setdiff1d(np.arange(len(band_psd)), cluster_indices)
    baseline_power = np.median(band_psd[outside_cluster]) if len(outside_cluster) > 0 else np.median(band_psd)
    cluster_total_power = cluster_power.sum()
    peak_confidence = cluster_total_power / baseline_power if baseline_power > 0 else 0.0

    logger.info(f"PSD: {len(cluster_indices)} peak(s) in cluster near {strongest_bpm:.1f} BPM, "
                f"centroid={mean_hr_scipy:.1f} BPM, confidence={peak_confidence:.2f}x baseline")
    logger.info("Successfully calculated heart rate using SciPy.")

    return freq_bpm_scipy, psd_scipy, freq_range_scipy, mean_hr_scipy, peak_confidence


def plot_heart_rate(
    freq_bpm_scipy: NDArray[floating[Any]],
    psd_scipy: NDArray[float64],
    freq_range_scipy: np.ndarray[Any, np.dtype[np.bool_]],
    mean_hr_scipy: float,
    hr_smooth_nk: NDArray[floating[Any]],
    mean_hr_nk: float,
    mode_hr_nk: float,
    final_hr: float,
    times_trimmed: NDArray[floating[Any]],
    overruled: bool,
    reconciliation_note: str,
    hr_window: int
) -> tuple[Figure, Figure]:
    """
    Generate plots comparing heart rate estimates from SciPy PSD and NeuroKit2.

    Parameters
    ----------
    freq_bpm_scipy : NDArray[floating[Any]]
        Frequencies in beats per minute from SciPy PSD analysis.
    psd_scipy : NDArray[float64]
        Power spectral density values corresponding to freq_bpm_scipy.
    freq_range_scipy : np.ndarray[Any, np.dtype[np.bool_]]
        Boolean mask indicating the heart rate frequency range used in PSD.
    mean_hr_scipy : float
        Heart rate estimated from SciPy PSD peak (cluster centroid).
    hr_smooth_nk : NDArray[floating[Any]]
        Smoothed instantaneous heart rate from NeuroKit2.
    mean_hr_nk : float
        Mean heart rate estimated from NeuroKit2's cleaned time series.
    mode_hr_nk : float
        Mode (most frequent value) of NeuroKit2's heart rate distribution -
        more robust than the mean to a residual minority of missed-beat dips.
    final_hr : float
        The final reconciled heart rate value, combining all three estimates
        (see reconcile_heart_rate_estimates).
    times_trimmed : NDArray[floating[Any]]
        Time points corresponding to hr_smooth_nk values.
    overruled : bool
        True if the final reconciled value differs from a simple average of
        all three (i.e. an outlier was discarded or a single estimate was
        trusted alone).
    reconciliation_note : str
        Human-readable explanation of which reconciliation path was taken -
        shown directly on the plot rather than a generic "was overruled" note.

    Returns
    -------
    tuple[Figure, Figure]
        - Figure showing the PSD and SciPy heart rate estimate.
        - Figure showing the time series comparison of heart rates.
    """

    # Create the first plot for the PSD. Add a yellow range to show what we will be filtering to.
    logger.info("Creating the figure...")
    fig1, ax1 = plt.subplots(figsize=(10, 5)) # type: ignore
    ax1.set_xlim(30, 300)
    ax1.plot(freq_bpm_scipy[freq_range_scipy], psd_scipy[freq_range_scipy]) # type: ignore
    ax1.axvline(x=mean_hr_scipy, color='red', linestyle='--', label=f'Mean HR: {mean_hr_scipy:.1f} BPM') # type: ignore
    ax1.axvspan(min(mean_hr_nk - hr_window, mean_hr_scipy - hr_window), max(mean_hr_nk + hr_window, mean_hr_scipy + hr_window), color='yellow', alpha=0.3, label=f'HR Range ±{hr_window} BPM') # type: ignore
    ax1.set_xlabel('Heart Rate (BPM)') # type: ignore
    ax1.set_ylabel('Power Spectral Density') # type: ignore
    ax1.set_title('PSD of fNIRS signal - Peak indicates Heart Rate') # type: ignore
    ax1.grid(True) # type: ignore

    # Was the value we reported here correct for the data on the graph or was it overruled?
    note_prefix = 'Reconciliation (outlier discarded):' if overruled else 'Reconciliation (estimates agreed):'
    note = f"\n{note_prefix}\n{reconciliation_note}"
    phantom = Line2D([0], [0], color='none', label=note)
    handles, _ = ax1.get_legend_handles_labels()
    ax1.legend(handles=handles + [phantom], fontsize=8)
    plt.close(fig1)

    # Create the second plot showing the rolling heart rate, as well as the two averages that were calculated
    logger.info("Creating the figure...")
    fig2, ax2 = plt.subplots(figsize=(14, 6)) # type: ignore
    ax2.plot(times_trimmed, hr_smooth_nk, label='Instantaneous HR (NeuroKit2)', color='blue', alpha=0.7)
    ax2.axhline(mean_hr_nk, color='steelblue', linestyle='--', alpha=0.7, label=f'NK Mean: {mean_hr_nk:.1f} BPM')
    ax2.axhline(mode_hr_nk, color='purple', linestyle='--', alpha=0.7, label=f'NK Mode: {mode_hr_nk:.1f} BPM')
    ax2.axhline(mean_hr_scipy, color='orange', linestyle=':', label=f'PSD Estimate: {mean_hr_scipy:.1f} BPM')
    ax2.axhline(final_hr, color='green', linestyle='-', linewidth=2, label=f'Final (Reconciled) HR: {final_hr:.1f} BPM')
    ax2.set_xlabel('Time (seconds)')
    ax2.set_ylabel('Heart Rate (BPM)')
    ax2.set_title('Heart Rate Estimates Comparison')
    ax2.legend(fontsize=9)
    ax2.grid(True)
    fig2.tight_layout()
    plt.close(fig2)

    return fig1, fig2



def detect_sensor_dropout(raw, threshold_ratio=0.05):
    """
    Identifies channels where signal variance drops significantly.
    Returns flagged channel names and a summary figure.
    """
    ch_names = raw.ch_names
    data = raw.get_data()
    n_samples = data.shape[1]
    quarter = n_samples // 4
    
    ratios = []
    dead_idx = []
    
    for i in range(len(ch_names)):
        start_var = np.var(data[i, :quarter])
        end_var = np.var(data[i, -quarter:])
        
        # Handle zero variance (dead from the start)
        if start_var == 0:
            ratio = 0.0
        else:
            ratio = end_var / start_var
            
        ratios.append(ratio)
        
        if ratio < threshold_ratio:
            dead_idx.append(i)
            print(f"Flagged {ch_names[i]}: Variance dropped to {ratio:.2%} of original.")

    # Pair-kill logic
    failed_bases = {ch_names[i].split(' ')[0] for i in dead_idx}
    bad_names = [ch for ch in ch_names if ch.split(' ')[0] in failed_bases]
    
    # --- Visualization ---
    fig_disp, ax = plt.subplots(figsize=(10, 5), constrained_layout=True)
    
    # Color logic: Coral for channels below the threshold
    colors = ['coral' if r < threshold_ratio else 'skyblue' for r in ratios]
    
    ax.bar(range(len(ratios)), ratios, color=colors)
    ax.axhline(threshold_ratio, color='red', linestyle='--', label=f'Threshold ({threshold_ratio:.0%})')
    
    ax.set_title("Sensor Dropout Check (Variance Stability)")
    ax.set_ylabel("Variance Ratio (End / Start)")
    ax.set_xlabel("Channel Index")
    ax.set_ylim(0, max(ratios + [threshold_ratio * 2])) # Scale to see the threshold clearly
    ax.legend()
    
    plt.close(fig_disp)
    
    print(f"Dropout Check: Flagged {len(failed_bases)} optode pairs.")
    return bad_names, fig_disp



def detect_spectral_noise_spike(raw, db_limit=-60, freq_div=4, min_freq=0.1, target_bandwith=0.2):
    """
    Identifies channels with excessive power at high frequencies 
    (sfreq/4), usually indicating electronic interference.
    """
    ch_names = raw.ch_names
    sfreq = raw.info['sfreq']
    target_freq = sfreq / freq_div
    
    # Compute PSD
    spectrum = raw.compute_psd(fmin=min_freq, fmax=sfreq/2)
    psd_data, freqs = spectrum.get_data(return_freqs=True)
    
    # Find power near the target frequency
    f_idx = np.where((freqs >= target_freq - target_bandwith) & (freqs <= target_freq + target_bandwith))[0]
    power_at_target = np.mean(psd_data[:, f_idx], axis=1)
    
    abs_threshold = 10 ** (db_limit / 10)
    noisy_idx = np.where(power_at_target > abs_threshold)[0]
    
    # Pair-kill logic
    failed_bases = {ch_names[i].split(' ')[0] for i in noisy_idx}
    bad_names = [ch for ch in ch_names if ch.split(' ')[0] in failed_bases]
    
    # --- Visualization ---
    fig, ax = plt.subplots(figsize=(8, 4))
    ax.plot(freqs, 10 * np.log10(psd_data.T), color='gray', alpha=0.2)
    if len(noisy_idx) > 0:
        ax.plot(freqs, 10 * np.log10(psd_data[noisy_idx].T), color='plum', label='Noisy Pairs')
    
    ax.axhline(db_limit, color='red', linestyle='--', label='Threshold')
    ax.set_title(f"PSD Noise Analysis (Target: {target_freq}Hz)")
    ax.set_ylabel("Power (dB)")
    ax.legend()
    plt.close(fig)
    
    print(f"Noise Check: Flagged {len(failed_bases)} optode pairs.")
    return bad_names, fig



def find_bad_channels_by_amplitude_range(raw, threshold=4):
    """Median absolute deviation"""
    picks = [ch for ch in raw.ch_names]
    data = raw.get_data(picks=picks)
    ranges = np.max(data, axis=1) - np.min(data, axis=1)
    
    # Calculate Z-Scores
    median_range = np.median(ranges)
    mad = np.median(np.abs(ranges - median_range))
    z_scores = 0.6745 * (ranges - median_range) / (mad if mad > 0 else 1e-15)
    
    # Identify failed bases
    failed_indices = np.where(np.abs(z_scores) > threshold)[0]
    failed_bases = {picks[i].split(' ')[0] for i in failed_indices}
    
    # Flag entire pairs
    bad_names = [ch for ch in picks if ch.split(' ')[0] in failed_bases]
    
    # --- Visualization ---
    fig_swing, ax = plt.subplots(figsize=(8, 4))
    # We color bars by the specific Z-score of that individual channel
    colors = ['coral' if np.abs(z) > threshold else 'skyblue' for z in z_scores]
    
    ax.bar(range(len(z_scores)), z_scores, color=colors)
    ax.axhline(threshold, color='red', linestyle='--', label='Outlier Threshold')
    ax.axhline(-threshold, color='red', linestyle='--')
    ax.set_title("Physiological Swing Analysis (Z-Scores)")
    ax.set_ylabel("Standardized Deviation")
    ax.set_xlabel("Channel Index")
    ax.legend()
    plt.close(fig_swing)
    
    return bad_names, fig_swing



def find_bad_channels_coeff_var(raw, coeff_var_threshold=25.0):
    """
    Identifies bad fNIRS channels using only the Coefficient of Variation (coeff_var).
    """
    print(f"\n--- Starting coeff_var-Only Quality Check on the channels ---")
    
    picks = [ch for ch in raw.ch_names]
    data = raw.get_data(picks=picks)
    
    # Calculate coeff_var (Coefficient of Variation)
    stds = np.std(data, axis=1)
    means = np.mean(data, axis=1)
    # Using a small epsilon (1e-15) to prevent division by zero
    coeff_var_scores = (stds / (means + 1e-15)) * 100
    
    # Find indices that exceed the threshold
    bad_coeff_var_indices = np.where(coeff_var_scores > coeff_var_threshold)[0]
    
    # Pair-kill logic: If one wavelength (HbO or HbR) fails, flag the pair
    failed_bases = set()
    for idx in bad_coeff_var_indices:
        base = picks[idx].split(' ')[0]
        failed_bases.add(base)
    
    bad_names = [ch for ch in picks if ch.split(' ')[0] in failed_bases]
    
    # Summary Prints
    print(f"coeff_var Check: Found {len(bad_coeff_var_indices)} channels exceeding {coeff_var_threshold}% noise threshold.")
    if failed_bases:
        print(f"Flagged {len(failed_bases)} optode pairs for removal:")
        for base in sorted(failed_bases):
            # Find the specific coeff_var for this base (using the first channel found for it)
            ch_idx = picks.index(next(p for p in picks if p.startswith(base)))
            print(f"  - {base}: coeff_var = {coeff_var_scores[ch_idx]:.2f}%")
    else:
        print("All channels passed the coeff_var check.")

    # --- Visualization ---
    fig_qc, ax = plt.subplots(figsize=(10, 5), constrained_layout=True)
    
    colors = ['coral' if c > coeff_var_threshold else 'skyblue' for c in coeff_var_scores]
    ax.bar(range(len(coeff_var_scores)), coeff_var_scores, color=colors)
    ax.axhline(coeff_var_threshold, color='red', linestyle='--', label=f'Threshold ({coeff_var_threshold}%)')
    
    ax.set_title("Coefficient of Variation (Relative Noise)")
    ax.set_ylabel("coeff_var %")
    ax.set_xlabel("Channel Index")
    ax.legend()
    
    plt.close(fig_qc)
    
    return bad_names, fig_qc



def detect_hbo_hbr_anticorrelation(raw_haemo, threshold: float = -0.2):
    """
    Flags channels where HbO and HbR are NOT showing the expected
    physiologically anti-correlated relationship. Real hemodynamic response
    typically shows HbO rising while HbR falls (and vice versa) - a channel
    pair with weak or positive correlation is a common signature of motion
    artifact or poor optode coupling that other QC metrics can miss, since
    it's checking signal SHAPE/relationship rather than amplitude, variance,
    or noise level.

    IMPORTANT: must be run on raw_haemo BEFORE enhance_negative_correlation
    (or any similar correction step) - that step actively forces HbO/HbR
    anti-correlation, so measuring this diagnostic afterward would just be
    checking whether the correction worked, not the underlying data quality.

    Parameters
    ----------
    raw_haemo : BaseRaw
        Haemoglobin-concentration data (post Beer-Lambert, pre-correction).
    threshold : float, default -0.2
        Channels with HbO/HbR correlation ABOVE this value are flagged as
        bad (i.e. not sufficiently anti-correlated). -0.2 is a permissive
        starting point - real channels often land well below this (-0.5 to
        -0.9), but very low SNR or task designs can naturally weaken the
        correlation without indicating artifact, so this shouldn't be set
        aggressively without checking against your own known-good data.

    Returns
    -------
    tuple[list[str], Figure]
        - list[str]: channel names (both hbo AND hbr for each flagged pair)
          below the anti-correlation threshold.
        - Figure: bar chart of correlation per channel pair, matching the
          visual style of detect_sensor_dropout.
    """
    ch_names = raw_haemo.ch_names
    data = raw_haemo.get_data()

    base_names = sorted({ch.split()[0] for ch in ch_names})
    correlations = {}
    bad_bases = []

    for base in base_names:
        try:
            hbo_idx = ch_names.index(f"{base} hbo")
            hbr_idx = ch_names.index(f"{base} hbr")
        except ValueError:
            continue  # channel doesn't have both chromophores present

        hbo_signal = data[hbo_idx]
        hbr_signal = data[hbr_idx]

        if np.std(hbo_signal) == 0 or np.std(hbr_signal) == 0:
            corr = 0.0  # flat channel - can't meaningfully correlate
        else:
            corr = float(np.corrcoef(hbo_signal, hbr_signal)[0, 1])

        correlations[base] = corr
        if corr > threshold:
            bad_bases.append(base)
            print(f"Flagged {base}: HbO/HbR correlation = {corr:.3f} (expected below {threshold})")

    bad_names = [ch for ch in ch_names if ch.split()[0] in bad_bases]

    fig, ax = plt.subplots(figsize=(10, 5), constrained_layout=True)
    bases_sorted = list(correlations.keys())
    corr_values = [correlations[b] for b in bases_sorted]
    colors = ['coral' if c > threshold else 'skyblue' for c in corr_values]

    ax.bar(range(len(corr_values)), corr_values, color=colors)
    ax.axhline(threshold, color='red', linestyle='--', label=f'Threshold ({threshold})')
    ax.axhline(0, color='black', linewidth=0.8)
    ax.set_title("HbO/HbR Anti-Correlation Check")
    ax.set_ylabel("Pearson r (HbO vs HbR)")
    ax.set_xlabel("Channel Pair Index")
    ax.legend()
    plt.close(fig)

    print(f"Anti-correlation check: flagged {len(bad_bases)} optode pair(s).")
    return bad_names, fig



def hr_calc(raw, seconds_to_strip_hr, l_freq, h_freq, search_min, search_max, max_low_hr, max_high_hr, smoothing_window_hr, hr_window, short_channels, short_channels_threshold, verbosity, psd_confidence_threshold: float = 3.0, band_halfwidth_hz: float = 0.3):
    if short_channels:
        short_chans = get_short_channels(raw, max_dist=short_channels_threshold)
    else:
        short_chans = None
    sfreq, signal_trimmed, times_trimmed = short_channel_processing_for_hr(raw, short_chans, seconds_to_strip_hr=seconds_to_strip_hr, verbosity=verbosity)
    hr_smooth_nk, mean_hr_nk, mode_hr_nk = calculate_heart_rate_neurokit(sfreq, signal_trimmed, hr_low_freq=l_freq, hr_high_freq=h_freq, max_low_hr=max_low_hr, max_high_hr=max_high_hr, smoothing_window_hr=smoothing_window_hr, short_channels=short_channels)
    freq_bpm_scipy, psd_scipy, freq_range_scipy, mean_hr_scipy, psd_confidence = calculate_heart_rate_scipy(sfreq, signal_trimmed, search_min=search_min, search_max=search_max)

    final_hr, overruled, reconciliation_note = reconcile_heart_rate_estimates(
        mean_hr_scipy, psd_confidence, mean_hr_nk, mode_hr_nk, agreement_tolerance_bpm=hr_window
    )
    logger.info(f"HR reconciliation: {reconciliation_note}")
    logger.info(f"Final heart rate: {final_hr:.1f} BPM")

    mean_hr_scipy = final_hr

    hr1, hr2 = plot_heart_rate(
        freq_bpm_scipy, psd_scipy, freq_range_scipy, mean_hr_scipy,
        hr_smooth_nk, mean_hr_nk, mode_hr_nk, final_hr,
        times_trimmed, overruled, reconciliation_note, hr_window=hr_window
    )
   
    fig = raw.compute_psd().plot(show=False)

    # Targeted frequency band for downstream calculations (SCI, PSP, etc.),
    # derived directly from the reconciled heart rate estimate above -
    # replaces a previous, disconnected all-channel median calculation that
    # never used the short-channel selection or reconciliation logic at all.
    hr_freq = final_hr / 60.0  # BPM -> Hz
    low = hr_freq - band_halfwidth_hz
    high = hr_freq + band_halfwidth_hz

    logger.info(f"SCI/PSP target band: {low:.3f}-{high:.3f} Hz "
                f"({final_hr:.1f} +/- {band_halfwidth_hz*60:.0f} BPM)")

    return fig, hr1, hr2, low, high, final_hr



def trim_participant_data(raw, seconds_to_keep: float):
    if hasattr(raw, 'annotations') and len(raw.annotations) > 0:
        # Get time of first event
        first_event_time = raw.annotations.onset[0]
        trim_time = max(0, first_event_time - seconds_to_keep)  # Ensure we don't go negative
        raw.crop(tmin=trim_time)
        # Shift annotation onsets to match new t=0

        ann = raw.annotations
        ann_shifted = Annotations(
            onset=ann.onset - trim_time,  # shift to start at zero
            duration=ann.duration,
            description=ann.description
        )
        data = raw.get_data()
        info = raw.info.copy()
        raw = RawArray(data, info)
        raw.set_annotations(ann_shifted)

        logger.info(f"Trimmed raw data: start at {trim_time}s (5s before first event), t=0 at new start")
    else:
        logger.warning("No events found, skipping trim step.")

    fig_trimmed = raw.plot(duration=raw.times[-1], n_channels=raw.info['nchan'], title="Trimmed Raw", show=False)
    return raw, fig_trimmed



def remove_bad_channels(raw, bad_channels, max_bad_channels: int):
    num_bad = len(bad_channels)
    
    # Check against the threshold
    if num_bad > max_bad_channels:
        raise Exception(
            f"Data Quality Error: {num_bad} channels flagged for removal, "
            f"which exceeds the limit of {max_bad_channels}. To avoid this, "
            f"either lower your filtering parameters or increase MAX_BAD_CHANNELS."
        )

    raw.pick_types(fnirs=True, exclude='bads')
    logger.info(f"Physically removed {len(bad_channels)} channels from the dataset.")
    return raw



def make_and_run_glm(raw_haemo, df_design_matrix, noise_model, bins, n_jobs, verbosity):

    glm_est = run_glm(raw_haemo, df_design_matrix, noise_model=noise_model, bins=bins, n_jobs=n_jobs, verbose=verbosity)
    fir_cols = [col for col in df_design_matrix.columns if "_delay_" in col]
    
    if fir_cols:
        # --- FIR MODEL HANDLING (Peak Delay Detection) ---
        logger.info("FIR model detected. Dynamically identifying peak delays...")
        
        # Extract base task conditions (e.g., "Tapping_Left", "Tapping_Right")
        base_conditions = list(set(col.split('_delay_')[0] for col in fir_cols))
        
        theta_list = glm_est.theta()
        columns = list(df_design_matrix.columns)
 
        peak_conditions = []
        for cond in base_conditions:
            cond_delays = [col for col in fir_cols if col.startswith(f"{cond}_delay_")]
 
            delay_impacts = {}
            for col in cond_delays:
                col_idx = columns.index(col)
                avg_absolute_theta = np.mean(np.abs([ch_theta[col_idx] for ch_theta in theta_list]))
                delay_impacts[col] = avg_absolute_theta
 
            peak_delay_col = max(delay_impacts, key=delay_impacts.get)
            logger.info(f"Condition '{cond}' peak response identified at delay column: {peak_delay_col}")
            peak_conditions.append(peak_delay_col)
            
        # Plot only the peak delays for a clean, single-column topomap per condition
        fig_glm_topo = glm_est.plot_topo(conditions=peak_conditions)
        
    else:
        # --- STANDARD HRF MODEL HANDLING ---
        # Extract only task conditions (ignore drifts, constants, and short channels)
        experimental_conditions = [
            col for col in df_design_matrix.columns 
            if not any(noise in col.lower() for noise in ['drift', 'constant', 'short'])
        ]
        fig_glm_topo = glm_est.plot_topo(conditions=experimental_conditions)

    plt.close(fig_glm_topo)
    return glm_est, fig_glm_topo



def _real_conditions(values, exclude_list=NUISANCE_EXCLUDE):
    """Filter out drift/constant/short-style nuisance regressor names,
    keeping only actual task conditions — same filtering logic already
    used for task_cols in generate_contrast_results, reused here so the
    plots only ever show things worth looking at."""
    return sorted({
        v for v in values
        if not any(ex in str(v).lower() for ex in exclude_list)
    })



def generate_channel_results(glm_est, file_path):
    df_cha = glm_est.to_dataframe()
    df_cha["ID"] = file_path

    df_cha = collapse_fir_condition_column(
        df_cha, value_col='theta',
        condition_col='Condition', group_cols=['ch_name', 'Chroma', 'ID']
    )
    
    return df_cha



def generate_roi_results(raw_haemo, df_design_matrix, glm_est, file_path, json_location):
    
    rois_formatted = {}
    try:
        with open(json_location, 'r') as f:
            roi_data = json.load(f)
    
        for region in roi_data.get("regions_of_interest", []):
            roi_name = region["name"]
            channels = region["channels"]
    
            mne_channels = []
            for ch in channels:
                mne_channels.append(f"{ch} hbo")
                mne_channels.append(f"{ch} hbr")
    
            valid_indices = []
            for ch in mne_channels:
                if ch in raw_haemo.ch_names:
                    idx = raw_haemo.ch_names.index(ch)
                    valid_indices.append(idx)
    
            if valid_indices:
                rois_formatted[roi_name] = valid_indices
            else:
                logger.warning(f"No channels from ROI '{roi_name}' found in raw_haemo.")
    
    except Exception as e:
        logger.error(f"Failed to load or parse ROI JSON: {e}")
        rois_formatted = {}
    
    # ---------------------------------------------------------------------
    # Tier 2: automatic geometric split — Left/Right, or Front/Back if
    # Left/Right isn't usable. Used whenever tier 1 produced nothing, whether
    # from a load/parse exception OR a file that loaded fine but matched zero
    # channels (the try/except alone doesn't catch that case, since no
    # exception is raised).
    # ---------------------------------------------------------------------
    if not rois_formatted:
        logger.error(
            f"'{json_location}' produced zero valid ROIs — attempting automatic "
            f"geometric fallback (Left/Right, then Front/Back)."
        )
        rois_formatted = _build_geometric_fallback_rois(raw_haemo)
    
    # ---------------------------------------------------------------------
    # Tier 3: per-channel (true last resort — only if neither geometric
    # split is usable, e.g. missing/degenerate location info).
    # ---------------------------------------------------------------------
    if not rois_formatted:
        logger.warning(
            "No usable geometric fallback — falling back to one ROI per "
            "channel. Statistics will run at single-channel resolution, not "
            "proper ROI aggregation, until regions.json or channel geometry "
            "is fixed."
        )
        rois_formatted = _build_per_channel_rois(raw_haemo)
    
    # 3. Calculate ROI results for all conditions
    conditions = df_design_matrix.columns
    
    # Compute output metrics by custom parsed ROIs (now passing lists of integers)
    df_roi = glm_est.to_dataframe_region_of_interest(rois_formatted, conditions)
    df_roi["ID"] = file_path


    df_roi = collapse_fir_condition_column(
        df_roi, value_col='theta',
        condition_col='Condition', group_cols=['ROI', 'Chroma', 'ID']
    )

    chroma = 'hbo'
    value_col = 'theta'

    real_conditions = _real_conditions(df_roi['Condition'].unique(), NUISANCE_EXCLUDE)
    sub = df_roi[(df_roi['Chroma'] == chroma) & (df_roi['Condition'].isin(real_conditions))]

    if sub.empty:
        print(f"No ROI data for chroma '{chroma}' after excluding nuisance conditions.")
        return
 
    subject_id = sub['ID'].iloc[0] if 'ID' in sub.columns else ''
    n_conditions = sub['Condition'].nunique()

    roi_channel_map = {
        raw_haemo.ch_names[idx]: roi_name
        for roi_name, indices in rois_formatted.items()
        for idx in indices
    }
 
    sns.set_theme(style="whitegrid")
    fig, ax = plt.subplots(figsize=(max(6, 1.5 * sub['ROI'].nunique()), 5))
 
    if n_conditions > 1:
        sns.barplot(data=sub, x='ROI', y=value_col, hue='Condition', ax=ax,
                    edgecolor='black', linewidth=1.2)
    else:
        sns.barplot(data=sub, x='ROI', y=value_col, ax=ax,
                    color='#2b5c8f', edgecolor='black', linewidth=1.2)
 
    ax.axhline(0, color='black', linewidth=1, linestyle='--')
    ax.set_ylabel(f'{value_col} ({chroma.upper()})', fontsize=12)
    ax.set_xlabel('Region of Interest (ROI)', fontsize=12)
    ax.set_title(f"Individual ROI Results ({chroma.upper()})\n{subject_id}",
                 fontsize=13, fontweight='bold')
    plt.tight_layout()
    plt.close(fig)

    return df_roi, roi_channel_map, fig



def generate_contrast_results(df_design_matrix, glm_est, file_path):

    contrast_results_dict = {}
    contrast_matrix = np.eye(df_design_matrix.shape[1])
    basic_conts = dict(
        [(column, contrast_matrix[i]) for i, column in enumerate(df_design_matrix.columns)]
    )
  
    if HRF_MODEL == "fir":
        all_delay_cols = [col for col in df_design_matrix.columns if "_delay_" in col]
        all_conditions = sorted({col.split("_delay_")[0] for col in all_delay_cols})
        if not all_conditions:
            raise ValueError("No FIR regressors found in the design matrix.")
 
        contrast_dict = {}
        for condition in all_conditions:
            delay_cols = [col for col in all_delay_cols if col.startswith(f"{condition}_delay_")]
            if not delay_cols:
                continue
            contrast_vector = np.mean([basic_conts[col] for col in delay_cols], axis=0)
            contrast_dict[condition] = contrast_vector
 
        for cond, contrast_vector in contrast_dict.items():
            contrast = glm_est.compute_contrast(contrast_vector)
            df = contrast.to_dataframe()
            df["ID"] = file_path
            contrast_results_dict[f"{cond}_vs_Zero"] = df
 
        for cond_a, cond_b in itertools.combinations(all_conditions, 2):
            if cond_a not in contrast_dict or cond_b not in contrast_dict:
                continue
            diff_vector = contrast_dict[cond_a] - contrast_dict[cond_b]
            contrast = glm_est.compute_contrast(diff_vector)
            df = contrast.to_dataframe()
            df["ID"] = file_path
            contrast_results_dict[f"{cond_a}_vs_{cond_b}"] = df
 
    else:
        # 0 is NOT a baseline.
        # AI Explaination:
        # When you regress a single condition (e.g., "Tapping_Right") against zero, the GLM asks:
        # "Is the signal during Tapping_Right significantly higher than the average signal across the entire run?"
        # Because of systemic physiology (the global blood pressure rise that happens during almost any active task),
        # the answer is almost always "Yes, the whole head is higher than the average."
 
        exclude_list = ["drift", "constant", "short"]
        task_cols = [c for c in df_design_matrix.columns if not any(ex in c.lower() for ex in exclude_list)]
 
        for cond in task_cols:
            vec = np.zeros(len(df_design_matrix.columns))
            vec[list(df_design_matrix.columns).index(cond)] = 1
            contrast = glm_est.compute_contrast(vec)
            df = contrast.to_dataframe()
            df["ID"] = file_path
            contrast_results_dict[f"{cond}_vs_Zero"] = df
 
        for cond_a, cond_b in itertools.combinations(task_cols, 2):
            vec = np.zeros(len(df_design_matrix.columns))
            vec[list(df_design_matrix.columns).index(cond_a)] = 1
            vec[list(df_design_matrix.columns).index(cond_b)] = -1
            contrast = glm_est.compute_contrast(vec)
            df = contrast.to_dataframe()
            df["ID"] = file_path
            contrast_results_dict[f"{cond_a}_vs_{cond_b}"] = df

    
    return contrast_results_dict



def haemoglobin_concentration(raw_od, file_path, override_ppf=False, ppf_lower_wavelength=6.0, ppf_upper_wavelength=6.0):
    if override_ppf:
        raw_haemo = beer_lambert_law(raw_od, ppf=(ppf_lower_wavelength, ppf_upper_wavelength))
    else:
        raw_haemo = beer_lambert_law(raw_od, ppf=calculate_dpf(file_path))
    return raw_haemo



def _png_worker(png_queue: Queue, fig_bytes_dict: dict, dpi: int = 100):
    """
    Runs on a single dedicated background thread for the lifetime of one
    process_participant() call. Consumes (label, fig) pairs and renders them
    to PNG bytes.
 
    Safety: figures are `plt.close()`-d by the producer (main thread) the
    instant they're created, BEFORE being queued. That means this thread
    never touches pyplot's global figure registry (Gcf) - it only calls
    fig.savefig(), which operates on the Figure/Canvas object directly.
    Since this is the only thread that ever calls savefig/draw, there's no
    concurrent-draw hazard, even while the main thread keeps creating and
    closing new figures elsewhere in the pipeline.
    """
    while True:
        item = png_queue.get()
        if item is None:  # sentinel - no more figures coming
            png_queue.task_done()
            break
        label, fig = item
        try:
            buf = BytesIO()
            fig.savefig(buf, format="png", dpi=dpi, pil_kwargs={"compress_level": 1})
            fig_bytes_dict[label] = buf.getvalue()
        finally:
            fig.clf()  # drop axes/artists now that we're done, frees memory sooner
            png_queue.task_done()



def initial_setup(file_path):
    timings = {}
    step_start = time.perf_counter()
    config_dict = {
        k: globals()[k]
        for k in __annotations__
        if k in globals()
    }

    fig_bytes_dict: dict[str, bytes] = {}
    qc: dict[str, Any] = {"file_path": file_path}
    png_queue: Queue = Queue()
    png_thread = threading.Thread(
        target=_png_worker, args=(png_queue, fig_bytes_dict), daemon=True
    )
    png_thread.start()

    return fig_bytes_dict, config_dict, png_queue, timings, step_start, qc



def _enqueue(label, fig, png_queue):
    if fig is None:
        return
    plt.close(fig)
    png_queue.put((label, fig))



def lap(start, timings, name):
    now = time.perf_counter()
    timings[name] = now - start
    return now



def process_participant(file_path, file_start, progress_callback=None):

    print(f"File was started with {time.time() - file_start:2f} seconds elapsed.")
    # Step 0: Setting up
    fig_bytes_dict, config_dict, png_queue, timings, step_start, qc = initial_setup(file_path)
    step_start = lap(step_start, timings, "Step 0")

    # Step 1: Preprocessing
    raw = load_snirf(file_path=file_path, downsample_frequency=DOWNSAMPLE_FREQUENCY, verbosity=VERBOSITY)
    qc["n_channels_loaded"] = raw.info['nchan']
    fig_raw = raw.plot(duration=raw.times[-1], n_channels=raw.info['nchan'], title="Loaded Raw", show=False)
    _enqueue("Loaded Raw Data", fig_raw, png_queue)
    if progress_callback: progress_callback(1)
    logger.info("Step 1 Completed.")
    step_start = lap(step_start, timings, "Step 1")

    # Step 2: Trimming
    if TRIM and not FOLDING_BYP:
        raw, fig_trimmed = trim_participant_data(raw, seconds_to_keep=SECONDS_TO_KEEP)
        _enqueue("Trimmed Raw Data", fig_trimmed, png_queue)
    if progress_callback: progress_callback(2)
    logger.info("Step 2 Completed.")
    step_start = lap(step_start, timings, "Step 2")

    # Step 3: Verify Optode Placement
    if OPTODE_PLACEMENT:
        fig_optodes = raw.plot_sensors(show_names=SHOW_OPTODE_NAMES, to_sphere=True, show=False, verbose=VERBOSITY) # type: ignore
        _enqueue("Plot Sensors", fig_optodes, png_queue)
    if progress_callback: progress_callback(3)
    logger.info("Step 3 Completed.")
    step_start = lap(step_start, timings, "Step 3")

    # Step 4: Short/Long Channels
    if SHORT_CHANNELS and not FOLDING_BYP:
        #NOTE: Have to split again later but since needed for heart rate, this will stay at step 4. Will split later again.
        _short_chans = get_short_channels(raw, max_dist=SHORT_CHANNELS_THRESHOLD) # JUST FOR PLOTTING THEM SEPERATELY
        fig_short_chans = _short_chans.plot(duration=raw.times[-1], n_channels=raw.info['nchan'], title="Short Channels Only", show=False, verbose=VERBOSITY)
        _enqueue("Short Channels Raw Data", fig_short_chans, png_queue)
    if LONG_CHANNELS:
        raw = get_long_channels(raw, min_dist=0, max_dist=LONG_CHANNELS_THRESHOLD)
    if progress_callback: progress_callback(4)
    logger.info("Step 4 Completed.")
    step_start = lap(step_start, timings, "Step 4")

    # Step 5: Heart Rate
    qc["heart_rate_ran"] = HEART_RATE and not FOLDING_BYP
    if HEART_RATE and not FOLDING_BYP:
        fig, hr1, hr2, low, high, final_hr = hr_calc(
            raw, 
            seconds_to_strip_hr=SECONDS_TO_STRIP_HR, 
            l_freq=HR_LOW_FREQ, 
            h_freq=HR_HIGH_FREQ, 
            search_min=HR_SEARCH_MIN, 
            search_max=HR_SEARCH_MAX, 
            max_low_hr=MAX_LOW_HR, 
            max_high_hr=MAX_HIGH_HR, 
            smoothing_window_hr=SMOOTHING_WINDOW_HR, 
            hr_window=HEART_RATE_WINDOW,
            short_channels=SHORT_CHANNELS, 
            short_channels_threshold=SHORT_CHANNELS_THRESHOLD,
            verbosity=VERBOSITY
        )
        qc["final_hr_bpm"] = round(final_hr, 1)
        _enqueue("Power Spectral Density", fig, png_queue)
        _enqueue('Heart Rate - PSD', hr1, png_queue)
        _enqueue('Heart Rate - Time', hr2, png_queue)
    if progress_callback: progress_callback(5)
    logger.info("Step 5 Completed.")    
    step_start = lap(step_start, timings, "Step 5")

    # Step 6: Scalp Coupling Index    
    bad_sci = []
    if SCI and not FOLDING_BYP:
        if HEART_RATE and SCI_USE_HEART_RATE_BAND:
            bad_sci, fig_sci_1, fig_sci_2 = calculate_scalp_coupling(raw, l_freq=low, h_freq=high, time_window=SCI_TIME_WINDOW, threshold=SCI_THRESHOLD)
        else:
            bad_sci, fig_sci_1, fig_sci_2 = calculate_scalp_coupling(raw, l_freq=SCI_LOW_FREQ, h_freq=SCI_HIGH_FREQ, time_window=SCI_TIME_WINDOW, threshold=SCI_THRESHOLD)
        _enqueue("Scalp Coupling Index Heatmap", fig_sci_1, png_queue)
        _enqueue("Scalp Coupling Index Binary Heatmap", fig_sci_2, png_queue)
    qc["n_bad_sci"] = len(bad_sci)
    if progress_callback: progress_callback(6)
    logger.info("Step 6 Completed.")
    step_start = lap(step_start, timings, "Step 6")

    # Step 7: Signal to Noise Ratio
    bad_snr = []
    if SNR and not FOLDING_BYP:
        bad_snr, fig_snr = calculate_signal_noise_ratio(raw)
        _enqueue("Signal To Noise Ratio", fig_snr, png_queue)
    qc["n_bad_snr"] = len(bad_snr)
    if progress_callback: progress_callback(7)
    logger.info("Step 7 Completed.")
    step_start = lap(step_start, timings, "Step 7")

    # Step 8: Peak Spectral Power
    bad_psp = []
    if PSP and not FOLDING_BYP:
        if HEART_RATE and PSP_USE_HEART_RATE_BAND:
            bad_psp, fig_psp1, fig_psp2 = calculate_peak_power(raw, time_window=PSP_TIME_WINDOW, threshold=PSP_THRESHOLD, l_freq=low, h_freq=high)
        else:
            bad_psp, fig_psp1, fig_psp2 = calculate_peak_power(raw, time_window=PSP_TIME_WINDOW, threshold=PSP_THRESHOLD, l_freq=PSP_LOW_FREQ, h_freq=PSP_HIGH_FREQ)
        _enqueue("Peak Spectral Power Heatmap", fig_psp1, png_queue)
        _enqueue("Peak Spectral Power Binary Heatmap", fig_psp2, png_queue)
    qc["n_bad_psp"] = len(bad_psp)
    if progress_callback: progress_callback(8)
    logger.info("Step 8 Completed.")
    step_start = lap(step_start, timings, "Step 8")

    # Step 9: Coefficient of Variation
    bad_coeff_var = []
    if COEFF_VAR and not FOLDING_BYP:
        bad_coeff_var, fig_coeff_var = find_bad_channels_coeff_var(raw, coeff_var_threshold=COEFF_VAR_THRESHOLD)
        _enqueue('Coefficient of Variation', fig_coeff_var, png_queue)
    qc["n_bad_coeff_var"] = len(bad_coeff_var)
    if progress_callback: progress_callback(9)
    logger.info("Step 9 Completed.")
    step_start = lap(step_start, timings, "Step 9")

    # Step 10: Median Absolute Deviation
    bad_amplitude_range = []
    if MAD and not FOLDING_BYP:
        bad_amplitude_range, fig_range = find_bad_channels_by_amplitude_range(raw, threshold=MAD_THRESHOLD)
        _enqueue('Median Absolute Deviation', fig_range, png_queue)
    qc["n_bad_mad"] = len(bad_amplitude_range)
    if progress_callback: progress_callback(10)
    logger.info("Step 10 Completed.")
    step_start = lap(step_start, timings, "Step 10")

    # Step 11: Power Spectral Density Noise
    bad_noise = []
    if PSD_NOISE and not FOLDING_BYP:
        bad_noise, fig_noise = detect_spectral_noise_spike(raw, db_limit=DB_LIMIT, freq_div=TARGET_FREQ_DIV, min_freq=PSD_MIN_FREQ, target_bandwith=PSD_TARGET_BANDWIDTH)
        _enqueue('Power Spectral Density Noise', fig_noise, png_queue)
    qc["n_bad_psd_noise"] = len(bad_noise)
    if progress_callback: progress_callback(11)
    logger.info("Step 11 Completed.")
    step_start = lap(step_start, timings, "Step 11")

    # Step 12: Channel Dropout
    bad_disp = []
    if SENSOR_DROPOUT and not FOLDING_BYP:
        bad_disp, fig_disp = detect_sensor_dropout(raw, threshold_ratio=SENSOR_DROPOUT_VARIANCE_THRESHOLD)
        _enqueue('Sensor Dropout', fig_disp, png_queue)
    qc["n_bad_dropout"] = len(bad_disp)
    if progress_callback: progress_callback(12)
    logger.info("Step 12 Completed.")
    step_start = lap(step_start, timings, "Step 12")

    # Step 13: Bad Channels Handling
    qc["n_bad_channels_total"] = 0
    qc["pct_bad_channels"] = 0.0
    qc["bad_channels_handling"] = "None"
    if BAD_CHANNELS_HANDLING != "None" and not FOLDING_BYP:
        raw, fig_dropped, fig_raw_before, bad_channels = mark_bads(raw, bad_sci, bad_snr, bad_psp, bad_coeff_var, bad_amplitude_range, bad_noise, bad_disp)
        qc["n_bad_channels_total"] = len(bad_channels)
        qc["pct_bad_channels"] = round(100 * len(bad_channels) / qc["n_channels_loaded"], 1) if qc["n_channels_loaded"] else 0.0
        qc["bad_channels_handling"] = BAD_CHANNELS_HANDLING
        if fig_dropped and fig_raw_before is not None:
            _enqueue("Bad Channels by Method", fig_dropped, png_queue)
            _enqueue("Bad Channels Data", fig_raw_before, png_queue)
        if bad_channels:
            if BAD_CHANNELS_HANDLING == "Interpolate":
                raw, fig_raw_after, fig_compare = interpolate_fNIRS_bads_weighted_average(raw, max_dist=MAX_DIST, min_neighbors=MIN_NEIGHBORS, short_channels_threshold=SHORT_CHANNELS_THRESHOLD)
                _enqueue("Data after Interpolating Bad Channels", fig_raw_after, png_queue)
                _enqueue("Bad Channels Interpolation Results", fig_compare, png_queue)
            elif BAD_CHANNELS_HANDLING == "Remove":
                raw = remove_bad_channels(raw, bad_channels, max_bad_channels=MAX_BAD_CHANNELS)
    if progress_callback: progress_callback(13)
    logger.info("Step 13 Completed.")
    step_start = lap(step_start, timings, "Step 13")

    # Step 14: Optical Density
    raw_od = optical_density(raw)
    fig_raw_od = raw_od.plot(duration=raw.times[-1], n_channels=raw.info['nchan'], title="Optical Density", show=False)
    _enqueue("Optical Density", fig_raw_od, png_queue)
    if progress_callback: progress_callback(14)
    logger.info("Step 14 Completed.")
    step_start = lap(step_start, timings, "Step 14")

    # Step 15: Temporal Derivative Distribution Repair Filtering
    if TDDR and not FOLDING_BYP:
        raw_od = temporal_derivative_distribution_repair(raw_od)
        fig_raw_od_tddr = raw_od.plot(duration=raw.times[-1], n_channels=raw.info['nchan'], title="After TDDR (Motion Correction)", show=False)
        _enqueue("Temporal Derivative Distribution Repair", fig_raw_od_tddr, png_queue)
    if progress_callback: progress_callback(15)
    logger.info("Step 15 Completed.")
    step_start = lap(step_start, timings, "Step 15")
    
    # Step 16: Wavelet Filtering
    if WAVELET and not FOLDING_BYP:
        raw_od, fig = calculate_and_apply_wavelet(data=raw_od, wavelet_type=WAVELET_TYPE, wavelet_level=WAVELET_LEVEL, iqr=IQR, verbosity=VERBOSITY)
        _enqueue("Wavelet", fig, png_queue)
    if progress_callback: progress_callback(16)
    logger.info("Step 16 Completed.")
    step_start = lap(step_start, timings, "Step 16")

    # Step 17: Haemoglobin Concentration
    raw_haemo = haemoglobin_concentration(raw_od, file_path, OVERRIDE_PPF, PPF_LOWER_WAVELENGTH, PPF_UPPER_WAVELENGTH)
    fig_raw_haemo_bll = raw_haemo.plot(duration=raw_haemo.times[-1], n_channels=raw_haemo.info['nchan'], title="HbO and HbR Signals", show=False)
    _enqueue("Modified Beer Lambert Law", fig_raw_haemo_bll, png_queue)
    if progress_callback: progress_callback(17)
    logger.info("Step 17 Completed.")
    step_start = lap(step_start, timings, "Step 17")

    bad_anticorr = []
    if FEATURE_2 and not FOLDING_BYP:
        bad_anticorr, fig_anticorr = detect_hbo_hbr_anticorrelation(raw_haemo, threshold=-0.2)
        _enqueue("HbO-HbR Anti-Correlation", fig_anticorr, png_queue)
    qc["n_bad_anticorrelation"] = len(bad_anticorr)

    # Step 18: Enhance Negative Correlation
    if ENHANCE_NEGATIVE_CORRELATION and not FOLDING_BYP:
        raw_haemo = enhance_negative_correlation(raw_haemo)
        fig_raw_haemo_enc = raw_haemo.plot(duration=raw_haemo.times[-1], n_channels=raw_haemo.info['nchan'], title="Enhance Negative Correlation", show=False)
        _enqueue("Enhance Negative Correlation", fig_raw_haemo_enc, png_queue)
    if progress_callback: progress_callback(18)
    logger.info("Step 18 Completed.")
    step_start = lap(step_start, timings, "Step 18")
    
    # Step 19: Filter
    if FILTER and not FOLDING_BYP:
        raw_haemo, fig_filter, fig_raw_haemo_filter = filter_the_data(
            raw_haemo,
            filter_algorithm=FILTER_ALGORITHM,
            l_freq=L_FREQ,
            h_freq=H_FREQ,
            l_trans_bandwidth=L_TRANS_BANDWIDTH,
            h_trans_bandwidth=H_TRANS_BANDWIDTH,
            # iir_type=IIR_TYPE,
            # iir_order=IIR_ORDER,
            filter_length=FILTER_LENGTH,
            filter_phase=FILTER_PHASE,
            fir_window=FIR_WINDOW,
            fir_design=FIR_DESIGN,
            # iir_output=IIR_OUTPUT,
            # passband_ripple=PASSBAND_RIPPLE,
            # stopband_attenuation=STOPBAND_ATTENUATION,
            filter_pad=FILTER_PAD,
            # skip_by_annotation=SKIP_BY_ANNOTATION,
            filter_n_jobs=FILTER_N_JOBS,
            verbosity=VERBOSITY
        )
        _enqueue("Filter_1", fig_filter, png_queue)
        _enqueue("Filter_2", fig_raw_haemo_filter, png_queue)
    if progress_callback: progress_callback(19)
    logger.info("Step 19 Completed.")
    step_start = lap(step_start, timings, "Step 19")

    # Step 20: Extracting Events
    if EVENTS and not FOLDING_BYP:
        if SHORT_CHANNELS:
            raw_haemo_evnt = get_long_channels(raw_haemo, min_dist=SHORT_CHANNELS_THRESHOLD, max_dist=LONG_CHANNELS_THRESHOLD)
        else:
            raw_haemo_evnt = raw_haemo
        events, event_dict = events_from_annotations(raw_haemo_evnt, event_id=EVENT_ID, regexp=EVENT_REGEX, verbose=VERBOSITY) #TODO: Implement the chunk duration
        fig_events = plot_events(events, event_id=event_dict, sfreq=raw_haemo_evnt.info["sfreq"], show=False)
        _enqueue("Events", fig_events, png_queue)
    if progress_callback: progress_callback(20)
    logger.info("Step 20 Completed.")
    step_start = lap(step_start, timings, "Step 20")

    # Step 21: Epoch Calculations
    epochs = None
    qc["n_epochs_final"] = 0
    if EPOCHS and EVENTS and not FOLDING_BYP:
        epochs = epochs_calculations(
            raw_haemo_evnt, 
            events, 
            event_dict, 
            epoch_handling=EPOCH_HANDLING,
            max_shift=MAX_SHIFT,
            t_min=T_MIN,
            t_max=T_MAX,
            baseline=(None,0), #TODO: Unhardcode this
            reject_epochs=REJECT_EPOCHS,
            reject_hbo_threshold=dict(hbo=REJECT_HBO_THRESHOLD),
            png_queue=png_queue
        )
        qc["n_epochs_final"] = len(epochs) if epochs is not None else 0
    if progress_callback: progress_callback(21)
    logger.info("Step 21 Completed.")
    step_start = lap(step_start, timings, "Step 21")

    # Step 22: Design Matrix
    raw_haemo, raw_haemo_dm, df_design_matrix, fig_design_matrix = make_design_matrix(
        raw_haemo=raw_haemo,
        resample=RESAMPLE,
        resample_freq=RESAMPLE_FREQ,
        stim_dur=STIM_DUR,
        hrf_model=HRF_MODEL,
        drift_model=DRIFT_MODEL,
        high_pass=HIGH_PASS,
        drift_order=DRIFT_ORDER,
        fir_delays=FIR_DELAYS,
        min_onset=MIN_ONSET,
        oversampling=OVERSAMPLING,
        short_channel_regression=SHORT_CHANNEL_REGRESSION,
        short_channels=SHORT_CHANNELS,
        long_channels=LONG_CHANNELS,
        short_channels_threshold=SHORT_CHANNELS_THRESHOLD,
        long_channels_threshold=LONG_CHANNELS_THRESHOLD,
        folding_bypass=FOLDING_BYP
    )
    _enqueue("Design Matrix", fig_design_matrix, png_queue)
    if progress_callback: progress_callback(22)
    logger.info("Step 22 Completed.")
    step_start = lap(step_start, timings, "Step 22")

    # Step 23: General Linear Model
    glm_est, fig_glm_topo = make_and_run_glm(raw_haemo_dm, df_design_matrix, noise_model=NOISE_MODEL, bins=BINS, n_jobs=N_JOBS, verbosity=VERBOSITY)
    _enqueue("GLM Topography", fig_glm_topo, png_queue)
    if progress_callback: progress_callback(23)
    logger.info("23")
    step_start = lap(step_start, timings, "Step 23")

    # Step 24: Generate GLM Results
    if "derivative" not in HRF_MODEL.lower():
        fig_glm_result = plot_glm_results(file_path, raw_haemo, glm_est, df_design_matrix)
        for name, fig in fig_glm_result:
            _enqueue(f"GLM {name}", fig, png_queue)
    if progress_callback: progress_callback(24)
    logger.info("24")
    step_start = lap(step_start, timings, "Step 24")

    # Step 25: Generate Channel Results
    df_cha = generate_channel_results(glm_est, file_path)
    if progress_callback: progress_callback(25)
    logger.info("25")
    step_start = lap(step_start, timings, "Step 25")

    # Step 26: Generate Region of Interest Results
    df_roi, roi_channel_map, fig_roi = generate_roi_results(raw_haemo, df_design_matrix, glm_est, file_path, json_location=JSON_LOCATION)
    _enqueue("Region of Interest", fig_roi, png_queue)
    if progress_callback: progress_callback(26)
    logger.info("26")
    step_start = lap(step_start, timings, "Step 26")

    # Step 27: Generate Contrast Results
    contrast_results_dict = generate_contrast_results(df_design_matrix, glm_est, file_path)
    if progress_callback: progress_callback(27)
    logger.info("27")
    step_start = lap(step_start, timings, "Step 27")


    # Step 27.5: Extract FIR Waveform Features & Enqueue Metric Plots
    fir_feature_dict = {'features': np.array([]), 'feature_names': [], 'feature_channels': []}

    if HRF_MODEL.lower() == "fir":
        try:
            fir_feature_dict = extract_fir_features_real_data(
                raw=raw_haemo,
                target_condition=None,  # e.g., 'reach'
                fir_delays=FIR_DELAYS,               # e.g., np.arange(0, 15)
                selected_metrics=tuple(METRIC_REGISTRY.keys()),   # e.g., ('Peak_Amp', 'TTP', 'AUC')
                roi_map=JSON_LOCATION,
                chromophores=('hbo', 'hbr', 'hbt'),
                glm_est=glm_est,
                df_design_matrix=df_design_matrix,
                png_queue=png_queue
            )
            logger.info("Step 27.5: FIR features successfully extracted and metric images enqueued.")
        except Exception as e:
            logger.warning(f"Step 27.5 Failed to extract FIR features: {e}")

     
    # Step 28: Finishing Up
    png_queue.put(None)   # sentinel
    png_queue.join()      # blocks only on remaining unfinished work
    if FOLDING_BYP:
        epochs = None
    sanitize_paths_for_pickle(raw_haemo, epochs)
    if progress_callback: progress_callback(28)
    logger.info("28")
    step_start = lap(step_start, timings, "Step 28")

    # Step 28.5: Return the results
    qc["total_processing_seconds"] = round(sum(timings.values()), 2)
    logger.info("Step timings:")
    for name, elapsed in timings.items():
        logger.info(f"  {name:<25} {elapsed:7.3f}s")

    logger.info(f"Total processing time: {sum(timings.values()):.3f}s")
    return raw_haemo, epochs, df_cha, df_roi, df_design_matrix, config_dict, fig_bytes_dict, contrast_results_dict, roi_channel_map, fir_feature_dict, qc, True



def sanitize_paths_for_pickle(raw_haemo, epochs):
    # Fix raw_haemo._filenames
    if hasattr(raw_haemo, '_filenames'):
        raw_haemo._filenames = [str(p) for p in raw_haemo._filenames]

    # Fix epochs._raw._filenames
    if hasattr(epochs, '_raw') and hasattr(epochs._raw, '_filenames'):
        epochs._raw._filenames = [str(p) for p in epochs._raw._filenames]



def _check_epoch_frequency_resolution(
    epoch_duration: float,
    fmin: float,
    min_cycles: float = 5.0,
    allow_unreliable: bool = False,
) -> None:
    required_duration = min_cycles / fmin
    if epoch_duration >= required_duration:
        return
    actual_cycles = epoch_duration * fmin
    message = (
        f"Epoch duration ({epoch_duration:.2f}s) gives only {actual_cycles:.2f} cycles "
        f"at fmin={fmin} Hz; need >= {min_cycles} cycles ({required_duration:.1f}s epochs) "
        f"for a reliable estimate. Below this threshold, the estimate is commonly dominated "
        f"by whatever frequency content IS well-resolved (often cardiac pulsation or motion "
        f"artifact), producing plausible-looking but spurious results rather than failing loudly."
    )
    if allow_unreliable:
        logger.warning(f"{message} Proceeding anyway because allow_unreliable=True.")
    else:
        raise ValueError(
            f"{message} Raise fmin to >= {min_cycles / epoch_duration:.3f} Hz for this epoch "
            f"length, use longer epochs, or pass allow_unreliable=True to proceed anyway."
        )
    

def functional_connectivity_spectral_epochs(
    epochs: Epochs,
    n_lines: int,
    vmin: float,
    fmin: float = 0.04,
    fmax: float = 0.2,
    method: str = "wpli2_debiased",
    allow_unreliable: bool = False,
    verbose: bool = False
) -> None:
    logger.info(f"[spectral_epochs] called (method={method}, fmin={fmin}, fmax={fmax})")

    epochs.load_data()
    hbo_epochs = epochs.copy().pick(picks="hbo")

    epoch_duration = hbo_epochs.times[-1] - hbo_epochs.times[0]
    _check_epoch_frequency_resolution(epoch_duration, fmin, allow_unreliable=allow_unreliable)

    con_coh = spectral_connectivity_epochs(
        hbo_epochs, method=method, mode="fourier", sfreq=hbo_epochs.info["sfreq"],
        fmin=fmin, fmax=fmax, faverage=True, verbose=verbose
    )
    coh = np.squeeze(con_coh.get_data(output="dense"))
    coh = coh + coh.T - np.diag(np.diag(coh))
    np.fill_diagonal(coh, 0)

    if verbose:
        logger.info(f"[{method}] matrix stats: min={coh.min():.6f}, max={coh.max():.6f}, "
                f"mean={coh.mean():.6f}, count above vmin({vmin})={np.sum(coh >= vmin)}")

    plot_connectivity_circle(
        coh, hbo_epochs.ch_names,
        title=f"fNIRS Functional Connectivity (HbO - {method}, {fmin}-{fmax} Hz)",
        n_lines=n_lines, vmin=vmin
    )
    logger.info("[spectral_epochs] finished")


def functional_connectivity_envelope(
    epochs: Epochs,
    n_lines: int,
    vmin: float,
    fmin: float = 0.04,
    fmax: float = 0.2,
    orthogonalize: bool = False,
    absolute: bool = True,
    allow_unreliable: bool = False,
    verbose: bool = True,
) -> None:
    """
    Compute and plot HbO functional connectivity via envelope correlation.
    See DESCRIPTION for full method explanation.
    """
    logger.info(f"[envelope] called (fmin={fmin}, fmax={fmax})")

    epochs.load_data()
    hbo_epochs = epochs.copy().pick(picks="hbo")

    epoch_duration = hbo_epochs.times[-1] - hbo_epochs.times[0]
    _check_epoch_frequency_resolution(epoch_duration, fmin, allow_unreliable=allow_unreliable)

    hbo_epochs.filter(l_freq=fmin, h_freq=fmax, verbose=True)  # MNE's filter() naming, translated internally

    data = hbo_epochs.get_data()

    env = envelope_correlation(data, orthogonalize=orthogonalize, absolute=absolute, verbose=verbose)
    env_data = env.get_data(output="dense")
    env_corr = np.squeeze(env_data.mean(axis=0))
    np.fill_diagonal(env_corr, 0)
    if verbose:
        _log_matrix_diagnostics("envelope", env_corr)

    plot_connectivity_circle(
        env_corr,
        hbo_epochs.ch_names,
        title=f"fNIRS HbO Envelope Correlation ({fmin}-{fmax} Hz, "
              f"{'orth' if orthogonalize else 'no orth'}, "
              f"{'abs' if absolute else 'signed'})",
        n_lines=n_lines,
        vmin=vmin
    )
    logger.info("[envelope] finished")


def functional_connectivity_spectral_time(
    epochs: Epochs,
    n_lines: int,
    vmin: float,
    fmin: float = 0.04,
    fmax: float = 0.2,
    n_freqs: int = 10,
    cycles_multiplier: float = 2.0,
    method: str = "wpli",
    allow_unreliable: bool = False,
    verbose: bool = True
) -> None:
    logger.info(f"[spectral_time] called (method={method}, fmin={fmin}, fmax={fmax})")

    epochs.load_data()
    hbo_epochs = epochs.copy().pick(picks="hbo")

    epoch_duration = hbo_epochs.times[-1] - hbo_epochs.times[0]
    _check_epoch_frequency_resolution(epoch_duration, fmin, allow_unreliable=allow_unreliable)

    data = hbo_epochs.get_data()
    names = hbo_epochs.ch_names
    sfreq = hbo_epochs.info["sfreq"]

    freqs = np.linspace(fmin, fmax, n_freqs)
    n_cycles = freqs * cycles_multiplier

    con_coh = spectral_connectivity_time(
        data, freqs=freqs, method=method, mode="multitaper", sfreq=sfreq,
        fmin=fmin, fmax=fmax, n_cycles=n_cycles, faverage=True, verbose=verbose
    )
    coh = con_coh.get_data(output="dense").squeeze()
    if coh.ndim == 3:
        coh = coh.mean(axis=0)
    coh = coh + coh.T - np.diag(np.diag(coh))
    np.fill_diagonal(coh, 0)

    if verbose:
        logger.info(f"[{method}] matrix stats: min={coh.min():.6f}, max={coh.max():.6f}, "
                f"mean={coh.mean():.6f}, count above vmin({vmin})={np.sum(coh >= vmin)}")
        
    plot_connectivity_circle(
        coh, names,
        title=f"fNIRS Functional Connectivity (HbO - {method}, Time-Resolved, {fmin}-{fmax} Hz, {n_freqs} bins)",
        n_lines=n_lines, vmin=vmin
    )
    logger.info("[spectral_time] finished")


def functional_connectivity_betas(
    raw_hbo: BaseRaw,
    n_lines: int,
    event_name: str | None = None,
    *,
    drift_model: str = "cosine",
    drift_order: int = 1,
    hrf_model: str = "glover",
    apply_gsr: bool = True,
    min_effect_size: float = 0.7,
    alpha: float = 0.05,
    resample_freq: float | None = 4.0,
    verbose: bool = True,
) -> None:
    logger.info(f"[betas] called (hrf_model={hrf_model}, event_name={event_name})")

    raw_hbo = raw_hbo.copy().pick(picks="hbo")

    if event_name is not None:
        keep_mask = [desc == event_name for desc in raw_hbo.annotations.description]
        raw_hbo.set_annotations(raw_hbo.annotations[keep_mask])

    if resample_freq is not None and raw_hbo.info["sfreq"] > resample_freq:
        raw_hbo.resample(resample_freq, npad="auto")

    raw_hbo.annotations.description = np.array([
        f"{desc}__trial_{i:03d}" for i, desc in enumerate(raw_hbo.annotations.description)
    ])

    design_kwargs = dict(raw=raw_hbo, drift_model=drift_model, drift_order=drift_order)
    if hrf_model == "fir":
        design_kwargs.update(hrf_model="fir", fir_delays=np.arange(0, 12, 1))
    else:
        design_kwargs.update(hrf_model=hrf_model)

    if verbose:
        logger.info(f"[betas] building design matrix, n_samples={len(raw_hbo.times)}...")
    design_matrix = make_first_level_design_matrix(**design_kwargs)
    if verbose:
        logger.info(f"[betas] design matrix built: shape={design_matrix.shape}")

    glm_results = run_glm(raw_hbo, design_matrix)
    betas = np.array(glm_results.theta())
    if betas.ndim == 3 and betas.shape[-1] == 1:
        betas = betas.squeeze(axis=-1)

    reg_names = list(design_matrix.columns)
    n_channels = betas.shape[0]
    assert betas.shape[1] == len(reg_names), (
        f"betas has {betas.shape[1]} columns but design matrix has "
        f"{len(reg_names)} regressors (betas.shape={betas.shape})"
    )

    trial_tags = sorted({
        (col.split("_delay")[0] if "_delay" in col else col)
        for col in reg_names
        if ("__trial_" in col)
        and (event_name is None or col.startswith(event_name + "__"))
    })

    if len(trial_tags) == 0:
        raise ValueError(f"No trials found for event_name={event_name}")
    if len(trial_tags) < 4:
        raise ValueError(
            f"Only {len(trial_tags)} trials found for event_name={event_name}; "
            "need at least 4 to compute correlation degrees of freedom."
        )
    
    if verbose:
        logger.info(f"[betas] trial_tags found: {len(trial_tags)}")

    beta_series = np.zeros((n_channels, len(trial_tags)))
    for t_idx, tag in enumerate(trial_tags):
        col_idx = [
            j for j, col in enumerate(reg_names)
            if (col.split("_delay")[0] if "_delay" in col else col) == tag
        ]
        beta_series[:, t_idx] = betas[:, col_idx].mean(axis=1)

    if apply_gsr:
        global_signal = np.mean(beta_series, axis=0)
        beta_series_clean = np.zeros_like(beta_series)
        for i in range(n_channels):
            slope, intercept = np.polyfit(global_signal, beta_series[i, :], 1)
            beta_series_clean[i, :] = beta_series[i, :] - (slope * global_signal + intercept)
        if verbose:
            logger.info("[betas] GSR applied")
    else:
        beta_series_clean = beta_series

    n_trials = beta_series_clean.shape[1]
    corr_matrix = np.corrcoef(beta_series_clean)

    with np.errstate(divide="ignore", invalid="ignore"):
        t_stats = corr_matrix * np.sqrt((n_trials - 2) / (1 - corr_matrix ** 2))
    p_matrix = 2 * t_dist.sf(np.abs(t_stats), df=n_trials - 2)
    np.fill_diagonal(p_matrix, 1.0)

    triu = np.triu_indices(n_channels, k=1)
    flat_p = p_matrix[triu]

    reject, _ = multipletests(flat_p, method="fdr_bh", alpha=alpha)[:2]
    sig_corr_matrix = np.zeros_like(corr_matrix)

    for idx, is_sig in enumerate(reject):
        r_val = corr_matrix[triu[0][idx], triu[1][idx]]
        if is_sig and abs(r_val) > min_effect_size:
            sig_corr_matrix[triu[0][idx], triu[1][idx]] = r_val
            sig_corr_matrix[triu[1][idx], triu[0][idx]] = r_val

    gsr_tag = "GSR" if apply_gsr else "no GSR"
    plot_connectivity_circle(
        sig_corr_matrix,
        raw_hbo.ch_names,
        title=f"Beta-Series Connectivity ({hrf_model}, FDR q<{alpha}, "
              f"|r|>{min_effect_size}, {gsr_tag})",
        n_lines=n_lines,
        vmin=min_effect_size,
        vmax=1.0,
        colormap="hot",
    )
    logger.info("[betas] finished")


def _log_matrix_diagnostics(name: str, mat: np.ndarray) -> None:
    """Temporary diagnostic: checks whether a connectivity matrix is fully
    populated and symmetric, or only has one triangle filled (the bug found
    in the group coherence path)."""
    n = mat.shape[0]
    total_offdiag = n * n - n
    nonzero = np.count_nonzero(mat)
    is_symmetric = np.allclose(mat, mat.T)
    upper_nonzero = np.count_nonzero(np.triu(mat, k=1))
    lower_nonzero = np.count_nonzero(np.tril(mat, k=-1))
    logger.info(
        f"[{name}] shape={mat.shape}, nonzero={nonzero}/{total_offdiag} off-diag cells, "
        f"symmetric={is_symmetric}, upper_tri_nonzero={upper_nonzero}, lower_tri_nonzero={lower_nonzero}"
    )








# ============================================================================
# Shared: channel alignment across participants
# ============================================================================

def _align_participant_matrices(
    subject_results: list[tuple[np.ndarray, list[str]]],
) -> tuple[np.ndarray, list[str], list[int]]:
    """
    Given a list of (corr_matrix, ch_names) per participant, find the
    channel set common to ALL participants, reindex every matrix to that
    common set (same order), and stack into one array.

    This is required because participants can legitimately end up with
    different channel counts/sets (bad-channel handling, interpolation
    failures, short/long trimming differences) - stacking raw matrices
    without this step either crashes or silently misaligns channel
    positions across participants.

    Returns
    -------
    stacked : (n_included_participants, n_common_channels, n_common_channels)
    common_names : list of channel names, in the order used for stacking
    dropped_indices : indices into subject_results that were excluded
        because they didn't have all common channels (shouldn't happen
        once common_names is the intersection, but guards against
        duplicate/malformed names)
    """
    if not subject_results:
        raise ValueError("No participant results to align - subject_results is empty.")

    name_sets = [set(names) for _, names in subject_results]
    common = set.intersection(*name_sets)

    if not common:
        raise ValueError(
            "No channels are common across all selected participants. "
            "Check that bad-channel handling / trimming produced consistent "
            "channel sets, or select a smaller/more homogeneous participant group."
        )

    dropped_channels_per_subject = {
        i: name_sets[i] - common for i in range(len(subject_results)) if name_sets[i] - common
    }
    if dropped_channels_per_subject:
        for i, dropped in dropped_channels_per_subject.items():
            logger.warning(
                f"Participant {i}: {len(dropped)} channel(s) not shared across "
                f"the group, excluded from group analysis: {sorted(dropped)}"
            )

    common_names = sorted(common)  # deterministic order
    stacked = []
    dropped_indices = []

    for i, (corr, names) in enumerate(subject_results):
        name_to_idx = {n: j for j, n in enumerate(names)}
        try:
            idx = [name_to_idx[n] for n in common_names]
        except KeyError:
            dropped_indices.append(i)
            continue
        reindexed = corr[np.ix_(idx, idx)]
        stacked.append(reindexed)

    if not stacked:
        raise ValueError("After channel alignment, no participants had usable data.")

    return np.array(stacked), common_names, dropped_indices


# ============================================================================
# Shared: group-level statistics (Fisher-Z average, t-test, FDR)
# ============================================================================

def _group_ttest_connectivity(
    corr_stack: np.ndarray,
    alpha: float,
    min_effect_size: float,
    min_participants: int = 3,
    channel_names: list[str] | None = None,
    n_top_pairs: int = 15,
) -> np.ndarray:
    n_participants = corr_stack.shape[0]
    if n_participants < min_participants:
        raise ValueError(
            f"Only {n_participants} participant(s) with usable data; need at "
            f"least {min_participants} for a group-level test to be meaningful."
        )

    n_channels = corr_stack.shape[1]
    z_stack = np.arctanh(np.clip(corr_stack, -0.999, 0.999))

    t_stats, p_values = ttest_1samp(z_stack, popmean=0, axis=0)

    triu = np.triu_indices(n_channels, k=1)
    flat_p = p_values[triu]

    nan_mask = np.isnan(flat_p)
    if nan_mask.any():
        logger.warning(f"{nan_mask.sum()} channel pair(s) had zero variance across participants.")
        flat_p = np.where(nan_mask, 1.0, flat_p)

    n_tests = len(flat_p)
    sorted_p = np.sort(flat_p)
    bh_thresholds = np.array([(k + 1) / n_tests * alpha for k in range(n_tests)])

    # --- Plot 1: histogram with BH threshold + expected-null reference ---
    fig1, ax1 = plt.subplots(figsize=(7, 5))
    counts, bins, _ = ax1.hist(flat_p, bins=20, range=(0, 1), color="steelblue", edgecolor="white")
    expected_uniform = n_tests / 20  # expected count per bin if p-values were uniform (null)
    ax1.axhline(expected_uniform, color="gray", linestyle="--", linewidth=1.5,
                label=f"expected under null ({expected_uniform:.1f}/bin)")
    ax1.axvline(alpha, color="red", linestyle=":", linewidth=1.5, label=f"alpha={alpha}")
    ax1.set_xlabel("raw p-value")
    ax1.set_ylabel("count")
    ax1.set_title(f"p-value distribution across {n_tests} channel pairs (before FDR)")
    ax1.legend()
    plt.show(block=False)

    # --- Plot 2: text report card of the smallest p-values, by actual pair ---
    avg_r = np.tanh(np.mean(z_stack, axis=0))
    order = np.argsort(flat_p)[:n_top_pairs]

    def _pair_label(flat_idx):
        r, c = triu[0][flat_idx], triu[1][flat_idx]
        if channel_names is not None:
            return f"{channel_names[r]} <-> {channel_names[c]}"
        return f"ch{r} <-> ch{c}"

    lines = [
        f"Group connectivity summary",
        f"n_participants={n_participants}   n_comparisons={n_tests}   alpha={alpha}",
        f"smallest p-value: {sorted_p[0]:.4f}   BH threshold at rank 1: {bh_thresholds[0]:.6f}",
        "",
        f"Top {n_top_pairs} pairs by raw p-value (uncorrected):",
        "-" * 60,
    ]
    for rank, idx in enumerate(order, start=1):
        lines.append(
            f"{rank:>2}. {_pair_label(idx):<28} p={flat_p[idx]:.4f}  avg_r={avg_r[triu[0][idx], triu[1][idx]]:+.3f}"
        )

    fig2, ax2 = plt.subplots(figsize=(8, 6))
    ax2.axis("off")
    ax2.text(0.02, 0.98, "\n".join(lines), va="top", ha="left", family="monospace", fontsize=9)
    plt.show(block=False)

    reject, _ = multipletests(flat_p, method="fdr_bh", alpha=alpha)[:2]
    logger.info(f"[group-stats] FDR pass: {reject.sum()}/{len(reject)} pairs significant at alpha={alpha}")

    sig_avg_r = np.zeros_like(avg_r)

    for i, is_sig in enumerate(reject):
        row, col = triu[0][i], triu[1][i]
        r_val = avg_r[row, col]
        if is_sig and abs(r_val) >= min_effect_size:
            sig_avg_r[row, col] = sig_avg_r[col, row] = r_val

    n_after_effect_size = np.count_nonzero(sig_avg_r) // 2  # matrix is symmetric, count unique pairs
    logger.info(
        f"[group-stats] after min_effect_size={min_effect_size} filter: "
        f"{n_after_effect_size}/{reject.sum()} FDR-significant pairs survive "
        f"(max avg_r in final matrix={sig_avg_r.max():.3f} if any)"
    )

    return sig_avg_r

# ============================================================================
# Per-subject extractors (raw, UNTHRESHOLDED correlation - group stats
# handle significance, not these)
# ============================================================================

def _single_subject_beta_corr(
    raw_haemo: BaseRaw,
    event_name: str | None,
    drift_model: str,
    drift_order: int,
    hrf_model: str,
    apply_gsr: bool,
    min_trials: int = 4,
    resample_freq: float | None = 4.0,
) -> tuple[np.ndarray | None, list[str] | None]:
    raw_hbo = raw_haemo.copy().pick(picks="hbo")

    if event_name is not None:
        keep_mask = [desc == event_name for desc in raw_hbo.annotations.description]
        raw_hbo.set_annotations(raw_hbo.annotations[keep_mask])
        logger.info(f"[betas] filtered to event '{event_name}': {len(raw_hbo.annotations)} annotations remain")

    if resample_freq is not None and raw_hbo.info["sfreq"] > resample_freq:
        logger.info(f"[betas] resampling {raw_hbo.info['sfreq']}Hz -> {resample_freq}Hz")
        raw_hbo.resample(resample_freq, npad="auto")

    raw_hbo.annotations.description = np.array([
        f"{desc}__trial_{i:03d}" for i, desc in enumerate(raw_hbo.annotations.description)
    ])
    logger.info(f"[betas] annotations rewritten: {len(raw_hbo.annotations)} trial regressors")

    design_kwargs = dict(raw=raw_hbo, drift_model=drift_model, drift_order=drift_order)
    if hrf_model == "fir":
        design_kwargs.update(hrf_model="fir", fir_delays=np.arange(0, 12, 1))
    else:
        design_kwargs.update(hrf_model=hrf_model)

    logger.info(f"[betas] building design matrix, hrf_model={hrf_model}, n_samples={len(raw_hbo.times)}...")
    design_matrix = make_first_level_design_matrix(**design_kwargs)
    logger.info(f"[betas] design matrix built: shape={design_matrix.shape}")

    glm_results = run_glm(raw_hbo, design_matrix)
    betas = np.array(glm_results.theta())
    if betas.ndim == 3 and betas.shape[-1] == 1:
        betas = betas.squeeze(axis=-1)

    reg_names = list(design_matrix.columns)
    n_channels = betas.shape[0]
    if betas.shape[1] != len(reg_names):
        logger.error(f"[betas] shape mismatch: {betas.shape[1]} vs {len(reg_names)}, skipping.")
        return None, None

    trial_tags = sorted({
        (col.split("_delay")[0] if "_delay" in col else col)
        for col in reg_names
        if ("__trial_" in col) and (event_name is None or col.startswith(event_name + "__"))
    })

    if len(trial_tags) < min_trials:
        logger.warning(f"[betas] only {len(trial_tags)} trials, need {min_trials}, skipping.")
        return None, None

    beta_series = np.zeros((n_channels, len(trial_tags)))
    for t_idx, tag in enumerate(trial_tags):
        col_idx = [j for j, col in enumerate(reg_names)
                   if (col.split("_delay")[0] if "_delay" in col else col) == tag]
        beta_series[:, t_idx] = betas[:, col_idx].mean(axis=1)

    if apply_gsr:
        global_signal = np.mean(beta_series, axis=0)
        beta_series_clean = np.zeros_like(beta_series)
        for i in range(n_channels):
            slope, intercept = np.polyfit(global_signal, beta_series[i, :], 1)
            beta_series_clean[i, :] = beta_series[i, :] - (slope * global_signal + intercept)
    else:
        beta_series_clean = beta_series

    corr_matrix = np.corrcoef(beta_series_clean)
    np.fill_diagonal(corr_matrix, 0)
    return corr_matrix, raw_hbo.ch_names


def _single_subject_epoch_coherence(
    epochs: Epochs,
    event_name: str | None = None,
    fmin: float = 0.04,
    fmax: float = 0.2,
    method: str = "wpli2_debiased",
    allow_unreliable: bool = False,
) -> tuple[np.ndarray | None, list[str] | None]:
    epochs.load_data()
    if event_name is not None:
        try:
            epochs = epochs[event_name]
        except KeyError:
            logger.warning(f"[coherence] event '{event_name}' not found, skipping.")
            return None, None

    logger.info(f"[coherence] n_epochs after event filter: {len(epochs)}")

    hbo_epochs = epochs.copy().pick(picks="hbo")

    epoch_duration = hbo_epochs.times[-1] - hbo_epochs.times[0]
    _check_epoch_frequency_resolution(epoch_duration, fmin, allow_unreliable=allow_unreliable)

    con_coh = spectral_connectivity_epochs(
        hbo_epochs, method=method, mode="fourier", sfreq=hbo_epochs.info["sfreq"],
        fmin=fmin, fmax=fmax, faverage=True, verbose=False,
    )
    coh = np.squeeze(con_coh.get_data(output="dense"))
    if coh.ndim != 2:
        logger.warning("Unexpected coherence shape for participant, skipping.")
        return None, None
    coh = coh + coh.T - np.diag(np.diag(coh))
    np.fill_diagonal(coh, 0)
    return coh, hbo_epochs.ch_names


# ============================================================================
# Group-level entry points
# ============================================================================

def run_group_functional_connectivity_betas(
    haemo_dict: dict[str, BaseRaw],
    selected_paths: List[str],
    event_name: Optional[str],
    n_lines: int,
    vmin: float,
    *,
    drift_model: str = "cosine",
    drift_order: int = 1,
    hrf_model: str = "glover",
    apply_gsr: bool = True,
    alpha: float = 0.05,
    min_participants: int = 3,
    resample_freq: float = 4.0,
) -> None:
    
    subject_results = []
    for path in selected_paths:
        raw = haemo_dict.get(path)
        logger.info(f"[group-betas] {path}: raw object id={id(raw)}")
        if raw is None:
            logger.warning(f"[group-betas] {path}: no haemo data, skipping.")
            continue
        corr, names = _single_subject_beta_corr(raw, event_name, drift_model, drift_order, hrf_model, apply_gsr, resample_freq=resample_freq)
        if corr is not None:
            subject_results.append((corr, names))
            logger.info(f"[group-betas] {path}: added, running total={len(subject_results)}")

    if not subject_results:
        logger.error("[group-betas] no usable participant data.")
        return

    stacked, common_names, dropped = _align_participant_matrices(subject_results)
    logger.info(f"[group-betas] aligned: stacked shape={stacked.shape}")
    if dropped:
        logger.warning(f"[group-betas] {len(dropped)} participant(s) excluded.")

    sig_avg_r = _group_ttest_connectivity(stacked, alpha=alpha, min_effect_size=vmin, min_participants=min_participants, channel_names=common_names)
    logger.info(f"[group-betas] stats done. nonzero significant entries={np.count_nonzero(sig_avg_r)}")

    plot_connectivity_circle(
        sig_avg_r, common_names, n_lines=n_lines,
        title=f"Group Betas Connectivity ({hrf_model}, n={stacked.shape[0]}, FDR q<{alpha}): "
              f"{event_name if event_name else 'All Events'}",
        vmin=vmin, vmax=1.0, colormap="hot",
    )


def run_group_functional_connectivity_epochs(
    epochs_dict: dict[str, Epochs],
    selected_paths: list[str],
    event_name: str | None,
    n_lines: int,
    vmin: float,
    *,
    fmin: float = 0.04,
    fmax: float = 0.2,
    method: str = "wpli2_debiased",
    alpha: float = 0.05,
    min_participants: int = 3,
    allow_unreliable: bool = False,
) -> None:
    logger.info(
        f"[group-coherence] START: method={method}, n_lines={n_lines}, vmin={vmin}, "
        f"fmin={fmin}, fmax={fmax}, alpha={alpha}, min_participants={min_participants}, "
        f"event_name={event_name}, n_selected_paths={len(selected_paths)}"
    )
    subject_results = []
    for path in selected_paths:
        epochs = epochs_dict.get(path)
        if epochs is None:
            continue
        corr, names = _single_subject_epoch_coherence(
            epochs, event_name=event_name, fmin=fmin, fmax=fmax,
            method=method, allow_unreliable=allow_unreliable
        )
        if corr is not None:
            subject_results.append((corr, names))

    if not subject_results:
        logger.error("No participants produced usable epoch coherence data for group analysis.")
        return

    stacked, common_names, dropped = _align_participant_matrices(subject_results)
    if dropped:
        logger.warning(f"{len(dropped)} participant(s) excluded during channel alignment.")

    sig_avg_r = _group_ttest_connectivity(
        stacked, alpha=alpha, min_effect_size=vmin, min_participants=min_participants, channel_names=common_names
    )

    plot_connectivity_circle(
        sig_avg_r, common_names, n_lines=n_lines,
        title=f"Group Connectivity ({method}, {fmin}-{fmax} Hz, n={stacked.shape[0]}, FDR q<{alpha})",
        vmin=vmin, vmax=1.0, colormap="hot",
    )










def sparks_csv_export(
    haemo_obj: BaseRaw,
    save_path: str,
) -> 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)

 
 
def peak_power_fast(
    raw,
    time_window=10,
    threshold=0.1,
    l_freq=0.7,
    h_freq=1.5,
    l_trans_bandwidth=0.3,
    h_trans_bandwidth=0.3,
    verbose=False,
):
    """
    Numerically-equivalent drop-in for mne_nirs' peak_power().
 
    Same algorithm, same math, same return shape/order. The only change:
    the periodogram() call is batched across ALL (channel-pair, window)
    combinations in a single vectorized call instead of one call per
    combination. np.correlate() is still called per (pair, window) - it's
    cheap relative to periodogram's per-call overhead - but if profiling
    shows it's now the bottleneck, that's the next thing to vectorize
    (via FFT-based batch cross-correlation), flagged separately since
    it's a bigger, riskier change to get bit-exact.
 
    Validate before trusting in production:
        _, scores_orig, times_orig = peak_power(raw, ...)
        _, scores_fast, times_fast = peak_power_fast(raw, ...)
        assert np.allclose(scores_orig, scores_fast)
    """
    raw = raw.copy().load_data()
    _validate_type(raw, BaseRaw, "raw")
 
    picks = _validate_nirs_info(raw.info)
    sfreq = raw.info["sfreq"]
 
    filtered_data = filter_data(
        raw._data,
        sfreq,
        l_freq,
        h_freq,
        picks=picks,
        verbose=verbose,
        l_trans_bandwidth=l_trans_bandwidth,
        h_trans_bandwidth=h_trans_bandwidth,
    )
 
    window_samples = int(np.ceil(time_window * sfreq))
    n_windows = int(np.floor(len(raw) / window_samples))
    n_pairs = len(picks) // 2
 
    scores = np.zeros((len(picks), n_windows))
    times = []
 
    if n_windows == 0:
        scores = scores[np.argsort(picks)]
        return raw, scores, times
 
    # All windows except the last are guaranteed to be exactly window_samples
    # long (the last window can be shorter due to the min() clamp below) -
    # batch-process those; handle the last window with the original scalar path.
    full_windows = n_windows - 1
    corr_len = 2 * window_samples - 1
 
    if full_windows > 0:
        c1_stack = np.empty((n_pairs, full_windows, window_samples))
        c2_stack = np.empty((n_pairs, full_windows, window_samples))
 
        for window in range(full_windows):
            start = window * window_samples
            end = start + window_samples
            for pi, ii in enumerate(range(0, len(picks), 2)):
                c1_stack[pi, window] = filtered_data[picks[ii]][start:end]
                c2_stack[pi, window] = filtered_data[picks[ii + 1]][start:end]
 
        std1 = c1_stack.std(axis=-1, keepdims=True)
        std1[std1 == 0] = 1
        std2 = c2_stack.std(axis=-1, keepdims=True)
        std2[std2 == 0] = 1
        c1_stack = c1_stack / std1
        c2_stack = c2_stack / std2
 
        corr_stack = np.empty((n_pairs, full_windows, corr_len))
        for pi in range(n_pairs):
            for window in range(full_windows):
                corr_stack[pi, window] = (
                    np.correlate(c1_stack[pi, window], c2_stack[pi, window], "full")
                    / window_samples
                )
 
        # single vectorized call replaces n_pairs * full_windows separate calls
        _, pxx = periodogram(corr_stack, fs=sfreq, window="hamming", axis=-1)
        window_scores = pxx.max(axis=-1)  # shape (n_pairs, full_windows)
 
        scores[0::2, :full_windows] = window_scores
        scores[1::2, :full_windows] = window_scores
 
        for window in range(full_windows):
            start = window * window_samples
            end = start + window_samples
            times.append((raw.times[start], raw.times[min(end, len(raw) - 1)]))
            if threshold is not None:
                for pi in np.where(window_scores[:, window] < threshold)[0]:
                    ii = pi * 2
                    raw.annotations.append(
                        raw.times[start],
                        time_window,
                        "BAD_PeakPower",
                        ch_names=[raw.ch_names[ii : ii + 2]],
                    )
 
    # last (possibly truncated) window - original scalar path, unchanged
    window = n_windows - 1
    start_sample = window * window_samples
    end_sample = int(np.min([start_sample + window_samples, len(raw) - 1]))
    t_start, t_stop = raw.times[start_sample], raw.times[end_sample]
    times.append((t_start, t_stop))
 
    for ii in range(0, len(picks), 2):
        c1 = filtered_data[picks[ii]][start_sample:end_sample]
        c2 = filtered_data[picks[ii + 1]][start_sample:end_sample]
        c1 = c1 / (np.std(c1) or 1)
        c2 = c2 / (np.std(c2) or 1)
        c = np.correlate(c1, c2, "full") / window_samples
        f, pxx = periodogram(c, fs=sfreq, window="hamming")
        scores[ii, window] = max(pxx)
        scores[ii + 1, window] = max(pxx)
        if (threshold is not None) and (max(pxx) < threshold):
            raw.annotations.append(
                t_start, time_window, "BAD_PeakPower",
                ch_names=[raw.ch_names[ii : ii + 2]],
            )
 
    scores = scores[np.argsort(picks)]
    return raw, scores, times


def write_qc_excel_summary(
    qc_rows: list[dict],
    output_path: str,
    population_flags: dict[str, list[dict]] | None = None,
) -> None:
    """
    Writes one Excel workbook summarizing QC metrics across all participants
    in a batch run - participants as columns, metrics as rows, each row
    color-scaled green (best) to red (worst) with direction-aware coloring,
    plus live summary formulas (mean/median/min/max/worst participant) per row.

    Failed participants are shown in a separate, clearly labeled block so
    they don't distort the color scale of successful participants' numbers.
    """
    successes = [r for r in qc_rows if r.get("status") == "success"]
    failures = [r for r in qc_rows if r.get("status") == "FAILED"]

    wb = Workbook()
    ws = wb.active
    ws.title = "QC Summary"

    header_font = Font(name="Arial", bold=True, size=11)
    label_font = Font(name="Arial", size=10)
    body_font = Font(name="Arial", size=10)
    fail_fill = PatternFill(start_color="FFC7CE", end_color="FFC7CE", fill_type="solid")

    metric_keys = [k for k in QC_METRIC_LABELS if k in QC_METRIC_DIRECTIONS or k in QC_METRIC_DEVIATION_BASED]
    n_participants = len(successes)

    # --- Header row ---
    ws.cell(row=1, column=1, value="Metric").font = header_font
    for i, row in enumerate(successes):
        col = i + 2
        c = ws.cell(row=1, column=col, value=row.get("file_path", f"P{i+1}"))
        c.font = header_font
        c.alignment = Alignment(horizontal="center", wrap_text=True)

    summary_start_col = n_participants + 3  # one blank column, then summary block
    for j, label in enumerate(["Mean", "Median", "Min", "Max", "Worst Participant"]):
        c = ws.cell(row=1, column=summary_start_col + j, value=label)
        c.font = header_font
        c.alignment = Alignment(horizontal="center", wrap_text=True)

    # --- Metric rows ---
    for r_idx, key in enumerate(metric_keys):
        row_num = r_idx + 2
        ws.cell(row=row_num, column=1, value=QC_METRIC_LABELS[key]).font = label_font

        for i, row in enumerate(successes):
            col = i + 2
            val = row.get(key)
            c = ws.cell(row=row_num, column=col, value=val)
            c.font = body_font

        if n_participants == 0:
            continue

        first_col_letter = get_column_letter(2)
        last_col_letter = get_column_letter(n_participants + 1)
        data_range = f"{first_col_letter}{row_num}:{last_col_letter}{row_num}"

        mean_col, median_col, min_col, max_col, worst_col = (
            summary_start_col, summary_start_col + 1, summary_start_col + 2,
            summary_start_col + 3, summary_start_col + 4,
        )

        ws.cell(row=row_num, column=mean_col, value=f"=AVERAGE({data_range})").font = body_font
        ws.cell(row=row_num, column=median_col, value=f"=MEDIAN({data_range})").font = body_font
        ws.cell(row=row_num, column=min_col, value=f"=MIN({data_range})").font = body_font
        ws.cell(row=row_num, column=max_col, value=f"=MAX({data_range})").font = body_font

        header_range = f"{first_col_letter}1:{last_col_letter}1"
        if key in QC_METRIC_DEVIATION_BASED:
            row_values = [row.get(key) for row in successes if row.get(key) is not None]
            row_paths = [row.get("file_path") for row in successes if row.get(key) is not None]
            if row_values:
                med = float(np.median(row_values))
                deviations = [abs(v - med) for v in row_values]
                worst_idx = int(np.argmax(deviations))
                worst_value = row_paths[worst_idx]
            else:
                worst_value = ""
            ws.cell(row=row_num, column=worst_col, value=worst_value).font = body_font
        else:
            # existing live-formula path for direction-based metrics, unchanged
            if QC_METRIC_DIRECTIONS[key]:
                worst_formula = f"=INDEX({header_range},MATCH(MAX({data_range}),{data_range},0))"
            else:
                worst_formula = f"=INDEX({header_range},MATCH(MIN({data_range}),{data_range},0))"
            ws.cell(row=row_num, column=worst_col, value=worst_formula).font = body_font
        # --- Color scale, direction-aware ---
        if key in QC_METRIC_DEVIATION_BASED:
            # color by |value - row median| via a helper column pattern isn't
            # natively supported by ColorScaleRule (it colors raw cell values,
            # not a derived formula) - approximate by centering the 3-color
            # scale on the row's own values, which still highlights the
            # extremes/outliers visually even without true deviation coloring.
            rule = ColorScaleRule(
                start_type="min", start_color="63BE7B",
                mid_type="percentile", mid_value=50, mid_color="FFEB84",
                end_type="max", end_color="F8696B",
            )
        elif QC_METRIC_DIRECTIONS[key]:  # lower is better: green=min, red=max
            rule = ColorScaleRule(
                start_type="min", start_color="63BE7B",
                mid_type="percentile", mid_value=50, mid_color="FFEB84",
                end_type="max", end_color="F8696B",
            )
        else:  # higher is better: green=max, red=min
            rule = ColorScaleRule(
                start_type="min", start_color="F8696B",
                mid_type="percentile", mid_value=50, mid_color="FFEB84",
                end_type="max", end_color="63BE7B",
            )
        ws.conditional_formatting.add(data_range, rule)

    if population_flags:
        outlier_row = len(metric_keys) + 2
        ws.cell(row=outlier_row, column=1, value="Population Outlier Flags").font = label_font

        any_flags = bool(population_flags)
        for i, row in enumerate(successes):
            col = i + 2
            path = row.get("file_path", "")
            flags = population_flags.get(path, []) if population_flags else []
            if flags:
                text = "; ".join(f"{f['metric']} (z={f['z']})" for f in flags)
                c = ws.cell(row=outlier_row, column=col, value=text)
                c.fill = PatternFill(
                    start_color="FFEB84" if len(flags) < 3 else "F8696B",
                    end_color="FFEB84" if len(flags) < 3 else "F8696B",
                    fill_type="solid",
                )
            else:
                c = ws.cell(row=outlier_row, column=col, value="No outliers detected")
                c.fill = PatternFill(start_color="C6EFCE", end_color="C6EFCE", fill_type="solid")
            c.font = body_font
            c.alignment = Alignment(horizontal="center")

    # --- Failed participants block, separate and clearly marked ---
    if failures:
        fail_row_start = len(metric_keys) + 5
        ws.cell(row=fail_row_start, column=1, value="FAILED PARTICIPANTS").font = Font(name="Arial", bold=True, size=12, color="CC0000")
        for i, row in enumerate(failures):
            r = fail_row_start + 1 + i
            path_cell = ws.cell(row=r, column=1, value=row.get("file_path", "unknown"))
            path_cell.font = body_font
            path_cell.fill = fail_fill
            err_cell = ws.cell(row=r, column=2, value=row.get("error", "unknown error"))
            err_cell.font = body_font
            err_cell.fill = fail_fill
    else:
        fail_row_start = len(metric_keys) + 5
        c = ws.cell(row=fail_row_start, column=1, value="All participants processed successfully - no failures.")
        c.font = Font(name="Arial", bold=True, size=11, color="006100")
        c.fill = PatternFill(start_color="C6EFCE", end_color="C6EFCE", fill_type="solid")

    ws.column_dimensions['A'].width = 32
    for i in range(n_participants):
        ws.column_dimensions[get_column_letter(i + 2)].width = 14
    ws.freeze_panes = "B2"

    wb.save(output_path)


METRIC_REGISTRY = {
    'Peak_Amp': 'Peak_Amp',
    'TTP': 'Time_to_Peak',
    'AUC': 'AUC',
    'Rising_Slope': 'Rising_Slope',
    'Recovery_Slope': 'Recovery_Slope',
    'FWHM': 'FWHM',
    'Onset_Latency': 'Onset_Latency',
    'P2P_Amp': 'Peak_to_Peak_Amp',
    'Signal_Std': 'Signal_Std',
    'RMS': 'RMS'
}


def plot_and_enqueue_waveform_metrics(
    roi_curves, 
    fir_delays, 
    selected_metrics, 
    target_condition='reach', 
    png_queue=None
):
    """
    Generates and enqueues visual plots for every calculated waveform metric
    across all Regions of Interest (ROIs).
    """
    # Structure metric values per ROI
    metric_data = {METRIC_REGISTRY[m]: {} for m in selected_metrics}

    for (chromo, roi_name), roi_fir_curve in roi_curves.items():
        metrics = compute_waveform_metrics(roi_fir_curve, fir_delays=fir_delays, selected_metrics=selected_metrics)
        for m_key, val in zip(selected_metrics, metrics):
            label = METRIC_REGISTRY[m_key]
            metric_data[label][f"{roi_name} ({chromo.upper()})"] = val

    # Generate an image for each metric across ROIs
    for metric_label, roi_dict in metric_data.items():
        if not roi_dict:
            continue
            
        fig, ax = plt.subplots(figsize=(8, 4.5))
        rois = list(roi_dict.keys())
        values = list(roi_dict.values())

        colors = ['#2b5c8f' if v >= 0 else '#d9534f' for v in values]
        ax.bar(rois, values, color=colors, alpha=0.85, edgecolor='black')
        
        ax.axhline(0, color='black', linewidth=0.8, linestyle='--')
        ax.set_title(f"FIR Waveform Metric: {metric_label} [{target_condition}]", fontsize=12, fontweight='bold')
        ax.set_xlabel("Region of Interest (ROI)", fontsize=10)
        ax.set_ylabel(metric_label, fontsize=10)
        plt.xticks(rotation=35, ha='right')
        plt.grid(axis='y', linestyle=':', alpha=0.6)
        plt.tight_layout()

        # Enqueue figure or close
        if png_queue is not None:
            _enqueue(f"FIR Waveform Metric - {metric_label}", fig, png_queue)
        else:
            plt.close(fig)


def _compute_roi_fir_curves(
    raw=None, 
    target_condition='reach', 
    fir_delays=np.arange(0, 15),
    roi_map={}, 
    chromophores=('hbo', 'hbr', 'hbt'),
    glm_est=None,
    df_design_matrix=None
):
    """
    Shared FIR-GLM curve extraction. If `glm_est` and `df_design_matrix` are supplied,
    it reuses pre-calculated GLM results directly to avoid duplicate processing.
    """
    roi_curves = {}
    active_roi_map = normalize_roi_map(roi_map)

    # --- SHORT-CIRCUIT: Reuse pre-calculated GLM estimation if available ---
    if glm_est is not None and df_design_matrix is not None:
        if hasattr(glm_est, 'to_dataframe'):
            glm_df = glm_est.to_dataframe().reset_index()
        elif isinstance(glm_est, pd.DataFrame):
            glm_df = glm_est.copy()
        else:
            raise ValueError("Unsupported format for precalculated glm_est.")

        glm_df.columns = [str(col).lower() for col in glm_df.columns]
        cond_col = 'condition' if 'condition' in glm_df.columns else 'regressor'
        ch_col = 'ch_name' if 'ch_name' in glm_df.columns else ('source' if 'source' in glm_df.columns else 'channel')

        print("Available conditions in GLM:", glm_df[cond_col].unique())
        fir_df = glm_df[glm_df[cond_col].astype(str).str.lower().str.contains(target_condition.lower())].copy()
        print(f"Matched rows for '{target_condition}': {len(fir_df)}")

        active_roi_map = normalize_roi_map(roi_map)

        if fir_df.empty:
            logger.warning(f"Condition '{target_condition}' not found in precalculated GLM estimates.")
            return roi_curves

        # 1. Extract base measured chromophores (e.g., hbo, hbr) directly from GLM data
        base_chromos = [c.lower() for c in chromophores if c.lower() != 'hbt']
        
        for chromo in base_chromos:
            chromo_df = fir_df[fir_df[ch_col].str.lower().str.contains(chromo)] if ch_col in fir_df.columns else fir_df
            
            ch_curves = {}
            for ch_name, ch_group in chromo_df.groupby(ch_col):
                pair = ch_name.split(' ')[0]
                ch_curves[pair] = ch_group['theta'].values if 'theta' in ch_group.columns else ch_group['beta'].values

            for roi_name, channels in active_roi_map.items():
                matching_curves = [ch_curves[ch] for ch in channels if ch in ch_curves]
                if matching_curves:
                    roi_curves[(chromo, roi_name)] = np.mean(matching_curves, axis=0)

        # 2. Derive HbT (HbO + HbR) dynamically if requested
        if 'hbt' in [c.lower() for c in chromophores]:
            for roi_name in active_roi_map.keys():
                hbo_key = ('hbo', roi_name)
                hbr_key = ('hbr', roi_name)
                if hbo_key in roi_curves and hbr_key in roi_curves:
                    roi_curves[('hbt', roi_name)] = roi_curves[hbo_key] + roi_curves[hbr_key]

        return roi_curves


def extract_fir_features_real_data(
    raw=None, 
    target_condition=None,
    fir_delays=np.arange(0, 15), 
    selected_metrics=('Peak_Amp',), 
    roi_map={},
    chromophores=('hbo', 'hbr', 'hbt'),
    glm_est=None,
    df_design_matrix=None,
    png_queue=None
):
    """
    Extracts FIR scalar waveform metrics, plots and enqueues figures for each
    waveform metric, and returns all outputs collapsed into a single dictionary variable.
    """

    raw_cols = [
        col for col in df_design_matrix.columns 
        if not any(k in col.lower() for k in ['drift', 'constant', 'short', 'nuisance'])
    ]
    
    # Strip '_delay_0', '_delay_1', etc. to get base condition names
    target_conditions = list(dict.fromkeys(
        col.split('_delay_')[0] if '_delay_' in col else col 
        for col in raw_cols
    ))

    collapsed_features = []
    feature_names = []
    feature_channels = []

    # 2. Iterate over EVERY condition
    for cond in target_conditions:
        try:
            roi_curves = _compute_roi_fir_curves(
                raw=raw, 
                target_condition=cond, 
                fir_delays=fir_delays,
                roi_map=roi_map, 
                chromophores=chromophores,
                glm_est=glm_est,
                df_design_matrix=df_design_matrix
            )

            if not roi_curves:
                print("999999999999999999999999999999999999999999999999999999999999999999999999999999999999999999999999")
                continue

            metric_labels = [METRIC_REGISTRY[m] for m in selected_metrics]

            # Extract metrics per ROI for this condition
            for (chromo, roi_name), roi_fir_curve in roi_curves.items():
                metrics = compute_waveform_metrics(roi_fir_curve, fir_delays=fir_delays, selected_metrics=selected_metrics)
                collapsed_features.extend(metrics)
                
                # Prefix feature names with condition AND chromophore
                feature_names.extend([f"{cond}_{chromo.upper()}_{roi_name}_{m}" for m in metric_labels])
                feature_channels.extend([f"{roi_name} ({chromo.upper()})"] * len(metric_labels))

            # Enqueue plots for this specific condition
            plot_and_enqueue_waveform_metrics(
                roi_curves=roi_curves,
                fir_delays=fir_delays,
                selected_metrics=selected_metrics,
                target_condition=cond,
                png_queue=png_queue
            )
        except Exception as e:
            logger.warning(f"Failed extracting FIR metrics for condition '{cond}': {e}")

    return {
        'features': np.array(collapsed_features),
        'feature_names': feature_names,
        'feature_channels': feature_channels
    }

def compute_waveform_metrics(fir_curve, fir_delays, selected_metrics=('Peak_Amp',)):
    peak_idx = np.argmax(fir_curve)
    peak_amp = fir_curve[peak_idx]
    ttp = fir_delays[peak_idx]

    calculated_metrics = {}

    if 'Peak_Amp' in selected_metrics:
        calculated_metrics['Peak_Amp'] = peak_amp

    if 'TTP' in selected_metrics:
        calculated_metrics['TTP'] = ttp

    if 'AUC' in selected_metrics:
        trapz_fn = getattr(np, 'trapezoid', getattr(np, 'trapz', None))
        calculated_metrics['AUC'] = trapz_fn(fir_curve, fir_delays)

    if 'Rising_Slope' in selected_metrics:
        calculated_metrics['Rising_Slope'] = (peak_amp - fir_curve[0]) / (ttp - fir_delays[0]) if ttp > fir_delays[0] else 0.0

    if 'Recovery_Slope' in selected_metrics:
        calculated_metrics['Recovery_Slope'] = (fir_curve[-1] - peak_amp) / (fir_delays[-1] - ttp) if fir_delays[-1] > ttp else 0.0

    if 'FWHM' in selected_metrics:
        half_max = peak_amp / 2.0
        above_half = np.where(fir_curve >= half_max)[0]
        calculated_metrics['FWHM'] = fir_delays[above_half[-1]] - fir_delays[above_half[0]] if len(above_half) > 1 else 0.0

    if 'Onset_Latency' in selected_metrics:
        onset_thresh = 0.2 * peak_amp
        above_onset = np.where(fir_curve >= onset_thresh)[0]
        calculated_metrics['Onset_Latency'] = fir_delays[above_onset[0]] if len(above_onset) > 0 else 0.0

    if 'P2P_Amp' in selected_metrics:
        calculated_metrics['P2P_Amp'] = peak_amp - np.min(fir_curve)

    if 'Signal_Std' in selected_metrics:
        calculated_metrics['Signal_Std'] = np.std(fir_curve)

    if 'RMS' in selected_metrics:
        calculated_metrics['RMS'] = np.sqrt(np.mean(fir_curve**2))

    return [calculated_metrics[m] for m in selected_metrics]


def normalize_roi_map(roi_map):
    """
    Normalizes ROI mapping inputs into a flat {roi_name: [channel_list]} dict.
    
    Accepts:
    1. String or Path pointing to a JSON file (JSON_LOCATION).
    2. Dict with 'regions_of_interest' list (loaded JSON).
    3. Standard flat dict {roi_name: [channels]}.
    """
    # 1. If roi_map is a file path string or Path, load JSON from disk
    if isinstance(roi_map, (str, Path)):
        json_path = Path(roi_map)
        if json_path.is_file():
            try:
                with open(json_path, 'r', encoding='utf-8') as f:
                    roi_map = json.load(f)
            except Exception as e:
                logger.error(f"Failed to load ROI JSON file from {json_path}: {e}")
                return {}
        else:
            logger.error(f"ROI JSON file path does not exist: {json_path}")
            return {}

    # 2. Handle nested JSON structure with 'regions_of_interest'
    if isinstance(roi_map, dict) and 'regions_of_interest' in roi_map:
        return {
            roi['name']: roi['channels'] 
            for roi in roi_map['regions_of_interest']
        }

    # 3. Fallback for flat dictionary {roi_name: [channels]}
    if isinstance(roi_map, dict):
        return roi_map

    return {}


if __name__ == "__main__":
    print("This file has no functionality when not used in tandem with the FLARES application.")
        
    # audit_log = logging.getLogger("memory_audit")
    # audit_log.setLevel(logging.INFO)
    # audit_log.propagate = False # This prevents it from talking to other loggers

    # # 2. Add a file handler specifically for this audit logger
    # if not audit_log.handlers:
    #     fh = logging.FileHandler('flares_memory_audit.log')
    #     fh.setFormatter(logging.Formatter('%(asctime)s | PID: %(process)d | %(message)s'))
    #     audit_log.addHandler(fh)

    # def get_mem_mb():
    #     return psutil.Process(os.getpid()).memory_info().rss / 1024 / 1024