# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/API/multi.ipynb.

# %% auto 0
__all__ = ['MultiContrast', 'combine', 'whorlmap']

# %% ../nbs/API/multi.ipynb 3
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
import warnings
from typing import List, Optional, Union, Tuple, Dict, Any


# %% ../nbs/API/multi.ipynb 6
class MultiContrast:
    """
    Unified multiple contrast object for forest plots and whorlmaps.
    
    Takes raw dabest objects and provides validated, processed data
    for downstream visualizations.
    """
    
    def __init__(self, 
                 dabest_objs: Union[List, List[List]], 
                 labels: Optional[List[str]] = None,
                 row_labels: Optional[List[str]] = None,
                 effect_size: str = "mean_diff",
                 ci_type: str = "bca"):
        """
        Initialize MultiContrast object with checking.
        
        Parameters
        ----------
        dabest_objs : Union[List, List[List]]
            Raw dabest objects. Can be:
            - 1D: [dabest_obj1, dabest_obj2, ...]  
            - 2D: [[dabest_obj1, dabest_obj2], [dabest_obj3, dabest_obj4]]
        labels : Optional[Union[List[str], List[List[str]]]], default=None
            Labels matching the contrast array structure. If None, defaults will be generated.
        effect_size : str, default="mean_diff"
            Effect size to extract from dabest objects
        ci_type : str, default="bca"
            Confidence interval type
        """
        # Store raw inputs for validation
        self._raw_dabest_objs = dabest_objs
        self._raw_labels = labels
        self._raw_row_labels = row_labels 

        # Validate and process inputs
        self.effect_size = self._validate_effect_size(effect_size)
        self.ci_type = self._validate_ci_type(ci_type)

        # Process structure (adapts forest_plot logic to handle 2D)
        self.structure = self._validate_and_parse_structure(dabest_objs, labels)

        # Validate all dabest objects consistency
        self.contrast_type = self._validate_contrast_consistency()

        # Extract data (adapts forest_plot's load_plot_data logic)
        self._bootstrap_data = None
        self._effect_size_data = None
        self._ci_data = None
        
    def _validate_effect_size(self, effect_size: str) -> str:
        """Validate effect size parameter (from forest_plot)."""
        possible_effect_sizes = [
            'mean_diff', 'median_diff', 'cohens_d', 
            'cohens_h', 'cliffs_delta', 'hedges_g', 'delta_g'
        ]
        
        if not isinstance(effect_size, str) or effect_size not in possible_effect_sizes:
            raise TypeError(
                f"effect_size must be one of: {possible_effect_sizes}"
            )
        return effect_size

    def _validate_ci_type(self, ci_type: str) -> str:
        """Validate CI type parameter (from forest_plot)."""
        if ci_type not in ('bca', 'pct'):
            raise TypeError("ci_type must be either 'bca' or 'pct'")
        return ci_type
    
    def _validate_and_parse_structure(self, dabest_objs, labels):
        """
        Validate and parse contrast structure, combining forest_plot 
        validation with whorlmap's 2D handling.
        """
        # Basic validation (from forest_plot)
        if not isinstance(dabest_objs, (list, tuple)) or len(dabest_objs) == 0:
            raise ValueError("dabest_objs must be a non-empty list")
        
        # Determine if 1D or 2D structure
        if isinstance(dabest_objs[0], (list, tuple)):
            # 2D structure (can be used to plot whorlmap or a stack of forest plots)
            structure_type = "2D"
            dabest_objs_2d = dabest_objs
            n_rows = len(dabest_objs)
            n_cols = len(dabest_objs[0])
            
            # Validate rectangular structure
            for i, row in enumerate(dabest_objs):
                if not isinstance(row, (list, tuple)):
                    raise TypeError(f"Row {i} must be a list/tuple in 2D structure")
                if len(row) != n_cols:
                    raise ValueError("All rows must have the same number of dabest_objs")
            
            # Handle 2D labels
            if labels:
                if not isinstance(labels, (list, tuple)):
                    raise TypeError("labels must be a list for 2D dabest_objs")
                if len(labels) != n_cols:
                    raise ValueError("Number of labels must match number of columns of dabest_objs")
                col_labels = labels
            else:
                col_labels = [f"Contrast {i+1}" for i in range(n_cols)]
            # Handle row_labels - use self._raw_row_labels if available
            if hasattr(self, '_raw_row_labels') and self._raw_row_labels:
                if not isinstance(self._raw_row_labels, (list, tuple)):
                    raise TypeError("row_labels must be a list for 2D dabest_objs")
                if len(self._raw_row_labels) != n_rows:
                    raise ValueError("Number of row_labels must match number of rows of dabest_objs")
                row_labels = self._raw_row_labels
            else:
                row_labels = [f"Row {i+1}" for i in range(n_rows)]
        else:
            # 1D structure (like forest_plot)
            structure_type = "1D"
            dabest_objs_2d = [dabest_objs]  # Wrap in single row for unified processing
            n_rows = 1
            n_cols = len(dabest_objs)
            
            # Handle 1D labels
            if labels:
                if not isinstance(labels, (list, tuple)):
                    raise TypeError("labels must be a list for 1D dabest_objs")
                if len(labels) != n_cols:
                    raise ValueError("Number of labels must match number of dabest_objs")
                col_labels = labels
            else:
                col_labels = [f"Contrast {i+1}" for i in range(n_cols)]
            row_labels = [""]  # Single empty row label
        
        return {
            'type': structure_type,
            'dabest_objs_2d': dabest_objs_2d,
            'n_rows': n_rows,
            'n_cols': n_cols,
            'col_labels': col_labels,
            'row_labels': row_labels,
            'total_dabest_objs': n_rows * n_cols
        }
    
    def _validate_contrast_consistency(self) -> Union[str, Dict]:
        """
        Validate contrast consistency with support for mixed types in whorlmap.
        
        Returns either:
        - str: Single contrast type for homogeneous data (forest_plot compatible)
        - dict: Row-wise contrast types for mixed data (whorlmap only)
        """
        all_dabest_objs = []
        for row in self.structure['dabest_objs_2d']:
            all_dabest_objs.extend(row)
        
        if not all_dabest_objs:
            raise ValueError("No valid dabest_objs found")
        
        # First, validate EACH contrast individually
        for i, dabest_obj in enumerate(all_dabest_objs):
            self._validate_individual_dabest_obj(dabest_obj, i)
        
        # Analyze contrast type structure
        contrast_types_by_row = []
        for row_idx, row in enumerate(self.structure['dabest_objs_2d']):
            row_types = []
            for contrast in row:
                contrast_type = ("delta2" if contrast.delta2 
                               else "mini_meta" if contrast.is_mini_meta
                               else "delta")
                row_types.append(contrast_type)
            contrast_types_by_row.append(row_types)
        
        # Check if all dabest_objs are the same type (forest_plot requirement)
        all_types_flat = [t for row_types in contrast_types_by_row for t in row_types]
        unique_types = set(all_types_flat)
        
        if len(unique_types) == 1:
            # Homogeneous: all same type (forest_plot compatible)
            contrast_type = list(unique_types)[0]
            self._validate_effect_size_compatibility(contrast_type)
            return contrast_type
        
        else:
            # Heterogeneous: mixed types (whorlmap only)
            if self.structure['type'] == '1D':
                raise ValueError(
                    "Mixed contrast types are only supported for 2D structures (whorlmaps). "
                    f"Found types: {unique_types}. For forest plots, all dabest_objs must be the same type."
                )
            
            # Validate within-row consistency for whorlmap
            for row_idx, row_types in enumerate(contrast_types_by_row):
                unique_row_types = set(row_types)
                if len(unique_row_types) > 1:
                    raise ValueError(
                        f"Within each row, all dabest_objs must be the same type. "
                        f"Row {row_idx} has mixed types: {unique_row_types}"
                    )
            
            # Validate effect size compatibility for each row type
            for row_types in contrast_types_by_row:
                row_type = row_types[0]  # All same within row
                self._validate_effect_size_compatibility(row_type)
            
            # Return row-wise type information
            return {
                'mixed': True,
                'by_row': [row_types[0] for row_types in contrast_types_by_row],
                'unique_types': list(unique_types)
            }
        
    def _validate_effect_size_compatibility(self, contrast_type: str):
        """Validate effect size compatibility with a specific contrast type."""
        if contrast_type == "mini_meta" and self.effect_size != 'mean_diff':
            raise ValueError("effect_size must be 'mean_diff' for mini-meta analyses")
        
        if contrast_type == "delta2" and self.effect_size not in ['mean_diff', 'hedges_g', 'delta_g']:
            raise ValueError(
                "effect_size must be 'mean_diff', 'hedges_g', or 'delta_g' for delta-delta analyses"
            )        
    
    def _validate_individual_dabest_obj(self, dabest_obj, position: int):
        """
        Validate individual dabest object.
        
        Parameters
        ----------
        dabest_obj : object
            Individual dabest object to validate
        position : int
            Position in the contrast list for error reporting
        """
        # Basic existence check
        if dabest_obj is None:
            raise ValueError(f"Dabest object at position {position} is None")
        
        # Required attributes for dabest objects
        required_attrs = ['delta2', 'is_mini_meta']
        for attr in required_attrs:
            if not hasattr(dabest_obj, attr):
                raise TypeError(
                    f"Object at position {position} is not a valid dabest object. "
                    f"Missing required attribute: '{attr}'"
                )
        
        # Validate effect size attribute exists
        effect_attr = "hedges_g" if self.effect_size == 'delta_g' else self.effect_size
        if not hasattr(dabest_obj, effect_attr):
            raise AttributeError(
                f"Dabest Object at position {position} does not have effect size '{self.effect_size}'. "
                f"Expected attribute: '{effect_attr}'"
            )
        
        # Test that we can actually access the effect size data
        try:
            effect_obj = getattr(dabest_obj, effect_attr)

            # For delta2/mini_meta, check the nested attributes exist
            if dabest_obj.delta2:
                if not hasattr(effect_obj, 'delta_delta'):
                    raise AttributeError(f"Delta-delta contrast at position {position} missing 'delta_delta' attribute")
            elif dabest_obj.is_mini_meta:
                if not hasattr(effect_obj, 'mini_meta'):
                    raise AttributeError(f"Mini-meta contrast at position {position} missing 'mini_meta' attribute")
            else:
                # Standard contrast - check results structure
                if not hasattr(effect_obj, 'results'):
                    raise AttributeError(f"Standard contrast at position {position} missing 'results' attribute")        
        except Exception as e:
            raise ValueError(
                f"Failed to access effect size data for dabest object at position {position}: {str(e)}"
            )
        
    def _extract_data(self) -> Tuple[List, List, List, List]:
        """
        Extract bootstrap, effect sizes, CI low bounds and CI high bounds.
        Handles mixed contrast types for whorlmap.
        """
        if self._bootstrap_data is not None:
            return self._bootstrap_data, self._effect_data, self._ci_lows, self._ci_highs
        
        # Process effect size attribute name
        effect_attr = "hedges_g" if self.effect_size == 'delta_g' else self.effect_size
        
        bootstraps = []
        differences = []
        ci_lows = []
        ci_highs = []
        
        if isinstance(self.contrast_type, dict) and self.contrast_type.get('mixed'):
            # Mixed types: process row by row
            for row_idx, row in enumerate(self.structure['dabest_objs_2d']):
                row_contrast_type = self.contrast_type['by_row'][row_idx]
                contrast_attr = {"delta2": "delta_delta", "mini_meta": "mini_meta"}.get(row_contrast_type)
                
                for contrast in row:
                    bootstrap, diff, ci_low, ci_high = self._extract_single_contrast(
                        contrast, effect_attr, row_contrast_type, contrast_attr
                    )
                    bootstraps.extend(bootstrap if isinstance(bootstrap, list) else [bootstrap])
                    differences.extend(diff if isinstance(diff, list) else [diff])
                    ci_lows.extend(ci_low if isinstance(ci_low, list) else [ci_low])
                    ci_highs.extend(ci_high if isinstance(ci_high, list) else [ci_high])
        
        else:
            # Homogeneous types: process all together (original logic)
            contrast_attr = {"delta2": "delta_delta", "mini_meta": "mini_meta"}.get(self.contrast_type)
            
            all_dabest_objs = []
            for row in self.structure['dabest_objs_2d']:
                all_dabest_objs.extend(row)
            
            for contrast in all_dabest_objs:
                bootstrap, diff, ci_low, ci_high = self._extract_single_contrast(
                    contrast, effect_attr, self.contrast_type, contrast_attr
                )
                bootstraps.extend(bootstrap if isinstance(bootstrap, list) else [bootstrap])
                differences.extend(diff if isinstance(diff, list) else [diff])
                ci_lows.extend(ci_low if isinstance(ci_low, list) else [ci_low])
                ci_highs.extend(ci_high if isinstance(ci_high, list) else [ci_high])
        
        # Cache results
        self._bootstrap_data = bootstraps
        self._effect_data = differences
        self._ci_lows = ci_lows
        self._ci_highs = ci_highs
        
        return bootstraps, differences, ci_lows, ci_highs
    
    def _extract_single_contrast(self, contrast, effect_attr, contrast_type, contrast_attr):
        """Extract data from a single contrast object."""
        if contrast_type == 'delta':
            # Standard dabest_objs - may have multiple comparisons
            effect_obj = getattr(contrast, effect_attr)
            boot_list = effect_obj.results.bootstraps.to_list()
            diff_list = effect_obj.results.difference.to_list()
            low_list = effect_obj.results.get(f'{self.ci_type}_low').to_list()
            high_list = effect_obj.results.get(f'{self.ci_type}_high').to_list()
            return boot_list, diff_list, low_list, high_list
            
        else:
            # Delta-delta or mini-meta - single value per contrast
            effect_obj = getattr(contrast, effect_attr)
            processed_obj = getattr(effect_obj, contrast_attr)
            
            if contrast_type == "delta2":
                bootstrap = processed_obj.bootstraps_delta_delta
                difference = processed_obj.difference
            else:  # mini_meta
                bootstrap = processed_obj.bootstraps_weighted_delta
                difference = processed_obj.difference
            
            ci_low = processed_obj.results.get(f'{self.ci_type}_low')[0]
            ci_high = processed_obj.results.get(f'{self.ci_type}_high')[0]
            
            return bootstrap, difference, ci_low, ci_high
    @property 
    def bootstraps(self) -> List:
        """Get bootstrap samples for all dabest_objs."""
        bootstraps, _, _, _ = self._extract_data()
        return bootstraps
    
    @property
    def effect_sizes(self) -> List:
        """Get effect sizes for all dabest_objs."""
        _, effects, _, _ = self._extract_data()
        return effects
    
    @property 
    def confidence_intervals(self) -> Tuple[List, List]:
        """Get confidence interval bounds."""
        _, _, ci_lows, ci_highs = self._extract_data()
        return ci_lows, ci_highs
    
    def forest_plot(self, forest_plot_title = None, forest_plot_kwargs = {}):
        """
        Create forest plot using validated data.
        
        This is a convenience method that calls the existing forest_plot function
        with validated dabest objects. # TODO: decide whether to
        migrate forest_plot to use MultiContrast data directly.
        """
        # Check compatibility with forest plot (mixed contrast types not supported)
        if isinstance(self.contrast_type, dict) and self.contrast_type.get('mixed'):
            raise ValueError(
                "Forest plots require all dabest_objs to be the same type. "
                f"This MultiContrast has mixed types: {self.contrast_type['unique_types']}. "
                "Consider creating separate MultiContrast objects for each type, "
                "or use whorlmap() which supports mixed types."
            )
        
        # Import forest_plot function
        from .forest_plot import forest_plot
        
        # # Get flattened contrast list for existing forest_plot function
        # all_dabest_objs = []
        # for row in self.structure['dabest_objs_2d']:
        #     all_dabest_objs.extend(row)
        
        # Call existing forest_plot with validated dabest objects

        f_forest, axes = plt.subplots(self.structure['n_rows'], 1, 
                               figsize=(8, 2 * self.structure['n_rows']), squeeze=False)
        for i, row in enumerate(self.structure['dabest_objs_2d']):
            # Set default parameters, allow kwargs to override
            forest_kwargs = {
                'effect_size': self.effect_size,
                'ci_type': self.ci_type,
                'ax': axes[i, 0],
                'labels': self.structure['col_labels'],
                'title': self.structure['row_labels'][i] if self.structure['n_rows'] > 1 else None,}
            forest_kwargs.update(forest_plot_kwargs)
            forest_plot(data = row, **forest_kwargs)
            if i == self.structure['n_rows'] - 1:
                axes[i, 0].set_xticks(axes[i, 0].get_xticks())
            else:
                axes[i, 0].set_xticks([])
        self.f_forest = f_forest
        if forest_plot_title:
            f_forest.suptitle(forest_plot_title)
        return f_forest, axes

    def whorlmap(self, **heatmap_kwargs):
        """
        Create whorlmap using validated data.
        
        This uses the whorlmap that can handle both homogeneous
        and mixed contrast types.
        """
        from .multi import whorlmap
        f_whorlmap = whorlmap(multi_contrast=self, **heatmap_kwargs)
        self.f_whorlmap = f_whorlmap
        # Call whorlmap with self as the multi_contrast object
        return f_whorlmap
    def get_bootstrap_by_position(self, row: int, col: int):
        """
        Get bootstrap data for a specific position in the grid.
        Useful for mixed-type whorlmaps.
        """
        if row >= self.structure['n_rows'] or col >= self.structure['n_cols']:
            raise IndexError(f"Position ({row}, {col}) out of bounds for {self.structure['n_rows']}×{self.structure['n_cols']} grid")
        
        contrast = self.structure['dabest_objs_2d'][row][col]
        effect_attr = "hedges_g" if self.effect_size == 'delta_g' else self.effect_size
        
        # Determine contrast type for this position
        if isinstance(self.contrast_type, dict) and self.contrast_type.get('mixed'):
            position_type = self.contrast_type['by_row'][row]
        else:
            position_type = self.contrast_type
        
        contrast_attr = {"delta2": "delta_delta", "mini_meta": "mini_meta"}.get(position_type)
        
        # Extract bootstrap for this specific contrast
        bootstrap, _, _, _ = self._extract_single_contrast(contrast, effect_attr, position_type, contrast_attr)
        
        # For standard dabest_objs, return first bootstrap (they may have multiple)
        if isinstance(bootstrap, list) and len(bootstrap) > 0:
            return bootstrap[0]
        return bootstrap
    
    def __repr__(self):
        if isinstance(self.contrast_type, dict) and self.contrast_type.get('mixed'):
            types_info = f"mixed({', '.join(self.contrast_type['unique_types'])})"
        else:
            types_info = self.contrast_type
            
        return (f"MultiContrast({self.structure['type']}: "
               f"{self.structure['n_rows']}x{self.structure['n_cols']}, "
               f"effect_size='{self.effect_size}', "
               f"contrast_type='{types_info}')")       

