MMN_MSfigs.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
MMN_MSfigs.py - Main analysis script for the Optimum-1 MMN study

Generates all figures and statistical analyses for the manuscript:
  - Figure 2: Dataset overview (unit counts, firing rates, participant distribution)
  - Figure 3: General deviance detection (combined deviant, epoch-based LME)
  - Figure 4: Feature-category deviance detection (epoch-based LME)
  - Figure 5: Time-resolved LME analysis (ΔAIC + t-values, Amy & HPC)
  - Figure 6: TTG & Ins individual unit responses
  - Supplementary figures: Individual deviants, carryover, SDF panels, heatmaps

Usage:
  1. Set `dataPath` to point to the Optimum1_ProcessedData/ directory
  2. Run sections sequentially (designed for interactive use in IDE with cell execution)
  3. Results are saved to Figures/ directory

Dependencies: MMN_funcs.py, MMN_stathelp.py
"""

# =============================================================================
# IMPORTS
# =============================================================================
from MMN_funcs import (
    buildSdfDF_mmn, create_mmn_colormap_for_plots, fit_glm_model,
    compute_model_metrics, remove_sdf_outlier_trials, prepare_mmn_data_for_glm,get_smoothed_data2
)


import numpy as np
import matplotlib.pyplot as plt
from scipy.io import loadmat
from matplotlib.patches import Rectangle
import pandas as pd
import seaborn as sns
import os
from statsmodels.stats.multitest import multipletests
import copy
from pathlib import Path
import time
from pymatreader import read_mat
from joblib import Parallel, delayed
import statsmodels.api as sm
import statsmodels.formula.api as smf
from scipy.stats import binomtest, combine_pvalues, mannwhitneyu, fisher_exact, kruskal
import pickle
import gc


def add_contrast_columns(data, contrast_type):
    """
    Add necessary contrast columns to the data based on contrast type.
    
    Parameters:
    -----------
    data : pandas.DataFrame
        DataFrame to modify in-place
    contrast_type : str
        Type of contrast to add columns for
    """
    # Define categories and deviants
    deviant_categories = {
        'Location': ['LocLeft', 'LocRight'],
        'Intensity': ['intensityDown', 'intensityUp'],
        'Frequency': ['Freq_450', 'Freq_550'],
        'Timing': ['Dur25', 'Gap']
    }
    all_deviants = ['Std', 'LocLeft', 'LocRight', 'intensityDown', 'intensityUp', 
                   'Freq_450', 'Freq_550', 'Dur25', 'Gap']
    
    # Add standard/deviant flag
    if 'isDeviant' not in data.columns:
        data['isDeviant'] = data['eventID'] != 'Std'
    
    # Add category flags
    for category, deviants in deviant_categories.items():
        column_name = f'is{category}'
        if column_name not in data.columns:
            data[column_name] = data['eventID'].isin(deviants)
    
    # Add individual deviant flags
    for deviant in all_deviants:
        column_name = f'is{deviant}'
        if column_name not in data.columns:
            data[column_name] = data['eventID'] == deviant
    
    # For 'categories' contrast, create DeviantCategory
    if contrast_type == 'categories':
        if 'DeviantCategory' not in data.columns:
            data['DeviantCategory'] = 'Standard'
            for category in deviant_categories:
                mask = data['eventID'].isin(deviant_categories[category])
                data.loc[mask, 'DeviantCategory'] = category
    
    # For 'individual' contrast, create DeviantType
    if contrast_type == 'individual':
        if 'DeviantType' not in data.columns:
            data['DeviantType'] = data['eventID']
            
def raincloud_plot(x, y, data, ax=None, palette=None, orient='v', 
                  violin_side='right', jitter=0.1, point_size=3,
                  show_box=True, box_width=0.15, point_style='contrast',
                  point_shift=0.15, point_hue=None, point_palette=None, 
                  point_alpha=0.7, **kwargs):
    """
    Create a complete raincloud plot combining half-violin, swarm, and boxplot
    
    Parameters:
    -----------
    x, y : str
        Column names for categorical and continuous variables
    data : DataFrame
        Input data
    ax : matplotlib axis
        Axis to plot on
    palette : str or list
        Color palette for different categories (violin colors)
    orient : str
        'v' for vertical, 'h' for horizontal
    violin_side : str
        Which side of violin to show ('left', 'right', 'top', 'bottom')
    jitter : float
        Amount of jitter for strip plot
    point_size : float
        Size of individual data points
    show_box : bool
        Whether to show boxplot
    box_width : float
        Width of boxplot
    point_style : str
        How to style points: 'contrast' (darker), 'edge' (black edges), 
        'shift' (moved away from violin), 'white' (white with edges), 'hue' (colored by point_hue)
    point_shift : float
        How much to shift points away from violin (when point_style='shift')
    point_hue : str or None
        Column name to use for coloring individual points (categorical or boolean)
    point_palette : str, list, or dict
        Color palette for point_hue. Can be seaborn palette name, list of colors,
        or dict mapping point_hue values to colors
    point_alpha : float
        Alpha (transparency) for data points
    **kwargs : additional arguments
        Passed to violin plot
    """
    if ax is None:
        fig, ax = plt.subplots(figsize=(10, 6))
    
    # Create violin plot with seaborn first
    violin_ax = sns.violinplot(x=x, y=y, data=data, ax=ax, 
                              inner=None, hue=x, palette=palette, legend=False, **kwargs)
    
    # Convert to half violins by clipping
    for item in violin_ax.collections:
        # Get the bounding box of each violin
        bbox = item.get_paths()[0].get_extents()
        x0, y0, width, height = bbox.x0, bbox.y0, bbox.width, bbox.height
        
        if orient == 'v':  # Vertical orientation
            if violin_side == 'right':
                # Keep right half
                clip_box = Rectangle((x0 + width/2, y0), width/2, height, 
                                   transform=ax.transData)
            elif violin_side == 'left':
                # Keep left half  
                clip_box = Rectangle((x0, y0), width/2, height, 
                                   transform=ax.transData)
        else:  # Horizontal orientation
            if violin_side == 'top':
                # Keep top half
                clip_box = Rectangle((x0, y0 + height/2), width, height/2, 
                                   transform=ax.transData)
            elif violin_side == 'bottom':
                # Keep bottom half
                clip_box = Rectangle((x0, y0), width, height/2, 
                                   transform=ax.transData)
        
        item.set_clip_path(clip_box)
    
    # Prepare point styling based on chosen style
    # Determine if we're using hue coloring (can be combined with any point_style)
    use_hue_coloring = point_hue is not None
    

    if use_hue_coloring:
        # Use point_hue for coloring - will be handled in plotting section
        point_edgecolor = 'black' if point_hue in data.columns and len(data[point_hue].unique()) > 2 else None
        
        # Set up point palette
        if point_palette is None:
            # Default palettes based on data type
            unique_vals = data[point_hue].unique()
            if len(unique_vals) == 2:
                # Binary/boolean - use contrasting colors
                if set(unique_vals).issubset({True, False, 0, 1, 'True', 'False'}):
                    final_point_palette = ['lightcoral', 'lightblue']
                else:
                    final_point_palette = ['orange', 'purple']
            else:
                # Multiple categories - use seaborn palette
                final_point_palette = sns.color_palette("Set1", len(unique_vals))
        elif isinstance(point_palette, str):
            # Seaborn palette name
            final_point_palette = sns.color_palette(point_palette, len(data[point_hue].unique()))
        elif isinstance(point_palette, dict):
            # Dictionary mapping
            final_point_palette = point_palette
        else:
            # List of colors
            final_point_palette = point_palette
    
    # Apply point_style-specific settings (these can work with or without hue coloring)
    if point_style == 'contrast' and not use_hue_coloring:
        # Make points darker than violins
        if isinstance(palette, str):
            temp_palette = sns.color_palette(palette, n_colors=len(data[x].unique()))
            final_point_palette = [tuple(np.array(color) * 0.7) for color in temp_palette]
        else:
            final_point_palette = palette
        point_edgecolor = None
        
    elif point_style == 'edge' and not use_hue_coloring:
        # Same color as violin but with black edges
        final_point_palette = palette
        point_edgecolor = 'black'
        
    elif point_style == 'white' and not use_hue_coloring:
        # White points with black edges
        final_point_palette = 'white'
        point_edgecolor = 'black'
        
    elif point_style == 'shift' and not use_hue_coloring:
        # Darker points, will be shifted
        final_point_palette = 'black'
        point_edgecolor = None
    
    # If no specific styling and no hue, use default
    if not use_hue_coloring and point_style not in ['contrast', 'edge', 'white', 'shift']:
        final_point_palette = 'black'
        point_edgecolor = None

    # Add strip plot (jittered points) with custom positioning and coloring
    if point_style == 'shift' or use_hue_coloring:
        # Manual positioning for shifted points or hue coloring (or both)
        categories = data[x].unique()
        
        for i, cat in enumerate(categories):
            cat_data = data[data[x] == cat]
            y_values = cat_data[y].values
            
            # Calculate shift direction based on violin side (only if point_style is 'shift')
            if orient == 'v':
                if violin_side == 'right' and point_style == 'shift':
                    x_pos = i - point_shift  # Shift left
                elif violin_side == 'left' and point_style == 'shift':
                    x_pos = i + point_shift  # Shift right
                else:
                    x_pos = i  # No shift
                
                # Add jitter
                x_jittered = x_pos + np.random.uniform(-jitter, jitter, len(y_values))
                
                if use_hue_coloring:
                    # Color points by hue variable
                    hue_values = cat_data[point_hue].values
                    unique_hue_vals = sorted(data[point_hue].unique())
                    
                    for hue_val in unique_hue_vals:
                        mask = hue_values == hue_val
                        if np.any(mask):
                            if isinstance(final_point_palette, dict):
                                color = final_point_palette[hue_val]
                            else:
                                hue_idx = unique_hue_vals.index(hue_val)
                                color = final_point_palette[hue_idx]
                            
                            ax.scatter(x_jittered[mask], y_values[mask], s=point_size**2, 
                                      color=color, alpha=point_alpha, label=f'{hue_val}' if i == 0 else "",
                                      edgecolors=point_edgecolor, linewidth=0.5)
                else:
                    # Single color (for shift style without hue)
                    if isinstance(final_point_palette, list) and len(final_point_palette) > i:
                        color = final_point_palette[i]
                    else:
                        color = final_point_palette
                        
                    ax.scatter(x_jittered, y_values, s=point_size**2, 
                              color=color, alpha=point_alpha,
                              edgecolors=point_edgecolor, linewidth=0.5)
                    
            else:  # Horizontal orientation
                if violin_side == 'top' and point_style == 'shift':
                    y_pos = i - point_shift  # Shift down
                elif violin_side == 'bottom' and point_style == 'shift':
                    y_pos = i + point_shift  # Shift up
                else:
                    y_pos = i  # No shift
                
                # Add jitter
                y_jittered = y_pos + np.random.uniform(-jitter, jitter, len(y_values))
                
                if use_hue_coloring:
                    # Color points by hue variable
                    hue_values = cat_data[point_hue].values
                    unique_hue_vals = sorted(data[point_hue].unique())
                    
                    for hue_val in unique_hue_vals:
                        mask = hue_values == hue_val
                        if np.any(mask):
                            if isinstance(final_point_palette, dict):
                                color = final_point_palette[hue_val]
                            else:
                                hue_idx = unique_hue_vals.index(hue_val)
                                color = final_point_palette[hue_idx]
                            
                            ax.scatter(y_values[mask], y_jittered[mask], s=point_size**2,
                                      color=color, alpha=point_alpha, label=f'{hue_val}' if i == 0 else "",
                                      edgecolors=point_edgecolor, linewidth=0.5)
                else:
                    # Single color (for shift style without hue)
                    if isinstance(final_point_palette, list) and len(final_point_palette) > i:
                        color = final_point_palette[i]
                    else:
                        color = final_point_palette
                        
                    ax.scatter(y_values, y_jittered, s=point_size**2,
                              color=color, alpha=point_alpha,
                              edgecolors=point_edgecolor, linewidth=0.5)
    else:
        # Use seaborn stripplot for other styles
        strip_kwargs = {'jitter': jitter, 'size': point_size, 'alpha': point_alpha}
        if point_edgecolor:
            strip_kwargs['edgecolor'] = point_edgecolor
            strip_kwargs['linewidth'] = 0.5
        
        if orient == 'v':
            sns.stripplot(x=x, y=y, data=data, ax=ax, 
                         palette=final_point_palette, **strip_kwargs)
        else:
            sns.stripplot(x=y, y=x, data=data, ax=ax, 
                         palette=final_point_palette, **strip_kwargs)
    
    # Add legend for point hue if used
    if use_hue_coloring and point_hue is not None:
        # Only add legend if we created labels (i.e., for the first category)
        handles, labels = ax.get_legend_handles_labels()
        if handles:
            ax.legend(handles, labels, title=point_hue, loc='best')
    
    # Add boxplot if requested
    if show_box:
        if orient == 'v':
            sns.boxplot(x=x, y=y, data=data, ax=ax, 
                       width=box_width, hue=x, palette=palette, legend=False,
                       boxprops={'facecolor': 'none', 'edgecolor': 'black'},
                       whiskerprops={'color': 'black'},
                       capprops={'color': 'black'},
                       medianprops={'color': 'red', 'linewidth': 2})
        else:
            sns.boxplot(x=y, y=x, data=data, ax=ax, 
                       width=box_width, hue=x, palette=palette, legend=False,
                       boxprops={'facecolor': 'none', 'edgecolor': 'black'},
                       whiskerprops={'color': 'black'},
                       capprops={'color': 'black'},
                       medianprops={'color': 'red', 'linewidth': 2})
    
    return ax


def unitOverview(df_work, save_fig=True, output_dir=None, figsize=(160/25.4, 100/25.4),didMap = None):
    """
    Generate comprehensive unit overview including counts, firing rate, and participant/session analysis.
    
    Parameters:
    -----------
    dfIn : pandas.DataFrame
        One of three DataFrame types:
        - DF_Units: Unit-level data (one row per unit) with meanFR column
        - DF_FRall: Trial-level FR data with meanFR column  
        - DF_sdfs: Trial-level SDF data with SpikeDensity column
        Should have: UniqueLabel, BrainRegion (or Bregion), SU, ID
    save_fig : bool
        Whether to save the figure
    output_dir : str or Path
        Directory to save figure
    figsize : tuple
        Figure size
        
    Returns:
    --------
    dict
        Dictionary containing summary tables and figure
    """
    
    # Initialize variables early to avoid undefined variable errors
    is_unit_level = False
    data_type = "Unknown"
    
    # Detect DataFrame type for appropriate processing
    if 'meanFR' in df_work.columns and len(df_work) == df_work['UniqueLabel'].nunique():
        is_unit_level = True
        data_type = "Unit-level (DF_Units)"
    elif 'meanFR' in df_work.columns:
        is_unit_level = False
        data_type = "Trial-level FR (DF_FRall)"
    elif 'SpikeDensity' in df_work.columns:
        is_unit_level = False
        data_type = "Trial-level SDF (DF_sdfs)"
    else:
        is_unit_level = False
        data_type = "Unknown - will attempt to detect firing rate column"
    
    print(f"Detected DataFrame type: {data_type}")
    
    # Reset index if ID is in index
    if 'ID' not in df_work.columns and 'ID' in df_work.index.names:
        df_work = df_work.reset_index()

    if didMap:
        df_work['orID'] = df_work['ID']
        df_work['ID'] = df_work['orID'].map(didMap)   
    
    # Handle brain region column naming variations
    if 'BrainRegion' not in df_work.columns:
        if 'Bregion' in df_work.columns:
            df_work['BrainRegion'] = df_work['Bregion']
        elif 'BrainRegion_Mapped' in df_work.columns:
            df_work['BrainRegion'] = df_work['BrainRegion_Mapped']
        else:
            raise ValueError("Cannot find brain region column (BrainRegion, Bregion, or BrainRegion_Mapped)")
    
    # Get unique units - different approach based on data type
    if is_unit_level:
        # For unit-level data, each row is already a unique unit
        unique_units = df_work.copy()
        print(f"Unit-level data: {len(unique_units)} units")
    else:
        # For trial-level data, get unique combinations
        unique_units = df_work.drop_duplicates(subset=['UniqueLabel'])
        print(f"Trial-level data: {len(unique_units)} unique units from {len(df_work)} trials")
    
    # Extract session from last character of UniqueLabel
    unique_units['SessionStr'] = unique_units['UniqueLabel'].str[-1]
    
    # Count units by brain region and unit type
    unit_counts = unique_units.groupby(['BrainRegion', 'SU']).size().unstack(fill_value=0)
    
    # Rename columns for clarity (False = MU, True = SU)
    unit_counts.columns = ['MU', 'SU']
    
    # Add total column
    unit_counts['Total'] = unit_counts['MU'] + unit_counts['SU']
    # Sort by brain region alphabetically for consistent ordering across plots
    unit_counts = unit_counts.sort_index()
    brain_region_order = unit_counts.index.tolist()
    
    # Print the summary
    print("Summary of Units by Brain Region")
    print("===============================")
    print(unit_counts)
    print(f"\nOverall Total: {unit_counts['Total'].sum()} units")
    print(f"Single Units: {unit_counts['SU'].sum()} ({unit_counts['SU'].sum()/unit_counts['Total'].sum()*100:.1f}%)")
    print(f"Multi Units: {unit_counts['MU'].sum()} ({unit_counts['MU'].sum()/unit_counts['Total'].sum()*100:.1f}%)")
    
    # Participant distribution analysis
    print("\nParticipant Distribution Analysis")
    print("================================")
    participant_summary = unique_units.groupby('ID').agg({
        'UniqueLabel': 'count',
        'SU': ['sum', lambda x: (x == False).sum()]  # Count SU and MU
    }).round(2)
    participant_summary.columns = ['Total_Units', 'Single_Units', 'Multi_Units']
    participant_summary['SU_Percentage'] = (participant_summary['Single_Units'] / 
                                           participant_summary['Total_Units'] * 100).round(1)
    print(participant_summary)
    
    # Session distribution analysis
    print("\nSession Distribution Analysis")
    print("============================")
    session_summary = unique_units.groupby('SessionStr').agg({
        'UniqueLabel': 'count',
        'SU': ['sum', lambda x: (x == False).sum()],
        'ID': 'nunique'  # Number of unique participants per session
    }).round(2)
    session_summary.columns = ['Total_Units', 'Single_Units', 'Multi_Units', 'Participants']
    session_summary['SU_Percentage'] = (session_summary['Single_Units'] / 
                                       session_summary['Total_Units'] * 100).round(1)
    print(session_summary)
    
    # Multi-session participant analysis
    participant_sessions = unique_units.groupby('ID')['SessionStr'].apply(list).apply(lambda x: sorted(list(set(x))))
    multi_session_participants = participant_sessions[participant_sessions.apply(len) > 1]
    
    print(f"\nMulti-session participants:")
    for pid, sessions in multi_session_participants.items():
        print(f"  {pid}: Sessions {', '.join(sessions)}")
    
    print(f"\nSession participation summary:")
    print(f"  Single session only: {len(participant_sessions[participant_sessions.apply(len) == 1])} participants")
    print(f"  Multiple sessions: {len(multi_session_participants)} participants")
    
    # For a more detailed breakdown including each participant's contribution
    print("\nDetailed Unit Count by Region and Participant")
    print("=============================================")
    participant_counts = unique_units.groupby(['BrainRegion', 'ID', 'SU']).size().unstack(fill_value=0)
    print(participant_counts)
    if participant_counts.shape[1] == 2:
        participant_counts.columns = ['MU', 'SU']
    elif participant_counts.shape[1] == 1:
        # Handle case where all units are same type
        if False in participant_counts.columns:
            participant_counts.columns = ['MU']
            participant_counts['SU'] = 0
        else:
            participant_counts.columns = ['SU']
            participant_counts['MU'] = 0
    participant_counts['Total'] = participant_counts['MU'] + participant_counts['SU']
    print(participant_counts)
    
    # Calculate mean firing rates for each unit
    # Handle different DataFrame types and firing rate columns
    if is_unit_level and 'meanFR' in df_work.columns:
        # For DF_Units - already has meanFR, just use it directly
        mean_frs = unique_units[['UniqueLabel', 'meanFR']].copy()
        mean_frs.columns = ['UniqueLabel', 'MeanFR']
        print("Using existing meanFR from unit-level data")
    elif 'SpikeDensity' in df_work.columns:
        # For DF_sdfs - calculate from trial data
        mean_frs = df_work.groupby('UniqueLabel')['SpikeDensity'].mean().reset_index()
        mean_frs.columns = ['UniqueLabel', 'MeanFR']
        print("Calculating mean firing rate from SpikeDensity (trial-level SDF data)")
    elif 'meanFR' in df_work.columns:
        # For DF_FRall - use existing meanFR column from trial data
        mean_frs = df_work.groupby('UniqueLabel')['meanFR'].mean().reset_index()
        mean_frs.columns = ['UniqueLabel', 'MeanFR']
        print("Using meanFR from trial-level FR data")
    else:
        # Fallback - look for any FR-related column
        fr_candidates = [col for col in df_work.columns if 'fr' in col.lower() or 'rate' in col.lower()]
        if fr_candidates:
            print(f"Warning: Using column '{fr_candidates[0]}' as firing rate")
            if is_unit_level:
                mean_frs = unique_units[['UniqueLabel', fr_candidates[0]]].copy()
            else:
                mean_frs = df_work.groupby('UniqueLabel')[fr_candidates[0]].mean().reset_index()
            mean_frs.columns = ['UniqueLabel', 'MeanFR']
        else:
            raise ValueError("Cannot find firing rate column (meanFR or SpikeDensity)")
    
    # Merge with unit info
    unit_fr_data = unique_units.merge(mean_frs, on='UniqueLabel')
    
    # Print SU/MU counts per region to console
    print("\nUnit counts by region (SU / MU):")
    for region, row in unit_counts.iterrows():
        print(f"  {region}: {row['SU']} single-units, {row['MU']} multi-units (total {row['Total']})")

    # Create the plots - 2 panels
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=figsize)

    # Font sizes (publication standard)
    _fs_label = 8   # axis labels
    _fs_tick  = 7   # tick labels
    _fs_title = 8   # subplot titles

    # Plot 1: Mean firing rates by brain region and unit type
    sns.swarmplot(x='BrainRegion', y='MeanFR', data=unit_fr_data, ax=ax1,
                  hue='SU', dodge=True, order=brain_region_order,size = 3)

    ax1.axhline(y=1, color='red', linestyle='--', alpha=0.7, linewidth=1, label='1 Hz threshold')
    ax1.set_ylabel('Mean Firing Rate (Hz)', fontsize=_fs_label)
    ax1.set_xlabel('Brain Region', fontsize=_fs_label)
    ax1.set_title('Firing Rates by Region and Unit Type', fontsize=_fs_title)
    ax1.tick_params(labelsize=_fs_tick)
    ax1.grid(True, alpha=0.3)

    # Fix legend labels: True→single-units, False→multi-units
    _handles, _labels = ax1.get_legend_handles_labels()
    _label_map = {'True': 'single-units', 'False': 'multi-units',
                  True: 'single-units', False: 'multi-units'}
    _labels = [_label_map.get(l, l) for l in _labels]
    ax1.legend(_handles, _labels, fontsize=_fs_tick, title_fontsize=_fs_tick)


    # Plot 2: Participant-Session contribution analysis
    # Create a heatmap showing units per participant-session per region
    # Each row represents a unique ID-Session combination
    
    # Create Session-ID combination identifier
    unique_units['Session_ID'] = 'S' + unique_units['SessionStr'] + '_' + unique_units['ID']
    
    # Create matrix with Session-ID combinations as rows, brain regions as columns
    session_participant_matrix = unique_units.groupby(['Session_ID', 'BrainRegion']).size().unstack(fill_value=0)
    
    # Sort rows to group by participant (sessions together)
    # Extract ID for sorting
    session_participant_matrix['sort_id'] = session_participant_matrix.index.str[3:]  # Remove 'S1_', 'S2_', etc.
    session_participant_matrix['sort_session'] = session_participant_matrix.index.str[1:2]  # Extract session number
    session_participant_matrix = session_participant_matrix.sort_values(['sort_id', 'sort_session'])
    
    # Remove sorting columns for the heatmap
    matrix_for_plot = session_participant_matrix.drop(['sort_id', 'sort_session'], axis=1)
    
    # Create row colors based on session
    def get_session_color(session_id):
        session = session_id[1]  # Extract session number (S1_, S2_, S3_)
        session_colors = {'1': 'lightblue', '2': 'lightgreen', '3': 'lightcoral'}
        return session_colors.get(session, 'gray')
    
    row_colors = [get_session_color(idx) for idx in matrix_for_plot.index]
    
    # Create the heatmap
    im = ax2.imshow(matrix_for_plot.values, aspect='auto', cmap='Blues')

    # Add session color bars on the left
    for i, color in enumerate(row_colors):
        ax2.add_patch(plt.Rectangle((-0.7, i-0.4), 0.3, 0.8, facecolor=color, edgecolor='black', linewidth=0.5))

    # Customize the heatmap
    ax2.set_xticks(range(len(matrix_for_plot.columns)))
    ax2.set_xticklabels(matrix_for_plot.columns, fontsize=_fs_tick)
    ax2.set_yticks(range(len(matrix_for_plot.index)))

    # Create cleaner y-tick labels (remove 'S1_', 'S2_', etc. prefix for display)
    clean_labels = [idx[3:] + f" (S{idx[1]})" for idx in matrix_for_plot.index]
    ax2.set_yticklabels(clean_labels, fontsize=_fs_tick)

    ax2.set_xlabel('Brain Region', fontsize=_fs_label)
    ax2.set_ylabel('Participant (Session)', fontsize=_fs_label)
    ax2.set_title('Units per Participant-Session per Region', fontsize=_fs_title)
    ax2.tick_params(labelsize=_fs_tick)

    # Add text annotations for non-zero values
    for i in range(len(matrix_for_plot.index)):
        for j in range(len(matrix_for_plot.columns)):
            value = matrix_for_plot.iloc[i, j]
            if value > 0:
                ax2.text(j, i, str(value), ha='center', va='center',
                        color='white' if value > matrix_for_plot.values.max()/2 else 'black',
                        fontweight='bold', fontsize=_fs_tick)

    # Add colorbar
    cbar = plt.colorbar(im, ax=ax2, shrink=0.8)
    cbar.set_label('Number of Units', fontsize=_fs_tick)
    cbar.ax.tick_params(labelsize=_fs_tick)

    # Add session legend with proper participation patterns
    from matplotlib.patches import Patch
    legend_elements = [
        Patch(facecolor='lightblue', label='Session 1'),
        Patch(facecolor='lightgreen', label='Session 2'),
        Patch(facecolor='lightcoral', label='Session 3')
    ]
    ax2.legend(handles=legend_elements, loc='upper left', bbox_to_anchor=(1.15, 1),
              title='Session', fontsize=_fs_tick, title_fontsize=_fs_tick)
    
    plt.tight_layout()
    
    # Print firing rate statistics
    print(f"\nFiring Rate Statistics")
    print("=====================")
    fr_stats = unit_fr_data.groupby(['BrainRegion', 'SU'])['MeanFR'].agg([
        'count', 'mean', 'std', 'min', 'max', 'median'
    ]).round(2)
    print(fr_stats)
    
    # Print participant statistics  
    print(f"\nParticipant Statistics")
    print("====================")
    print(f"Total participants: {unique_units['ID'].nunique()}")
    print(f"Units per participant (mean ± std): {participant_summary['Total_Units'].mean():.1f} ± {participant_summary['Total_Units'].std():.1f}")
    print(f"Range: {participant_summary['Total_Units'].min()}-{participant_summary['Total_Units'].max()} units")
    
    # Brain region contribution by participant
    region_by_participant = unique_units.groupby(['ID', 'BrainRegion']).size().unstack(fill_value=0)
    print(f"\nBrain region sampling across participants:")
    for region in region_by_participant.columns:
        n_participants = (region_by_participant[region] > 0).sum()
        print(f"  {region}: {n_participants}/{len(region_by_participant)} participants ({n_participants/len(region_by_participant)*100:.1f}%)")
    
    # Units below threshold analysis
    threshold = 1.0  # Hz
    below_threshold = unit_fr_data[unit_fr_data['MeanFR'] < threshold]
    print(f"\nUnits below {threshold} Hz threshold:")
    print(f"Total: {len(below_threshold)}/{len(unit_fr_data)} ({len(below_threshold)/len(unit_fr_data)*100:.1f}%)")
    if len(below_threshold) > 0:
        threshold_breakdown = below_threshold.groupby(['BrainRegion', 'SU']).size()
        print("Breakdown by region and type:")
        print(threshold_breakdown)
    
    # Save figure if requested
    if save_fig and output_dir:
        output_path = Path(output_dir)
        output_path.mkdir(parents=True, exist_ok=True)
        fig_path = output_path / 'unit_overview.png'
        plt.savefig(fig_path, dpi=300, bbox_inches='tight')
        print(f"\nFigure saved to: {fig_path}")
        plt.savefig(os.path.join(output_path, 'unit_overview.svg'), 
                   format='svg', bbox_inches='tight')
    
    # Return summary data
    results = {
        'unit_counts': unit_counts,
        'participant_counts': participant_counts,
        'participant_summary': participant_summary,
        'session_summary': session_summary,
        'firing_rate_stats': fr_stats,
        'unit_fr_data': unit_fr_data,
        'below_threshold': below_threshold,
        'session_participant_matrix': matrix_for_plot,  # New session-participant matrix
        'participant_region_matrix': unique_units.groupby(['ID', 'BrainRegion']).size().unstack(fill_value=0),  # Keep original for compatibility
        'figure': fig,
        'data_type': data_type  # Include detected data type
    }
    
    return results
def add_deidentified_unit_id(df_units):
    """Add dUnitID column: P01_M1, P01_S1, etc."""
    df = df_units.copy()
    df['SU'] = df['SU'].astype(bool)
    df['UnitType'] = df['SU'].map({True: 'S', False: 'M'})
    df = df.sort_values(['ID', 'UnitType', 'UniqueLabel']).reset_index(drop=True)
    df['UnitNum'] = df.groupby(['ID', 'UnitType']).cumcount() + 1
    df['dUnitID'] = df['ID'] + '_' + df['UnitType'] + df['UnitNum'].astype(str)
    df = df.drop(columns=['UnitNum'])
    return df

#%% Main execution - define parameters, load and preprocess data
# =============================================================================
# CONFIGURATION
# =============================================================================

MMNsessions = ['1', '2', '3']

# De-identified participant IDs per session
allIDs = {
    'Session1': ['P01', 'P02', 'P03', 'P04', 'P05', 'P06', 'P07', 'P08',
                 'P09', 'P10', 'P11', 'P12', 'P13'],
    'Session2': ['P02', 'P08'],
    'Session3': ['P06'],
}

# Analysis parameters
sdf_norm = 'RA'              # 'block' or 'RA' - both divide by standard trials
normPeriod = (-0.1, 0.45)    # Normalization period (seconds)
rDec = 10                     # Decimation factor for time dimension
minFRcrit = 1                 # Minimum firing rate (Hz) for unit inclusion
outlier_threshold = 4         # IQR multiplier for outlier trial removal
verbose = True

# Data path - set this to the location of the de-identified data
# The folder should contain Session1/, Session2/, Session3/ subdirectories
dataPath = Path('../Optimum1_ProcessedData')

# Configure normalization settings
if sdf_norm == 'block':
    sdf_pos2 = '_Norm'
elif sdf_norm == 'RA':
    sdf_pos2 = '_NormRA'
else:
    sdf_pos2 = ''

# Define brain regions dictionary (mapping raw electrode labels to regions)
target_regions = {
    'HPC': ['HCzD', 'HCzI', 'HcaI', 'HCa', 'hkaI', 'hkaD', 'HkD', 'HkI',
            'HipI', 'HCuI', 'HCpD', 'HCu', 'HcuD', 'HCpI', 'hquD', 'HIPI','HD'],
    'Amy': ['AMD', 'Am', 'AmD', 'AMII', 'AMID', 'AI', 'AmI'],
    'Ins': ['InAnt'],
    'TTG': ['TTG', 'chan']
}



# =============================================================================
# DATA LOADING
# =============================================================================

# Storage for SDF data from each subject
sdf_parsing = []
DF_Units = pd.DataFrame()
DF_FRall = pd.DataFrame()
resultsFolder = Path(f'Figures/Session{"".join(MMNsessions)}')

for selMMNSession in MMNsessions:
    # Process each subject
    for subj in allIDs[f'Session{selMMNSession}']:
        if verbose:
            print(f"Processing SDF data for {subj}")
        
        # Define filenames (de-identified data structure)
        sdfFN = dataPath / f'Session{selMMNSession}' / 'SDFs' / f'{subj}_analysis.mat'
        FRtab_FN = dataPath / f'Session{selMMNSession}' / 'FRTabs' / f'{subj}_FRtable'
        
        try:
            sdfMAT = read_mat(str(sdfFN))
            spkMAT = loadmat(f'{FRtab_FN}.mat', struct_as_record=False, squeeze_me=True)
            DF_FR = pd.read_csv(f'{FRtab_FN}.csv')
            spkMAT['nsEventsALL'] = sdfMAT['nsEventsALL'].copy()
        except Exception as e:
            if verbose:
                print(f"  Error loading data for {subj}: {e}")
            
        # Extract unit information
        unitLabels = spkMAT['ALLspikes'].label
        unitBreg = spkMAT['ALLspikesInfo'].Bregion
        unitisSUAs = spkMAT['ALLspikesInfo'].isSUA
        meanFRs = spkMAT['ALLspikesInfo'].meanFR
        
        # Handle single unit case
        
        if isinstance(unitLabels, str):
            unitLabels = [unitLabels]
            unitBreg = np.array([unitBreg])
            unitisSUAs = [unitisSUAs]
            meanFRs = np.array([meanFRs])
            
        # --- Probe label corrections ---
        if subj == 'P01':
            unitBreg[(spkMAT['ALLspikesInfo'].MW == 1)] = 'AmI'#Left HPC probe is actually Amy
        if subj == 'P02':
            unitBreg[(spkMAT['ALLspikesInfo'].MW == 2)] = 'HD' #PIns probe targeted HPC
        if subj == 'P12':
            unitBreg[(spkMAT['ALLspikesInfo'].MW == 2)] = 'AmI'# HPC probe targeted Amy

            
        unitBregionOrig = unitBreg.copy()
        # Map brain regions for units
        for targReg in target_regions:
            unitBreg[np.isin(unitBreg, target_regions[targReg])] = targReg
            
        # CREATE UNIT-LEVEL DATAFRAME
        # Build unit-level DataFrame for this subject
        n_units = len(unitLabels)
        
        unit_data = {
            'ID': [subj] * n_units,
            'SessionStr': [selMMNSession] * n_units,
            'UnitLabel': unitLabels,
            'UniqueLabel': [f"{subj}_{label}S{selMMNSession}" for label in unitLabels],
            'BregionOrig': unitBregionOrig,
            'BrainRegion': unitBreg,
            'SU': unitisSUAs,
            'meanFR': meanFRs
        }
        
        # Add any additional fields from ALLspikesInfo if they exist
        additional_fields = ['SNR', 'amplitude', 'templateWidth', 'waveformDuration', 'firingRateStability']
        for field in additional_fields:
            if hasattr(spkMAT['ALLspikesInfo'], field):
                field_data = getattr(spkMAT['ALLspikesInfo'], field)
                if isinstance(field_data, (int, float, str)):
                    field_data = [field_data] * n_units
                elif len(field_data) != n_units:
                    continue  # Skip if length doesn't match
                unit_data[field] = field_data
        
        # Convert to DataFrame and append
        DF_subject_units = pd.DataFrame(unit_data)
        DF_Units = pd.concat([DF_Units, DF_subject_units], ignore_index=True)

        # Accumulate trial-level FR data (all units, including low-FR) for supplementary analysis
        label_to_unique = {label: f"{subj}_{label}S{selMMNSession}" for label in unitLabels}
        DF_FR_subj = DF_FR.copy()
        DF_FR_subj['UniqueLabel'] = DF_FR_subj['UnitLabel'].map(label_to_unique)
        DF_FR_subj['ID'] = subj
        DF_FR_subj['SessionStr'] = selMMNSession
        DF_FRall = pd.concat([DF_FRall, DF_FR_subj], ignore_index=True)

        # Get trial info for this subject
        trlInfoTemp = DF_FR.loc[DF_FR.UnitLabel == unitLabels[0]].copy()
        
        # Check data consistency
        if trlInfoTemp.shape[0] != sdfMAT['sdf']['trial'].shape[0]:
            print(f'{subj}: Inconsistent trial numbers between FRtab and trl')
            continue
        
        if sdfMAT['sdf']['label'] != spkMAT['ALLspikes'].label.tolist():
            print(f'{subj}: SDF and ALLspikes labels do not match')
            continue
        
        # Skip if all units are below minimum firing rate
        if np.all(meanFRs <= minFRcrit):
            print(f'{subj}: all {len(meanFRs)} units below minimum firing rate ({minFRcrit}) Hz')
            continue
        # Build SDF DataFrame using our custom function
        try:
            DF_subj = buildSdfDF_mmn(sdfMAT, spkMAT['ALLspikesInfo'], subj, trlInfoTemp, 
                                     target_regions = target_regions,r=rDec, minFR=minFRcrit, normBy=sdf_norm,normPeriod = normPeriod, sessionStr = selMMNSession)
            # Drop NaN values in SpikeDensity
            DF_subj = DF_subj.dropna(subset=['SpikeDensity'])
            # If no valid units, skip to next subject
            if DF_subj.empty:
                if verbose:
                    print(f"  No valid units found for {subj}")
                
            
            # Remove outliers if threshold is provided
            if outlier_threshold:
                # Function to remove outlier trials based on SDF values
                # After outlier removal
                DF_subj2, outlier_info = remove_sdf_outlier_trials(
                    DF_subj,
                    epoch=(-0.45, 0.45),
                    iqr_multiplier=outlier_threshold,
                    sdfStr=f'SpikeDensity{sdf_pos2}',
                    verbose = verbose
                )              

                    
                    # save the outlier information
                # Create directory for outlier info
                outlier_dir = resultsFolder / 'OutlierInfo' / f'Session{selMMNSession}'
                Path(outlier_dir).mkdir(parents=True, exist_ok=True)
                
                # Save outlier information for this subject
                outlier_file = outlier_dir / f'{subj}_outliers_{sdf_pos2}_{outlier_threshold}.pkl'
                
                # Create a simpler version of outlier_info that only contains what we need for raster plots
                raster_outlier_info = {
                    'unit_labels': outlier_info['unit_labels'],
                    'outlier_trials': outlier_info['outlier_trials'],
                    'subject': subj,
                    'threshold': outlier_threshold,
                    'epoch': (-0.45, 0.45),
                    'sdf_type': sdf_pos2
                }
                # Save outlier info to file
                with open(outlier_file, 'wb') as f:
                    pickle.dump(raster_outlier_info, f)
                    
            # Add to list of SDF dataframes
            sdf_parsing.append(DF_subj2)
            
            if verbose:
                print(f"  Added {len(DF_subj)} rows from {subj}")
                
        except Exception as e:
            print(f"Error processing SDF data for {subj}: {e}")

# Data is already de-identified - IDs in files are P01, P02, etc.
# No runtime de-identification needed

# Combine all SDFs
if sdf_parsing:
    DF_sdfs = pd.concat(sdf_parsing)
    if verbose:
        print(f"Combined SDF data contains {len(DF_sdfs)} rows")
    del(sdf_parsing)
    del(sdfMAT)

else:
    if verbose:
        print("No SDF data was successfully processed")
DF_sdfs = DF_sdfs.reset_index(['BrainRegion','SU'])        
DF_sdfs = DF_sdfs[~DF_sdfs.index.get_level_values('habituation').astype(bool)]     

#remove low-firing units again
# check for minimum firing rate after removing outlier trials
mFRs = DF_sdfs.groupby('UniqueLabel')['SpikeDensity'].mean()
print(mFRs[mFRs<=minFRcrit].index)

if len(mFRs[mFRs<=minFRcrit])>0:
    DF_sdfs = DF_sdfs.loc[~DF_sdfs['UniqueLabel'].isin(mFRs[mFRs<=minFRcrit].index.values)]
        
    
    
#last removal of bad units
# These units were excluded based on visual inspection of raster plots
# (inconsistent spiking, cross-talk, sudden FR changes)
rmUnitLabels = [
    'P03_NSX84-1S1',              # P03
    'P03_NSX84-2S1',              # P03
    'P04_NSX111-1S1',             # P04
    'P04_NSX124-2S1',             # P04 - questionable
    'P06_NSX80-1S1',              # P06 - inconsistent spiking
    'P07_NSX80-1S1',              # P07 - FR increase on last block
    'P08_NSX123-1S1',             # P08
    'P08_NSX124-2S1',             # P08
    'P08_NSX127-1S1',             # P08
    'P11_NSX101-1S1',             # P11
    'P12_NSX73-1S1',              # P12 - cross-talk
    'P12_NSX104-1S1',             # P12 - cross-talk for intensity-up
    'P12_NSX111-2S1',             # P12 - inconsistent spiking
    'P08_NSX66-1S2',              # P08 S2 - sudden FR increase
    'P08_NSX87-2S2',              # P08 S2 - sudden FR increase
    'P08_NSX87-4S2',              # P08 S2 - sudden FR increase
    'P08_NSX125-1S2',             # P08 S2 - sudden FR increase
    'P06_NSX80-1S3',              # P06 S3 - inconsistent spiking
    'P06_NSX103-1S3',             # P06 S3 - inconsistent spiking
    'P06_NSX104-1S3',             # P06 S3 - inconsistent spiking
]


DF_sdfs = DF_sdfs.loc[~DF_sdfs['UniqueLabel'].isin(rmUnitLabels)]


#optimizing  DF_sdfs
DF_sdfs['eventID'] = DF_sdfs['eventID'].astype('category')

for col in ['trlNum']:
    DF_sdfs[col] = pd.to_numeric(DF_sdfs[col], downcast='integer')

# Downcast SpikeDensity columns to float32 to reduce memory
sdf_cols = [c for c in DF_sdfs.columns if c.startswith('SpikeDensity')]
for col in sdf_cols:
    DF_sdfs[col] = DF_sdfs[col].astype(np.float32)


DF_sdfs['SU'] = DF_sdfs['SU'].astype('bool')

#% Define mean sdf values across predefined epochs
# Epoch definitions (seconds, relative to stimulus onset)
defIntervals = {
    'Baseline': [-0.1, 0.0],
    'Early':    [0.02, 0.1],
    'Middle':   [0.1, 0.25],
    'Late':     [0.25, 0.4],
    'Trial':    [0, 0.4],
    'EarlyNext': [0.5, 0.7]
}




DF_UnitsAnalyzed = DF_Units.loc[DF_Units['UniqueLabel'].isin(DF_sdfs['UniqueLabel'].unique())]
DF_UnitsAnalyzed = add_deidentified_unit_id(DF_UnitsAnalyzed)


# Initialize an empty dictionary to store mean/epoch responses for each interval as DataFrames
mean_dict = {}
for SpikeVar in ['SpikeDensity',f'SpikeDensity{sdf_pos2}']:
    # Loop over each defined interval to calculate and store means
    for perStr, perDef in defIntervals.items():
        print(f'{perStr}, {SpikeVar}')
        # Filter the DataFrame for the current interval
        df_sdfTemp = DF_sdfs[(DF_sdfs['time'] >= perDef[0]) & (DF_sdfs['time'] <= perDef[1]) &
                               (DF_sdfs.index.get_level_values('habituation') == False)]
        
        # Group by metadata and calculate the mean of SpikeDensity for the interval
        meanResponse = df_sdfTemp.groupby(
            ['trlNum','UniqueLabel']
        )[SpikeVar].mean().rename(f'{SpikeVar}{perStr}')
        
        # Add the resulting series to the dictionary
        mean_dict[f'{SpikeVar}{perStr}'] = meanResponse

# Concatenate all the interval columns into a single DataFrame on the index
DF_SDFmeans = pd.concat(mean_dict.values(), axis=1)
# Reset index to make it easier to work with
DF_SDFmeans = DF_SDFmeans.reset_index()


# Add metadata columns from DF_sdfs to DF_SDFmeans
metadata_mapping = DF_sdfs.reset_index().groupby(['trlNum', 'UniqueLabel']).agg({
    'BrainRegion':'first',
    'SU': 'first',
    'Block': 'first',
    'trlNumInBlock': 'first',
    'eventID': 'first'
}).reset_index()

DF_SDFmeans = DF_SDFmeans.merge(
    metadata_mapping,
    on=['trlNum', 'UniqueLabel'],
    how='left'
)

# Get unique combinations of BrainRegion, SU status, and UniqueLabel
unique_unitsTmp = DF_SDFmeans.drop_duplicates(subset=['BrainRegion', 'SU', 'UniqueLabel'])

# Count units by brain region and unit type
unit_counts = unique_unitsTmp.groupby(['BrainRegion', 'SU']).size().unstack(fill_value=0)

plt.rcParams['svg.fonttype'] = 'none'
plt.rcParams['pdf.fonttype'] = 'truetype'

resultsSorted = unitOverview(DF_Units, save_fig=True, output_dir=resultsFolder)  


resultsAnalyzed= unitOverview(DF_sdfs, save_fig=True, output_dir=f'{resultsFolder}/SDF')     



#%%

    
# 2. ADD HPC LOCATION COLUMN
HPC_HEAD = ['HCzD', 'HCzI', 'HkI', 'HcaI', 'HCa', 'hkaI', 'hkaD']
HPC_BODY = ['HCuI', 'HCpD', 'HCpI', 'HCu', 'hquD', 'HcuD']
def get_hpc_segment(region):
    if pd.isna(region):
        return np.nan
    if region in HPC_HEAD:
        return 'Head'
    elif region in HPC_BODY:
        return 'Body'
    else:
        return np.nan


# Add HPC Location column  
DF_Units['HPC_Segment'] = DF_Units['BregionOrig'].apply(get_hpc_segment)
DF_Units.loc[DF_Units['ID'].isin(['P02','P07','P13','P10']),'RHemisphere'] = True
DF_Units.loc[DF_Units['ID'].isin(['P09']),'RHemisphere'] = False

# add missing HPC field column
DF_Units.loc[DF_Units['ID'].isin(['P01']),'HPC_Segment'] = 'Head'
DF_Units.loc[DF_Units['ID'].isin(['P11']),'HPC_Segment'] = 'Head'

# Check which units need manual HPC assignment
DF_Units[DF_Units['HPC_Segment'].isna() & (DF_Units['BrainRegion'] == 'HPC')][['UniqueLabel', 'BregionOrig']]

# Add HPC Location column  
DF_UnitsAnalyzed['HPC_Segment'] = DF_UnitsAnalyzed['BregionOrig'].apply(get_hpc_segment)

# Check which units need manual HPC assignment
DF_UnitsAnalyzed[DF_UnitsAnalyzed['HPC_Segment'].isna() & (DF_UnitsAnalyzed['BrainRegion'] == 'HPC')][['UniqueLabel', 'BregionOrig']]

#%%


def epoch_permutation_test(formula, data, factor_of_interest,family_config, n_permutations=1000, 
                          model_type='GLM', random_state=42):
    """
    Perform permutation test for a single factor within epoch data for one unit.
    
    Parameters:
    -----------
    formula : str
        Model formula (e.g., 'SpikeDensity_NormRAEarly ~ isDeviant')
    data : pandas.DataFrame
        Trial-level data for one unit and one epoch
    factor_of_interest : str
        Factor to test (e.g., 'isDeviant')
    n_permutations : int
        Number of permutations
    model_type : str
        'GLM' or 'LME'
    random_state : int
        Random seed for reproducibility
        
    Returns:
    --------
    dict
        Permutation test results
    """
    np.random.seed(random_state)
    
    if len(data) < 5:  # Minimum trials needed
        return None
    
    # Fit original model
    try:
        if model_type == 'GLM':
            original_model = smf.glm(formula, data=data, family=family_config)
            original_result = original_model.fit()
        else:  # LME (though less relevant for single unit)
            # For single unit, LME doesn't make much sense, but keeping for consistency
            data_temp = data.copy()
            original_model = smf.mixedlm(formula, data=data_temp, 
                                       groups=data_temp['Block'], re_formula="1")
            original_result = original_model.fit(reml=False)
        
        # Get coefficient for factor of interest
        param_name = f"{factor_of_interest}[T.True]"
        if param_name in original_result.params:
            original_coef = original_result.params[param_name]
        elif factor_of_interest in original_result.params:
            original_coef = original_result.params[factor_of_interest]
            param_name = factor_of_interest
        else:
            return None
            
    except Exception as e:
        print(f"Original model failed: {str(e)}")
        return None
    
    # Get response variable name
    response_var = formula.split('~')[0].strip()
    
    # Run permutations
    permuted_coefs = []
    
    for perm_i in range(n_permutations):
        # Permute response variable
        data_perm = data.copy()
        data_perm[response_var] = np.random.permutation(data[response_var].values)
        
        try:
            if model_type == 'GLM':
                perm_model = smf.glm(formula, data=data_perm, family=family_config)
                perm_result = perm_model.fit()
            else:
                perm_model = smf.mixedlm(formula, data=data_perm, 
                                       groups=data_perm['Block'], re_formula="1")
                perm_result = perm_model.fit(reml=False)
            
            # Extract coefficient
            if param_name in perm_result.params:
                permuted_coefs.append(perm_result.params[param_name])
            elif factor_of_interest in perm_result.params:
                permuted_coefs.append(perm_result.params[factor_of_interest])
                
        except:
            # Skip failed permutations
            continue
    
    if len(permuted_coefs) < n_permutations * 0.8:  # Need at least 80% success rate
        return None
    
    # Calculate statistics
    permuted_coefs = np.array(permuted_coefs)
    n_perms_actual = len(permuted_coefs)
    # Two-tailed p-value: proportion of permutations as extreme as observed
    n_as_extreme = np.sum(np.abs(permuted_coefs) >= np.abs(original_coef))
    if n_as_extreme == 0:
        # Conservative bound: p <= 1/(n+1)
        p_value = 1.0 / (n_perms_actual + 1)
    else:
        # Standard calculation
        p_value = n_as_extreme / n_perms_actual
    
    ci_lower = np.percentile(permuted_coefs, 2.5)
    ci_upper = np.percentile(permuted_coefs, 97.5)
    
    return {
        'original_coefficient': original_coef,
        'permuted_coefficients': permuted_coefs,
        'p_value': p_value,
        'ci_lower': ci_lower,
        'ci_upper': ci_upper,
        'n_permutations': n_perms_actual,
        'param_name': param_name,
        'p_exact_zero': (n_as_extreme == 0)  # Flag for reporting
    }


def run_unit_epoch_models_with_permutation(DF_SDFmeans, epochs, model_name = 'null',
                                          baseEquation = 'SpikeDensity ~ 1',
                                          normPos='',
                                          model_type='GLM', modelFamily = [],
                                          n_permutations=1000, n_jobs=-1, FRmin2=None, factors = None):
    """
    Run epoch-based unit-level models with permutation testing.
    
    Parameters:
    -----------
    DF_SDFmeans : pandas.DataFrame
        Epoch-averaged spike density data
    epochs : list
        List of epochs to analyze (e.g., ['Early', 'Middle', 'Late'])
    contrast_type : str
        Type of contrast ('individual', 'categories', 'combined', etc.)
    normPos : str
        Normalization suffix ('_NormRA', '_Norm', '')
    model_type : str
        'GLM' or 'LME'
    n_permutations : int
        Number of permutations per test
    n_jobs : int
        Number of parallel jobs
    FRmin2 : float or None
        Minimum firing rate threshold
        
    Returns:
    --------
    dict
        Results organized by unit and epoch
    """

    
    print(f"Starting epoch-based unit analysis with {n_permutations} permutations...")


    # Prepare data with contrast columns
    data_prepared = DF_SDFmeans.copy()
    if 'isDeviant' in baseEquation and not any(x in baseEquation for x in ['isLocation', 'isFrequency', 'isIntensity', 'isTiming']):
        contrast_type = 'combined'
    elif any(x in baseEquation for x in ['isLocLeft', 'isLocRight', 'isDur25', 'isGap', 'isFreq_450', 'isFreq_550', 'isintensityDown', 'isintensityUp']):
        contrast_type = 'individual'
    elif any(x in baseEquation for x in ['isLocation', 'isFrequency', 'isIntensity', 'isTiming']):
        contrast_type = 'categories'
    elif 'StdTime' in model_name:
        contrast_type = 'StdTime'
    else:
        contrast_type = 'custom'
        
    add_contrast_columns(data_prepared, contrast_type)
    
    # Apply FR filter if specified - use Trial epoch for filtering
    if FRmin2 is not None:
        trial_col = 'SpikeDensityTrial'
        if trial_col in data_prepared.columns:
            mean_frs = data_prepared.groupby('UniqueLabel')[trial_col].mean()
            valid_units = mean_frs[mean_frs >= FRmin2].index
            data_prepared = data_prepared[data_prepared['UniqueLabel'].isin(valid_units)]
            print(f"Filtered to {len(valid_units)} units with FR >= {FRmin2} Hz")
        else:
            print(f"Warning: {trial_col} not found, skipping FR filter")
    
    # Get unique units
    units = data_prepared['UniqueLabel'].unique()
    print(f"Processing {len(units)} units across {len(epochs)} epochs")
    
    # Prepare tasks for parallel processing
    tasks = []
    
    for unit in units:
        unit_data = data_prepared[data_prepared['UniqueLabel'] == unit]
        brain_region = unit_data['BrainRegion'].iloc[0] if 'BrainRegion' in unit_data.columns else 'Unknown'
        
        for epoch in epochs:
            # Get response variable for this epoch
            response_var = f'SpikeDensity{normPos}{epoch}'

            
            if response_var not in unit_data.columns:
                continue
                
            # Update equations for this epoch
            epoch_equations = baseEquation.replace('responseVar', response_var) 
            
            # Extract factors from equations
            if factors is None:
                factors = set()
                rhs = epoch_equations.split('~')[1].strip()
                for term in rhs.replace('*', '+').replace(':', '+').split('+'):
                    term = term.strip()
                    if term not in ['1', '+', '-', '*', '/', '(', ')', ','] and 'C(' not in term:
                        factors.add(term)
            
            # Create task for each model and factor combination
            for factor in factors:
                if factor in epoch_equations:
                    tasks.append({
                        'unit': unit,
                        'brain_region': brain_region,
                        'epoch': epoch,
                        'model_name': model_name,
                        'formula': epoch_equations,
                        'factor': factor,
                        'data': unit_data,
                        'n_permutations': n_permutations,
                        'model_type': model_type,
                        'model_family': modelFamily
                    })
    
    print(f"Created {len(tasks)} analysis tasks")
    
    # Run parallel processing
    def process_single_task(task):
        try:
            result = epoch_permutation_test(
                task['formula'], task['data'], task['factor'], task['model_family']['function'],
                task['n_permutations'], task['model_type']
            )
            
            # Also fit original model for additional metrics
            if model_type == 'GLM':
                original_model = smf.glm(task['formula'], data=task['data'], 
                                       family=task['model_family']['function'])
                fitted = original_model.fit()
            else:
                data_temp = task['data'].copy()
                original_model = smf.mixedlm(task['formula'], data=data_temp, 
                                           groups=data_temp['Block'], re_formula="1")
                fitted = original_model.fit(reml=False)
            
            # Calculate t-value
            param_name = result['param_name'] if result else None
            if result and param_name in fitted.tvalues:
                t_value = fitted.tvalues[param_name]
            else:
                t_value = np.nan
                
            return {
                'unit': task['unit'],
                'brain_region': task['brain_region'],
                'epoch': task['epoch'],
                'model_name': task['model_name'],
                'model_family_name': task['model_family']['name'],
                'factor': task['factor'],
                'permutation_result': result,
                't_value': t_value,
                'original_fit': fitted if result else None
            }
            
        except Exception as e:
            print(f"Task failed for {task['unit']}, {task['epoch']}, {task['factor']}: {str(e)}")
            return None
    
    # Execute in parallel
    start_time = time.time()
    results = Parallel(n_jobs=n_jobs, verbose=1)(
        delayed(process_single_task)(task) for task in tasks
    )
    
    print(f"Completed in {time.time() - start_time:.1f} seconds")
    
    # Organize results
    organized_results = {}
    
    for result in results:
        if result is None or result['permutation_result'] is None:
            continue
            
        unit = result['unit']
        epoch = result['epoch']
        model = result['model_name']
        famly = result['model_family_name']
        factor = result['factor']
        
        # Initialize nested structure
        if unit not in organized_results:
            organized_results[unit] = {
                'brain_region': result['brain_region'],
                'epochs': {}
            }
        
        if epoch not in organized_results[unit]['epochs']:
            organized_results[unit]['epochs'][epoch] = {}
            
        if model not in organized_results[unit]['epochs'][epoch]:
            organized_results[unit]['epochs'][epoch][model] = {}
            
        # Store results
        organized_results[unit]['epochs'][epoch][model][factor] = {
            'permutation_result': result['permutation_result'],
            't_value': result['t_value'],
            'original_fit': result['original_fit']
        }
    
    print(f"Organized results for {len(organized_results)} units")
    
    return organized_results

def create_epoch_unit_summary_table(unit_results, epochs, models, factors, 
                                   unitsInfo, significance_threshold=0.05):
    """
    Create summary table of results across all units, epochs, models, and factors.
    """
    # Create lookup dictionary for unit info
    unit_info_dict = {}
    if 'UniqueLabel' in unitsInfo.columns:
        for _, row in unitsInfo.iterrows():
            unit_info_dict[row['UniqueLabel']] = {
                'SU': row['SU']
            }
    summary_rows = []
    
    for unit, unit_data in unit_results.items():
        brain_region = unit_data['brain_region']
        # Get unit info from lookup
        unit_info = unit_info_dict.get(unit, {'SU': None})
        for epoch in epochs:
            for model in models:
                for factor in factors:
                    if (epoch in unit_data['epochs'] and 
                        model in unit_data['epochs'][epoch] and 
                        factor in unit_data['epochs'][epoch][model]):
                        
                        factor_result = unit_data['epochs'][epoch][model][factor]
                        perm_result = factor_result['permutation_result']
                        
                        summary_rows.append({
                            'Unit': unit,
                            'BrainRegion': brain_region,
                            'SU': unit_info['SU'],  
                            'Epoch': epoch,
                            'Model': model,
                            'Factor': factor,
                            'Coefficient': perm_result['original_coefficient'],
                            'T_value': factor_result['t_value'],
                            'P_value': perm_result['p_value'],
                            'CI_lower': perm_result['ci_lower'],
                            'CI_upper': perm_result['ci_upper'],
                            'Significant': perm_result['p_value'] < significance_threshold,
                            'N_permutations': perm_result['n_permutations']
                        })
    
    return pd.DataFrame(summary_rows)

def combined_region_summary_heatmap(summary_table, model_name, output_dir, regions_to_plot, 
                                   population_test_results=None, significance_threshold=0.05,
                                   colormap_limits=None, factor_name='Combined',
                                   show_directionality=True, test_type='fisher'):
    """
    Create heatmap visualization for combined case showing regions × epochs.
    
    Parameters:
    -----------
    summary_table : pandas.DataFrame
        Summary table containing results
    model_name : str
        Name of the model
    output_dir : str
        Directory to save the output
    regions_to_plot : list
        List of brain regions to plot
    fisher_results : pandas.DataFrame, optional
        Results from Fisher's combined test containing 'Significant_FDR' column
    significance_threshold : float
        Threshold for individual unit significance (default 0.05)
    colormap_limits : tuple, optional
        (vmin, vmax) limits for colormap
    factor_name : str
        Name of the factor being analyzed
    show_directionality : bool
        If True, show enhancement/suppression counts in annotation
    """
    
    # Create figure
    fig, ax = plt.subplots(1, 1, figsize=(max(5, len(summary_table.Epoch.unique())*0.8), 
                                          max(6, len(regions_to_plot)*0.8)))
    
    # Initialize data structures
    epochs = summary_table.Epoch.unique()
    regions = [r for r in regions_to_plot if r in summary_table.BrainRegion.unique()]
    
    # Matrices for percentages and counts
    percentage_matrix = np.zeros((len(regions), len(epochs)))
    count_matrix = np.zeros((len(regions), len(epochs)), dtype=object)
    
    # Fill the matrices
    for region_idx, region in enumerate(regions):
        for epoch_idx, epoch in enumerate(epochs):
            # Filter data for this epoch/region combination
            epoch_region_data = summary_table[
                (summary_table.Epoch == epoch) & 
                (summary_table.BrainRegion == region)
            ]
            
            if len(epoch_region_data) > 0:
                n_total = len(epoch_region_data)
                
                # Determine direction column
                if 'T_value' in epoch_region_data.columns:
                    direction_col = 'T_value'
                elif 'MeanDifference' in epoch_region_data.columns:
                    direction_col = 'MeanDifference'
                elif 'MedianDifference' in epoch_region_data.columns:
                    direction_col = 'MedianDifference'
                elif 'Coefficient' in epoch_region_data.columns:
                    direction_col = 'Coefficient'
                else:
                    direction_col = None
                
                # Count enhanced and suppressed responses
                sig_mask = epoch_region_data['P_value'] < significance_threshold
                
                if direction_col is not None:
                    n_enhanced = (sig_mask & (epoch_region_data[direction_col] > 0)).sum()
                    n_suppressed = (sig_mask & (epoch_region_data[direction_col] < 0)).sum()
                else:
                    # Fallback: just count significant without direction
                    n_enhanced = sig_mask.sum()
                    n_suppressed = 0
                
                n_responsive = n_enhanced + n_suppressed
                percentage = (n_responsive / n_total) * 100 if n_total > 0 else 0
                
                percentage_matrix[region_idx, epoch_idx] = percentage
                
                # Create text annotation with directionality
                if show_directionality and direction_col is not None and (n_enhanced > 0 or n_suppressed > 0):
                    text = f"{n_enhanced}↑/{n_suppressed}↓\n({percentage:.1f}%)"
                else:
                    text = f"{n_responsive}/{n_total}\n({percentage:.1f}%)"
                
                count_matrix[region_idx, epoch_idx] = text
            else:
                percentage_matrix[region_idx, epoch_idx] = 0
                count_matrix[region_idx, epoch_idx] = "0/0"
    
    # Set colormap limits
    if colormap_limits is None:
        vmin, vmax = 0, percentage_matrix.max() if percentage_matrix.max() > 0 else 100
    else:
        vmin, vmax = colormap_limits
    
    # Create heatmap
    im = ax.imshow(percentage_matrix, cmap='Reds', aspect='auto', 
                   vmin=vmin, vmax=vmax, interpolation='nearest')
    
    # Add text annotations with asterisks
    for region_idx, region in enumerate(regions):
        for epoch_idx, epoch in enumerate(epochs):
            # Base text
            text = count_matrix[region_idx, epoch_idx]
            color = 'white' if percentage_matrix[region_idx, epoch_idx] > (vmax * 0.5) else 'black'
        
            # Add asterisk if population test is significant
            if population_test_results is not None:
                # Find the corresponding factor in test results
                possible_factors = [factor_name, f'is{factor_name}', f'{factor_name}[T.True]']
                
                for possible_factor in possible_factors:
                    test_match = population_test_results[
                        (population_test_results.BrainRegion == region) &
                        (population_test_results.Epoch == epoch) &
                        (population_test_results.Factor == possible_factor)
                    ]
                    
                    if len(test_match) > 0 and test_match.iloc[0]['Significant_FDR']:
                        text += "*"
                        break
            
            # Plot text
            ax.text(epoch_idx, region_idx, text, ha='center', va='center', 
                   color=color, fontsize=11, fontweight='bold')
    
    # Set labels and ticks
    ax.set_xticks(range(len(epochs)))
    ax.set_xticklabels(epochs)
    ax.set_yticks(range(len(regions)))
    ax.set_yticklabels(regions)
    
    ax.set_xlabel('Epoch', fontweight='bold')
    ax.set_ylabel('Brain Region', fontweight='bold')
    
    # Add colorbar
    cbar = plt.colorbar(im, ax=ax, shrink=0.8)
    cbar.set_label('% Responsive Units', fontweight='bold')
    
    # Title
    title = f'Combined Analysis: {factor_name} Factor ({model_name} model)\n'
    if show_directionality:
        title += 'Numbers = Enhanced↑/Suppressed↓ (% Total Responsive)'
    else:
        title += 'Numbers = Responding/Total Units'
    if population_test_results is not None:
        if test_type == 'fisher':
            title += '\n* = FDR-corrected Fisher\'s combined test'
        elif test_type == 'binomial':
            title += '\n* = FDR-corrected binomial test'
    ax.set_title(title, fontsize=14, pad=20)
    
    # Adjust layout
    plt.tight_layout()
    
    # Save figure
    filename_base = f'{model_name}_{factor_name}_combined_epoch_region_heatmap'
    plt.savefig(os.path.join(output_dir, f'{filename_base}.png'), 
               dpi=300, bbox_inches='tight')
    plt.savefig(os.path.join(output_dir, f'{filename_base}.svg'), 
               format='svg', bbox_inches='tight')
    
    return fig, percentage_matrix, count_matrix

def region_summary_heatmaps(summary_table, model_name, output_dir, regions_to_plot, 
                           population_test_results=None, significance_threshold=0.05,
                           combined_mode=False, colormap_limits=None,
                           show_directionality=True, test_type='fisher'):
    """
    Create heatmap visualizations with population-level test results.
    
    NOW INCLUDES: Enhanced↑/Suppressed↓ annotations for all modes (not just combined).
    
    Parameters:
    -----------
    summary_table : DataFrame
        Results from create_epoch_unit_summary_table
    model_name : str
        Name of the model
    output_dir : str
        Directory to save outputs
    regions_to_plot : list
        List of brain regions to include
    population_test_results : DataFrame, optional
        Results from Fisher's or binomial test
    significance_threshold : float
        Threshold for individual unit significance (default 0.05)
    combined_mode : bool
        If True, use combined_region_summary_heatmap
    colormap_limits : tuple, optional
        (vmin, vmax) for colormap
    show_directionality : bool
        If True, include direction info in annotations
    test_type : str
        Type of test: 'fisher', 'binomial', or 'none'
    """
    
    # Check if single factor case
    factors = summary_table.Factor.unique()
    
    if len(factors) == 1 or combined_mode:
        factor_name = factors[0] if len(factors) == 1 else 'Combined'
        clean_factor = factor_name.replace('is', '').replace('[T.True]', '')
        
        return combined_region_summary_heatmap(
            summary_table, model_name, output_dir, regions_to_plot,
            population_test_results, significance_threshold, colormap_limits, 
            clean_factor, show_directionality, test_type
        )

    # Calculate dimensions based on data
    n_factors = len(summary_table['Factor'].unique())
    n_regions = len(regions_to_plot)
    
    # Scale: more space needed for directionality annotations
    if n_regions <= 2:  # Side by side
        fig_width = n_factors * 1.5 * 2  # Increased width for annotations
        fig_height = 5
        n_rows, n_cols = 1, 2
    else:  # Grid layout
        fig_width = n_factors * 1.5 * 2 *0.65
        fig_height = 3.5 * 3 *0.85# 3 rows
        n_rows, n_cols = 3, 2
    
    fig, axes = plt.subplots(n_rows, n_cols, figsize=(fig_width, fig_height))
    
    if n_regions == 1:
        axes = [axes]
    elif n_rows == 1 or n_cols == 1:
        axes = axes.flatten()
    else:
        axes = axes.flatten()
    
    # Plot each region
    for i, region in enumerate(regions_to_plot):
        ax = axes[i]
        
        region_data = summary_table[summary_table.BrainRegion == region]
        
        if len(region_data) > 0:
            # Determine direction column
            if 'T_value' in region_data.columns:
                direction_col = 'T_value'
            elif 'MeanDifference' in region_data.columns:
                direction_col = 'MeanDifference'
            elif 'MedianDifference' in region_data.columns:
                direction_col = 'MedianDifference'
            elif 'Coefficient' in region_data.columns:
                direction_col = 'Coefficient'
            else:
                direction_col = None
            
            # Get unique epochs and factors
            epochs = region_data['Epoch'].unique()
            factors_clean = [f.replace('is', '').replace('[T.True]', '') 
                           for f in region_data['Factor'].unique()]
            factors_original = region_data['Factor'].unique()
            
            # Create matrices for heatmap
            percentage_matrix = np.zeros((len(epochs), len(factors_original)))
            annotation_matrix = np.empty((len(epochs), len(factors_original)), dtype=object)
            
            for epoch_idx, epoch in enumerate(epochs):
                for factor_idx, factor in enumerate(factors_original):
                    # Filter data for this epoch/factor combination
                    cell_data = region_data[
                        (region_data.Epoch == epoch) & 
                        (region_data.Factor == factor)
                    ]
                    
                    if len(cell_data) > 0:
                        n_total = len(cell_data)
                        sig_mask = cell_data['P_value'] < significance_threshold
                        
                        if direction_col is not None and show_directionality:
                            n_enhanced = (sig_mask & (cell_data[direction_col] > 0)).sum()
                            n_suppressed = (sig_mask & (cell_data[direction_col] < 0)).sum()
                        else:
                            n_enhanced = sig_mask.sum()
                            n_suppressed = 0
                        
                        n_responsive = n_enhanced + n_suppressed
                        percentage = (n_responsive / n_total) * 100 if n_total > 0 else 0
                        
                        percentage_matrix[epoch_idx, factor_idx] = percentage
                        
                        # Create annotation with directionality
                        if show_directionality and direction_col is not None:
                            annotation_matrix[epoch_idx, factor_idx] = f"{n_enhanced}↑/{n_suppressed}↓\n({percentage:.1f}%)"
                        else:
                            annotation_matrix[epoch_idx, factor_idx] = f"{n_responsive}/{n_total}\n({percentage:.1f}%)"
                    else:
                        percentage_matrix[epoch_idx, factor_idx] = 0
                        annotation_matrix[epoch_idx, factor_idx] = "0/0"
            
            # Set colormap limits
            if colormap_limits is None:
                vmin = 0
                vmax = percentage_matrix.max() if percentage_matrix.max() > 0 else 100
            else:
                vmin, vmax = colormap_limits
            
            # Create heatmap
            im = ax.imshow(percentage_matrix, cmap='Reds', aspect='auto',
                          vmin=vmin, vmax=vmax, interpolation='nearest')
            
            # Add text annotations with asterisks for population significance
            for epoch_idx, epoch in enumerate(epochs):
                for factor_idx, factor in enumerate(factors_original):
                    text = annotation_matrix[epoch_idx, factor_idx]
                    color = 'white' if percentage_matrix[epoch_idx, factor_idx] > (vmax * 0.5) else 'black'
                    
                    # Add asterisk if population test is significant
                    if population_test_results is not None:
                        possible_factors = [factor, factor.replace('is', ''), f'{factor}[T.True]']
                        
                        for possible_factor in possible_factors:
                            test_match = population_test_results[
                                (population_test_results.BrainRegion == region) &
                                (population_test_results.Epoch == epoch) &
                                (population_test_results.Factor == possible_factor)
                            ]
                            
                            if len(test_match) > 0 and test_match.iloc[0]['Significant_FDR']:
                                text += "*"
                                break
                    
                    ax.text(factor_idx, epoch_idx, text, ha='center', va='center',
                           color=color, fontsize=9, fontweight='bold')
            
            # Set labels and ticks
            ax.set_xticks(range(len(factors_clean)))
            ax.set_xticklabels(factors_clean, rotation=45, ha='right')
            ax.set_yticks(range(len(epochs)))
            ax.set_yticklabels(epochs)
            
            ax.set_title(f'{region}', fontweight='bold')
            ax.set_xlabel('Factor/Deviant Type')
            if i % n_cols == 0:
                ax.set_ylabel('Epoch')
            
        else:
            ax.text(0.5, 0.5, f'No data for {region}', 
                   transform=ax.transAxes, ha='center', va='center', 
                   fontsize=12, color='gray')
            ax.set_title(f'{region}')
            ax.set_xticks([])
            ax.set_yticks([])
    
    # Hide unused subplots
    for i in range(n_regions, len(axes)):
        axes[i].set_visible(False)
    
    # Add colorbar to the figure
    cbar_ax = fig.add_axes([0.92, 0.15, 0.02, 0.7])
    cbar = fig.colorbar(im, cax=cbar_ax)
    cbar.set_label('% Responsive Units', fontweight='bold')
    
    # Add overall title
    title = f'Responsive Units: Epoch × Factor ({model_name} model)'
    if population_test_results is not None:
        if test_type == 'fisher':
            title += '\n* = FDR-corrected Fisher\'s combined test'
        elif test_type == 'binomial':
            title += '\n* = FDR-corrected binomial test'
    if show_directionality:
        title += '\n(Numbers = Enhanced↑/Suppressed↓)'
    plt.suptitle(title, fontsize=14, y=1.02)
    
    plt.tight_layout(rect=[0, 0, 0.9, 0.95])
    
    # Save
    plt.savefig(os.path.join(output_dir, f'{model_name}_region_heatmaps_combined.png'), 
               dpi=300, bbox_inches='tight')
    plt.savefig(os.path.join(output_dir, f'{model_name}_region_heatmaps_combined.svg'), 
               format='svg', bbox_inches='tight')
    
    return fig

def run_binomial_significance_test(summary_table, alpha=0.05):
    """
    Run binomial tests on significance rates and apply FDR correction.
    
    Parameters:
    -----------
    summary_table : pandas.DataFrame
        Summary table from create_epoch_unit_summary_table
    alpha : float
        Expected significance rate under null hypothesis (default 0.05)
    
    Returns:
    --------
    pandas.DataFrame
        Results with binomial p-values and FDR corrections
    """
    
    binomial_results = []
    
    # Group by Region, Factor, and Epoch
    for (region, factor, epoch), group in summary_table.groupby(['BrainRegion', 'Factor', 'Epoch']):
        n_total = len(group)
        n_significant = (group['P_value'] < alpha).sum()
        
        if n_total > 0:
            # Binomial test: is significance rate > chance (alpha)?
            # Using 'greater' because we test if observed rate > expected rate
            binom_temp = binomtest(n_significant, n_total, p=alpha, alternative='greater')
            p_binomial = binom_temp.pvalue
            
            binomial_results.append({
                'BrainRegion': region,
                'Factor': factor,
                'Epoch': epoch,
                'N_total': n_total,
                'N_significant': n_significant,
                'Significance_rate': n_significant / n_total,
                'Expected_rate': alpha,
                'P_binomial': p_binomial
            })
    
    binomial_df = pd.DataFrame(binomial_results)
    
    # Apply FDR correction within each epoch
    binomial_df['P_binomial_FDR'] = np.nan
    binomial_df['Significant_FDR'] = False
    
    for epoch in binomial_df['Epoch'].unique():
        for region in binomial_df['BrainRegion'].unique():
            # Select this region×epoch combination
            mask = (binomial_df['Epoch'] == epoch) & (binomial_df['BrainRegion'] == region)
            region_epoch_pvals = binomial_df.loc[mask, 'P_binomial'].values
            
            if len(region_epoch_pvals) > 0:
                # FDR correction across factors for this region×epoch
                _, p_corrected, _, _ = multipletests(region_epoch_pvals, alpha=alpha, method='fdr_bh')
                binomial_df.loc[mask, 'P_binomial_FDR'] = p_corrected
                binomial_df.loc[mask, 'Significant_FDR'] = p_corrected < alpha
    
    return binomial_df


def create_population_test_summary_heatmap(test_results, output_dir, model_name, test_type='fisher'):
    """
    Create heatmap showing FDR-corrected population-level test results.
    
    Parameters:
    -----------
    test_results : pandas.DataFrame
        Results from run_binomial_significance_test or run_fishers_combined_test
    output_dir : str
        Directory to save the output
    model_name : str
        Name of the model
    test_type : str
        Type of test: 'fisher' or 'binomial'
    """
    # Determine which p-value column to use
    if test_type == 'fisher':
        pval_column = 'P_fisher_FDR'
        test_name = "Fisher's combined test"
    elif test_type == 'binomial':
        pval_column = 'P_binomial_FDR'
        test_name = "binomial test"
    else:
        raise ValueError(f"test_type must be 'fisher' or 'binomial', got '{test_type}'")
    
    # Create separate heatmaps for each epoch
    epochs = test_results['Epoch'].unique()
    n_epochs = len(epochs)
    
    fig, axes = plt.subplots(1, n_epochs, figsize=(5*n_epochs, 4))
    if n_epochs == 1:
        axes = [axes]
    
    for i, epoch in enumerate(epochs):
        epoch_data = test_results[test_results['Epoch'] == epoch]
        
        # Create pivot for heatmap
        pivot_data = epoch_data.pivot(index='BrainRegion', columns='Factor', values=pval_column)
        
        # Create significance mask
        sig_mask = epoch_data.pivot(index='BrainRegion', columns='Factor', values='Significant_FDR')
        
        # Plot heatmap
        sns.heatmap(pivot_data, annot=True, fmt='.3f', cmap='Blues_r', 
                   ax=axes[i], cbar_kws={'label': 'FDR-corrected p-value'},
                   mask=~sig_mask, linewidths=0.5, vmin=0, vmax=0.05)
        
        axes[i].set_title(f'{epoch}')
        axes[i].set_xlabel('Factor')
        if i == 0:
            axes[i].set_ylabel('Brain Region')
        else:
            axes[i].set_ylabel('')
    
    plt.suptitle(f'Population-level Significance: {model_name} Model\n(FDR-corrected {test_name})', 
                 fontsize=14)
    plt.tight_layout()
    
    # Save
    plt.savefig(os.path.join(output_dir, f'{model_name}_population_test_heatmap.png'), 
               dpi=300, bbox_inches='tight')
    plt.savefig(os.path.join(output_dir, f'{model_name}_population_test_heatmap.svg'), 
               format='svg', bbox_inches='tight')
    
    return fig
def run_fishers_combined_test(summary_table, alpha=0.05):
    """
    Fisher's combined probability test for population-level significance.
    Handles bidirectional responses (enhancement and suppression).
    
    Parameters:
    -----------
    summary_table : pandas.DataFrame
        Summary table from create_epoch_unit_summary_table or 
        create_epoch_unit_summary_table_simple
    alpha : float
        Significance level for FDR correction
        
    Returns:
    --------
    pandas.DataFrame
        Results with Fisher's statistics and directionality info
    """
    from scipy.stats import combine_pvalues
    
    results = []
    
    # Group by Region, Factor, and Epoch
    for (region, factor, epoch), group in summary_table.groupby(['BrainRegion', 'Factor', 'Epoch']):
        
        # Get p-values (should be two-tailed from your permutation tests)
        p_values = group['P_value'].values
        
        # Handle different column names for effect direction
        if 'T_value' in group.columns:
            direction = group['T_value'].values
        elif 'MeanDifference' in group.columns:
            direction = group['MeanDifference'].values
        elif 'MedianDifference' in group.columns:
            direction = group['MedianDifference'].values
        elif 'Coefficient' in group.columns:
            direction = group['Coefficient'].values
        else:
            # Fallback: use sign of any available metric
            direction = np.zeros(len(group))
        
        # Remove NaNs
        valid_mask = ~np.isnan(p_values)
        p_values_clean = p_values[valid_mask]
        direction_clean = direction[valid_mask]
        
        if len(p_values_clean) < 3:
            continue
        
        # Run Fisher's combined test
        try:
            chi2_stat, combined_p = combine_pvalues(p_values_clean, method='fisher')
        except:
            # Handle edge cases (e.g., all p-values = 1)
            chi2_stat = 0
            combined_p = 1.0
        
        # Count directionality for units that are individually significant
        sig_mask = p_values_clean < 0.05  # or whatever threshold you use
        n_enhanced = np.sum((direction_clean > 0) & sig_mask)
        n_suppressed = np.sum((direction_clean < 0) & sig_mask)
        n_nonresponsive = len(p_values_clean) - n_enhanced - n_suppressed
        
        # Also get total directionality (regardless of significance)
        n_positive_trend = np.sum(direction_clean > 0)
        n_negative_trend = np.sum(direction_clean < 0)
        
        results.append({
            'BrainRegion': region,
            'Factor': factor,
            'Epoch': epoch,
            'N_total': len(p_values_clean),
            'N_enhanced': n_enhanced,
            'N_suppressed': n_suppressed,
            'N_nonresponsive': n_nonresponsive,
            'N_positive_trend': n_positive_trend,
            'N_negative_trend': n_negative_trend,
            'Percent_responsive': 100 * (n_enhanced + n_suppressed) / len(p_values_clean),
            'Chi2_statistic': chi2_stat,
            'P_fisher': combined_p,
            'Mean_abs_effect': np.mean(np.abs(direction_clean))
        })
    
    results_df = pd.DataFrame(results)
    
    # Apply FDR correction within each region×epoch
    results_df['P_fisher_FDR'] = np.nan
    results_df['Significant_FDR'] = False
    
    for (region, epoch) in results_df.groupby(['BrainRegion', 'Epoch']).groups.keys():
        mask = (results_df['BrainRegion'] == region) & (results_df['Epoch'] == epoch)
        p_vals = results_df.loc[mask, 'P_fisher'].values
        
        if len(p_vals) > 0 and not np.all(np.isnan(p_vals)):
            _, p_corrected, _, _ = multipletests(p_vals, alpha=alpha, method='fdr_bh')
            results_df.loc[mask, 'P_fisher_FDR'] = p_corrected
            results_df.loc[mask, 'Significant_FDR'] = p_corrected < alpha
    
    return results_df
plt.rcParams['font.family'] = 'Liberation Sans'  # Best Arial alternative
plt.rcParams['svg.fonttype'] = 'none'
plt.rcParams['pdf.fonttype'] = 'truetype'
plt.rcParams['font.size'] = 11
plt.rcParams['axes.labelsize'] = 10
plt.rcParams['axes.titlesize'] = 10
plt.rcParams['xtick.labelsize'] = 7
plt.rcParams['ytick.labelsize'] = 7
plt.rcParams['legend.fontsize'] = 9
plt.rcParams['figure.titlesize'] = 14
ylim_configEpoch = {
    'combined': (-3.5, 3.5),
    'category': (-4, 4),
    'individual': (-4, 4)
}

catPalette, mmnPalette = create_mmn_colormap_for_plots()
     
# Configuration parameters
epochs_to_analyze = ['Early', 'Middle','Late']

modelsToTest = {
    # Individual Deviant Models
    'individual': 'responseVar ~ isLocLeft + isLocRight + isintensityDown + isintensityUp + isFreq_450 + isFreq_550 + isDur25 + isGap',
    
    # Category-based Models
    'category': 'responseVar ~ isLocation + isFrequency + isIntensity + isTiming',
    
    # Simple Models
    'combined': 'responseVar ~ isDeviant'
}


modelFamily = {'name': 'gaussian_identity',
               'function':sm.families.Gaussian()}

normPos = '_NormRA'  # Use the same normalization as before
model_type = 'GLM'  # within-unit, so keep GLM
n_permutations = 2000  # Number of permutations per test
significance_test = 'fisher'  # Options: 'fisher', 'binomial', or None
FRmin2 = 1  # Minimum firing rate threshold
n_jobs = -1  # Use all available cores

loadResultsEpoch = True
# Use spike density data
DF_analysis = DF_SDFmeans.copy()
normPos = normPos  # or '' for unnormalized
analysis_name = 'SDF'


# Set up output directory

epoch_output_dir = resultsFolder / f'EpochUnitAnalysis_{analysis_name}' / f'{model_type}{normPos}{modelFamily["name"]}'
if FRmin2:
    epoch_output_dir = epoch_output_dir / f'minFR{FRmin2}'

epoch_output_dir.mkdir(parents=True, exist_ok=True)

print("="*60)
print(f"EPOCH-BASED UNIT ANALYSIS - {analysis_name.upper()} DATA")
print("="*60)
print(f"Output directory: {epoch_output_dir}")
print(f"Epochs: {epochs_to_analyze}")
print(f"Model type: {model_type}")
print(f"Model family: {modelFamily['name']}")
print(f"Permutations: {n_permutations}")
print(f"Min FR: {FRmin2}")
print("="*60)

# Verify data
print(f"\nData shape: {DF_analysis.shape}")
print(f"Unique units: {DF_analysis['UniqueLabel'].nunique()}")
print(f"Brain regions: {DF_analysis['BrainRegion'].unique()}")

# Check epoch columns
epoch_columns = [f'SpikeDensity{normPos}{epoch}' for epoch in epochs_to_analyze]
missing_columns = [col for col in epoch_columns if col not in DF_analysis.columns]

if missing_columns:
    print(f"WARNING: Missing epoch columns: {missing_columns}")
    available_epochs = [epoch for epoch in epochs_to_analyze 
                       if f'SpikeDensity{normPos}{epoch}' in DF_analysis.columns]
    print(f"Using available epochs: {available_epochs}")
    epochs_to_analyze = available_epochs

if not epochs_to_analyze:
    raise ValueError("ERROR: No valid epoch columns found!")

# Store all results
all_epoch_results = {}
all_summary_tables = {}
all_population_results = {}
for model_name, baseEquation in modelsToTest.items():
    print(f"\n{'='*20} {model_name.upper()} models {'='*20}")
    
    start_time = time.time()
    
    try:
        # Create subfolder for this contrast type
        contrast_output_dir = epoch_output_dir / model_name
        contrast_output_dir.mkdir(exist_ok=True)
        
        # Check if results already exist
        results_file = contrast_output_dir / 'epoch_results.pkl'
        
        formula_rhs = baseEquation.split('~')[1].strip()
        factors = []
        for term in formula_rhs.replace('*', '+').replace(':', '+').split('+'):
            term = term.strip()
            if term not in ['1', '+', '-', '*', '/', '(', ')', ','] and 'C(' not in term:
                factors.append(term)    
                
        if results_file.exists() & loadResultsEpoch:
            print(f"Loading existing results from {results_file}")
            with open(results_file, 'rb') as f:
                epoch_results = pickle.load(f)
        else:
            print(f"Running new analysis for {model_name}...")
            
            # Run the analysis
            epoch_results = run_unit_epoch_models_with_permutation(
                DF_SDFmeans=DF_analysis,
                epochs=epochs_to_analyze,
                model_name = model_name, 
                baseEquation = baseEquation,
                normPos=normPos,
                model_type=model_type,
                modelFamily = modelFamily,
                n_permutations=n_permutations,
                n_jobs=n_jobs,
                FRmin2=FRmin2,
                factors = factors
            )
            
            # Save results
            with open(results_file, 'wb') as f:
                pickle.dump(epoch_results, f)
            
            print(f"Results saved to {results_file}")
        
        # Store in master results dictionary
        all_epoch_results[model_name] = epoch_results
        
        elapsed_time = time.time() - start_time
        print(f"Analysis completed in {elapsed_time:.1f} seconds")
            
        
        # Generate visualizations and summaries
        if epoch_results:
            print(f"Creating visualizations for {len(epoch_results)} units...")
            
            # Get all models and factors from results
            all_models = set()
            all_factors = []
            for f_i in factors:
                all_factors.append(f_i) 
            all_models.add(model_name)

            print(f"Found {len(all_factors)} factors: {list(all_factors)}")
            
            # Create summary table
            print("Creating summary table...")
            summary_table = create_epoch_unit_summary_table(
                epoch_results, 
                epochs_to_analyze, 
                list(all_models), 
                list(all_factors),
                unique_unitsTmp,  
            )
            # Run population-level significance test if specified
            population_results = None
            if significance_test == 'binomial':
                print("Running population-level binomial tests...")
                population_results = run_binomial_significance_test(summary_table, alpha=0.05)
                
                # Save results
                test_file = contrast_output_dir / 'binomial_results.csv'
                population_results.to_csv(test_file, index=False)
                
                # Print summary
                print(f"\nPOPULATION-LEVEL SIGNIFICANCE RESULTS (Binomial) for {model_name}:")
                print(f"  Total region×factor×epoch combinations: {len(population_results)}")
                sig_count = (population_results['P_binomial_FDR'] < 0.05).sum()
                print(f"  Significant after FDR correction: {sig_count}")
                if sig_count > 0:
                    print("  Significant combinations:")
                    sig_combos = population_results[population_results['P_binomial_FDR'] < 0.05]
                    for _, row in sig_combos.iterrows():
                        print(f"    {row['BrainRegion']}-{row['Factor']}-{row['Epoch']}: " +
                              f"{row['N_significant']}/{row['N_total']} units " +
                              f"(p_FDR = {row['P_binomial_FDR']:.3f})")
            
            elif significance_test == 'fisher':
                print("Running Fisher's combined test...")
                population_results = run_fishers_combined_test(summary_table, alpha=0.05)
                
                # Save results
                test_file = contrast_output_dir / 'fisher_results.csv'
                population_results.to_csv(test_file, index=False)
                
                # Print summary
                print(f"\nFISHER'S COMBINED TEST RESULTS for {model_name}:")
                print(f"  Total region×factor×epoch combinations: {len(population_results)}")
                sig_count = (population_results['P_fisher_FDR'] < 0.05).sum()
                print(f"  Significant after FDR correction: {sig_count}")
                if sig_count > 0:
                    print("  Significant combinations:")
                    sig_combos = population_results[population_results['P_fisher_FDR'] < 0.05]
                    for _, row in sig_combos.iterrows():
                        print(f"    {row['BrainRegion']}-{row['Factor']}-{row['Epoch']}: " +
                              f"{row['N_enhanced']}↑/{row['N_suppressed']}↓ " +
                              f"(χ²={row['Chi2_statistic']:.1f}, p_FDR={row['P_fisher_FDR']:.4f})")
            
            else:
                print("No population-level significance test selected.")

            # Store for combined manuscript figures (built after the loop)
            all_summary_tables[model_name] = summary_table
            all_population_results[model_name] = population_results

            for br_i in ['Amy', 'HPC']:
                # Get number of unique factors for this model
                n_factors = len(summary_table['Factor'].unique()) if model_name != 'combined' else 1
                
                # Scale width: 1 inch per factor, minimum 3 inches
                fig_width = max(3, n_factors * 0.85 + 1.8)
                fig_height = len(epochs_to_analyze) * 3.5
                
                fig, ax = plt.subplots(len(epochs_to_analyze), 1, figsize=(fig_width, fig_height))

            
                pltYvar = 'T_value'
                for ep_i, epoch in enumerate(epochs_to_analyze):
                    dfPlt =  summary_table.loc[(summary_table.BrainRegion == br_i) &
                                              (summary_table.Epoch == epoch)].copy()
                    if model_name == 'individual':
                        rcPalette = mmnPalette
                        dfPlt['Factor'] = dfPlt['Factor'].str.replace('^is', '', regex=True)
                    elif model_name == 'category':
                        rcPalette = catPalette
                    else:
                        rcPalette = 'Set2'
                    raincloud_plot(x='Factor', y=pltYvar, data=dfPlt,
                                  ax = ax[ep_i], palette=rcPalette, violin_side='right',
                                  point_style='shift', point_shift=0.2,point_size=5,
                                  point_hue='Significant', point_palette=['black', 'red'])
                    ax[ep_i].legend().set_visible(False)  # Turn off legend
                    if ep_i < len(epochs_to_analyze)-1:
                        ax[ep_i].set_xticklabels([])
                        ax[ep_i].set_xlabel('')
                    ax[ep_i].set_title(epoch)
                    ax[ep_i].axhline(y=0, color='k', linestyle='--', alpha=0.3)
                    ax[ep_i].spines[['right', 'top']].set_visible(False)
                    if model_name in ylim_configEpoch:
                        ax[ep_i].set_ylim(ylim_configEpoch[model_name])
                plt.suptitle(f'{model_name} {model_type} model ({br_i}) ')
                plt.tight_layout()
                plt.savefig(os.path.join(contrast_output_dir, f"rainPlots_{pltYvar}_{br_i}.png"), 
                           dpi=300, bbox_inches='tight')
                plt.savefig(os.path.join(contrast_output_dir, f"rainPlots_{pltYvar}_{br_i}.svg"), 
                           format='svg', bbox_inches='tight')
            # Save summary table
            summary_file = contrast_output_dir / 'epoch_summary_table.csv'
            summary_table.to_csv(summary_file, index=False)
            
            # Print summary statistics
            print(f"\nSUMMARY STATISTICS for {model_name}:")
            print(f"  Total tests: {len(summary_table)}")
            print(f"  Significant tests (p<0.05): {(summary_table['P_value'] < 0.05).sum()}")
            print(f"  Percentage significant: {(summary_table['P_value'] < 0.05).mean()*100:.1f}%")
            
            # By brain region
            if 'BrainRegion' in summary_table.columns:
                region_summary = summary_table.groupby('BrainRegion').agg({
                    'P_value': lambda x: (x < 0.05).sum(),
                    'Coefficient': 'count'
                }).rename(columns={'P_value': 'Significant', 'Coefficient': 'Total'})
                region_summary['Percentage'] = (region_summary['Significant'] / region_summary['Total'] * 100).round(1)
                print(f"\n  By brain region:")
                print(region_summary)
            
            # By epoch
            epoch_summary = summary_table.groupby('Epoch').agg({
                'P_value': lambda x: (x < 0.05).sum(),
                'Coefficient': 'count'
            }).rename(columns={'P_value': 'Significant', 'Coefficient': 'Total'})
            epoch_summary['Percentage'] = (epoch_summary['Significant'] / epoch_summary['Total'] * 100).round(1)
            print(f"\n  By epoch:")
            print(epoch_summary)
            
            regions_to_plot = ['Amy', 'HPC', 'TTG', 'Ins']
            
            # Create heatmaps with population test results if available
            if model_name == 'combined':
                figtemp = region_summary_heatmaps(
                    summary_table, 
                    model_name, 
                    str(contrast_output_dir), 
                    regions_to_plot,
                    population_test_results=population_results,
                    combined_mode=True,
                    colormap_limits=(0, 30),
                    show_directionality=True,
                    test_type=significance_test if significance_test else 'none'
                )
            else:
                figtemp = region_summary_heatmaps(
                    summary_table, 
                    model_name, 
                    str(contrast_output_dir), 
                    regions_to_plot, 
                    population_test_results=population_results,
                    combined_mode=False, 
                    colormap_limits=(0, 30),
                    show_directionality=True,
                    test_type=significance_test if significance_test else 'none'
                )
            # After running the population test, if you want to create this heatmap:
            if population_results is not None:
                create_population_test_summary_heatmap(
                    population_results, 
                    str(contrast_output_dir), 
                    model_name,
                    test_type=significance_test
                )

        else:
            print(f"No results obtained for {model_name}")
            
    except Exception as e:
        print(f"ERROR in {model_name} analysis: {str(e)}")
        import traceback
        traceback.print_exc()

#%% Supp Table: per-patient × brain-region summary

def build_patient_region_summary_table(df_units_analyzed, modulated_units_set,
                                       output_dir,outputFN = 'patient_region_summary_table', region_order=None):
    """
    Build a per-patient × brain-region summary table.

    Parameters
    ----------
    df_units_analyzed : DataFrame
        DF_UnitsAnalyzed — one row per unit; requires columns
        ID, SessionStr, BrainRegion, SU (bool), meanFR, UniqueLabel.
    modulated_units_set : set
        UniqueLabels of units significant in ≥1 epoch × factor combination.
    output_dir : Path
        Directory in which to save the CSV and PDF/SVG figures.
    region_order : list, optional
        Regions in display order. Default: ['Amy', 'HPC', 'TTG', 'Ins'].

    Returns
    -------
    table_df : DataFrame
        Long-form table (patient rows + subtotal rows).
    fig : matplotlib Figure
    """
    import matplotlib.colors as mcolors

    if region_order is None:
        region_order = ['Amy', 'HPC', 'TTG', 'Ins']

    df = df_units_analyzed.copy()
    df['SU'] = df['SU'].astype(bool)
    df['is_modulated'] = df['UniqueLabel'].isin(modulated_units_set)

    # ── build row list ─────────────────────────────────────────────────────────
    rows = []
    for region in region_order:
        reg_df = df[df['BrainRegion'] == region]
        if reg_df.empty:
            continue
        first_in_group = True                          # ← ADD THIS
        for patient_id, pat_df in reg_df.groupby('ID', sort=True):
            n_sessions = pat_df['SessionStr'].nunique()
            n_total    = len(pat_df)
            n_su       = int(pat_df['SU'].sum())
            n_mu       = int((~pat_df['SU']).sum())
            med_fr     = pat_df['meanFR'].median()
            fr_min     = pat_df['meanFR'].min()
            fr_max     = pat_df['meanFR'].max()
            n_mod      = int(pat_df['is_modulated'].sum())
            pct_mod    = n_mod / n_total * 100 if n_total > 0 else 0.0

            rows.append({
                'Region':         region,
                'BRegion':        region if first_in_group else '',   # ← ADD THIS
                'Patient':        patient_id,
                'Sessions':       n_sessions,
                'N total':        n_total,
                'N SU':           n_su,
                'N MU':           n_mu,
                'Median FR (Hz)': f'{med_fr:.1f}',
                'FR range':       f'{fr_min:.1f}\u2013{fr_max:.1f}',
                'N modulated':    n_mod,
                '% modulated':    f'{pct_mod:.1f}',
                '_is_subtotal':   False,
            })
            first_in_group = False                     # ← ADD THIS

        # ── region subtotal row ────────────────────────────────────────────────
        n_pts          = reg_df['ID'].nunique()
        n_sessions_reg = int(reg_df.groupby('ID')['SessionStr'].nunique().sum())
        n_total_r      = len(reg_df)
        n_su_r         = int(reg_df['SU'].sum())
        n_mu_r         = int((~reg_df['SU']).sum())
        med_fr_r       = reg_df['meanFR'].median()
        fr_min_r       = reg_df['meanFR'].min()
        fr_max_r       = reg_df['meanFR'].max()
        n_mod_r        = int(reg_df['is_modulated'].sum())
        pct_mod_r      = n_mod_r / n_total_r * 100 if n_total_r > 0 else 0.0

        rows.append({
            'Region':         region,
            'BRegion':  region,
            'Patient':        f'Total ({n_pts} patients)',
            'Sessions':       n_sessions_reg,
            'N total':        n_total_r,
            'N SU':           n_su_r,
            'N MU':           n_mu_r,
            'Median FR (Hz)': f'{med_fr_r:.1f}',
            'FR range':       f'{fr_min_r:.1f}\u2013{fr_max_r:.1f}',
            'N modulated':    n_mod_r,
            '% modulated':    f'{pct_mod_r:.1f}',
            '_is_subtotal':   True,
        })

    table_df = pd.DataFrame(rows).reset_index(drop=True)

    # ── save CSV (drop internal styling flag) ─────────────────────────────────
    output_dir = Path(output_dir)
    csv_path = output_dir / f'{outputFN}.csv'
    table_df.drop(columns=['_is_subtotal']).to_csv(csv_path, index=False,encoding='utf-8-sig')
    print(f'Saved summary table CSV → {csv_path}')

    # ── matplotlib figure ──────────────────────────────────────────────────────
    display_cols = ['BRegion', 'Patient', 'Sessions', 'N total', 'N SU', 'N MU',
                    'Median FR (Hz)', 'FR range', 'N modulated', '% modulated']
    col_widths   = [0.085, 0.155, 0.065, 0.075, 0.06, 0.06,
                    0.105, 0.125, 0.105, 0.095]   # sum ≈ 0.93

    # Subtle region accent colours
    region_accent = {'Amy': '#D6EAF8', 'HPC': '#D5F5E3',
                     'TTG': '#F5CBA7', 'Ins': '#E8DAEF'}

    n_data_rows  = len(table_df)
    row_h_in     = 0.21   # inches per data row
    header_h_in  = 0.28   # inches for header
    fig_h        = header_h_in + n_data_rows * row_h_in + 0.35
    fig_w        = 7.0    # ≈ 18 cm

    fig, ax = plt.subplots(figsize=(fig_w, fig_h))
    ax.axis('off')

    cell_data  = table_df[display_cols].values.tolist()
    the_table  = ax.table(
        cellText  = cell_data,
        colLabels = display_cols,
        cellLoc   = 'center',
        loc       = 'center',
        colWidths = col_widths,
    )
    the_table.auto_set_font_size(False)
    the_table.set_fontsize(7)

    # Header row (row index 0)
    for j in range(len(display_cols)):
        cell = the_table[0, j]
        cell.set_facecolor('#D8D8D8')
        cell.set_text_props(fontweight='bold')
        cell.set_edgecolor('#AAAAAA')
        cell.set_height(header_h_in / fig_h)

    # Data rows
    for i, row_ser in table_df.iterrows():
        is_sub = row_ser['_is_subtotal']
        region = row_ser['Region']
        base   = mcolors.to_rgba(region_accent.get(region, '#FFFFFF'))

        for j in range(len(display_cols)):
            cell = the_table[i + 1, j]   # +1 offsets past header
            cell.set_edgecolor('#CCCCCC')
            cell.set_height(row_h_in / fig_h)

            if is_sub:
                # Darker tint + bold for subtotal rows
                face = tuple(min(1.0, c * 0.55 + 0.45) for c in base[:3]) + (1.0,)
                cell.set_facecolor(face)
                cell.set_text_props(fontweight='bold')
            else:
                # Alternating light tint within each region
                tint = 0.28 if (i % 2 == 0) else 0.12
                face = tuple(min(1.0, c * tint + (1 - tint)) for c in base[:3]) + (1.0,)
                cell.set_facecolor(face)
                

    fig.tight_layout(pad=0.2)

    for ext in ('pdf', 'svg'):
        fig_path = output_dir / f'{outputFN}.{ext}'
        fig.savefig(fig_path, bbox_inches='tight', dpi=300)
        print(f'Saved figure → {fig_path}')

    return table_df, fig


# ── Config: which model(s) contribute to "modulated" status ───────────────────
# A unit is modulated if significant in ≥1 row in ANY of the listed models.
models_for_modulation = ['combined']
#models_for_modulation = ['category']
#models_for_modulation = ['individual']

modulated_units_R2_4 = set()
for _mname in models_for_modulation:
    _csv = epoch_output_dir / _mname / 'epoch_summary_table.csv'
    if _csv.exists():
        _st = pd.read_csv(_csv)
    elif _mname in all_summary_tables:
        _st = all_summary_tables[_mname]
    else:
        print(f'Warning: no epoch summary table found for model "{_mname}", skipping.')
        continue
    _st['BrainRegion'] = _st['BrainRegion'].replace('AIns', 'Ins')
    modulated_units_R2_4 |= set(_st.loc[_st['Significant'], 'Unit'].unique())

print(f'Modulated units ({models_for_modulation}): {len(modulated_units_R2_4)}')
outputFN = 'patient_region_summary_table_' + "".join([s[:3] for s in models_for_modulation])
tbl_R2_4, fig_R2_4 = build_patient_region_summary_table(
    DF_UnitsAnalyzed, modulated_units_R2_4, epoch_output_dir,outputFN
)
plt.show()

#%% Build 3 manuscript epoch figures (heatmap + raincloud combined)

def build_ms_epoch_figure(summary_tbl, pop_results, model_name, regions, epochs,
                          ylim_cfg, cat_palette, mmn_palette, output_dir):
    """
    Build a combined manuscript-style figure: heatmap panel(s) + raincloud panel(s).

    For 'category' and 'individual': one column per region (Amy, HPC),
    with a heatmap row on top and one raincloud row per epoch below.
    For 'combined': single heatmap (regions × epochs) + raincloud rows for Amy & HPC.
    """
    pltYvar = 'T_value'
    epochs_list = list(epochs)
    n_epochs = len(epochs_list)
    vmin, vmax = 0, 30  # colormap limits
    sig_threshold = 0.05

    def _draw_heatmap(ax, region_data, factors_original, factors_clean, epoch_list,
                      pop_res, region, factor_name=None):
        """Draw a single heatmap into ax; returns the imshow object."""
        n_e = len(epoch_list)
        n_f = len(factors_original)
        pct_mat = np.zeros((n_e, n_f))
        ann_mat = np.empty((n_e, n_f), dtype=object)

        for ei, ep in enumerate(epoch_list):
            for fi, fac in enumerate(factors_original):
                cell = region_data[(region_data.Epoch == ep) & (region_data.Factor == fac)]
                if len(cell) > 0:
                    n_tot = len(cell)
                    sig_mask = cell['P_value'] < sig_threshold
                    n_enh = (sig_mask & (cell['T_value'] > 0)).sum() if 'T_value' in cell.columns else sig_mask.sum()
                    n_sup = (sig_mask & (cell['T_value'] < 0)).sum() if 'T_value' in cell.columns else 0
                    n_resp = n_enh + n_sup
                    pct = (n_resp / n_tot) * 100
                    pct_mat[ei, fi] = pct
                    ann_txt = f"{n_enh}↑/{n_sup}↓\n({pct:.1f}%)"
                else:
                    pct_mat[ei, fi] = 0
                    ann_txt = "0/0"

                # Asterisk for FDR significance
                if pop_res is not None:
                    for pf in [fac, fac.replace('is', ''), fac + '[T.True]']:
                        match = pop_res[(pop_res.BrainRegion == region) &
                                        (pop_res.Epoch == ep) &
                                        (pop_res.Factor == pf)]
                        if len(match) > 0 and match.iloc[0].get('Significant_FDR', False):
                            ann_txt += '*'
                            break
                ann_mat[ei, fi] = ann_txt

        im = ax.imshow(pct_mat, cmap='Reds', aspect='auto', vmin=vmin, vmax=vmax,
                       interpolation='nearest')
        for ei, ep in enumerate(epoch_list):
            for fi in range(n_f):
                txt = ann_mat[ei, fi]
                col = 'white' if pct_mat[ei, fi] > vmax * 0.5 else 'black'
                ax.text(fi, ei, txt, ha='center', va='center',
                        color=col, fontsize=11, fontweight='bold')
        ax.set_xticks(range(n_f))
        ax.set_xticklabels(factors_clean, rotation=45, ha='right', fontsize=9)
        ax.set_yticks(range(n_e))
        ax.set_yticklabels(epoch_list, fontsize=7)
        ax.set_xlabel('Factor/Deviant Type', fontsize=7)
        ax.set_ylabel('Epoch', fontsize=7)
        return im

    # -------------------------------------------------------------------------
    # Figure 1: Category model — Amy (a,b) and HPC (c,d)
    # -------------------------------------------------------------------------
    if 'category' in all_summary_tables:
        st = all_summary_tables['category']
        pr = all_population_results.get('category')
        factors_orig = st['Factor'].unique().tolist()
        factors_cln = [f.replace('is', '').replace('[T.True]', '') for f in factors_orig]
        fig_regions = ['Amy', 'HPC']
        n_cols = len(fig_regions)
        n_rows_fig = 1 + n_epochs   # 1 heatmap row + n_epochs raincloud rows
        fig_w = max(8, n_cols * len(factors_orig) * 0.9 + 2)
        fig_h = 3.5 + n_epochs * 3.0
        figMS1, axes1 = plt.subplots(n_rows_fig, n_cols,
                                     figsize=(fig_w, fig_h))
        axes1 = np.array(axes1).reshape(n_rows_fig, n_cols)
        panel_labels = ['a', 'b', 'c', 'd']
        pl_idx = 0
        for ci, region in enumerate(fig_regions):
            rd = st[st.BrainRegion == region]
            # Heatmap row
            ax_hm = axes1[0, ci]
            im = _draw_heatmap(ax_hm, rd, factors_orig, factors_cln, epochs_list, pr, region)
            ax_hm.set_title(f'{region} (N = {rd["Unit"].nunique()})',
                             fontweight='bold', fontsize=9)
            ax_hm.text(-0.15, 1.05, f'({panel_labels[pl_idx]})', transform=ax_hm.transAxes,
                       fontsize=12, fontweight='bold', va='top')
            pl_idx += 1
            # Raincloud rows (one per epoch)
            for ri, epoch in enumerate(epochs_list):
                ax_rc = axes1[ri + 1, ci]
                dfPlt = rd[rd.Epoch == epoch].copy()
                dfPlt['Factor'] = dfPlt['Factor'].str.replace('^is', '', regex=True)
                raincloud_plot(x='Factor', y=pltYvar, data=dfPlt, ax=ax_rc,
                               palette=cat_palette, violin_side='right',
                               point_style='shift', point_shift=0.2, point_size=4,
                               point_hue='Significant', point_palette=['black', 'red'])
                ax_rc.legend().set_visible(False)
                ax_rc.axhline(y=0, color='k', linestyle='--', alpha=0.3)
                
                ax_rc.spines[['right', 'top']].set_visible(False)
                if 'category' in ylim_cfg:
                    ax_rc.set_ylim(ylim_cfg['category'])
                if ri < n_epochs - 1:
                    ax_rc.set_xticklabels([])
                    ax_rc.set_xticks([])
                    ax_rc.set_xlabel('')
                    ax_rc.spines['bottom'].set_visible(False)
                else:
                    ax_rc.set_xlabel('Factor', fontsize=9)
                # ── epoch label (left column only, as subplot title) ──────────
                if ci == 0:
                    ax_rc.set_title(epoch, fontsize=7, loc='right', pad=2)
                    ax_rc.set_ylabel('T value', fontsize=7)
                    ax_rc.tick_params(axis='y', labelsize=6)
                else:
                    # right column: strip entire y-axis
                    ax_rc.set_ylabel('')
                    ax_rc.set_title('')
                    ax_rc.tick_params(axis='y', left=False, labelleft=False)
                    ax_rc.spines['left'].set_visible(False)
                
                if ri == 0:
                    ax_rc.text(-0.15, 1.05, f'({panel_labels[pl_idx]})',
                               transform=ax_rc.transAxes, fontsize=12,
                               fontweight='bold', va='top')
                    pl_idx += 1

        # Shared colorbar
        cbar = figMS1.colorbar(im, ax=axes1[0, :].tolist(), shrink=0.8, pad=0.02)
        cbar.set_label('% Responsive Units', fontsize=9)
        #figMS1.suptitle('Category model — Epoch × Factor', fontsize=11, y=1.01)
        plt.tight_layout()
        figMS1.savefig(os.path.join(output_dir, 'MS_Fig_category_heatmap_raincloud.png'),
                       dpi=300, bbox_inches='tight')
        figMS1.savefig(os.path.join(output_dir, 'MS_Fig_category_heatmap_raincloud.svg'),
                       format='svg', bbox_inches='tight')
        print("Saved MS_Fig_category_heatmap_raincloud")

        # Figure 2: Combined deviant model — regions × epochs heatmap + raincloud
        # -------------------------------------------------------------------------
        if 'combined' in all_summary_tables:
            st = all_summary_tables['combined']
            pr = all_population_results.get('combined')
            fig_regions_comb = ['Amy', 'HPC']
            factors_orig = st['Factor'].unique().tolist()
            n_reg = len(fig_regions_comb)
            n_ep  = len(epochs_list)
    
            # ── figure: 80 × 170 mm ──────────────────────────────────────────────
            fig_w2 = 80  / 25.4   # ~3.15 in
            fig_h2 = 170 / 25.4   # ~6.69 in
    
            figMS2 = plt.figure(figsize=(fig_w2, fig_h2))
    
            # GridSpec: row 0 = full-width heatmap, rows 1..n_epochs = raincloud pairs
            from matplotlib.gridspec import GridSpec
            gs = GridSpec(
                1 + n_epochs, 2,
                figure=figMS2,
                height_ratios=[1.0] + [1.6] * n_epochs,
                hspace=0.45,
                wspace=0.08,
            )
    
            # ── Panel (a): heatmap spanning BOTH columns ──────────────────────────
            ax_hm = figMS2.add_subplot(gs[0, :])   # <-- spans col 0 and col 1
    
            pct_mat_comb = np.zeros((n_reg, n_ep))
            ann_mat_comb = np.empty((n_reg, n_ep), dtype=object)
    
            for ri2, region in enumerate(fig_regions_comb):
                for ei2, ep in enumerate(epochs_list):
                    cell = st[(st.BrainRegion == region) & (st.Epoch == ep)]
                    if len(cell) > 0:
                        n_tot = len(cell)
                        sig_mask = cell['P_value'] < sig_threshold
                        n_enh = (sig_mask & (cell['T_value'] > 0)).sum() \
                                if 'T_value' in cell.columns else sig_mask.sum()
                        n_sup = (sig_mask & (cell['T_value'] < 0)).sum() \
                                if 'T_value' in cell.columns else 0
                        n_resp = n_enh + n_sup
                        pct = (n_resp / n_tot) * 100
                        pct_mat_comb[ri2, ei2] = pct
                        ann_txt = f"{n_enh}↑/{n_sup}↓\n({pct:.1f}%)"
                        if pr is not None:
                            for fac in factors_orig:
                                for pf in [fac, fac.replace('is', ''), fac + '[T.True]']:
                                    match = pr[(pr.BrainRegion == region) &
                                               (pr.Epoch == ep) &
                                               (pr.Factor == pf)]
                                    if len(match) > 0 and match.iloc[0].get('Significant_FDR', False):
                                        ann_txt += '*'
                                        break
                    else:
                        pct_mat_comb[ri2, ei2] = 0
                        ann_txt = "0/0"
                    ann_mat_comb[ri2, ei2] = ann_txt
    
            im2 = ax_hm.imshow(pct_mat_comb, cmap='Reds', aspect='auto',
                                vmin=vmin, vmax=vmax, interpolation='nearest')
            for ri2 in range(n_reg):
                for ei2 in range(n_ep):
                    txt = ann_mat_comb[ri2, ei2]
                    col = 'white' if pct_mat_comb[ri2, ei2] > vmax * 0.5 else 'black'
                    ax_hm.text(ei2, ri2, txt, ha='center', va='center',
                               color=col, fontsize=7, fontweight='bold')
    
            ax_hm.set_xticks(range(n_ep))
            ax_hm.set_xticklabels(epochs_list, fontsize=7)
            ax_hm.set_yticks(range(n_reg))
            ax_hm.set_yticklabels(fig_regions_comb, fontsize=7)
            ax_hm.set_xlabel('')        # no x-label — tick labels are self-explanatory
            ax_hm.set_title('')         # no title on heatmap
            ax_hm.spines[['right', 'top']].set_visible(False)
            ax_hm.text(-0.10, 1.14, '(a)', transform=ax_hm.transAxes,
                       fontsize=8, fontweight='bold', va='top')
    
            cbar2 = figMS2.colorbar(im2, ax=ax_hm, shrink=0.85, pad=0.02)
            cbar2.set_label('% Responsive Units', fontsize=6)
            cbar2.ax.tick_params(labelsize=6)
    
            # ── Column headers: "Amy" and "HPC" centered above each raincloud col ─
            # Use the GridSpec subplot spec positions to place text in figure coords
            # We add them as titles of the first raincloud row (ri3==0), centred
            _col_header_added = {0: False, 1: False}
    
            # ── Panel (b): raincloud rows ─────────────────────────────────────────
            rc_axes = {}
            for ri3, ep in enumerate(epochs_list):
                for ci3, region in enumerate(['Amy', 'HPC']):
                    ax_rc = figMS2.add_subplot(gs[ri3 + 1, ci3])
                    rc_axes[(ri3, ci3)] = ax_rc
    
                    dfPlt = st[(st.BrainRegion == region) & (st.Epoch == ep)].copy()
                    dfPlt['Factor'] = dfPlt['Factor'].str.replace('^is', '', regex=True)
                    raincloud_plot(x='Factor', y=pltYvar, data=dfPlt, ax=ax_rc,
                                   palette='Set2', violin_side='right',
                                   point_style='shift', point_shift=0.2, point_size=3,
                                   point_hue='Significant', point_palette=['black', 'red'])
                    ax_rc.legend().set_visible(False)
                    ax_rc.axhline(y=0, color='k', linestyle='--', alpha=0.3, linewidth=0.8)
                    ax_rc.spines[['right', 'top']].set_visible(False)
                    if 'combined' in ylim_cfg:
                        ax_rc.set_ylim(ylim_cfg['combined'])
    
                    # ── epoch label (left column only, as subplot title) ──────────
                    if ci3 == 0:
                        ax_rc.set_title(ep, fontsize=7, loc='left', pad=2)
                        ax_rc.set_ylabel('T value', fontsize=7)
                        ax_rc.tick_params(axis='y', labelsize=6)
                    else:
                        # right column: strip entire y-axis
                        ax_rc.set_ylabel('')
                        ax_rc.set_title('')
                        ax_rc.tick_params(axis='y', left=False, labelleft=False)
                        ax_rc.spines['left'].set_visible(False)
    
                    # ── x-axis: only bottom row, both columns ─────────────────────
                    if ri3 < n_epochs - 1:
                        ax_rc.set_xlabel('')
                        ax_rc.set_xticklabels([])
                        ax_rc.tick_params(axis='x', bottom=False)
                        ax_rc.spines['bottom'].set_visible(False)
                    else:
                        ax_rc.set_xlabel('Factor', fontsize=7)
                        ax_rc.tick_params(axis='x', labelsize=7)
    
                    # ── (b) panel label on very first raincloud subplot ───────────
                    if ri3 == 0 and ci3 == 0:
                        ax_rc.text(-0.10, 1.30, '(b)', transform=ax_rc.transAxes,
                                   fontsize=8, fontweight='bold', va='top')
    
            # ── Centered column headers "Amy" / "HPC" above the first raincloud row
            for ci3, region in enumerate(['Amy', 'HPC']):
                ax_top = rc_axes[(0, ci3)]
                ax_top.set_title(region, fontsize=8, fontweight='bold',
                                 loc='center', pad=14)
    
            # no suptitle
            plt.savefig(os.path.join(output_dir, 'MS_Fig_combined_heatmap_raincloud.png'),
                        dpi=300, bbox_inches='tight')
            plt.savefig(os.path.join(output_dir, 'MS_Fig_combined_heatmap_raincloud.svg'),
                        format='svg', bbox_inches='tight')
            plt.close(figMS2)
            print("Saved MS_Fig_combined_heatmap_raincloud")

    # -------------------------------------------------------------------------
    # Figure 3: Individual GLM model — Amy only, 8 factors, heatmap (a) + raincloud (b)
    # -------------------------------------------------------------------------
    if 'individual' in all_summary_tables:
        st = all_summary_tables['individual']
        pr = all_population_results.get('individual')
        for region in ['Amy','HPC']:
            rd = st[st.BrainRegion == region]
            factors_orig = rd['Factor'].unique().tolist()
            factors_cln = [f.replace('is', '').replace('[T.True]', '') for f in factors_orig]
            n_rows_fig = 1 + n_epochs
            fig_w = max(6, len(factors_orig) * 0.85 + 2)
            fig_h = 3.5 + n_epochs * 3.0
            figMS3, axes3 = plt.subplots(n_rows_fig, 1, figsize=(fig_w, fig_h))
            axes3 = np.array(axes3).flatten()
            # Heatmap
            im3 = _draw_heatmap(axes3[0], rd, factors_orig, factors_cln, epochs_list, pr, region)
            n_units = rd['Unit'].nunique()
            axes3[0].set_title(f'individual GLM model ({region}) — N = {n_units}',
                                fontsize=9, fontweight='bold')
            axes3[0].text(-0.05, 1.05, '(a)', transform=axes3[0].transAxes,
                          fontsize=10, fontweight='bold', va='top')
            cbar3 = figMS3.colorbar(im3, ax=axes3[0], shrink=0.8)
            cbar3.set_label('% Responsive Units', fontsize=7)
            # Raincloud rows per epoch
            for ri4, ep in enumerate(epochs_list):
                ax_rc = axes3[ri4 + 1]
                dfPlt = rd[rd.Epoch == ep].copy()
                dfPlt['Factor'] = dfPlt['Factor'].str.replace('^is', '', regex=True)
                raincloud_plot(x='Factor', y=pltYvar, data=dfPlt, ax=ax_rc,
                               palette=mmn_palette, violin_side='right',
                               point_style='shift', point_shift=0.2, point_size=4,
                               point_hue='Significant', point_palette=['black', 'red'])
                ax_rc.legend().set_visible(False)
                ax_rc.axhline(y=0, color='k', linestyle='--', alpha=0.3)
                ax_rc.set_title(ep, fontsize=8)
                ax_rc.spines[['right', 'top']].set_visible(False)
                if 'individual' in ylim_cfg:
                    ax_rc.set_ylim(ylim_cfg['individual'])
                if ri4 < n_epochs - 1:
                    ax_rc.set_xticklabels([])
                    ax_rc.set_xlabel('')
                    ax_rc.set_xticks([])
                    ax_rc.set_xlabel('')
                    ax_rc.spines['bottom'].set_visible(False)
                else:
                    ax_rc.set_xlabel('Factor', fontsize=7)
                ax_rc.set_ylabel('T_value', fontsize=7)
                
                if ri4 == 0:
                    ax_rc.text(-0.05, 1.05, '(b)', transform=ax_rc.transAxes,
                               fontsize=10, fontweight='bold', va='top')
            figMS3.suptitle(f'individual GLM model ({region})', fontsize=11, y=1.01)
            plt.tight_layout()
            figMS3.savefig(os.path.join(output_dir, f'MS_Fig_individual_{region}_heatmap_raincloud.png'),
                           dpi=300, bbox_inches='tight')
            figMS3.savefig(os.path.join(output_dir, f'MS_Fig_individual_{region}_heatmap_raincloud.svg'),
                           format='svg', bbox_inches='tight')
            print(f"Saved MS_Fig_individual_{region}_heatmap_raincloud")


# Build the 3 manuscript figures using stored summary tables
if all_summary_tables:
    ms_fig_output = str(epoch_output_dir)
    build_ms_epoch_figure(
        summary_tbl=None,  # unused, reads from all_summary_tables directly
        pop_results=None,
        model_name=None,
        regions=['Amy', 'HPC', 'TTG', 'Ins'],
        epochs=epochs_to_analyze,
        ylim_cfg=ylim_configEpoch,
        cat_palette=catPalette,
        mmn_palette=mmnPalette,
        output_dir=ms_fig_output
    )



sigWhat = summary_table.loc[summary_table['Significant'] & (summary_table['Factor']=='isGap')]

#%% Per-Unit Heatmap Figure — Category + Combined, REFINED 
# ============================================================================
# Heatmap-only figure, refined layout:
#   Subplot grid: rows = regions (Amy, HPC), columns = epochs (Early, Mid, Late)
#   Within each subplot:
#     - rows = individual units, grouped by participant
#     - columns = 4 categories + 1 combined (5 total)
#     - color = t-value (diverging RdBu_r, shared scale)
#     - '*' centered in significant cells
#     - horizontal lines separate participants
#     - colored squares below x-axis for category identity
#     - short x-tick labels (F, I, L, T, All), no rotation
# ============================================================================

def build_per_unit_heatmap_refined(DF_UnitsAnalyzed, all_summary_tables,
                                   epochs_to_analyze, epoch_output_dir,
                                   regions_to_plot=('Amy', 'HPC'),
                                   sig_threshold=0.05,
                                   model_key='category',
                                   cat_palette=None):
    """
    Refined per-unit heatmap with combined column and colored category markers.
    model_key: 'category' or 'individual' — selects which model's factors to plot.
    """
    from matplotlib.gridspec import GridSpec
    from matplotlib.colors import TwoSlopeNorm
    import matplotlib.patches as mpatches

    st_cat = all_summary_tables.get(model_key)
    st_comb = all_summary_tables.get('combined')
    if st_cat is None:
        print("ERROR: 'category' model not found.")
        return None

    df_units = DF_UnitsAnalyzed.copy()
    df_units['SU'] = df_units['SU'].astype(bool)

    # Factor setup — categories
    factors_orig = st_cat['Factor'].unique().tolist()
    factor_clean = {f: f.replace('is', '').replace('[T.True]', '') for f in factors_orig}
    factors_display_full = [factor_clean[f] for f in factors_orig]

    # Short labels: first 1-2 letters
    short_labels_map = {
        'Frequency': 'F', 'Intensity': 'I', 'Location': 'L', 'Timing': 'T',
        'LocLeft': 'L←',  'LocRight': 'L→',
        'intensityDown': 'I↓', 'intensityUp': 'I↑',
        'Freq_450': 'F↓', 'Freq_550': 'F↑',
        'Dur25': 'D',   'Gap': 'G'
    }
    factors_short = [short_labels_map.get(fc, fc[:2]) for fc in factors_display_full]

    # Add combined as last column
    has_combined = st_comb is not None
    if has_combined:
        factors_short_all = factors_short + ['All']
        factors_display_all = factors_display_full + ['Combined']
    else:
        factors_short_all = factors_short
        factors_display_all = factors_display_full

    n_factors_cat = len(factors_orig)
    n_factors_all = len(factors_short_all)
    n_epochs = len(epochs_to_analyze)
    n_regions = len(regions_to_plot)


    # Category colors — use the paper's standard palette
    category_palette, condition_colors = create_mmn_colormap_for_plots()
    cat_colors = {
        'Frequency': category_palette['Frequency'],   # soft mint
        'Intensity': category_palette['Intensity'],   # soft apricot
        'Location':  category_palette['Location'],    # mid-blue
        'Timing':    category_palette['Timing'],      # soft lavender
        'Combined':  category_palette['isDeviant'],   # gray for pooled deviants
        # individual deviant colors — inherit from parent category
        'LocLeft':       condition_colors['LocLeft'],
        'LocRight':      condition_colors['LocRight'],
        'intensityDown': condition_colors['intensityDown'],
        'intensityUp':   condition_colors['intensityUp'],
        'Freq_450':      condition_colors['Freq_450'],
        'Freq_550':      condition_colors['Freq_550'],
        'Dur25':         condition_colors['Dur25'],
        'Gap':           condition_colors['Gap'],
    }
    unit_to_pid = df_units.set_index('UniqueLabel')['ID'].to_dict()

    # ---- Global t-value range ---------------------------------------------
    abs_max_global = 0
    for region in regions_to_plot:
        for st in [st_cat, st_comb]:
            if st is not None:
                st_reg = st[st['BrainRegion'] == region]
                if len(st_reg) > 0 and 'T_value' in st_reg.columns:
                    rmax = st_reg['T_value'].abs().max()
                    if not np.isnan(rmax):
                        abs_max_global = max(abs_max_global, rmax)
    if abs_max_global == 0:
        abs_max_global = 1.0

    # ---- Pre-compute unit orders per region -------------------------------
    region_unit_data = {}
    for region in regions_to_plot:
        st_reg = st_cat[st_cat['BrainRegion'] == region].copy()
        if len(st_reg) == 0:
            region_unit_data[region] = None
            continue

        st_reg['PID'] = st_reg['Unit'].map(unit_to_pid)
        units_in_region = st_reg['Unit'].unique()

        unit_pid_df = pd.DataFrame({
            'Unit': units_in_region,
            'PID': [unit_to_pid.get(u, 'ZZZ') for u in units_in_region]
        })
        pid_unit_counts = unit_pid_df.groupby('PID').size().sort_values(ascending=False)
        pid_order_region = pid_unit_counts.index.tolist()
        unit_pid_df['PID_rank'] = unit_pid_df['PID'].map(
            {p: i for i, p in enumerate(pid_order_region)})
        unit_pid_df = unit_pid_df.sort_values(['PID_rank', 'Unit']).reset_index(drop=True)

        ordered_units = unit_pid_df['Unit'].tolist()
        ordered_pids = unit_pid_df['PID'].tolist()
        n_units = len(ordered_units)

        boundaries = []
        pid_midpoints = {}
        current_pid = ordered_pids[0]
        start_idx = 0
        for i, pid in enumerate(ordered_pids):
            if pid != current_pid:
                boundaries.append(i)
                pid_midpoints[current_pid] = (start_idx + i - 1) / 2.0
                current_pid = pid
                start_idx = i
        pid_midpoints[current_pid] = (start_idx + n_units - 1) / 2.0

        region_unit_data[region] = {
            'st_reg': st_reg,
            'ordered_units': ordered_units,
            'ordered_pids': ordered_pids,
            'n_units': n_units,
            'boundaries': boundaries,
            'pid_midpoints': pid_midpoints,
            'pid_order': pid_order_region
        }

    # ---- Figure dimensions ------------------------------------------------
    unit_counts_per_region = []
    for region in regions_to_plot:
        rd = region_unit_data.get(region)
        unit_counts_per_region.append(rd['n_units'] if rd else 5)

    row_height_per_unit = 0.12
    fig_heights = [max(n * row_height_per_unit, 2.0) for n in unit_counts_per_region]
    total_height = sum(fig_heights)*0.6
    fig_width = len(factors_display_all)*0.6  # narrower

    fig = plt.figure(figsize=(fig_width, total_height))
    gs = GridSpec(n_regions, n_epochs + 1, figure=fig,
                  height_ratios=unit_counts_per_region,
                  width_ratios=[1.0] * n_epochs + [0.05],
                  hspace=0.05, wspace=0.35)

    last_im = None

    for reg_idx, region in enumerate(regions_to_plot):
        rd = region_unit_data.get(region)
        if rd is None:
            for ep_idx in range(n_epochs):
                ax = fig.add_subplot(gs[reg_idx, ep_idx])
                ax.text(0.5, 0.5, f'No data\n{region}',
                       transform=ax.transAxes, ha='center', va='center')
                ax.set_xticks([]); ax.set_yticks([])
            continue

        ordered_units = rd['ordered_units']
        ordered_pids = rd['ordered_pids']
        n_units = rd['n_units']
        boundaries = rd['boundaries']
        pid_midpoints = rd['pid_midpoints']
        pid_order_reg = rd['pid_order']
        st_reg = rd['st_reg']

        # Also get combined data for this region
        st_comb_reg = None
        if has_combined:
            st_comb_reg = st_comb[st_comb['BrainRegion'] == region].copy()
            if len(st_comb_reg) > 0:
                st_comb_reg['PID'] = st_comb_reg['Unit'].map(unit_to_pid)

        for ep_idx, epoch in enumerate(epochs_to_analyze):
            ax = fig.add_subplot(gs[reg_idx, ep_idx])

            # Build matrix: n_units × (n_factors_cat + 1 for combined)
            tval_mat = np.full((n_units, n_factors_all), np.nan)
            sig_mat = np.zeros((n_units, n_factors_all), dtype=bool)

            for ui, unit in enumerate(ordered_units):
                # Category factors
                for fi, factor in enumerate(factors_orig):
                    cell = st_reg[(st_reg['Unit'] == unit) &
                                  (st_reg['Epoch'] == epoch) &
                                  (st_reg['Factor'] == factor)]
                    if len(cell) > 0:
                        tval_mat[ui, fi] = cell['T_value'].values[0]
                        sig_mat[ui, fi] = cell['P_value'].values[0] < sig_threshold

                # Combined factor (last column)
                if has_combined and st_comb_reg is not None:
                    cell_c = st_comb_reg[(st_comb_reg['Unit'] == unit) &
                                         (st_comb_reg['Epoch'] == epoch)]
                    if len(cell_c) > 0:
                        tval_mat[ui, n_factors_cat] = cell_c['T_value'].values[0]
                        sig_mat[ui, n_factors_cat] = cell_c['P_value'].values[0] < sig_threshold

            # Build annotation matrix: '*' for significant, '' otherwise
            annot_mat = np.where(sig_mat, r"$\ast$", '')

            # Plot with sns.heatmap — handles centering of annotations
            tval_df = pd.DataFrame(tval_mat, columns=factors_short_all)
            annot_df = pd.DataFrame(annot_mat, columns=factors_short_all)

            sns.heatmap(tval_df, ax=ax, cmap='RdBu_r',
                       vmin=-abs_max_global, vmax=abs_max_global,
                       annot=annot_df, fmt='s', annot_kws={'fontsize': 6, 'fontweight': 'bold',
                                                            'color': 'black','ha':'center','va':'center'},
                       linewidths=0, cbar=False,
                       xticklabels=False, yticklabels=False)

            last_im = ax.collections[0]  # for colorbar reference

            # Participant boundary lines (sns.heatmap uses 0-based y from top)
            for b in boundaries:
                ax.axhline(y=b, color='black', linewidth=0.8)

            # Separator line before combined column
            if has_combined:
                ax.axvline(x=n_factors_cat, color='black',
                          linewidth=0.8, linestyle='--', alpha=0.5)

            # X-axis: short labels, no rotation, with colored squares below
            ax.set_xticks([i + 0.5 for i in range(n_factors_all)])
            if reg_idx == n_regions - 1:
                ax.set_xticklabels(factors_short_all, fontsize=8, rotation=0, ha='center')
                trans = ax.get_xaxis_transform()
                for fi, fc in enumerate(factors_display_all):
                    color = cat_colors.get(fc, 'gray')
                    rect = mpatches.Rectangle(
                        (fi + 0.1, -0.75), width=0.85, height=0.55,
                        linewidth=0, facecolor=color,
                        transform=trans, clip_on=False
                    )
                    ax.add_patch(rect)
               

                ax.tick_params(axis='x', pad=12)
            else:
                ax.set_xticklabels([])
                ax.set_xticks([]) 

            # Y-axis: participant labels on first column only
            if ep_idx == 0:
                ytick_positions = []
                ytick_labels = []
                for pid in pid_order_reg:
                    if pid in pid_midpoints:
                        ytick_positions.append(pid_midpoints[pid] + 0.5)  # +0.5 for heatmap coords
                        n_pid_units = sum(1 for p in ordered_pids if p == pid)
                        ytick_labels.append(f'{pid} ({n_pid_units})')
                ax.set_yticks(ytick_positions)
                ax.set_yticklabels(ytick_labels, fontsize=7, rotation=0)
                ax.set_ylabel(f'{region}', fontsize=11, fontweight='bold')
            else:
                ax.set_yticks([])

            # Epoch title on top row
            if reg_idx == 0:
                ax.set_title(epoch, fontsize=11, fontweight='bold')

    # ---- Shared colorbar --------------------------------------------------
    if last_im is not None:
        cbar_ax = fig.add_subplot(gs[:, -1])
        cbar = fig.colorbar(last_im, cax=cbar_ax)
        cbar.set_label('t-value', fontsize=9)
        cbar.ax.tick_params(labelsize=8)

    # ---- Legend for colored squares ---------------------------------------
    legend_patches = [mpatches.Patch(facecolor=cat_colors.get(fc, 'gray'),linewidth=0, label=fc)
                      for fc in factors_display_all]
    fig.legend(handles=legend_patches, loc='lower center',
              ncol=len(legend_patches), fontsize=7, framealpha=0.9,
              bbox_to_anchor=(0.45, -0.03),
              handlelength=3.0,handleheight=0.5, handletextpad=0.4, columnspacing=1.0)

    # ---- Save -------------------------------------------------------------

    out_dir = Path(epoch_output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)
    fname_base = f'MS_Fig_per_unit_{model_key}_heatmap_v2'
    for ext in ['png', 'svg']:
        fig.savefig(out_dir / f'{fname_base}.{ext}',
                   dpi=300, bbox_inches='tight', format=ext)
    print(f"Saved per-unit {model_key} heatmap to {out_dir}")

    plt.show()
    return fig


# ---- EXECUTE ----
if all_summary_tables and 'category' in all_summary_tables:
    fig_unit_hm_v2 = build_per_unit_heatmap_refined(
        DF_UnitsAnalyzed=DF_UnitsAnalyzed,
        all_summary_tables=all_summary_tables,
        epochs_to_analyze=epochs_to_analyze,
        epoch_output_dir=epoch_output_dir,
        regions_to_plot=('Amy', 'HPC', 'TTG', 'Ins'),
        sig_threshold=0.05,
        cat_palette=catPalette if isinstance(catPalette, dict) else None
    )
else:
    print("WARNING: 'category' model not available.")

# ---- Individual-deviant figure (separate) ----
if all_summary_tables and 'individual' in all_summary_tables:
    fig_unit_hm_indiv = build_per_unit_heatmap_refined(
        DF_UnitsAnalyzed=DF_UnitsAnalyzed,
        all_summary_tables=all_summary_tables,
        epochs_to_analyze=epochs_to_analyze,
        epoch_output_dir=epoch_output_dir,
        regions_to_plot=('Amy', 'HPC', 'TTG', 'Ins'),
        sig_threshold=0.05,
        model_key='individual',
        cat_palette=catPalette if isinstance(catPalette, dict) else None
    )
else:
    print("WARNING: 'individual' model not available.")

#%% Time-resolved LMEs/GLMs

def modelEqs_Figs(qlmode, normPos=''):
    """
    Generate model equations for MMN task based on specified mode.
    
    Parameters:
    -----------
    qlmode : str
        Type of analysis: 'combined', 'categories', 'individual', or 'hierarchical'
    normPos : str
        String suffix for the dependent variable (e.g., 'Norm', 'NormSTD')
    
    Returns:
    --------
    dict
        Dictionary of model equations
    """
    if qlmode == 'combined':
        # Simple approach: Standard vs all deviant types combined
        lmeEqs = {
            'null': f'SpikeDensity{normPos} ~ 1',
            'isDeviant': f'SpikeDensity{normPos} ~ isDeviant'
        }

    elif qlmode == 'individual':
        # Each deviant type separately
        lmeEqs = {
            'null': f'SpikeDensity{normPos} ~ 1',
            'allDeviants': f'SpikeDensity{normPos} ~ isLocLeft + isLocRight + isintensityDown + isintensityUp + isFreq_450 + isFreq_550 + isDur25 + isGap'
        }

    elif qlmode == 'categories':
        # Use categorical variables (for factorial ANOVA-style analysis)
        lmeEqs = {
            'null': f'SpikeDensity{normPos} ~ 1',
            'deviantType': f'SpikeDensity{normPos} ~ isLocation + isIntensity + isFrequency + isTiming'
        }
    elif qlmode == 'CAT+IND':
        # Use categorical variables (for factorial ANOVA-style analysis)
        lmeEqs = {
            'null': f'SpikeDensity{normPos} ~ 1',
            'cat': f'SpikeDensity{normPos} ~ isLocation + isIntensity + isFrequency + isTiming',
            'ind': f'SpikeDensity{normPos} ~ isLocLeft + isLocRight + isintensityDown + isintensityUp + isFreq_450 + isFreq_550 + isDur25 + isGap'
        }
    elif qlmode == 'prev_CAT+IND':
        # Use categorical variables (for factorial ANOVA-style analysis)
        lmeEqs = {
            'null': f'SpikeDensity{normPos} ~ 1',
            'cat': f'SpikeDensity{normPos} ~ isLocation + isIntensity  + isTiming',  # isFrequency as reference
            'ind': f'SpikeDensity{normPos} ~ isLocLeft + isLocRight + isintensityDown + isintensityUp + isFreq_450 + isDur25 + isGap'  # isFreq_550 as reference
        }   
    
    return lmeEqs


def prepare_carryover_data(DFprepIn, time_window=None, respVar='SpikeDensity', FRmin2=None):
    """
    Prepare data for carry-over analysis using only standard trials,
    coded by the PREVIOUS trial's deviant type.
    
    Uses a hierarchical approach:
    1. First, create a trial-level dictionary for trials with consistent eventID across units
    2. For inconsistent trials, use unit-specific lookup
    3. Only fill as NaN if previous trial doesn't exist for that unit
    
    Optimized version using vectorized operations.
    """
    
    # Filter for standard trials only
    DF_standards = DFprepIn[DFprepIn['eventID'] == 'Std'].copy()
    
    # Track original counts
    n_total_standards = DF_standards['trlNum'].nunique()
    n_total_timepoints = len(DF_standards)
    
    print(f"\n  === Building Previous Trial Lookup ===")
    
    # Step 1: Check which trials have consistent eventID across units
    DF_check = DFprepIn.drop_duplicates(subset=['trlNum', 'UniqueLabel'])
    
    # Check if eventID is consistent for each trial
    eventID_per_trial = DF_check.groupby('trlNum')['eventID'].agg(['first', 'nunique'])
    eventID_per_trial.columns = ['EventID', 'NumUnique']
    eventID_per_trial['Consistent'] = eventID_per_trial['NumUnique'] == 1
    
    n_consistent = eventID_per_trial['Consistent'].sum()
    n_inconsistent = (~eventID_per_trial['Consistent']).sum()
    pct_consistent = (n_consistent / len(eventID_per_trial)) * 100
    
    print(f"  Total unique trials: {len(eventID_per_trial)}")
    print(f"  Consistent trials (same eventID across units): {n_consistent} ({pct_consistent:.1f}%)")
    print(f"  Inconsistent trials: {n_inconsistent} ({100-pct_consistent:.1f}%)")
    
    # Step 2: Create trial-level dictionary for consistent trials
    consistent_trials = eventID_per_trial[eventID_per_trial['Consistent']].index
    trial_eventID_dict = eventID_per_trial.loc[consistent_trials, 'EventID'].to_dict()
    
    print(f"\n  Building previous trial dictionary for consistent trials...")
    
    # Create lookup: trlNum -> previous eventID (for consistent trials only)
    prev_trial_lookup = {}
    for trlNum in sorted(trial_eventID_dict.keys()):
        prev_trlNum = trlNum - 1
        if prev_trlNum in trial_eventID_dict:
            prev_trial_lookup[trlNum] = trial_eventID_dict[prev_trlNum]
    
    print(f"  Dictionary entries created: {len(prev_trial_lookup)}")
    
    # Step 3: For inconsistent trials, create unit-specific lookup
    inconsistent_trials = eventID_per_trial[~eventID_per_trial['Consistent']].index
    
    if len(inconsistent_trials) > 0:
        print(f"\n  Building unit-specific lookup for {len(inconsistent_trials)} inconsistent trials...")
        
        # Get trial info per unit
        trial_info = DFprepIn.groupby(['UniqueLabel', 'trlNum']).first().reset_index()
        trial_info = trial_info.sort_values(['UniqueLabel', 'trlNum'])
        
        # Create unit-specific lookup only for inconsistent trials
        unit_specific_lookup = {}
        for unit in trial_info['UniqueLabel'].unique():
            unit_trials = trial_info[trial_info['UniqueLabel'] == unit]
            unit_eventID_dict = dict(zip(unit_trials['trlNum'], unit_trials['eventID']))
            
            for trlNum in inconsistent_trials:
                if trlNum in unit_eventID_dict:
                    prev_trlNum = trlNum - 1
                    if prev_trlNum in unit_eventID_dict:
                        if unit not in unit_specific_lookup:
                            unit_specific_lookup[unit] = {}
                        unit_specific_lookup[unit][trlNum] = unit_eventID_dict[prev_trlNum]
        
        print(f"  Unit-specific lookup created for {len(unit_specific_lookup)} units")
    else:
        unit_specific_lookup = {}
    
    # Step 4: Apply lookups to standard trials (VECTORIZED)
    print(f"\n  Applying previous trial lookups to standard trials...")
    
    # Initialize column
    DF_standards['prev_eventID'] = None
    
    # Vectorized approach: map using trial dictionary first
    DF_standards['prev_eventID'] = DF_standards['trlNum'].map(prev_trial_lookup)
    
    # Count how many came from dictionary
    n_from_dict_mask = DF_standards['prev_eventID'].notna()
    n_from_dict = n_from_dict_mask.sum()
    
    # For rows still missing, try unit-specific lookup
    if len(unit_specific_lookup) > 0:
        missing_mask = DF_standards['prev_eventID'].isna()
        
        # Apply unit-specific lookup to missing values
        def lookup_unit_specific(row):
            if pd.isna(row['prev_eventID']):
                unit = row['UniqueLabel']
                trlNum = row['trlNum']
                if unit in unit_specific_lookup and trlNum in unit_specific_lookup[unit]:
                    return unit_specific_lookup[unit][trlNum]
            return row['prev_eventID']
        
        DF_standards.loc[missing_mask, 'prev_eventID'] = DF_standards[missing_mask].apply(
            lookup_unit_specific, axis=1
        )
    
    # Count final statistics
    n_from_unit = (DF_standards['prev_eventID'].notna() & ~n_from_dict_mask).sum()
    n_missing = DF_standards['prev_eventID'].isna().sum()
    
    print(f"\n  Lookup Statistics (per timepoint):")
    print(f"    From trial dictionary: {n_from_dict}")
    print(f"    From unit-specific lookup: {n_from_unit}")
    print(f"    Missing (no previous trial): {n_missing}")
    
    # Create isDeviant flag for previous trial
    DF_standards['prev_isDeviant'] = DF_standards['prev_eventID'].apply(
        lambda x: False if x == 'Std' else (True if pd.notna(x) else None)
    )
    
    # Remove trials with no previous trial data
    DF_carryover = DF_standards[DF_standards['prev_eventID'].notna()].copy()
    
    # Calculate removal statistics
    n_kept_trials = DF_carryover['trlNum'].nunique()
    n_removed_trials = n_total_standards - n_kept_trials
    pct_removed_trials = (n_removed_trials / n_total_standards) * 100 if n_total_standards > 0 else 0
    
    n_kept_timepoints = len(DF_carryover)
    n_removed_timepoints = n_total_timepoints - n_kept_timepoints
    pct_removed_timepoints = (n_removed_timepoints / n_total_timepoints) * 100 if n_total_timepoints > 0 else 0
    
    # Print detailed report
    print(f"\n  === Carry-Over Data Preparation ===")
    print(f"  Total standard trials (original): {n_total_standards}")
    print(f"  Standard trials removed (no previous trial): {n_removed_trials} ({pct_removed_trials:.1f}%)")
    print(f"  Standard trials retained: {n_kept_trials} ({100-pct_removed_trials:.1f}%)")
    print(f"\n  Total timepoints (original): {n_total_timepoints}")
    print(f"  Timepoints removed: {n_removed_timepoints} ({pct_removed_timepoints:.1f}%)")
    print(f"  Timepoints retained: {n_kept_timepoints} ({100-pct_removed_timepoints:.1f}%)")
    
    # Report breakdown of previous trial types
    if len(DF_carryover) > 0:
        trial_data = DF_carryover.groupby(['UniqueLabel', 'trlNum']).first()
        n_prev_std = (trial_data['prev_isDeviant'] == False).sum()
        n_prev_dev = (trial_data['prev_isDeviant'] == True).sum()
        pct_prev_std = (n_prev_std / n_kept_trials) * 100 if n_kept_trials > 0 else 0
        pct_prev_dev = (n_prev_dev / n_kept_trials) * 100 if n_kept_trials > 0 else 0
        
        print(f"\n  === Previous Trial Type Distribution ===")
        print(f"  Previous trial was STANDARD: {n_prev_std} trials ({pct_prev_std:.1f}%)")
        print(f"  Previous trial was DEVIANT: {n_prev_dev} trials ({pct_prev_dev:.1f}%)")
        
        if n_prev_dev > 0:
            print(f"\n  === Previous Deviant Type Breakdown ===")
            prev_dev_counts = trial_data[trial_data['prev_isDeviant'] == True]['prev_eventID'].value_counts()
            for dev_type, count in prev_dev_counts.sort_index().items():
                print(f"    {dev_type}: {count} ({count/n_prev_dev*100:.1f}%)")
    
    # KEY STEP: Replace eventID with prev_eventID so that prepare_mmn_data_for_glm
    # creates the factor columns based on PREVIOUS trial type
    DF_carryover['eventID'] = DF_carryover['prev_eventID']
    
    # Drop the prev_ columns as they're no longer needed
    DF_carryover = DF_carryover.drop(columns=['prev_eventID', 'prev_isDeviant'])
    
    print(f"\n  === Final Dataset (before prepare_mmn_data_for_glm) ===")
    print(f"  Total trials: {DF_carryover['trlNum'].nunique()}")
    print(f"  Total units: {DF_carryover['UniqueLabel'].nunique()}")
    print(f"  Total timepoints: {len(DF_carryover)}")
    
    return DF_carryover


def process_timepoint_mmn(t, lmeEqs, data_t, respVar='SpikeDensity', groupVar='UniqueLabel', 
                     model_type='LME', model_info=None):
    """
    Process a single time point for GLM/LME analysis.
    
    Parameters:
    -----------
    t : float
        The time point being processed
    lmeEqs : dict
        Dictionary of model equations to fit
    data_t : pandas.DataFrame
        DataFrame containing only the data for this specific time point
    respVar : str
        Name of the response variable (default: 'SpikeDensity')
    groupVar : str
        Name of the grouping variable for mixed models (default: 'UniqueLabel')
    model_type : str
        Type of model to fit: 'LME' or 'GLM' (default: 'LME')
    model_info : dict or None
        Additional model information for GLM fitting
        
    Returns:
    --------
    tuple
        (t, t_results) where t_results is a dictionary with model results
    """
    t_results = {}
    
    for model_name, eq in lmeEqs.items():
        try:
            # Fit appropriate model type
            if model_type == 'LME':
                model = smf.mixedlm(
                    eq,
                    data=data_t,
                    groups=data_t[groupVar],
                    re_formula="1"
                )
                result = model.fit(reml=False,maxiter=300,   gtol=1e-3)
            else:  # GLM
                result = fit_glm_model(eq, data_t, model_info)
            
            if result is not None:
                metrics_temp = compute_model_metrics(result, data_t, respVar, 
                                             model_type=model_type, group_var=groupVar)
                # Store results
                t_results[model_name] = {
                    'tvalues': result.tvalues,
                    'pvalues': result.pvalues,
                    'coef': result.params,
                    'converged': result.converged,
                    'metrics': metrics_temp
                }
                
                # Add LME-specific diagnostics
                if model_type == 'LME':
                    # Extract random effects variance
                    try:
                        re_var = result.cov_re.iloc[0, 0] if hasattr(result, 'cov_re') else np.nan
                        re_sd = np.sqrt(re_var) if not np.isnan(re_var) else np.nan
                    except (AttributeError, IndexError):
                        re_var = np.nan
                        re_sd = np.nan
                    
                    # Extract number of groups
                    try:
                        n_groups = len(result.groups) if hasattr(result, 'groups') else data_t[groupVar].nunique()
                    except:
                        n_groups = np.nan
                    
                    # Add to results dictionary
                    t_results[model_name].update({
                        'random_effects_var': re_var,
                        'random_effects_sd': re_sd,
                        'n_groups': n_groups,
                        'group_var': groupVar
                    })
                del result
        except Exception as e:
            print(f"Failed to fit model {model_name} at time {t}: {str(e)}")
            t_results[model_name] = None
    
    return t, t_results


def run_mmn_glm_analysis(DF_sdfs, model_type='GLM', whatorun='combined', selRegs=None, 
                        normPos='', smoothWindowSize=None, FDRpcrit=0.05, 
                        currFolder=None, loadResults=False, time_window=None, FRmin2=None,
                        xlim_override=None, ylim_dict=None):
    """
    Run GLM/LME analysis on MMN data with consolidated figure output.
    """
    import gc
    # Convert whatorun to list if it's a string
    if isinstance(whatorun, str):
        whatorun = [whatorun]
    
    # Set default regions if not provided
    if selRegs is None:
        selRegs = ['Amy', 'HPC', 'Ins', 'TTG']
    
    # Set model-specific parameters
    if model_type == 'LME':
        groupVar = 'ID'
        model_info = None
    else:  # GLM
        groupVar = None
        if normPos in ('_Norm', '_NormRA'):
            model_info = {
                'family': sm.families.Gaussian(),
                'link': None
            }
        else:
            model_info = {
                'family': 'zip',
                'link': None
            }
    
    # Create output folder
    if currFolder is None:
        currFolder = Path(f'Figures/{model_type}s_mmn_timeresolved')
    Path(currFolder).mkdir(parents=True, exist_ok=True)
    
    # Get factor colors using existing function
    category_palette, condition_colors = create_mmn_colormap_for_plots()
    factor_colors = (category_palette|condition_colors)
    
    # Run analysis for each contrast type and region
    for run_i in whatorun:
        # Create subfolder for this contrast type
        contrast_folder = Path(currFolder) / run_i
        Path(contrast_folder).mkdir(parents=True, exist_ok=True)
        
        # Get model equations for this contrast type
        lmeEqs = modelEqs_Figs(run_i, normPos=normPos)
        print(f"Processing {run_i} contrast type")
        
        # Prepare data for GLM/LME
        if run_i == 'prev_CAT+IND':
            # Prepare carry-over data
            DF_lme1 = prepare_carryover_data(
                DF_sdfs, 
                time_window=None,
                respVar=f'SpikeDensity{normPos}',
                FRmin2=None
            )
            DF_lme = prepare_mmn_data_for_glm(DF_lme1, deviant_contrasts=run_i,  
                                              time_window=time_window,
                                              lmeEqs=lmeEqs, respVar=f'SpikeDensity{normPos}', 
                                              FRmin2=FRmin2)
            DF_lme = DF_lme.loc[DF_lme['eventID'] != 'Std']
            
        else:
            DF_lme = prepare_mmn_data_for_glm(DF_sdfs, deviant_contrasts=run_i,  
                                              time_window=time_window,
                                              lmeEqs=lmeEqs, respVar=f'SpikeDensity{normPos}', 
                                              FRmin2=FRmin2)
        DF_lme['time'] = DF_lme['time']
        # Define time markers for MMN task (in ms)
        time_markers_all = [-500, 0, 500]

        # Apply smoothing if requested
        if smoothWindowSize:
            DF_lme = get_smoothed_data2(DF_lme,
                                         sdf_col=f'SpikeDensity{normPos}',
                                         window=smoothWindowSize)

        for BReg in selRegs:
            print(f"  Analyzing region {BReg}")
            DF_modinBReg = DF_lme.loc[DF_lme['BrainRegion'] == BReg].copy()
            DF_modinBReg['time'] = DF_modinBReg['time']*1000
            # Get time points for this region (sorted, in ms) — must match time_data_dict keys exactly
            lme_times = np.sort(DF_modinBReg['time'].unique())
            # Filter time markers to those within this region's time range
            time_markers = [m for m in time_markers_all if lme_times.min() <= m <= lme_times.max()]
            # Define output file
            results_file = f'{contrast_folder}/results_{BReg}_{normPos}.pkl'

            # Pre-filter data by time point
            time_data_dict = {}
            for t, tGroup in DF_modinBReg.groupby('time'):
                time_data_dict[t] = tGroup
            
            if loadResults:
                # Load previously saved results
                try:
                    with open(results_file, 'rb') as f:
                        results = pickle.load(f)
                    print(f"    Loaded saved results from {results_file}")
                    # Collect all factors that appear in any model at any timepoint
                    all_factors = set()
                    for model_name, model_result in results.items():
                        if model_result is not None:
                            all_factors.update(model_result['tvalues'].columns)
                except FileNotFoundError:
                    print(f"    No saved results found, computing new results")
                    loadResults = False
            
            if not loadResults:

                # Initialize storage for results
                results = {
                    model: {
                        'tvalues': pd.DataFrame(index=lme_times),
                        'pvalues': pd.DataFrame(index=lme_times),
                        'coef': pd.DataFrame(index=lme_times),
                        'aic': [],
                        'bic': [],
                        'converged': [],
                        'r2_marginal': [],
                        'r2_conditional': []
                    } for model in lmeEqs.keys()
                }
                # Add LME-specific storage
                if model_type == 'LME':
                    for model in lmeEqs.keys():
                        results[model].update({
                            'random_effects_var': [],
                            'random_effects_sd': [],
                            'n_groups': [],
                            'group_var': []
                        })
              
                # Run parallel computation for all time points
                time_results = Parallel(n_jobs=-1)(
                    delayed(process_timepoint_mmn)(t, lmeEqs, time_data_dict[t], f'SpikeDensity{normPos}',
                                                groupVar, model_type, model_info) 
                    for t in lme_times
                )
                
                # Collect all factors that appear in any model at any timepoint
                all_factors = set()
                for t, t_results in time_results:
                    for model_name, model_result in t_results.items():
                        if model_result is not None:
                            all_factors.update(model_result['tvalues'].index)
                
                # Process results
                for t, t_results in time_results:
                    for model_name, model_result in t_results.items():
                        if model_result is not None:
                            # Store t-values, p-values, and coefficients
                            for factor in all_factors:
                                if factor in model_result['tvalues']:
                                    results[model_name]['tvalues'].loc[t, factor] = model_result['tvalues'][factor]
                                    results[model_name]['pvalues'].loc[t, factor] = model_result['pvalues'][factor]
                                    results[model_name]['coef'].loc[t, factor] = model_result['coef'][factor]
                            
                            # Store model comparison metrics
                            results[model_name]['aic'].append(model_result['metrics']['AIC'])
                            results[model_name]['bic'].append(model_result['metrics']['BIC'])
                            results[model_name]['converged'].append(model_result['converged'])
                            results[model_name]['r2_marginal'].append(model_result['metrics']['R2_marginal'])
                            results[model_name]['r2_conditional'].append(model_result['metrics']['R2_conditional'])
                            # Store LME-specific metrics if present
                            if model_type == 'LME':
                                results[model_name]['random_effects_var'].append(
                                    model_result.get('random_effects_var', np.nan))
                                results[model_name]['random_effects_sd'].append(
                                    model_result.get('random_effects_sd', np.nan))
                                results[model_name]['n_groups'].append(
                                    model_result.get('n_groups', np.nan))
                                results[model_name]['group_var'].append(
                                    model_result.get('group_var', groupVar))
                        
                        else:
                            # Fill with NaN if model failed
                            for factor in all_factors:
                                results[model_name]['tvalues'].loc[t, factor] = np.nan
                                results[model_name]['pvalues'].loc[t, factor] = np.nan
                                results[model_name]['coef'].loc[t, factor] = np.nan
                            
                            # Fill metrics with NaN
                            results[model_name]['aic'].append(np.nan)
                            results[model_name]['bic'].append(np.nan)
                            results[model_name]['r2_marginal'].append(np.nan)
                            results[model_name]['r2_conditional'].append(np.nan)
                            results[model_name]['converged'].append(False)
                            # Fill LME-specific with NaN
                            if model_type == 'LME':
                                results[model_name]['random_effects_var'].append(np.nan)
                                results[model_name]['random_effects_sd'].append(np.nan)
                                results[model_name]['n_groups'].append(np.nan)
                                results[model_name]['group_var'].append(groupVar)
                
                # Save results
                with open(results_file, 'wb') as f:
                    pickle.dump(results, f)
            
            # === CREATE CONSOLIDATED FIGURE ===
            NBRunits = len(DF_modinBReg.UniqueLabel.unique())
            
            # Remove 'Intercept' from factors if present
            plotting_factors = [f.replace('[T.True]', '') for f in sorted(all_factors) if f != 'Intercept']
            
            # Define factor groupings for individual model
            deviant_categories = {
                'Location': ['isLocLeft', 'isLocRight'],
                'Intensity': ['isintensityDown', 'isintensityUp'],
                'Frequency': ['isFreq_450', 'isFreq_550'],
                'Timing': ['isDur25', 'isGap']
            }
            
            # Determine subplot organization based on contrast type
            if run_i in ['CAT+IND','prev_CAT+IND']:
                models_to_plot = ['cat', 'ind']
                # Create subplot groups combining categorical + individual factors
                subplot_groups = []
                for category, individual_factors in deviant_categories.items():
                    cat_factor = f'is{category}'  # e.g., 'isLocation'
                    # Combine categorical + individual factors for this group
                    factors_in_group = [cat_factor] + individual_factors
                    subplot_groups.append((category, factors_in_group))
                n_rows = 1 + len(subplot_groups)  # ΔAIC + 4 category groups
            elif run_i == 'individual':
                models_to_plot = ['allDeviants']
                # Group factors by category - 4 subplots with 2 traces each
                subplot_groups = [(category, factors) for category, factors in deviant_categories.items()]
                n_rows = 1 + len(subplot_groups)  # 1 for ΔAIC + category groups
            else:
                models_to_plot = [m for m in lmeEqs.keys() if m != 'null']
                # For combined/categorical: one subplot per factor
                n_rows = 1 + len(plotting_factors)  # 1 for ΔAIC + individual factors
                subplot_groups = [(f.replace('is', '').replace('[T.True]', ''), [f]) for f in plotting_factors]
            
            
            # Create figure with single column, multiple rows
            fig = plt.figure(figsize=(8.25, 3 * n_rows))
            
            # Get the non-null model name (should be only one since we have null + 1 other)
            analysis_model = [m for m in lmeEqs.keys() if m != 'null'][0]
            
            # ===============================
            # TOP SUBPLOT: ΔAIC
            # ===============================
            ax_aic = plt.subplot(n_rows, 1, 1)
            
            if run_i in ['CAT+IND','prev_CAT+IND']:
                # Plot ΔAIC for both models
                colors = ['#2c3e50', '#e74c3c']  # Dark blue for categorical, red for individual
                labels = ['Categorical', 'Individual']
                
                for i, (model_name, color, label) in enumerate(zip(models_to_plot, colors, labels)):
                    if 'null' in results and model_name in results:
                        delta_aic = np.array(results[model_name]['aic']) - np.array(results['null']['aic'])
                        convergence_issues = ~np.array(results[model_name]['converged'])
                        
                        # Plot ΔAIC
                        ax_aic.plot(lme_times, delta_aic, color=color, linewidth=2, 
                                   label=f'{label} vs null', alpha=0.8)
                        
                        # Mark convergence issues
                        if np.any(convergence_issues):
                            ax_aic.scatter(lme_times[convergence_issues], delta_aic[convergence_issues], 
                                         color=color, marker='s', s=30, alpha=0.8)
            else:
                # Single model comparison (existing logic)
                analysis_model = models_to_plot[0]
                if 'null' in results and analysis_model in results:
                    delta_aic = np.array(results[analysis_model]['aic']) - np.array(results['null']['aic'])
                    convergence_issues = ~np.array(results[analysis_model]['converged'])
                    
                    # Plot ΔAIC
                    ax_aic.plot(lme_times, delta_aic, color='#2c3e50', linewidth=2, 
                               label=f'{analysis_model} vs null')
                    
                    # Mark convergence issues
                    if np.any(convergence_issues):
                        ax_aic.scatter(lme_times[convergence_issues], delta_aic[convergence_issues], 
                                     color='red', marker='s', s=30, alpha=0.8, label='Non-convergent')

            # Add reference lines
            ax_aic.axhline(y=0, color='k', linestyle='-', alpha=0.3)
            ax_aic.axhline(y=-2, color='gray', linestyle='--', alpha=0.7, label='ΔAIC = -2')
            ax_aic.axhline(y=-4, color='gray', linestyle='--', alpha=0.7, label='ΔAIC = -4')
            ax_aic.axhline(y=-10, color='gray', linestyle=':', alpha=0.7, label='ΔAIC = -10')
            
            # Add time markers
            for tt in time_markers:
                ax_aic.axvline(x=tt, color='grey', linestyle='--', alpha=0.5)
            
            ax_aic.set_ylabel('ΔAIC')
            if run_i == 'CAT+IND':
                ax_aic.set_title(f'{BReg} ({NBRunits} units) - Model Comparison: Categorical vs Individual vs Null')
            elif run_i == 'prev_CAT+IND':
                ax_aic.set_title(f'{BReg} ({NBRunits} units) - Model Comparison: Previous deviant: Categorical vs Individual vs Null')
            else:
                ax_aic.set_title(f'{BReg} ({NBRunits} units) - Model Comparison: {models_to_plot[0]} vs null')
            ax_aic.legend(loc='best')
            ax_aic.spines['top'].set_visible(False)
            ax_aic.spines['right'].set_visible(False)
            if xlim_override is not None:
                ax_aic.set_xlim(xlim_override)
            else:
                ax_aic.set_xlim(lme_times[0], lme_times[-1])
            if ylim_dict is not None and 'deltaAIC' in ylim_dict:
                ax_aic.set_ylim(ylim_dict['deltaAIC'])
            
            # ===============================
            # FACTOR SUBPLOTS: t-values
            # ===============================
            for i, (group_name, factors_in_group) in enumerate(subplot_groups):
                ax = plt.subplot(n_rows, 1, i + 2)  # +2 because first row is ΔAIC
                
                # Collect all t-values for this subplot to determine y-limits
                all_t_vals_subplot = []

                # Plot each factor in this group
                for factor in factors_in_group:
                    factor_key = factor
                    if factor_key not in results[models_to_plot[0]]['tvalues'].columns:
                        factor_key = factor + '[T.True]'
                    
                    # Determine which model to use and styling
                    if run_i in ['CAT+IND','prev_CAT+IND']:
                        # Categorical factors use 'cat' model, individual factors use 'ind' model
                        if factor.startswith('is') and not any(factor in cat_factors for cat_factors in deviant_categories.values()):
                            # This is a categorical factor (e.g., isLocation)
                            model_to_use = 'cat'
                            line_style = '-'      # Solid line
                            line_width = 2.5
                            alpha = 0.9
                            factor_label = factor.replace('is', '') + ' (Cat)'
                        else:
                            # This is an individual factor (e.g., isLocLeft)
                            model_to_use = 'ind'
                            line_style = '--'     # Dashed line
                            line_width = 1.5
                            alpha = 0.7
                            factor_label = factor.replace('is', '') + ' (Ind)'
                    else:
                        # Single model case
                        model_to_use = models_to_plot[0]
                        line_style = '-'
                        line_width = 1.5
                        alpha = 1.0
                        factor_label = factor.replace('is', '').replace('[T.True]', '')
                    
                    # Skip if factor not found in results
                    if factor_key not in results[model_to_use]['tvalues'].columns:
                        continue
                        
                    t_vals = results[model_to_use]['tvalues'][factor_key]
                    p_vals = results[model_to_use]['pvalues'][factor_key]
                    
                    # Collect t-values for y-limit calculation
                    all_t_vals_subplot.extend(t_vals[~np.isnan(t_vals)])
                    
                    # Apply FDR correction
                    valid_p = ~np.isnan(p_vals)
                    if np.any(valid_p):
                        _, p_vals_fdr_all, _, _ = multipletests(p_vals[valid_p], alpha=0.05, method='fdr_bh')
                        p_vals_fdr = np.full_like(p_vals, np.nan)
                        p_vals_fdr[valid_p] = p_vals_fdr_all
                    else:
                        p_vals_fdr = p_vals
            
                    # Get color for this factor (remove suffix for color lookup)
                    base_factor = factor_label.replace(' (Cat)', '').replace(' (Ind)', '')
                    factor_color = factor_colors.get(base_factor, '#34495e')
                    
                    # Plot t-values
                    ax.plot(lme_times, t_vals, color=factor_color, alpha=alpha, 
                           linewidth=line_width, label=factor_label, linestyle=line_style)
                    
                    # Highlight significant timepoints with thicker lines
                    sig_mask = p_vals_fdr < FDRpcrit
                    if np.any(sig_mask):
                        thick_line_width = line_width + 2.5
                        ax.plot(lme_times, np.where(sig_mask, t_vals, np.nan), 
                                color=factor_color, alpha=alpha,
                                linewidth=thick_line_width, linestyle=line_style)
                
                
                # Set y-limits: use ylim_dict if provided, otherwise calculate dynamically
                if ylim_dict is not None and group_name in ylim_dict:
                    ax.set_ylim(ylim_dict[group_name])
                elif all_t_vals_subplot:
                    t_min = min(all_t_vals_subplot)
                    t_max = max(all_t_vals_subplot)
                    y_min = min(-2, t_min)
                    y_max = max(2, t_max)
                    
                    # Add some padding (10% of range)
                    y_range = y_max - y_min
                    y_padding = y_range * 0.1
                    y_min -= y_padding
                    y_max += y_padding
                    
                    ax.set_ylim(y_min, y_max)
                else:
                    # Fallback if no data
                    ax.set_ylim(-3, 3)
                
                # Add reference lines
                ax.axhline(y=0, color='k', linestyle='-', alpha=0.3)
                
                # Add time markers
                for tt in time_markers:
                    ax.axvline(x=tt, color='grey', linestyle='--', alpha=0.5)
                
                # Customize subplot
                ax.set_ylabel('t-value')
                ax.set_title(f'{group_name}')
                
                # Add legend if multiple factors in group
                if len(factors_in_group) > 1:
                    ax.legend(loc='upper left', fontsize=9)
                
                ax.spines['top'].set_visible(False)
                ax.spines['right'].set_visible(False)
                if xlim_override is not None:
                    ax.set_xlim(xlim_override)
                else:
                    ax.set_xlim(lme_times[0], lme_times[-1])
                
                # Only add x-label to bottom subplot
                if i == len(subplot_groups) - 1:
                    ax.set_xlabel('Time (ms)')
            
            # Adjust layout
            plt.tight_layout()
            
            # Save consolidated figure with appropriate filename
            if run_i in ['CAT+IND','prev_CAT+IND']:
                fig_suffix = 'CatInd'
            else:
                fig_suffix = run_i
            
            # Save consolidated figure
            fig.savefig(f'{contrast_folder}/consolidated_{BReg}{normPos}_{fig_suffix}.png', 
                       dpi=300, bbox_inches='tight')
            fig.savefig(f'{contrast_folder}/consolidated_{BReg}{normPos}_{fig_suffix}.svg', 
                       format="svg", bbox_inches='tight')     
            gc.collect()
            #del time_results
            #del results


def plot_MS_timeresolved(
    region_list,
    run_i,
    contrast_folder,
    normPos='_NormRA',
    xlim_override=None,
    ylim_dict=None,
    FDRpcrit=0.05,
    unit_counts=None,
    fig_width=140/25.4,  # 180 mm — double-column figure width
):
    """
    Create a combined side-by-side figure for two brain regions (2-column layout).

    Layout rules:
      - Left column:  y-axis with labels/ticks.
      - Right column: no y spine, no y ticks/ticklabels (shared scale via sharey).
      - Top row(s):   no x spine; tick marks only (no labels).
      - Bottom row:   x-axis with label.
      - Single legend placed centrally between the top and second rows.
    """
    import gc as _gc
    assert len(region_list) == 2, "Exactly two regions required"

    # Load results for both regions
    region_results = {}
    region_times   = {}
    for BReg in region_list:
        pkl_path = f'{contrast_folder}/results_{BReg}_{normPos}.pkl'
        with open(pkl_path, 'rb') as fh:
            region_results[BReg] = pickle.load(fh)
        first_model = list(region_results[BReg].keys())[0]
        region_times[BReg] = region_results[BReg][first_model]['tvalues'].index.values

    # Subplot structure (mirrors run_mmn_glm_analysis logic)
    lmeEqs = modelEqs_Figs(run_i, normPos=normPos)
    deviant_categories = {
        'Location':  ['isLocLeft',      'isLocRight'],
        'Intensity': ['isintensityDown', 'isintensityUp'],
        'Frequency': ['isFreq_450',      'isFreq_550'],
        'Timing':    ['isDur25',         'isGap'],
    }

    if run_i in ['CAT+IND', 'prev_CAT+IND']:
        models_to_plot = ['cat', 'ind']
        subplot_groups = [
            (cat, [f'is{cat}'] + inds)
            for cat, inds in deviant_categories.items()
        ]
    elif run_i == 'individual':
        models_to_plot = ['allDeviants']
        subplot_groups = list(deviant_categories.items())
    else:
        models_to_plot = [m for m in lmeEqs.keys() if m != 'null']
        all_factors = set()
        for BReg in region_list:
            for mname, mres in region_results[BReg].items():
                if mname != 'null':
                    all_factors.update(mres['tvalues'].columns)
        plotting_factors = [f for f in sorted(all_factors) if f != 'Intercept']
        subplot_groups = [
            (f.replace('is', '').replace('[T.True]', ''), [f])
            for f in plotting_factors
        ]

    n_rows = 1 + len(subplot_groups)   # row 0 = ΔAIC, rows 1+ = factor t-values

    # Colors
    category_palette, condition_colors = create_mmn_colormap_for_plots()
    fcolors = category_palette | condition_colors

    # Figure
    row_height = fig_width * 0.23   # maintain aspect ratio
    fig, axes = plt.subplots(
        n_rows, 2,
        figsize=(fig_width, row_height * n_rows),
        sharey='row',
        sharex=True,
    )
    if n_rows == 1:
        axes = axes[np.newaxis, :]

    legend_handles = []
    legend_labels  = []
    time_markers_all = [-500, 0, 500]
    for col_idx, BReg in enumerate(region_list):
        results   = region_results[BReg]
        lme_times = region_times[BReg]
        time_markers = [t for t in time_markers_all if lme_times.min() <= t <= lme_times.max()]

        n_units_str = ''
        if unit_counts and BReg in unit_counts:
            n_units_str = f' ({unit_counts[BReg]} units)'

        # --- Row 0: ΔAIC ---
        ax_aic = axes[0, col_idx]

        if run_i in ['CAT+IND', 'prev_CAT+IND']:
            aic_colors = ['#2c3e50', '#e74c3c']
            aic_labels = ['Categorical', 'Individual']
            for mname, color, label in zip(models_to_plot, aic_colors, aic_labels):
                if 'null' in results and mname in results:
                    delta_aic = (np.array(results[mname]['aic']) -
                                 np.array(results['null']['aic']))
                    h, = ax_aic.plot(lme_times, delta_aic, color=color,
                                     linewidth=1.4, alpha=0.8)
                    if col_idx == 0:
                        legend_handles.append(h)
                        legend_labels.append(f'{label} vs null')
        else:
            amodel = models_to_plot[0]
            if 'null' in results and amodel in results:
                delta_aic = (np.array(results[amodel]['aic']) -
                             np.array(results['null']['aic']))
                h, = ax_aic.plot(lme_times, delta_aic, color='#2c3e50', linewidth=1.4)
                if col_idx == 0:
                    legend_handles.append(h)
                    legend_labels.append(f'{amodel} vs null')

        ax_aic.axhline(y=0,   color='k',    linestyle='-',  alpha=0.3)
        h2  = ax_aic.axhline(y=-2,  color='gray', linestyle='--', alpha=0.7)
        h4  = ax_aic.axhline(y=-4,  color='gray', linestyle='--', alpha=0.7)
        h10 = ax_aic.axhline(y=-10, color='gray', linestyle=':',  alpha=0.7)
        if col_idx == 0:
            legend_handles += [h2, h4, h10]
            legend_labels  += ['ΔAIC = \u22122', 'ΔAIC = \u22124', 'ΔAIC = \u221210']

        for tt in time_markers:
            ax_aic.axvline(x=tt, color='grey', linestyle='--', alpha=0.5)

        ax_aic.set_title(f'{BReg}{n_units_str}', fontsize=9, fontweight='bold')
        if col_idx == 0:
            ax_aic.set_ylabel('ΔAIC', fontsize=8)
        if ylim_dict and 'deltaAIC' in ylim_dict:
            ax_aic.set_ylim(ylim_dict['deltaAIC'])
        if xlim_override:
            ax_aic.set_xlim(xlim_override)
        ax_aic.tick_params(axis='x', labelsize=8)

        # --- Rows 1+: t-value panels ---
        for i, (group_name, factors_in_group) in enumerate(subplot_groups):
            ax = axes[i + 1, col_idx]
            all_t_vals_subplot = []

            for factor in factors_in_group:
                factor_key = factor
                ref_model = ('cat' if run_i in ['CAT+IND', 'prev_CAT+IND']
                             else models_to_plot[0])
                if factor_key not in results[ref_model]['tvalues'].columns:
                    factor_key = factor + '[T.True]'

                if run_i in ['CAT+IND', 'prev_CAT+IND']:
                    is_cat = (factor.startswith('is') and
                              not any(factor in v for v in deviant_categories.values()))
                    model_to_use = 'cat' if is_cat else 'ind'
                    line_style   = '-'  if is_cat else '--'
                    line_width   = 1.25  if is_cat else 1
                    alpha        = 0.9  if is_cat else 0.7
                    factor_label = factor.replace('is', '') + (' (Cat)' if is_cat else ' (Ind)')
                else:
                    model_to_use = models_to_plot[0]
                    line_style   = '-'
                    line_width   = 1.2
                    alpha        = 1.0
                    factor_label = factor.replace('is', '').replace('[T.True]', '')

                if factor_key not in results[model_to_use]['tvalues'].columns:
                    continue

                t_vals = results[model_to_use]['tvalues'][factor_key]
                p_vals = results[model_to_use]['pvalues'][factor_key]
                all_t_vals_subplot.extend(t_vals[~np.isnan(t_vals)])

                valid_p = ~np.isnan(p_vals)
                if np.any(valid_p):
                    _, p_fdr_all, _, _ = multipletests(
                        p_vals[valid_p], alpha=0.05, method='fdr_bh'
                    )
                    p_vals_fdr = np.full_like(p_vals, np.nan)
                    p_vals_fdr[valid_p] = p_fdr_all
                else:
                    p_vals_fdr = p_vals

                base_factor  = factor_label.replace(' (Cat)', '').replace(' (Ind)', '')
                factor_color = fcolors.get(base_factor, '#34495e')

                ax.plot(lme_times, t_vals, color=factor_color, alpha=alpha,
                        linewidth=line_width, linestyle=line_style)
                sig_mask = p_vals_fdr < FDRpcrit
                if np.any(sig_mask):
                    ax.plot(lme_times, np.where(sig_mask, t_vals, np.nan),
                            color=factor_color, alpha=alpha,
                            linewidth=line_width + 1.3, linestyle=line_style)

            if ylim_dict and group_name in ylim_dict:
                ax.set_ylim(ylim_dict[group_name])
            elif all_t_vals_subplot:
                t_min = min(all_t_vals_subplot)
                t_max = max(all_t_vals_subplot)
                ylo   = min(-2, t_min)
                yhi   = max( 2, t_max)
                pad   = (yhi - ylo) * 0.1
                ax.set_ylim(ylo - pad, yhi + pad)
            else:
                ax.set_ylim(-3, 3)

            ax.axhline(y=0, color='k', linestyle='-', alpha=0.3)
            for tt in time_markers:
                ax.axvline(x=tt, color='grey', linestyle='--', alpha=0.5)

            if col_idx == 0:
                ax.set_ylabel('t-value', fontsize=8)
            if i == len(subplot_groups) - 1:
                ax.set_xlabel('Time (ms)', fontsize=8)
            if xlim_override:
                ax.set_xlim(xlim_override)
            ax.tick_params(axis='x', labelsize=8)

    # Spine / tick cleanup
    for row_idx in range(n_rows):
        for col_idx in range(2):
            ax = axes[row_idx, col_idx]
            ax.spines['top'].set_visible(False)
            if col_idx == 1:
                # Right column: strip entire y axis
                ax.spines['left'].set_visible(False)
                ax.spines['right'].set_visible(False)
                ax.tick_params(axis='y', which='both',
                               left=False, right=False, labelleft=False)
            else:
                ax.spines['right'].set_visible(False)
            if row_idx < n_rows - 1:
                # Non-bottom rows: remove x spine, keep tick marks only
                ax.spines['bottom'].set_visible(False)
                ax.tick_params(axis='x', which='both',
                               bottom=True, labelbottom=False)

    # Single horizontal legend centred between the top and second rows
    legend_y = 1.0 - 0.5 / n_rows
    fig.legend(
        legend_handles, legend_labels,
        loc='center',
        bbox_to_anchor=(0.5, legend_y),
        bbox_transform=fig.transFigure,
        frameon=True,
        fontsize=7.5,
        ncol=1,
        handlelength=2,
        columnspacing=1.2,
    )

    plt.tight_layout()
    plt.subplots_adjust(hspace=0.25, wspace=0.08)

    # Save
    pair_tag = '_'.join(region_list)
    out_base = f'{contrast_folder}/MS_{pair_tag}{normPos}_{run_i}'
    fig.savefig(out_base + '.png', dpi=300, bbox_inches='tight')
    fig.savefig(out_base + '.svg', format='svg', bbox_inches='tight')
    print(f"Saved combined figure: {out_base}")
    _gc.collect()
    return fig


 #%% Time-resolved LME analysis
model_type = 'LME'  # LME with participant as random intercept
smoothWindowSize = 0.1
time_window = [-0.1, 0.5] # [-0.1, 0.7], or [-0.1, 0.5] IF COMBINED
whatorun = ['combined'] # ['combined', 'categories', 'individual','CAT+IND']
FRmin2 = 1

resultFolder = f'{resultsFolder}/{model_type}_timeresolved'

normPos = '_NormRA' #'_NormRA', '_Norm', ''

if smoothWindowSize:
    resultFolder += '_smooth'
if FRmin2:
    resultFolder += f'minFR{FRmin2}'
selRegions = ['Amy','HPC']
ylim_shared = {
    'deltaAIC': (-11.7, 2.3),
    'Deviant': (-4, 3),        # for 'combined' analysis
}
# After creating DF_sdfs with buildSdfDF_mmn
run_mmn_glm_analysis(
    DF_sdfs,
    model_type= model_type,  # or 'LME'
    whatorun=whatorun,  # Run multiple approaches
    selRegs= selRegions,  # Analyze these regions
    normPos=normPos,  # Use block-normalized data
    smoothWindowSize=smoothWindowSize,  # 100ms smoothing window
    FDRpcrit=0.05,  # 5% FDR threshold
    currFolder=resultFolder,
    loadResults=True,
    time_window = time_window,
    FRmin2 = FRmin2,
    xlim_override=(-100, 700), 
    ylim_dict=ylim_shared
)

selRegions = ['TTG','Ins']###

# After creating DF_sdfs with buildSdfDF_mmn
run_mmn_glm_analysis(
    DF_sdfs,
    model_type= model_type,  # or 'LME'
    whatorun=whatorun,  # Run multiple approaches
    selRegs= selRegions,  # Analyze these regions
    normPos=normPos,  # Use block-normalized data
    smoothWindowSize=smoothWindowSize,  # 100ms smoothing window
    FDRpcrit=0.05,  # 5% FDR threshold
    currFolder=resultFolder,
    loadResults=False,
    time_window = time_window,
    FRmin2 = FRmin2,
    xlim_override=(-100, 700)
)


time_window = [-0.1, 0.7] #
whatorun = ['CAT+IND'] # ['combined', 'categories', 'individual','CAT+IND']
selRegions = ['Amy','HPC']
ylim_shared = {
    'deltaAIC': (-15, 15),
    'Location': (-4.5, 4.5),       # for 'CAT+IND' analysis
    'Intensity': (-4.5, 4.5),
    'Frequency': (-4.5, 4.5),
    'Timing': (-4.5, 4.5)
}
# After creating DF_sdfs with buildSdfDF_mmn
run_mmn_glm_analysis(
    DF_sdfs,
    model_type= model_type,  # or 'LME'
    whatorun=whatorun,  # Run multiple approaches
    selRegs= selRegions,  # Analyze these regions
    normPos=normPos,  # Use block-normalized data
    smoothWindowSize=smoothWindowSize,  # 100ms smoothing window
    FDRpcrit=0.05,  # 5% FDR threshold
    currFolder=resultFolder,
    loadResults=True,
    time_window = time_window,
    FRmin2 = FRmin2,
    ylim_dict=ylim_shared
)


time_window = [-0.1, 0.7] #
whatorun = ['CAT+IND'] # ['combined', 'categories', 'individual','CAT+IND']
selRegions = ['TTG','Ins']
ylim_shared = {
    'deltaAIC': (-40, 15),
    'Location': (-3.5, 3),       # for 'CAT+IND' analysis
    'Intensity': (-4, 6.5),
    'Frequency': (-3, 5),
    'Timing': (-5, 4)
}
# After creating DF_sdfs with buildSdfDF_mmn
run_mmn_glm_analysis(
    DF_sdfs,
    model_type= model_type,  # or 'LME'
    whatorun=whatorun,  # Run multiple approaches
    selRegs= selRegions,  # Analyze these regions
    normPos=normPos,  # Use block-normalized data
    smoothWindowSize=smoothWindowSize,  # 100ms smoothing window
    FDRpcrit=0.05,  # 5% FDR threshold
    currFolder=resultFolder,
    loadResults=True,
    time_window = time_window,
    FRmin2 = FRmin2,
    ylim_dict=ylim_shared
)

#%% MS region-pair figures (time-resolved ensemble analysis)
# Amy + HPC side-by-side
ylim_MS = {
    'deltaAIC': (-11.7, 2.3),
    'Deviant':  (-4, 3),
}

plot_MS_timeresolved(
    region_list=['Amy', 'HPC'],
    run_i='combined',
    contrast_folder=f'{resultFolder}/combined',
    normPos=normPos,
    xlim_override=(-100, 700),
    ylim_dict=ylim_MS,
    FDRpcrit=0.05,
)

ylim_MS = {
    'deltaAIC': (-35, 4),
    'Deviant':  (-6, 6),
}
# TTG + Ins side-by-side
plot_MS_timeresolved(
    region_list=['TTG', 'Ins'],
    run_i='combined',
    contrast_folder=f'{resultFolder}/combined',
    normPos=normPos,
    xlim_override=(-100, 700),
    ylim_dict=ylim_MS,
    FDRpcrit=0.05,
)


ylim_MS = {
    'deltaAIC': (-20, 15),
    'Location': (-5, 5),       # for 'CAT+IND' analysis
    'Intensity': (-5, 5),
    'Frequency': (-5, 5),
    'Timing': (-5, 5)
}


plot_MS_timeresolved(
    region_list=['Amy', 'HPC'],
    run_i='CAT+IND',
    contrast_folder=f'{resultFolder}/CAT+IND',
    normPos=normPos,
    xlim_override=(-100, 700),
    ylim_dict=ylim_MS,
    FDRpcrit=0.05,
)

ylim_MS = {
    'deltaAIC': (-40, 15),
    'Location': (-3.5, 4.5),       # for 'CAT+IND' analysis
    'Intensity': (-4, 6.5),
    'Frequency': (-3, 5),
    'Timing': (-5, 4)
}
# TTG + Ins side-by-side
plot_MS_timeresolved(
    region_list=['TTG', 'Ins'],
    run_i='CAT+IND',
    contrast_folder=f'{resultFolder}/CAT+IND',
    normPos=normPos,
    xlim_override=(-100, 700),
    ylim_dict=ylim_MS,
    FDRpcrit=0.05,
)

# NOTE: The duplicate "LME instead of GLMs" section was removed.
# The primary analysis above already uses LME as the default model type.

#%% Separate SU/MU analyses 


model_type = 'LME' 
time_window = [-0.1, 0.5] # [-0.1, 0.7], or [-0.1, 0.5] IF COMBINED
whatorun = ['combined'] # ['combined', 'categories', 'individual','CAT+IND']

resultFolder = f'{resultsFolder}/{model_type}_timeresolved_SU'


if smoothWindowSize:
    resultFolder += '_smooth'
if FRmin2:
    resultFolder += f'minFR{FRmin2}'
selRegions = ['Amy','HPC']
ylim_shared = {
    'deltaAIC': (-11.7, 2.3),
    'Deviant': (-4, 3),        # for 'combined' analysis
}
# After creating DF_sdfs with buildSdfDF_mmn
run_mmn_glm_analysis(
    DF_sdfs.loc[DF_sdfs['SU']],
    model_type= model_type,  # or 'LME'
    whatorun=whatorun,  # Run multiple approaches
    selRegs= selRegions,  # Analyze these regions
    normPos=normPos,  # Use block-normalized data
    smoothWindowSize=smoothWindowSize,  # 100ms smoothing window
    FDRpcrit=0.05,  # 5% FDR threshold
    currFolder=resultFolder,
    loadResults=False,
    time_window = time_window,
    FRmin2 = FRmin2,
    xlim_override=(-100, 700), 
    ylim_dict=None# ylim_shared# 
)


time_window = [-0.1, 0.7] #
whatorun = ['CAT+IND'] # ['combined', 'categories', 'individual','CAT+IND']
selRegions = ['Amy','HPC']
ylim_shared = {
    'deltaAIC': (-15, 15),
    'Location': (-4.5, 4.5),       # for 'CAT+IND' analysis
    'Intensity': (-4.5, 4.5),
    'Frequency': (-4.5, 4.5),
    'Timing': (-4.5, 4.5)
}
# After creating DF_sdfs with buildSdfDF_mmn
run_mmn_glm_analysis(
    DF_sdfs.loc[DF_sdfs['SU']],
    model_type= model_type,  # or 'LME'
    whatorun=whatorun,  # Run multiple approaches
    selRegs= selRegions,  # Analyze these regions
    normPos=normPos,  # Use block-normalized data
    smoothWindowSize=smoothWindowSize,  # 100ms smoothing window
    FDRpcrit=0.05,  # 5% FDR threshold
    currFolder=resultFolder,
    loadResults=False,
    time_window = time_window,
    FRmin2 = FRmin2,
    ylim_dict=None #ylim_shared
)


time_window = [-0.1, 0.5] # [-0.1, 0.7], or [-0.1, 0.5] IF COMBINED
whatorun = ['combined'] # ['combined', 'categories', 'individual','CAT+IND']
resultFolder = f'{resultsFolder}/{model_type}_timeresolved_MU'

if smoothWindowSize:
    resultFolder += '_smooth'
if FRmin2:
    resultFolder += f'minFR{FRmin2}'
selRegions = ['Amy','HPC']
ylim_shared = {
    'deltaAIC': (-11.7, 2.3),
    'Deviant': (-4, 3),        # for 'combined' analysis
}
# After creating DF_sdfs with buildSdfDF_mmn
run_mmn_glm_analysis(
    DF_sdfs.loc[~DF_sdfs['SU']],
    model_type= model_type,  # or 'LME'
    whatorun=whatorun,  # Run multiple approaches
    selRegs= selRegions,  # Analyze these regions
    normPos=normPos,  # Use block-normalized data
    smoothWindowSize=smoothWindowSize,  # 100ms smoothing window
    FDRpcrit=0.05,  # 5% FDR threshold
    currFolder=resultFolder,
    loadResults=False,
    time_window = time_window,
    FRmin2 = FRmin2,
    xlim_override=(-100, 700), 
    ylim_dict=None# ylim_shared# 
)


time_window = [-0.1, 0.7] #
whatorun = ['CAT+IND'] # ['combined', 'categories', 'individual','CAT+IND']
selRegions = ['Amy','HPC']
ylim_shared = {
    'deltaAIC': (-15, 15),
    'Location': (-4.5, 4.5),       # for 'CAT+IND' analysis
    'Intensity': (-4.5, 4.5),
    'Frequency': (-4.5, 4.5),
    'Timing': (-4.5, 4.5)
}
# After creating DF_sdfs with buildSdfDF_mmn
run_mmn_glm_analysis(
    DF_sdfs.loc[~DF_sdfs['SU']],
    model_type= model_type,  # or 'LME'
    whatorun=whatorun,  # Run multiple approaches
    selRegs= selRegions,  # Analyze these regions
    normPos=normPos,  # Use block-normalized data
    smoothWindowSize=smoothWindowSize,  # 100ms smoothing window
    FDRpcrit=0.05,  # 5% FDR threshold
    currFolder=resultFolder,
    loadResults=False,
    time_window = time_window,
    FRmin2 = FRmin2,
    ylim_dict=None #ylim_shared
)


#%% MS region-pair figures (time-resolved ensemble analysis), SU/MU
# Amy + HPC side-by-side

resultFolder = f'{resultsFolder}/{model_type}_timeresolved_SU'
if smoothWindowSize:
    resultFolder += '_smooth'
if FRmin2:
    resultFolder += f'minFR{FRmin2}'
ylim_MS = {
    'deltaAIC': (-15, 15),
    'Location': (-4.5, 4.5),       # for 'CAT+IND' analysis
    'Intensity': (-4.5, 4.5),
    'Frequency': (-4.5, 4.5),
    'Timing': (-4.5, 4.5)
}


plot_MS_timeresolved(
    region_list=['Amy', 'HPC'],
    run_i='CAT+IND',
    contrast_folder=f'{resultFolder}/CAT+IND',
    normPos=normPos,
    xlim_override=(-100, 700),
    ylim_dict=ylim_MS,
    FDRpcrit=0.05,
)


resultFolder = f'{resultsFolder}/{model_type}_timeresolved_MU'
if smoothWindowSize:
    resultFolder += '_smooth'
if FRmin2:
    resultFolder += f'minFR{FRmin2}'

ylim_MS = {
    'deltaAIC': (-15, 15),
    'Location': (-4.5, 4.5),       # for 'CAT+IND' analysis
    'Intensity': (-4.5, 4.5),
    'Frequency': (-4.5, 4.5),
    'Timing': (-4.5, 4.5)
}


plot_MS_timeresolved(
    region_list=['Amy', 'HPC'],
    run_i='CAT+IND',
    contrast_folder=f'{resultFolder}/CAT+IND',
    normPos=normPos,
    xlim_override=(-100, 700),
    ylim_dict=ylim_MS,
    FDRpcrit=0.05,
)


#%% Carry-over time-resolved analysis


def generate_pairwise_models(analysis_type='individual', normPos=''):
    """
    Generate model equations for pairwise carryover comparisons.
    
    Each model compares one deviant type (as previous trial) against standard
    (as previous trial) in their effects on current standard trials.
    
    Parameters:
    -----------
    analysis_type : str
        'individual' : One model per individual deviant
        'category'   : One model per deviant category
        'both'       : Both individual and categorical models
    normPos : str
        Suffix for response variable (e.g., '_Norm', '_NormRA')
    
    Returns:
    --------
    dict
        Dictionary organized as:
        {
            'prev_LocLeft': {
                'null': 'SpikeDensity ~ 1',
                'effect': 'SpikeDensity ~ isLocLeft'
            },
            'prev_LocRight': {...},
            ...
        }
    """
    
    models = {}
    
    if analysis_type in ['individual', 'both']:
        # Individual deviant types
        individual_deviants = [
            'LocLeft', 'LocRight', 
            'intensityDown', 'intensityUp',
            'Freq_450', 'Freq_550', 
            'Dur25', 'Gap'
        ]
        
        for dev in individual_deviants:
            models[f'prev_{dev}'] = {
                'null': f'SpikeDensity{normPos} ~ 1',
                'effect': f'SpikeDensity{normPos} ~ is{dev}'
            }
    
    if analysis_type in ['category', 'both']:
        # Deviant categories
        categories = ['Location', 'Intensity', 'Frequency', 'Timing']
        
        for cat in categories:
            models[f'prev_{cat}'] = {
                'null': f'SpikeDensity{normPos} ~ 1',
                'effect': f'SpikeDensity{normPos} ~ is{cat}'
            }
    
    return models


def filter_data_for_pairwise(DF_carryover, deviant_name, is_category=False):
    """
    Filter carryover data for a specific pairwise comparison.
    
    Keeps only standards that were preceded by:
    1. The specific deviant of interest, OR
    2. OTHER deviants (NOT standards)
    
    This creates a clean pairwise comparison: 
    - "previous deviant X" vs. "previous other deviants"
    
    Note: Since the task alternates Std-Dev-Std-Dev, there are no Std→Std sequences
    (except in habituation period which should be excluded).
    
    Parameters:
    -----------
    DF_carryover : pandas.DataFrame
        Full carryover dataset (standards coded by previous trial type)
    deviant_name : str
        Name of the deviant (e.g., 'LocLeft') or category (e.g., 'Location')
    is_category : bool
        Whether deviant_name refers to a category vs. individual deviant
    
    Returns:
    --------
    pandas.DataFrame
        Filtered dataset containing only relevant trials
    """
    
    # All deviant types
    all_deviants = ['LocLeft', 'LocRight', 'intensityDown', 'intensityUp',
                    'Freq_450', 'Freq_550', 'Dur25', 'Gap']
    
    if is_category:
        # For categories, define which deviants belong to this category
        category_mapping = {
            'Location': ['LocLeft', 'LocRight'],
            'Intensity': ['intensityDown', 'intensityUp'],
            'Frequency': ['Freq_450', 'Freq_550'],
            'Timing': ['Dur25', 'Gap']
        }
        
        # This category's deviants
        this_category_deviants = category_mapping[deviant_name]
        
        # Other deviants (NOT in this category, NOT standard)
        other_deviants = [d for d in all_deviants if d not in this_category_deviants]
        
        # Keep: this category OR other deviants
        mask = DF_carryover['eventID'].isin(this_category_deviants + other_deviants)
        
    else:
        # For individual deviants: keep this deviant OR all other deviants
        other_deviants = [d for d in all_deviants if d != deviant_name]
        mask = DF_carryover['eventID'].isin([deviant_name] + other_deviants)
    
    # Exclude any standards that might have slipped through (Std→Std sequences)
    filtered = DF_carryover[mask].copy()
    filtered = filtered[filtered['eventID'] != 'Std']
    
    return filtered


def process_timepoint_pairwise(t, model_eqs, data_t, respVar, groupVar, 
                               model_type, model_info, factor_name):
    """
    Process a single time point for pairwise GLM/LME analysis.
    
    Parameters:
    -----------
    t : float
        Time point being processed
    model_eqs : dict
        Dictionary with 'null' and 'effect' model equations
    data_t : pandas.DataFrame
        Data for this time point
    respVar : str
        Response variable name
    groupVar : str or None
        Grouping variable for mixed models
    model_type : str
        'GLM' or 'LME'
    model_info : dict or None
        GLM family/link info
    factor_name : str
        Name of the factor being tested (e.g., 'isLocLeft')
    
    Returns:
    --------
    tuple
        (t, t_results) where t_results contains model fit statistics
    """
    
    t_results = {}
    
    for model_name, eq in model_eqs.items():
        try:
            # Fit model
            if model_type == 'LME':
                model = smf.mixedlm(
                    eq,
                    data=data_t,
                    groups=data_t[groupVar],
                    re_formula="1"
                )
                result = model.fit(reml=False)
            else:  # GLM
                result = fit_glm_model(eq, data_t, model_info)
            
            if result is not None:
                metrics = compute_model_metrics(
                    result, data_t, respVar,
                    model_type=model_type, 
                    group_var=groupVar
                )
                
                t_results[model_name] = {
                    'tvalues': result.tvalues,
                    'pvalues': result.pvalues,
                    'coef': result.params,
                    'converged': result.converged,
                    'metrics': metrics
                }
                del result
        except Exception as e:
            print(f"Failed to fit {model_name} model at time {t}: {str(e)}")
            t_results[model_name] = None
    
    return t, t_results


def run_pairwise_timeresolved(DF_comparison, model_eqs, lme_times, 
                              model_type, normPos, groupVar, model_info,
                              comparison_name):
    """
    Run time-resolved analysis for a single pairwise comparison.
    
    Parameters:
    -----------
    DF_comparison : pandas.DataFrame
        Filtered data for this pairwise comparison
    model_eqs : dict
        Model equations ('null' and 'effect')
    lme_times : array
        Time points to analyze
    model_type : str
        'GLM' or 'LME'
    normPos : str
        Normalization suffix
    groupVar : str or None
        Grouping variable
    model_info : dict or None
        GLM configuration
    comparison_name : str
        Name of comparison (e.g., 'prev_LocLeft')
    
    Returns:
    --------
    dict
        Results with tvalues, pvalues, coef, and metrics for null and effect models
    """
    
    # Pre-filter data by time point
    time_data_dict = {}
    for t, tGroup in DF_comparison.groupby('time'):
        time_data_dict[t] = tGroup
    
    # Extract factor name from comparison
    factor_name = comparison_name.replace('prev_', '')
    factor_column = f'is{factor_name}'
    
    # Initialize results storage
    results = {
        'null': {
            'tvalues': pd.DataFrame(index=lme_times),
            'pvalues': pd.DataFrame(index=lme_times),
            'coef': pd.DataFrame(index=lme_times),
            'aic': [],
            'bic': [],
            'converged': [],
            'r2_marginal': [],
            'r2_conditional': []
        },
        'effect': {
            'tvalues': pd.DataFrame(index=lme_times),
            'pvalues': pd.DataFrame(index=lme_times),
            'coef': pd.DataFrame(index=lme_times),
            'aic': [],
            'bic': [],
            'converged': [],
            'r2_marginal': [],
            'r2_conditional': []
        }
    }
    
    # Run parallel computation
    time_results = Parallel(n_jobs=-1)(
        delayed(process_timepoint_pairwise)(
            t, model_eqs, time_data_dict[t], 
            f'SpikeDensity{normPos}', groupVar,
            model_type, model_info, factor_column
        ) 
        for t in lme_times
    )
    
    # Collect results
    for t, t_results in time_results:
        for model_name in ['null', 'effect']:
            if t_results[model_name] is not None:
                model_result = t_results[model_name]
                
                # Store coefficients, t-values, p-values for all parameters
                for param in model_result['tvalues'].index:
                    results[model_name]['tvalues'].loc[t, param] = model_result['tvalues'][param]
                    results[model_name]['pvalues'].loc[t, param] = model_result['pvalues'][param]
                    results[model_name]['coef'].loc[t, param] = model_result['coef'][param]
                
                # Store metrics
                results[model_name]['aic'].append(model_result['metrics']['AIC'])
                results[model_name]['bic'].append(model_result['metrics']['BIC'])
                results[model_name]['r2_marginal'].append(model_result['metrics']['R2_marginal'])
                results[model_name]['r2_conditional'].append(model_result['metrics']['R2_conditional'])
                results[model_name]['converged'].append(model_result['converged'])
            else:
                # Fill with NaN if model failed
                results[model_name]['aic'].append(np.nan)
                results[model_name]['bic'].append(np.nan)
                results[model_name]['r2_marginal'].append(np.nan)
                results[model_name]['r2_conditional'].append(np.nan)
                results[model_name]['converged'].append(False)
    
    return results


def create_pairwise_plots(all_results, BReg, analysis_type, contrast_folder, 
                         normPos, lme_times, time_markers, FDRpcrit=0.05):
    """
    Create comprehensive visualization for pairwise carryover comparisons.
    
    Organization:
    - Row 1: ΔAIC comparison across all pairwise tests
    - Rows 2-5: Feature-specific t-values grouped by category
    
    Parameters:
    -----------
    all_results : dict
        Complete results for all pairwise comparisons in this region
    BReg : str
        Brain region name
    analysis_type : str
        'individual', 'category', or 'both'
    contrast_folder : Path
        Output folder
    normPos : str
        Normalization suffix
    lme_times : array
        Time points
    time_markers : list
        Time markers for vertical lines
    FDRpcrit : float
        FDR significance threshold
    """
    
    # Get color mapping
    category_palette, condition_colors = create_mmn_colormap_for_plots()
    factor_colors = {**category_palette, **condition_colors}
    
    # Define category groupings
    category_groups = {
        'Location': ['LocLeft', 'LocRight'],
        'Intensity': ['intensityDown', 'intensityUp'],
        'Frequency': ['Freq_450', 'Freq_550'],
        'Timing': ['Dur25', 'Gap']
    }
    
    # Count number of units
    NBRunits = len(all_results)  # This needs to be passed differently
    
    # Determine what to plot based on analysis type
    if analysis_type == 'individual':
        subplot_groups = list(category_groups.items())
        n_rows = 1 + len(subplot_groups)  # ΔAIC + 4 categories
    elif analysis_type == 'category':
        subplot_groups = [(cat, [cat]) for cat in category_groups.keys()]
        n_rows = 1 + len(subplot_groups)  # ΔAIC + 4 categories
    else:  # 'both'
        subplot_groups = list(category_groups.items())
        n_rows = 1 + len(subplot_groups)  # ΔAIC + 4 categories
    
    # Create figure
    fig = plt.figure(figsize=(12, 4 * n_rows))
    
    # ========================================
    # ROW 1: ΔAIC COMPARISON
    # ========================================
    ax_aic = plt.subplot(n_rows, 1, 1)
    
    for comparison_name, comp_results in all_results.items():
        # Calculate ΔAIC
        delta_aic = np.array(comp_results['effect']['aic']) - np.array(comp_results['null']['aic'])
        convergence_issues = ~np.array(comp_results['effect']['converged'])
        
        # Extract deviant name for labeling and coloring
        dev_name = comparison_name.replace('prev_', '')
        color = factor_colors.get(dev_name, '#34495e')
        
        # Plot ΔAIC
        ax_aic.plot(lme_times, delta_aic, color=color, linewidth=1.5, 
                   label=f'{dev_name}', alpha=0.8)
        
        # Mark convergence issues
        if np.any(convergence_issues):
            ax_aic.scatter(lme_times[convergence_issues], 
                          delta_aic[convergence_issues],
                          color=color, marker='s', s=30, alpha=0.5)
    
    # Add reference lines
    ax_aic.axhline(y=0, color='k', linestyle='-', alpha=0.3)
    ax_aic.axhline(y=-2, color='gray', linestyle='--', alpha=0.5, label='ΔAIC = -2')
    ax_aic.axhline(y=-4, color='gray', linestyle='--', alpha=0.5, label='ΔAIC = -4')
    ax_aic.axhline(y=-10, color='gray', linestyle=':', alpha=0.5, label='ΔAIC = -10')
    ax_aic.set_xlim(lme_times[0],lme_times[-1])
    # Add time markers
    for tt in time_markers:
        ax_aic.axvline(x=tt, color='grey', linestyle='--', alpha=0.5)
    
    ax_aic.set_ylabel('ΔAIC (vs null)')
    ax_aic.set_title(f'{BReg} - Pairwise Carryover Analysis: ΔAIC for All Comparisons')
    ax_aic.legend(loc='upper left', ncol=4, fontsize=8)
    ax_aic.spines['top'].set_visible(False)
    ax_aic.spines['right'].set_visible(False)
    
    # ========================================
    # ROWS 2+: CATEGORY-SPECIFIC T-VALUES
    # ========================================
    for i, (category_name, deviants_in_category) in enumerate(subplot_groups):
        ax = plt.subplot(n_rows, 1, i + 2)
        
        # Collect all t-values for y-axis scaling
        all_t_vals = []
        
        # Plot each deviant in this category
        for dev in deviants_in_category:
            comparison_name = f'prev_{dev}'
            
            if comparison_name not in all_results:
                continue
            
            comp_results = all_results[comparison_name]
            factor_column = f'is{dev}'
            
            # Check if factor is in results
            if factor_column not in comp_results['effect']['tvalues'].columns:
                # Try with [T.True] suffix
                factor_column_alt = f'{factor_column}[T.True]'
                if factor_column_alt in comp_results['effect']['tvalues'].columns:
                    factor_column = factor_column_alt
                else:
                    continue
            
            t_vals = comp_results['effect']['tvalues'][factor_column].values
            p_vals = comp_results['effect']['pvalues'][factor_column].values
            
            # Collect for y-limit calculation
            all_t_vals.extend(t_vals[~np.isnan(t_vals)])
            
            # Apply FDR correction
            valid_p = ~np.isnan(p_vals)
            if np.any(valid_p):
                _, p_vals_fdr_all, _, _ = multipletests(
                    p_vals[valid_p], 
                    alpha=FDRpcrit, 
                    method='fdr_bh'
                )
                p_vals_fdr = np.full_like(p_vals, np.nan)
                p_vals_fdr[valid_p] = p_vals_fdr_all
            else:
                p_vals_fdr = p_vals
            
            # Get color
            color = factor_colors.get(dev, '#34495e')
            
            # Plot t-values
            ax.plot(lme_times, t_vals, color=color, alpha=0.8, 
                   linewidth=1.5, label=f'{dev}')
            
            # Highlight significant periods
            sig_mask = p_vals_fdr < FDRpcrit
            if np.any(sig_mask):
                ax.plot(lme_times, np.where(sig_mask, t_vals, np.nan),
                       color=color, alpha=0.9, linewidth=4)
        
        # Calculate dynamic y-limits
        if all_t_vals:
            t_min = min(all_t_vals)
            t_max = max(all_t_vals)
            y_min = min(-2, t_min)
            y_max = max(2, t_max)
            y_range = y_max - y_min
            y_padding = y_range * 0.1
            ax.set_ylim(y_min - y_padding, y_max + y_padding)
        else:
            ax.set_ylim(-3, 3)
        
        # Add reference lines
        ax.axhline(y=0, color='k', linestyle='-', alpha=0.3)
        
        # Add time markers
        for tt in time_markers:
            ax.axvline(x=tt, color='grey', linestyle='--', alpha=0.5)
        
        # Labels and legend
        ax.set_ylabel('t-value')
        ax.set_title(f'Previous {category_name} → Current Standard')
        # Set x-axis limits to match time range
        ax.set_xlim(lme_times[0],lme_times[-1])
        if len(deviants_in_category) > 1:
            ax.legend(loc='upper left', fontsize=9)
        
        ax.spines['top'].set_visible(False)
        ax.spines['right'].set_visible(False)
        
        # Only add x-label to bottom subplot
        if i == len(subplot_groups) - 1:
            ax.set_xlabel('Time (ms)')
    
    # Adjust layout and save
    plt.tight_layout()
    
    fig.savefig(f'{contrast_folder}/pairwise_carryover_{BReg}{normPos}_{analysis_type}.png',
               dpi=300, bbox_inches='tight')
    fig.savefig(f'{contrast_folder}/pairwise_carryover_{BReg}{normPos}_{analysis_type}.svg',
               format='svg', bbox_inches='tight')
    plt.close(fig)
    gc.collect()



def create_MS_pairwise_plots(
    region_list,
    analysis_type,
    contrast_folder,
    normPos='_NormRA',
    xlim_override=None,
    ylim_dict=None,
    FDRpcrit=0.05,
    fig_width=140/25.4,  # 180 mm — double-column figure width
):
    """
    Create a combined side-by-side carryover figure for two brain regions.

    Layout rules (same as plot_MS_timeresolved):
      - Left column:  y-axis labels/ticks.
      - Right column: no y spine, no y ticks/ticklabels (shared scale).
      - Top row(s):   no x spine; tick marks only (no labels).
      - Bottom row:   x-axis with label.
      - Single legend placed centrally between the top and second rows.

    Loads pkl files saved by run_carryover_pairwise_analysis.
    """
    import gc as _gc
    assert len(region_list) == 2, "Exactly two regions required"

    # Load saved results for both regions
    region_results = {}
    for BReg in region_list:
        pkl_path = f'{contrast_folder}/pairwise_results_{BReg}{normPos}.pkl'
        with open(pkl_path, 'rb') as fh:
            region_results[BReg] = pickle.load(fh)

    # Infer lme_times from first available comparison in first region
    first_reg_res = region_results[region_list[0]]
    first_comp    = list(first_reg_res.keys())[0]
    lme_times     = first_reg_res[first_comp]['effect']['tvalues'].index.values
    time_markers  = [t for t in [0] if lme_times.min() <= t <= lme_times.max()]

    # Category groupings (same as create_pairwise_plots)
    category_groups = {
        'Location':  ['LocLeft',      'LocRight'],
        'Intensity': ['intensityDown', 'intensityUp'],
        'Frequency': ['Freq_450',      'Freq_550'],
        'Timing':    ['Dur25',         'Gap'],
    }

    if analysis_type == 'category':
        subplot_groups = [(cat, [cat]) for cat in category_groups]
    else:  # 'individual' or 'both'
        subplot_groups = list(category_groups.items())

    n_rows = 1 + len(subplot_groups)  # row 0 = ΔAIC, rows 1+ = t-value panels

    # Colors
    category_palette, condition_colors = create_mmn_colormap_for_plots()
    fcolors = {**category_palette, **condition_colors}

    # Figure dimensions (maintain aspect ratio from create_pairwise_plots: 4/12)
    row_height = fig_width * (4 / 16)
    fig, axes = plt.subplots(
        n_rows, 2,
        figsize=(fig_width, row_height * n_rows),
        sharey='row',
        sharex=True,
    )
    if n_rows == 1:
        axes = axes[np.newaxis, :]

    legend_handles = []
    legend_labels  = []

    for col_idx, BReg in enumerate(region_list):
        results     = region_results[BReg]
        lme_t       = lme_times      # same time axis for all regions
        t_markers   = time_markers

        # --- Row 0: ΔAIC (one line per comparison) ---
        ax_aic = axes[0, col_idx]

        for comparison_name, comp_results in results.items():
            delta_aic = (np.array(comp_results['effect']['aic']) -
                         np.array(comp_results['null']['aic']))
            dev_name  = comparison_name.replace('prev_', '')
            color     = fcolors.get(dev_name, '#34495e')
            h, = ax_aic.plot(lme_t, delta_aic, color=color, linewidth=1.5, alpha=0.8)
            if col_idx == 0:
                legend_handles.append(h)
                legend_labels.append(dev_name)

        ax_aic.axhline(y=0,   color='k',    linestyle='-',  alpha=0.3)
        h2  = ax_aic.axhline(y=-2,  color='gray', linestyle='--', alpha=0.5)
        h4  = ax_aic.axhline(y=-4,  color='gray', linestyle='--', alpha=0.5)
        h10 = ax_aic.axhline(y=-10, color='gray', linestyle=':',  alpha=0.5)
        if col_idx == 0:
            legend_handles += [h2, h4, h10]
            legend_labels  += ['ΔAIC = −2', 'ΔAIC = −4', 'ΔAIC = −10']

        for tt in t_markers:
            ax_aic.axvline(x=tt, color='grey', linestyle='--', alpha=0.5)

        ax_aic.set_title(BReg, fontsize=9, fontweight='bold')
        if col_idx == 0:
            ax_aic.set_ylabel('ΔAIC', fontsize=8)
        if ylim_dict and 'deltaAIC' in ylim_dict:
            ax_aic.set_ylim(ylim_dict['deltaAIC'])
        if xlim_override:
            ax_aic.set_xlim(xlim_override)

        # --- Rows 1+: t-value panels (one row per category) ---
        for i, (category_name, deviants_in_category) in enumerate(subplot_groups):
            ax = axes[i + 1, col_idx]
            all_t_vals = []

            for dev in deviants_in_category:
                comparison_name = f'prev_{dev}'
                if comparison_name not in results:
                    continue
                comp_results  = results[comparison_name]
                factor_column = f'is{dev}'
                if factor_column not in comp_results['effect']['tvalues'].columns:
                    factor_column_alt = f'{factor_column}[T.True]'
                    if factor_column_alt in comp_results['effect']['tvalues'].columns:
                        factor_column = factor_column_alt
                    else:
                        continue

                t_vals = comp_results['effect']['tvalues'][factor_column].values
                p_vals = comp_results['effect']['pvalues'][factor_column].values
                all_t_vals.extend(t_vals[~np.isnan(t_vals)])

                valid_p = ~np.isnan(p_vals)
                if np.any(valid_p):
                    _, p_fdr_all, _, _ = multipletests(
                        p_vals[valid_p], alpha=FDRpcrit, method='fdr_bh'
                    )
                    p_vals_fdr = np.full_like(p_vals, np.nan)
                    p_vals_fdr[valid_p] = p_fdr_all
                else:
                    p_vals_fdr = p_vals

                color = fcolors.get(dev, '#34495e')
                ax.plot(lme_t, t_vals, color=color, alpha=0.8, linewidth=1.5)
                sig_mask = p_vals_fdr < FDRpcrit
                if np.any(sig_mask):
                    ax.plot(lme_t, np.where(sig_mask, t_vals, np.nan),
                            color=color, alpha=0.9, linewidth=4)

            if ylim_dict and category_name in ylim_dict:
                ax.set_ylim(ylim_dict[category_name])
            elif all_t_vals:
                t_min = min(all_t_vals)
                t_max = max(all_t_vals)
                ylo   = min(-2, t_min)
                yhi   = max( 2, t_max)
                pad   = (yhi - ylo) * 0.1
                ax.set_ylim(ylo - pad, yhi + pad)
            else:
                ax.set_ylim(-3, 3)

            ax.axhline(y=0, color='k', linestyle='-', alpha=0.3)
            for tt in t_markers:
                ax.axvline(x=tt, color='grey', linestyle='--', alpha=0.5)

            if col_idx == 0:
                ax.set_ylabel('t-value', fontsize=8)
            if i == len(subplot_groups) - 1:
                ax.set_xlabel('Time (ms)', fontsize=8)
            if xlim_override:
                ax.set_xlim(xlim_override)

    # Spine / tick cleanup
    for row_idx in range(n_rows):
        for col_idx in range(2):
            ax = axes[row_idx, col_idx]
            ax.spines['top'].set_visible(False)
            if col_idx == 1:
                ax.spines['left'].set_visible(False)
                ax.spines['right'].set_visible(False)
                ax.tick_params(axis='y', which='both',
                               left=False, right=False, labelleft=False)
            else:
                ax.spines['right'].set_visible(False)
            if row_idx < n_rows - 1:
                ax.spines['bottom'].set_visible(False)
                ax.tick_params(axis='x', which='both',
                               bottom=True, labelbottom=False)

    # Single horizontal legend centred between the top and second rows
    legend_y = 1.0 - 0.5 / n_rows
    # Split legend into two rows if many entries (≥6)
    ncols_legend = min(len(legend_handles), 3)
    fig.legend(
        legend_handles, legend_labels,
        loc='center',
        bbox_to_anchor=(0.5, legend_y),
        bbox_transform=fig.transFigure,
        frameon=True,
        fontsize=7.5,
        ncol=ncols_legend,
        handlelength=2,
        columnspacing=1.0,
    )

    plt.tight_layout()
    plt.subplots_adjust(hspace=0.2, wspace=0.08)

    # Save
    pair_tag = '_'.join(region_list)
    out_base = f'{contrast_folder}/MS_{pair_tag}{normPos}_{analysis_type}'
    fig.savefig(out_base + '.png', dpi=300, bbox_inches='tight')
    fig.savefig(out_base + '.svg', format='svg', bbox_inches='tight')
    print(f"Saved MS carryover figure: {out_base}")
    _gc.collect()
    return fig


def run_carryover_pairwise_analysis(DF_in, model_type='GLM', analysis_type='individual',
                                    selRegs=None, normPos='', smoothWindowSize=None,
                                    FDRpcrit=0.05, currFolder=None, loadResults=False,
                                    time_window=None, FRmin2=None):
    """
    Run pairwise carryover analysis: each previous deviant type vs. previous standard.
    
    This function performs separate analyses for each deviant type, comparing standards
    that were preceded by that deviant to standards preceded by other standards.
    
    Parameters:
    -----------
    DF_sdfs : pandas.DataFrame
        Full spike density dataframe
    model_type : str
        'GLM' or 'LME'
    analysis_type : str
        'individual' : One model per individual deviant (8 comparisons)
        'category'   : One model per deviant category (4 comparisons)
        'both'       : Both individual and categorical analyses
    selRegs : list or None
        Brain regions to analyze
    normPos : str
        Normalization suffix (e.g., '_NormRA')
    smoothWindowSize : float or None
        Smoothing window in seconds
    FDRpcrit : float
        FDR significance threshold
    currFolder : str or Path or None
        Output directory
    loadResults : bool
        Whether to load previously saved results
    time_window : list or None
        Time window to analyze [start, end]
    FRmin2 : float or None
        Minimum firing rate threshold
    
    Returns:
    --------
    dict
        Complete results organized by region and comparison
    """
    
    # Import the carryover data preparation function    
    # Set default regions
    if selRegs is None:
        selRegs = ['Amy', 'HPC', 'Ins', 'TTG']
    
    # Set model-specific parameters
    if model_type == 'LME':
        groupVar = 'ID'
        model_info = None
    else:  # GLM
        groupVar = None
        if normPos in ('_Norm', '_NormRA'):
            model_info = {
                'family': sm.families.Gaussian(),
                'link': None
            }
        else:
            model_info = {
                'family': 'zip',
                'link': None
            }
    
    # Create output folder
    if currFolder is None:
        currFolder = Path(f'Figures/{model_type}s_mmn_carryover_pairwise')
    Path(currFolder).mkdir(parents=True, exist_ok=True)
    
    # Create subfolder for this analysis type
    contrast_folder = Path(currFolder) / analysis_type
    Path(contrast_folder).mkdir(parents=True, exist_ok=True)
    
    print(f"\n{'='*70}")
    print(f"PAIRWISE CARRYOVER ANALYSIS: {analysis_type}")
    print(f"{'='*70}")
    
    # Step 1: Prepare carryover data
    print("\nStep 1: Preparing carryover data...")
    DF_carryover = prepare_carryover_data(
        DF_in,
        time_window=time_window,
        respVar=f'SpikeDensity{normPos}',
        FRmin2=FRmin2
    )

    # Create dummy equations that include ALL possible factors
    dummy_lmeEqs = {
        'dummy': f'SpikeDensity{normPos} ~ isLocLeft + isLocRight + isintensityDown + isintensityUp + isFreq_450 + isFreq_550 + isDur25 + isGap + isLocation + isIntensity + isFrequency + isTiming'
    }
    # Add contrast columns
    DF_carryover = prepare_mmn_data_for_glm(
        DF_carryover,
        deviant_contrasts='individual',
        time_window=time_window,
        lmeEqs=dummy_lmeEqs,  
        respVar=f'SpikeDensity{normPos}',
        FRmin2=FRmin2
    )
    

    
    # Apply smoothing if requested
    if smoothWindowSize:
        print(f"Applying smoothing (window={smoothWindowSize}s)...")
        DF_carryover = get_smoothed_data2(
            DF_carryover,
            sdf_col=f'SpikeDensity{normPos}',
            window=smoothWindowSize
        )
    
    # Step 2: Generate pairwise models
    print("\nStep 2: Generating pairwise model equations...")
    pairwise_models = generate_pairwise_models(
        analysis_type=analysis_type,
        normPos=normPos
    )
    print(f"Generated {len(pairwise_models)} pairwise comparisons")
    # Get time points
    DF_carryover['time'] = DF_carryover['time']*1000
    lme_times = DF_carryover['time'].unique()
    time_markers = [-500, 0, 500]
    time_markers = [m for m in time_markers if lme_times.min() <= m <= lme_times.max()]
    # Step 3: Run analysis for each region
    all_region_results = {}
    
    for BReg in selRegs:
        print(f"\n{'='*70}")
        print(f"REGION: {BReg}")
        print(f"{'='*70}")
        
        DF_region = DF_carryover[DF_carryover['BrainRegion'] == BReg].copy()
        if BReg in ['Ins']:
            DF_region = DF_region.loc[DF_region['UniqueLabel'].str[-2:]=='S1'] # S2 for P02 has SOA = 1s
        if len(DF_region) == 0:
            print(f"No data for region {BReg}, skipping...")
            continue
        
        # Count units
        n_units = DF_region['UniqueLabel'].nunique()
        print(f"Number of units: {n_units}")
        
        # Define results file
        results_file = f'{contrast_folder}/pairwise_results_{BReg}{normPos}.pkl'
        
        if loadResults and Path(results_file).exists():
            print(f"Loading saved results from {results_file}")
            with open(results_file, 'rb') as f:
                region_results = pickle.load(f)
        else:
            print("Running pairwise analyses...")
            region_results = {}
            
            # Run each pairwise comparison
            for comparison_name, model_eqs in pairwise_models.items():
                print(f"\n  Processing: {comparison_name}")
                
                # Extract deviant name
                dev_name = comparison_name.replace('prev_', '')
                is_category = dev_name in ['Location', 'Intensity', 'Frequency', 'Timing']
                
                # Filter data for this comparison
                DF_comparison = filter_data_for_pairwise(
                    DF_region, 
                    dev_name, 
                    is_category=is_category
                )
                
                n_trials = DF_comparison['trlNum'].nunique()
                print(f"    Trials in comparison: {n_trials}")
                
                if n_trials < 10:
                    print(f"    Too few trials, skipping...")
                    continue
                
                # Run time-resolved analysis
                comparison_results = run_pairwise_timeresolved(
                    DF_comparison=DF_comparison,
                    model_eqs=model_eqs,
                    lme_times=lme_times,
                    model_type=model_type,
                    normPos=normPos,
                    groupVar=groupVar,
                    model_info=model_info,
                    comparison_name=comparison_name
                )
                
                region_results[comparison_name] = comparison_results
            
            # Save results
            print(f"\nSaving results to {results_file}")
            with open(results_file, 'wb') as f:
                pickle.dump(region_results, f)
        
        # Store for final output
        all_region_results[BReg] = region_results
        
        # Step 4: Create visualizations
        print(f"\nCreating visualizations for {BReg}...")
        create_pairwise_plots(
            all_results=region_results,
            BReg=BReg,
            analysis_type=analysis_type,
            contrast_folder=contrast_folder,
            normPos=normPos,
            lme_times=lme_times,
            time_markers=time_markers,
            FDRpcrit=FDRpcrit
        )
        
        print(f"Done: Completed analysis for {BReg}")
        gc.collect()
    
    print(f"\n{'='*70}")
    print("PAIRWISE CARRYOVER ANALYSIS COMPLETE")
    print(f"Results saved to: {contrast_folder}")
    print(f"{'='*70}\n")
    
    return all_region_results 

# Configuration
model_type = 'LME'  # LME with participant as random intercept
analysis_type = 'individual'  # Run both individual and categorical
normPos = '_NormRA'
smoothWindowSize = 0.1
time_window = [-0.4, 0.4]
FRmin2 = 1
FDRpcrit = 0.05
# Select regions
selRegions = ['Amy', 'HPC', 'Ins', 'TTG']

# Output folder
resultFolder = f'{resultsFolder}/Carryover_Pairwise_{model_type}_timeresolved'
if smoothWindowSize:
    resultFolder += '_smooth'
if FRmin2:
    resultFolder += f'_minFR{FRmin2}'

# Run analysis
results = run_carryover_pairwise_analysis(
    DF_sdfs,
    model_type=model_type,
    analysis_type=analysis_type,
    selRegs=selRegions,
    normPos=normPos,
    smoothWindowSize=smoothWindowSize,
    FDRpcrit=FDRpcrit,
    currFolder=resultFolder,
    loadResults=True,
    time_window=time_window,
    FRmin2=FRmin2
)
analysis_type = 'category'  # Run both individual and categorical
# Run analysis
results = run_carryover_pairwise_analysis(
    DF_sdfs,
    model_type=model_type,
    analysis_type=analysis_type,
    selRegs=selRegions,
    normPos=normPos,
    smoothWindowSize=smoothWindowSize,
    FDRpcrit=FDRpcrit,
    currFolder=resultFolder,
    loadResults=True,
    time_window=time_window,
    FRmin2=FRmin2
)

#%% MS region-pair figures (carryover analysis)
ylim_shared = {
    'deltaAIC': (-15, 5),
    'Location': (-4, 4),       # for 'CAT+IND' analysis
    'Intensity': (-4, 4),
    'Frequency': (-4, 4),
    'Timing': (-4, 4)
}

# Amy + HPC side-by-side
create_MS_pairwise_plots(
    region_list=['Amy', 'HPC'],
    analysis_type='individual',
    contrast_folder=f'{resultFolder}/individual',
    normPos=normPos,
    xlim_override=(-400, 400),
    ylim_dict=ylim_shared,
    FDRpcrit=0.05,
)
ylim_shared = {
    'deltaAIC': (-20, 5),
    'Location': (-4.5, 4.5),       # for 'CAT+IND' analysis
    'Intensity': (-4.5, 4.5),
    'Frequency': (-4.5, 4.5),
    'Timing': (-4.5, 4.5)
}
# TTG + AIns side-by-side
create_MS_pairwise_plots(
    region_list=['TTG', 'Ins'],
    analysis_type='individual',
    contrast_folder=f'{resultFolder}/individual',
    normPos=normPos,
    xlim_override=(-400, 400),
    ylim_dict=ylim_shared,
    FDRpcrit=0.05,
)

#%% =============================================================================
# SUPPLEMENTARY ANALYSIS: LOW FIRING RATE (LowFR) UNIT DEVIANCE DETECTION
# =============================================================================
# Units excluded from main analysis due to low FR (<1 Hz, minFRcrit) are
# characterised here using non-parametric tests on raw spike counts (FR_Total).
#
# Stage 1 (all low-FR units):
#   - Mann-Whitney U: spike counts (Std vs all deviants pooled)
#   - Fisher's exact: proportion of trials with >=1 spike (Std vs deviants)
#   - Population binomial test vs 5% chance level
#
# Stage 2 (Stage 1 passers only):
#   - Kruskal-Wallis across conditions
#   - Post-hoc MWU (Std vs each condition) with FDR correction
#
# Outputs:
#   LowFR_Supplementary/LowFR_unit_summary.csv
#   LowFR_Supplementary/{UniqueLabel}_raster_dist.png/svg  (Stage 1 passers)
# =============================================================================

# ---------------------------------------------------------------------------
# Helper: raster + spike-count distribution figure for one low-FR unit
# ---------------------------------------------------------------------------
def plot_lowfr_unit(unit_label, df_unit_counts, unit_info_row,
                    spike_times_all, spike_trials_all,
                    condition_colors, output_dir):
    """
    Two-column figure: raster (left) and per-condition spike-count
    distribution (right) for a single low-FR unit.

    Parameters
    ----------
    unit_label : str
    df_unit_counts : pd.DataFrame
        Rows from DF_lowfr_counts for this unit, non-habituation trials only.
        Must have columns: _seq (absolute 1-based trial index matching spkMAT),
                           trlNum, eventID, FR_Total.
    unit_info_row : pd.Series   (row from DF_LowFR)
    spike_times_all : np.ndarray  (spike times in seconds, relative to onset)
    spike_trials_all : np.ndarray  (1-based sequential trial index per spike)
    condition_colors : dict  (from create_mmn_colormap_for_plots)
    output_dir : Path
    """
    from matplotlib.gridspec import GridSpec

    brain_region = unit_info_row['BrainRegion']
    su_str  = 'SU' if bool(unit_info_row['SU']) else 'MU'
    mean_fr = unit_info_row['meanFR']

    cond_order = ['Std', 'LocLeft', 'LocRight',
                  'intensityDown', 'intensityUp',
                  'Freq_450', 'Freq_550', 'Dur25', 'Gap']
    cond_labels = {
        'Std': 'Standard',
        'LocLeft': 'Loc Left', 'LocRight': 'Loc Right',
        'intensityDown': 'Int Down', 'intensityUp': 'Int Up',
        'Freq_450': 'Freq 450', 'Freq_550': 'Freq 550',
        'Dur25': 'Dur 25', 'Gap': 'Gap',
    }

    fig = plt.figure(figsize=(9, 5))
    gs  = GridSpec(1, 2, width_ratios=[1.5, 1.0],
                   wspace=0.38, left=0.09, right=0.97,
                   top=0.89, bottom=0.15)
    ax_raster = fig.add_subplot(gs[0, 0])
    ax_dist   = fig.add_subplot(gs[0, 1])

    # ---- Left panel: raster ------------------------------------------------
    current_y = 0
    tick_pos  = []
    tick_labs = []

    for cond in cond_order:
        cond_rows = df_unit_counts.loc[
            df_unit_counts['eventID'] == cond
        ].sort_values('trlNum')
        n_trials = len(cond_rows)
        if n_trials == 0:
            continue

        color = condition_colors.get(cond, '#888888')
        ax_raster.axhspan(current_y, current_y + n_trials,
                          color=color, alpha=0.18, linewidth=0)

        events = []
        for _, row in cond_rows.iterrows():
            seq = int(row['_seq'])
            mask = (spike_trials_all == seq)
            spk_ms = spike_times_all[mask] * 1000
            events.append(spk_ms)

        if events:
            ax_raster.eventplot(
                events,
                lineoffsets=np.arange(current_y + 0.5, current_y + n_trials + 0.5),
                linelengths=1.2,
                linewidths=1.5,
                colors='black',
                alpha=0.8
            )

        tick_pos.append(current_y + n_trials / 2)
        tick_labs.append(cond_labels.get(cond, cond))
        current_y += n_trials

    ax_raster.axvline(x=0,   linestyle='--', color='firebrick', alpha=0.7, linewidth=0.9)
    ax_raster.axvline(x=500, linestyle='--', color='gray',      alpha=0.5, linewidth=0.7)
    ax_raster.set_xlim([-100, 620])
    ax_raster.set_ylim([0, current_y])
    ax_raster.set_yticks(tick_pos)
    ax_raster.set_yticklabels(tick_labs, fontsize=7)
    ax_raster.set_xlabel('Time (ms)', fontsize=7)
    ax_raster.set_ylabel('Condition', fontsize=7)
    ax_raster.set_title('A', loc='left', fontsize=9, fontweight='bold')
    ax_raster.tick_params(axis='x', labelsize=6)
    ax_raster.spines[['right', 'top']].set_visible(False)

    # ---- Right panel: spike-count distribution -----------------------------
    import seaborn as sns
    plot_rows = []
    for cond in cond_order:
        cond_rows = df_unit_counts.loc[df_unit_counts['eventID'] == cond]
        for _, row in cond_rows.iterrows():
            plot_rows.append({
                'Condition':   cond_labels.get(cond, cond),
                'Spike count': row['FR_Total'],
            })
    df_plot = pd.DataFrame(plot_rows)

    ordered_labels = [cond_labels[c] for c in cond_order
                      if cond_labels.get(c, '') in df_plot['Condition'].values]
    palette = {cond_labels[c]: condition_colors.get(c, '#888888') for c in cond_order}

    sns.boxplot(data=df_plot, x='Condition', y='Spike count',
                order=ordered_labels, palette=palette,
                width=0.45, fliersize=0, linewidth=0.8, ax=ax_dist)
    sns.stripplot(data=df_plot, x='Condition', y='Spike count',
                  order=ordered_labels, palette=palette,
                  size=3, jitter=True, alpha=0.55, ax=ax_dist)

    ax_dist.set_xticklabels(ordered_labels, rotation=45, ha='right', fontsize=6)
    ax_dist.set_ylabel('Spike count (0–500 ms)', fontsize=7)
    ax_dist.set_xlabel('', fontsize=7)
    ax_dist.set_title('B', loc='left', fontsize=9, fontweight='bold')
    ax_dist.axhline(y=0, color='k', linestyle='--', alpha=0.2, linewidth=0.7)
    ax_dist.tick_params(axis='y', labelsize=6)
    ax_dist.spines[['right', 'top']].set_visible(False)

    fig.suptitle(
        f'{unit_label}  |  {su_str}, {brain_region}  |  meanFR = {mean_fr:.3f} Hz',
        fontsize=8, fontweight='bold'
    )

    safe_label = unit_label.replace('/', '_').replace(' ', '_')
    fig.savefig(output_dir / f'{safe_label}_raster_dist.png', dpi=300, bbox_inches='tight')
    fig.savefig(output_dir / f'{safe_label}_raster_dist.svg', format='svg', bbox_inches='tight')
    plt.close(fig)


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
lowfr_alpha_s1   = 0.05   # Stage 1 threshold (uncorrected)
lowfr_alpha_s2   = 0.05   # Stage 2 FDR threshold
lowfr_min_trials = 5      # minimum trials per condition for KW
lowfr_cat_thresh = 15     # use feature categories when median deviant-type trial count < this
lowfr_output_dir = resultsFolder / 'LowFR_Supplementary'
lowfr_output_dir.mkdir(parents=True, exist_ok=True)

deviant_categories = {
    'Location':  ['LocLeft', 'LocRight'],
    'Intensity': ['intensityDown', 'intensityUp'],
    'Frequency': ['Freq_450', 'Freq_550'],
    'Timing':    ['Dur25', 'Gap'],
}
individual_deviants = ['LocLeft', 'LocRight', 'intensityDown', 'intensityUp',
                       'Freq_450', 'Freq_550', 'Dur25', 'Gap']

# ---------------------------------------------------------------------------
# Identify low-FR units
# ---------------------------------------------------------------------------
# DF_Units holds ALL recorded units (built unconditionally in the loading loop).
# Units absent from DF_sdfs were excluded by minFRcrit. Also remove
# quality-flagged rmUnitLabels from this supplementary set.
DF_LowFR = DF_Units.loc[
    ~DF_Units['UniqueLabel'].isin(DF_sdfs['UniqueLabel'].unique()) &
    ~DF_Units['UniqueLabel'].isin(rmUnitLabels)
].copy().reset_index(drop=True)

print(f"\nLow-FR units (meanFR < {minFRcrit} Hz, after quality exclusions): {len(DF_LowFR)}")
print(DF_LowFR.groupby('BrainRegion')['UniqueLabel'].count().rename('N_units').to_string())

# ---------------------------------------------------------------------------
# Filter DF_FRall to low-FR units, exclude habituation trials
# ---------------------------------------------------------------------------
DF_lowfr_counts = DF_FRall.loc[
    DF_FRall['UniqueLabel'].isin(DF_LowFR['UniqueLabel']) &
    ~DF_FRall['habituation'].astype(bool)
].copy()

# ---------------------------------------------------------------------------
# Stage 1: Mann-Whitney U + Fisher's exact (all low-FR units)
# ---------------------------------------------------------------------------
lowfr_results = []

for _, unit_row in DF_LowFR.iterrows():
    unit_label = unit_row['UniqueLabel']
    subset = DF_lowfr_counts.loc[DF_lowfr_counts['UniqueLabel'] == unit_label]

    if subset.empty:
        continue

    std_counts = subset.loc[subset['eventID'] == 'Std', 'FR_Total'].values.astype(float)
    dev_counts = subset.loc[subset['eventID'] != 'Std', 'FR_Total'].values.astype(float)

    # Skip units with zero spikes in all trials
    if std_counts.sum() == 0 and dev_counts.sum() == 0:
        print(f"  {unit_label}: zero spikes in all trials, skipping")
        continue

    # Mann-Whitney U (spike counts, two-sided)
    if len(std_counts) >= 2 and len(dev_counts) >= 2:
        mwu_stat, mwu_p = mannwhitneyu(std_counts, dev_counts, alternative='two-sided')
    else:
        mwu_stat, mwu_p = np.nan, np.nan

    # Fisher's exact (proportion of trials with >=1 spike)
    std_has = int((std_counts >= 1).sum())
    std_not = int((std_counts <  1).sum())
    dev_has = int((dev_counts >= 1).sum())
    dev_not = int((dev_counts <  1).sum())

    if std_has + std_not > 0 and dev_has + dev_not > 0:
        fisher_stat, fisher_p = fisher_exact(
            [[std_has, std_not], [dev_has, dev_not]], alternative='two-sided'
        )
    else:
        fisher_stat, fisher_p = np.nan, np.nan

    lowfr_results.append({
        'UniqueLabel':       unit_label,
        'ID':                unit_row['ID'],
        'SessionStr':        str(unit_row['SessionStr']),
        'BrainRegion':       unit_row['BrainRegion'],
        'SU':                bool(unit_row['SU']),
        'meanFR':            unit_row['meanFR'],
        'n_std_trials':      len(std_counts),
        'n_dev_trials':      len(dev_counts),
        'MWU_stat':          mwu_stat,
        'MWU_p':             mwu_p,
        'Fisher_stat':       fisher_stat,
        'Fisher_p':          fisher_p,
        'KW_p':              np.nan,
        'KW_strategy':       '',
        'KW_sig_conditions': '',
    })

DF_lowfr_res = pd.DataFrame(lowfr_results)

# Stage 1 population summary
print(f"\nStage 1 complete: {len(DF_lowfr_res)} units processed")
for test_col, test_label in [('MWU_p', 'Mann-Whitney U'), ('Fisher_p', "Fisher's exact")]:
    valid   = DF_lowfr_res[test_col].dropna()
    n_valid = len(valid)
    n_sig   = int((valid < lowfr_alpha_s1).sum())
    pct_sig = 100 * n_sig / n_valid if n_valid > 0 else 0
    binom_res = binomtest(n_sig, n_valid, p=lowfr_alpha_s1, alternative='greater')
    sig_marker = ('***' if binom_res.pvalue < 0.001 else
                  '**'  if binom_res.pvalue < 0.01  else
                  '*'   if binom_res.pvalue < 0.05  else 'ns')
    print(f"  {test_label}: {n_sig}/{n_valid} sig ({pct_sig:.1f}%) vs 5% chance "
          f"[binom p={binom_res.pvalue:.4f} {sig_marker}]")

# ---------------------------------------------------------------------------
# Stage 2: Kruskal-Wallis + post-hoc MWU with FDR (Stage 1 passers only)
# ---------------------------------------------------------------------------
stage1_pass_mask = (
    (DF_lowfr_res['MWU_p']    < lowfr_alpha_s1) |
    (DF_lowfr_res['Fisher_p'] < lowfr_alpha_s1)
)
stage1_passers = DF_lowfr_res.loc[stage1_pass_mask, 'UniqueLabel'].tolist()
print(f"\nStage 1 passers (either test p<{lowfr_alpha_s1}): {len(stage1_passers)}")

for unit_label in stage1_passers:
    idx    = DF_lowfr_res.index[DF_lowfr_res['UniqueLabel'] == unit_label][0]
    subset = DF_lowfr_counts.loc[DF_lowfr_counts['UniqueLabel'] == unit_label]

    # Choose strategy based on median deviant-type trial count
    dev_trial_counts = [
        (subset['eventID'] == dev).sum()
        for dev in individual_deviants
        if (subset['eventID'] == dev).sum() > 0
    ]
    median_dev_n = np.median(dev_trial_counts) if dev_trial_counts else 0

    if median_dev_n < lowfr_cat_thresh:
        strategy = 'categories'
        groups_for_kw = {}
        for cat_name, cat_devs in deviant_categories.items():
            cat_counts = subset.loc[
                subset['eventID'].isin(cat_devs), 'FR_Total'
            ].values.astype(float)
            if len(cat_counts) >= lowfr_min_trials:
                groups_for_kw[cat_name] = cat_counts
    else:
        strategy = 'individual'
        groups_for_kw = {}
        for dev in individual_deviants:
            dev_counts = subset.loc[
                subset['eventID'] == dev, 'FR_Total'
            ].values.astype(float)
            if len(dev_counts) >= lowfr_min_trials:
                groups_for_kw[dev] = dev_counts

    std_counts_kw = subset.loc[subset['eventID'] == 'Std', 'FR_Total'].values.astype(float)
    if len(std_counts_kw) >= lowfr_min_trials:
        groups_for_kw['Std'] = std_counts_kw

    DF_lowfr_res.at[idx, 'KW_strategy'] = strategy

    if len(groups_for_kw) < 2:
        print(f"  {unit_label}: insufficient conditions for KW, skipping Stage 2")
        continue

    try:
        kw_stat, kw_p = kruskal(*groups_for_kw.values())
    except ValueError as e:
        print(f"  {unit_label}: KW failed ({e})")
        continue

    DF_lowfr_res.at[idx, 'KW_p'] = kw_p

    # Post-hoc MWU with FDR (only if KW significant)
    if kw_p < lowfr_alpha_s2:
        std_g          = groups_for_kw.get('Std', np.array([]))
        posthoc_pvals  = []
        posthoc_labels = []
        for cond_name, cond_g in groups_for_kw.items():
            if cond_name == 'Std':
                continue
            if len(std_g) < 2 or len(cond_g) < 2:
                continue
            _, ph_p = mannwhitneyu(std_g, cond_g, alternative='two-sided')
            posthoc_pvals.append(ph_p)
            posthoc_labels.append(cond_name)

        if posthoc_pvals:
            _, pvals_fdr, _, _ = multipletests(posthoc_pvals, method='fdr_bh')
            sig_conds = [lbl for lbl, pfdr in zip(posthoc_labels, pvals_fdr)
                         if pfdr < lowfr_alpha_s2]
            DF_lowfr_res.at[idx, 'KW_sig_conditions'] = ', '.join(sig_conds)

# ---------------------------------------------------------------------------
# Save summary CSV
# ---------------------------------------------------------------------------
DF_lowfr_res.to_csv(lowfr_output_dir / 'LowFR_unit_summary.csv', index=False)
print(f"\nSummary saved: {lowfr_output_dir / 'LowFR_unit_summary.csv'}")

# ---------------------------------------------------------------------------
# Figures: raster + distribution for Stage 1 passers
# ---------------------------------------------------------------------------
if stage1_passers:
    print(f"\nGenerating figures for {len(stage1_passers)} Stage 1 passer(s)...")
    _, condition_colors_lowfr = create_mmn_colormap_for_plots()

    passers_info = DF_LowFR.loc[DF_LowFR['UniqueLabel'].isin(stage1_passers)]

    for (session_str, subj_id), grp in passers_info.groupby(['SessionStr', 'ID']):
        FRtab_FN_fig = dataPath / f'Session{session_str}' / 'FRTabs' / f'{subj_id}_FRtable'
        try:
            spkMAT_fig = loadmat(
                f'{FRtab_FN_fig}.mat', struct_as_record=False, squeeze_me=True
            )
        except Exception as e:
            print(f"  Cannot load {subj_id} S{session_str} for figures: {e}")
            continue

        unitLabels_fig = spkMAT_fig['ALLspikes'].label
        if isinstance(unitLabels_fig, str):
            unitLabels_fig = [unitLabels_fig]
        else:
            unitLabels_fig = list(unitLabels_fig)

        for _, unit_row in grp.iterrows():
            unit_label = unit_row['UniqueLabel']
            raw_label  = unit_row['UnitLabel']

            if raw_label not in unitLabels_fig:
                print(f"  {unit_label}: label not found in spkMAT, skipping figure")
                continue
            unit_idx_fig = unitLabels_fig.index(raw_label)

            # Extract spike arrays (handle single-unit edge case)
            spike_data = spkMAT_fig['spikeTrials']
            if isinstance(spike_data.label, str):
                spk_times  = np.atleast_1d(np.asarray(spike_data.time,  dtype=float))
                spk_trials = np.atleast_1d(np.asarray(spike_data.trial, dtype=int))
            else:
                spk_times  = np.atleast_1d(np.asarray(spike_data.time[unit_idx_fig],  dtype=float))
                spk_trials = np.atleast_1d(np.asarray(spike_data.trial[unit_idx_fig], dtype=int))

            # Build df_unit_counts with absolute sequential trial index (_seq)
            # matching the 1-based indexing used by spkMAT['spikeTrials'].trial
            all_unit_fr = DF_FRall.loc[
                DF_FRall['UniqueLabel'] == unit_label
            ].sort_values('trlNum').reset_index(drop=True)
            all_unit_fr['_abs_seq'] = all_unit_fr.index + 1
            seq_map = all_unit_fr.set_index('trlNum')['_abs_seq'].to_dict()

            df_unit_counts = DF_lowfr_counts.loc[
                DF_lowfr_counts['UniqueLabel'] == unit_label
            ].copy()
            df_unit_counts['_seq'] = df_unit_counts['trlNum'].map(seq_map)

            plot_lowfr_unit(
                unit_label       = unit_label,
                df_unit_counts   = df_unit_counts,
                unit_info_row    = unit_row,
                spike_times_all  = spk_times,
                spike_trials_all = spk_trials,
                condition_colors = condition_colors_lowfr,
                output_dir       = lowfr_output_dir,
            )
            print(f"  Saved: {unit_label}")

# ---------------------------------------------------------------------------
# Final printed summary
# ---------------------------------------------------------------------------
kw_run = DF_lowfr_res.loc[
    DF_lowfr_res['UniqueLabel'].isin(stage1_passers), 'KW_p'
].dropna()
kw_sig = int((kw_run < lowfr_alpha_s2).sum())

print("\n" + "=" * 60)
print("LOW-FR SUPPLEMENTARY ANALYSIS — SUMMARY")
print("=" * 60)
print(f"Total low-FR units identified : {len(DF_LowFR)}")
print(f"Units with non-zero spike data: {len(DF_lowfr_res)}")
print()
for test_col, test_label in [('MWU_p', 'Mann-Whitney U'), ('Fisher_p', "Fisher's exact")]:
    valid   = DF_lowfr_res[test_col].dropna()
    n_valid = len(valid)
    n_sig   = int((valid < lowfr_alpha_s1).sum())
    pct     = 100 * n_sig / n_valid if n_valid else 0
    binom_r = binomtest(n_sig, n_valid, p=0.05, alternative='greater')
    sig_m   = ('***' if binom_r.pvalue < 0.001 else
               '**'  if binom_r.pvalue < 0.01  else
               '*'   if binom_r.pvalue < 0.05  else 'ns')
    print(f"{test_label}: {n_sig}/{n_valid} significant ({pct:.1f}%) "
          f"vs 5% chance  [binom p={binom_r.pvalue:.4f} {sig_m}]")
print()
print(f"Stage 1 passers (either test) : {len(stage1_passers)}")
print(f"Stage 2 KW significant        : {kw_sig}/{len(kw_run)}")
print()
print(f"Output: {lowfr_output_dir}")
print("=" * 60)