# %% ../nbs/API/multi.ipynb 8
def combine(dabest_objs: Union[List, List[List]], 
           labels: Optional[List[str]] = None,
           row_labels: Optional[List[str]] = None,
           effect_size: str = "mean_diff",
           ci_type: str = "bca",
           allow_mixed_types: bool = False) -> MultiContrast:
    """
    Create a MultiContrast object from raw dabest objects.
    
    This is the main entry point that users should use to create
    multi-contrast visualizations.
    
    Parameters
    ----------
    dabest_objs : Union[List, List[List]]
        Raw dabest objects in 1D or 2D structure
    labels : Optional[Union[List[str], List[List[str]]]], default=None
        Labels for dabest_objs
    effect_size : str, default="mean_diff"  
        Effect size to extract
    ci_type : str, default="bca"
        Confidence interval type
    allow_mixed_types : bool, default=False
        If True, allows different contrast types in different rows (whorlmap only)
        If False, enforces homogeneous types (forest_plot compatible)
        
    Returns
    -------
    MultiContrast
        Validated multi-contrast object ready for visualization
        
    Examples
    --------
    # Homogeneous 1D structure (forest_plot and whorlmap compatible)
    mc = combine([dabest1, dabest2, dabest3], 
                labels=['Treatment A', 'Treatment B', 'Treatment C'])
    mc.forest_plot()
    mc.whorlmap()  # Will arrange in single row
    
    # Homogeneous 2D structure (forest_plot flattens, whorlmap uses grid)
    mc = combine([[dabest1, dabest2], [dabest3, dabest4]], 
                labels=[['Dose Low', 'Dose High'], ['Time 1', 'Time 2']])
    mc.whorlmap()  # 2x2 grid
    mc.forest_plot()  # Flattened to 1D
    
    # Mixed types 2D structure (whorlmap only!)
    mc = combine([[standard_dabest1, standard_dabest2], 
                  [delta2_dabest1, delta2_dabest2]],
                 labels=[['Standard A', 'Standard B'], 
                         ['Delta2 A', 'Delta2 B']],
                 allow_mixed_types=True)
    mc.whorlmap()  # Works: mixed spiral types per row
    # mc.forest_plot()  # Raises error: incompatible with mixed types
    
    # Mini-meta + Delta2 mixed example
    mc = combine([[mini_meta1, mini_meta2], 
                  [delta2_obj1, delta2_obj2]],
                 allow_mixed_types=True)
    mc.whorlmap()  # Top row: mini-meta spirals, bottom row: delta2 spirals
    """
    mc = MultiContrast(dabest_objs, labels, row_labels, effect_size, ci_type)
    
    # Check mixed types policy
    if isinstance(mc.contrast_type, dict) and mc.contrast_type.get('mixed'):
        if not allow_mixed_types:
            raise ValueError(
                f"Mixed contrast types detected: {mc.contrast_type['unique_types']}. "
                "Set allow_mixed_types=True to enable mixed-type whorlmaps, "
                "or ensure all dabest_objs are the same type for forest_plot compatibility."
            )
    
    return mc

# %% ../nbs/API/multi.ipynb 10
def _sample_bootstrap(bootstrap, m, n, reverse_neg, abs_rank, chop_tail):
    """Sample bootstrap values and prepare for spiral visualization."""
    bootstrap_sorted = sorted(bootstrap)
    chop_tail_int = int(np.ceil(len(bootstrap_sorted) * chop_tail / 100))
    bootstrap_sorted = bootstrap_sorted[chop_tail_int : len(bootstrap_sorted) - chop_tail_int]
    
    ranks_to_look = np.linspace(0, len(bootstrap_sorted), m * n, dtype=int)   
    ranks_to_look[0] = 1
    
    if np.sum(np.array(bootstrap_sorted) > 0) < len(bootstrap_sorted) / 2:
        if reverse_neg:
            bootstrap_sorted = bootstrap_sorted[::-1]
    
    if abs_rank:
        bootstrap_sorted = sorted(bootstrap_sorted, key=abs)
    
    long_ranks = [bootstrap_sorted[r - 1] for r in ranks_to_look]
    return long_ranks

# %% ../nbs/API/multi.ipynb 11
def _spiralize(fill, m, n):
    """Convert linear array into spiral pattern."""
    i = 0
    j = 0
    k = 0
    array = np.zeros((m, n))
    
    while m > 0 and k < len(fill):
        jj = j
        ii = i
        
        # Right
        for j in range(j, n):
            if k >= len(fill):
                break
            array[i, j] = fill[k]
            k += 1
        
        # Down
        for i in range(ii + 1, m):
            if k >= len(fill):
                break
            array[i, j] = fill[k]
            k += 1
        
        # Left
        for j in range(n - 2, jj - 1, -1):
            if k >= len(fill):
                break
            array[i, j] = fill[k]
            k += 1
        
        # Up
        for i in range(m - 2, ii, -1):
            if k >= len(fill):
                break
            array[i, j] = fill[k]
            k += 1
        
        m -= 1
        n -= 1
        j += 1
        
    return array

# %% ../nbs/API/multi.ipynb 12
def whorlmap(multi_contrast, n=21, sort_by=None, cmap = 'vlag', vmax = None, vmin = None, reverse_neg=True, 
             abs_rank=False, chop_tail=0, ax=None, fig_size=None, title = None, heatmap_kwargs=None, plot_kwargs=None):
    """
    Create a whorlmap visualization of multiple contrasts.
    
    Parameters
    ----------
    multi_contrast : MultiContrast
        Object containing multiple dabest objects
    n : int, default 21
        Size of each spiral (n x n grid per contrast)
    sort_by : list, optional
        Order to sort contrasts by
    vmax, vmin : float, default None, None
        Color scale limits
    reverse_neg : bool, default True
        Whether to reverse negative values
    abs_rank : bool, default False
        Whether to rank by absolute value
    chop_tail : float, default 0
        Percentage of extreme values to exclude
    ax : matplotlib.Axes, optional
        Existing axes to plot on
    fig_size : tuple, optional
        Figure size (width, height) in inches
    title : str, optional
        Plot title
    heatmap_kwargs : dict, optional
        Additional keyword arguments passed to sns.heatmap().
        Common options include:
        - 'cmap': colormap (overrides direct cmap parameter)
        - 'vmin', 'vmax': color scale limits (override direct parameters)
        - 'center': center value for colormap
        - 'annot': whether to annotate cells with values
        - 'fmt': format string for annotations
        - 'linewidths': width of lines between cells
        - 'linecolor': color of lines between cells
        - 'cbar': whether to show colorbar
        - 'cbar_kws': colorbar customization dict
        - 'square': whether to make cells square
        - 'xticklabels', 'yticklabels': tick label control
        - 'mask': boolean array to mask cells
    plot_kwargs : dict, optional
        Additional keyword arguments for plot styling and layout.
        Available options (WIP):
        - 'title': plot title
        - 'xlabel', 'ylabel': axis labels
        - 'xticklabels', 'yticklabels': tick labels
        - 'xticklabels_rotation', 'yticklabels_rotation': tick label rotation angles
        - 'xticklabels_ha', 'yticklabels_ha': horizontal alignment 
    Returns
    -------
    tuple
        (figure, axes, mean_delta_dataframe) if ax is None, 
        else (axes, mean_delta_dataframe)
    """
    from .misc_tools import merge_two_dicts

    structure = multi_contrast.structure
    n_rows = structure['n_rows']
    n_cols = structure['n_cols']
    col_labels = structure['col_labels'] 
    row_labels = structure['row_labels']
    was_1d = (structure['type'] == '1D')

    # Initialize spirals and mean_delta DataFrames
    spirals = pd.DataFrame(np.zeros((n_rows * n, n_cols * n)))
    
    mean_delta = pd.DataFrame(np.zeros((n_rows, n_cols)), 
                             columns=col_labels, 
                             index=row_labels)
    
    # Get all bootstrap data from MultiContrast
    all_bootstraps = multi_contrast.bootstraps
    bootstrap_idx = 0

    for i in range(n_rows):
        for j in range(n_cols):
            contrast_idx = sort_by[j] if sort_by is not None else j
            
            # For mixed types, get bootstrap for specific position
            if isinstance(multi_contrast.contrast_type, dict) and multi_contrast.contrast_type.get('mixed'):
                bootstrap = multi_contrast.get_bootstrap_by_position(i, contrast_idx)
            else:
                # For homogeneous types, use the flattened bootstrap list
                flat_idx = i * n_cols + contrast_idx
                if flat_idx < len(all_bootstraps):
                    bootstrap = all_bootstraps[flat_idx]
                else:
                    # Handle case where we have fewer bootstraps than expected
                    bootstrap = all_bootstraps[bootstrap_idx]
                    bootstrap_idx += 1
            
            long_ranks = _sample_bootstrap(bootstrap, n, n, reverse_neg, abs_rank, chop_tail)
            spiral = _spiralize(long_ranks, n, n)
            spirals.iloc[i*n:i*n+n, j*n:j*n+n] = spiral
            mean_delta.iloc[i, j] = np.mean(long_ranks)
    
    if ax is None:
        f, a = plt.subplots(1, 1)
    else:
        a = ax
    if was_1d:
        cbar_orientation, cbar_location = 'horizontal', 'top'
    else:
        cbar_orientation, cbar_location = 'vertical', 'right'

    # heatmap kwargs
    default_heatmap_kwargs = {
        "cmap": cmap,
        "vmax": np.max(spirals.values) if vmax is None else vmax,
        "vmin": np.min(spirals.values) if vmin is None else vmin,
        "center": 0,
        "cbar_kws": {"shrink": 1, "pad": .05, "orientation": cbar_orientation, "location": cbar_location},
    }
    if heatmap_kwargs is None:
        heatmap_kwargs = default_heatmap_kwargs
    else:
        heatmap_kwargs = merge_two_dicts(
            default_heatmap_kwargs, heatmap_kwargs
        )

    # Create heatmap
    sns.heatmap(spirals, ax=a, **heatmap_kwargs)


    # Plot kwargs
    default_plot_kwargs = {
        "title": title,
        "xticklabels": col_labels,
        "xticklabels_rotation": 45,
        "xticklabels_ha":'right',
        "yticklabels": row_labels if not was_1d else [],
        "yticklabels_rotation": 0,
        "yticklabels_ha": 'right',
    }
    if plot_kwargs is None:
        plot_kwargs = default_plot_kwargs
    else:
        plot_kwargs = merge_two_dicts(
            default_plot_kwargs, plot_kwargs
        )

    # Set title
    if plot_kwargs.get('title') is not None:
        if ax is None:
            f.suptitle(plot_kwargs.get('title'))
        else:
            a.set_title(plot_kwargs.get('title'))

    # Set labels
    if plot_kwargs.get('xlabel') is not None:
        a.set_xlabel(plot_kwargs.get('xlabel'))
    if plot_kwargs.get('ylabel') is not None:
        a.set_ylabel(plot_kwargs.get('ylabel'))

    # Set labels
    a.set_xticks(np.linspace(n/2, n_cols*n-n/2, n_cols))
    a.set_xticklabels(plot_kwargs.get("xticklabels"), rotation=plot_kwargs.get("xticklabels_rotation"), ha=plot_kwargs.get("xticklabels_ha"))

    a.set_yticks([] if was_1d else np.linspace(n/2, n_rows*n-n/2, n_rows))
    a.set_yticklabels(plot_kwargs.get("yticklabels"), rotation=plot_kwargs.get("yticklabels_rotation"), ha=plot_kwargs.get("yticklabels_ha"))

    if ax is None:
        f.gca().set_aspect('equal')
        if fig_size is None:
            f.set_size_inches(n_cols/3, n_rows/3)
        else:
            f.set_size_inches(fig_size)
        return f, a, mean_delta
    else:
        return a, mean_delta

# %% ../nbs/API/multi.ipynb 13
__all__ = ['MultiContrast', 'combine', 'whorlmap']

