model_selection.py

import os
import shutil
import sys
import tempfile
from itertools import combinations
from typing import Dict
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from scipy import stats
from augment_omics.style import Style, set_pub_style, wrap_labels

import warnings
warnings.filterwarnings('ignore')

# Set publication style
set_pub_style()

PROT_VS_MET_PATH = 'results/data/survival_metrics/paired_proteomics_vs_metabolomics_with_cov.csv'
THRESHOLD = 'mid'

STYLE_COLORS = [Style.blue, Style.teal, Style.purple, Style.orange, Style.green,
                Style.red, Style.yellow, Style.brown, Style.olive, Style.silver]

# Models with a fixed color; everything else falls back to STYLE_COLORS by position.
FIXED_COLORS = {
    'CLIP':          Style.blue,
    'FedCoder':      Style.teal,
    'Transformer':   Style.purple,
    'MLP':           Style.orange,
    'Proteomics':    Style.green,
    'Ridge':         Style.red,
    'Metabolomics':  Style.yellow,
    'CLIP (τ=100)':  Style.brown,
}


def build_color_map(models) -> dict:
    """Assign a plot color per model, honouring FIXED_COLORS."""
    models = list(models) + ['Metabolomics']
    return {m: FIXED_COLORS.get(m, STYLE_COLORS[i % len(STYLE_COLORS)])
            for i, m in enumerate(models)}

def compare_models(model_dict: Dict[str, str], output_path: str = "model_comparison_results.csv"):
    """
    Compute per-model disease counts needed for plotting.

    Parameters:
    -----------
    model_dict : Dict[str, str]
        Dictionary where keys are model names and values are paths to CSV files
    output_path : str
        Path to save the results CSV

    Returns:
    --------
    pd.DataFrame
        DataFrame with Clinically_Predictable_Count per model
    """
    results = []

    for model_name, csv_path in model_dict.items():
        df = pd.read_csv(csv_path)
        print(f"Loaded {model_name}: {len(df)} diseases")
        results.append({
            'Model': model_name,
            'Total_Diseases': len(df),
            'Clinically_Predictable_Count': (df['cindex_1_mean'] > 0.75).sum(),
            'Metabolomics_Clinically_Predictable_Count': (df['cindex_2_mean'] > 0.75).sum(),
        })

    results_df = pd.DataFrame(results)
    results_df.to_csv(output_path, index=False)
    print(f"\nResults saved to: {output_path}")
    return results_df





def get_disease_subsets(prot_vs_met_path: str = PROT_VS_MET_PATH,
                        threshold: str = THRESHOLD):
    """
    Return two disease code arrays based on the proteomics-vs-metabolomics comparison.

    potential : diseases where proteomics is better than metabolomics (per threshold)
    rest      : all other diseases (not in potential)
    """
    df = pd.read_csv(prot_vs_met_path)

    if threshold == 'soft':
        mask = df['cindex_diff_mean'] > 0
    elif threshold == 'mid':
        mask = df['cindex_diff_mean'] > 0.01
    elif threshold == 'hard':
        mask = (df['cindex_diff_mean'] > 0.01) & (df['cindex_diff_pval'] < 0.05)
    else:
        raise ValueError(f"Unknown threshold: {threshold!r}. Choose 'soft', 'mid', or 'hard'.")

    potential = df[mask]['code'].values
    rest = df[~mask]['code'].values
    print(f"Disease subsets (threshold='{threshold}'): "
          f"potential={len(potential)}, rest={len(rest)}")
    return potential, rest


def load_all_models(model_dict: Dict[str, str],
                    disease_codes=None) -> Dict[str, pd.DataFrame]:
    """Load all model data from CSV files, optionally filtering to a set of disease codes."""
    all_data = {}
    for model_name, csv_path in model_dict.items():
        df = pd.read_csv(csv_path)
        if disease_codes is not None:
            df = df[df['code'].isin(disease_codes)].reset_index(drop=True)
        all_data[model_name] = df
        print(f"Loaded {model_name}: {len(df)} diseases")
    return all_data


def extract_metrics(all_data: Dict[str, pd.DataFrame]) -> pd.DataFrame:
    """
    Extract C-index, DR, and LR metrics for all models and baseline.
    """
    metrics_list = []
    
    for model_name, df in all_data.items():
        for idx, row in df.iterrows():
            metrics_list.append({
                'disease': row['disease'],
                'model_name': model_name,
                'cindex': row['cindex_1_mean'],
                'dr': row['dr_1_mean'],
                'lr': row['lr_1_mean'],
                'metric_type': 'model'
            })
            
            metrics_list.append({
                'disease': row['disease'],
                'model_name': f"{model_name}_baseline",
                'cindex': row['cindex_2_mean'],
                'dr': row['dr_2_mean'],
                'lr': row['lr_2_mean'],
                'metric_type': 'baseline'
            })
    
    return pd.DataFrame(metrics_list)


def perform_pairwise_tests(metrics_df: pd.DataFrame, metric_name: str) -> pd.DataFrame:
    """
    Perform pairwise statistical tests between all models.
    """
    models = [m for m in metrics_df['model_name'].unique() if not m.endswith('_baseline')]
    baseline_models = [m for m in metrics_df['model_name'].unique() if m.endswith('_baseline')]
    if baseline_models:
        models.append(baseline_models[0])

    results = []

    for model1, model2 in combinations(models, 2):
        data1 = metrics_df[metrics_df['model_name'] == model1].sort_values('disease')[metric_name].values
        data2 = metrics_df[metrics_df['model_name'] == model2].sort_values('disease')[metric_name].values
        t_stat, t_pval = stats.ttest_rel(data1, data2)
        w_stat, w_pval = stats.wilcoxon(data1, data2)

        diff = data1 - data2
        cohens_d = np.mean(diff) / np.std(diff)
        mean_diff = np.mean(data1) - np.mean(data2)

        if model1.endswith('_baseline'):
            model1 = 'Metabolomics'
        if model2.endswith('_baseline'):
            model2 = 'Metabolomics'
            
        if mean_diff > 0 and t_pval < 0.05:
            winner = model1
        elif mean_diff < 0 and t_pval < 0.05:
            winner = model2
        else:
            winner = "No significant difference"
        
        results.append({
            'Model_1': model1, 'Model_2': model2,
            'Metric': metric_name.upper(),
            'Mean_1': np.mean(data1), 'Mean_2': np.mean(data2),
            'Mean_Difference': mean_diff,
            'T_statistic': t_stat, 'T_pvalue': t_pval,
            'Wilcoxon_statistic': w_stat, 'Wilcoxon_pvalue': w_pval,
            'Cohens_d': cohens_d, 'Winner': winner,
            'Significant_t': 'Yes' if t_pval < 0.05 else 'No',
            'Significant_w': 'Yes' if w_pval < 0.05 else 'No'
        })

    out = pd.DataFrame(results)

    # All-pairs model comparison is a screen, so correct for it. The family is
    # the pairs among architectures for one metric in one disease subset; the
    # subsets are separate runs and the metrics are corrected separately.
    # Reported alongside the raw p rather than replacing it, and Winner and
    # Significant_t are deliberately left on the uncorrected p so that figures
    # and counts elsewhere in this script are unaffected.
    out['T_qvalue'] = stats.false_discovery_control(out.T_pvalue.values, method='bh')
    out['Wilcoxon_qvalue'] = stats.false_discovery_control(out.Wilcoxon_pvalue.values,
                                                          method='bh')
    out['Significant_t_fdr'] = np.where(out.T_qvalue < 0.05, 'Yes', 'No')
    out['Significant_w_fdr'] = np.where(out.Wilcoxon_qvalue < 0.05, 'Yes', 'No')

    return out


def create_summary_statistics(metrics_df: pd.DataFrame) -> pd.DataFrame:
    """Create summary statistics table for all models and metrics."""
    model_data = metrics_df[~metrics_df['model_name'].str.endswith('_baseline')]
    
    summary_list = []
    
    for model in model_data['model_name'].unique():
        model_subset = model_data[model_data['model_name'] == model]
        for metric in ['cindex', 'dr', 'lr']:
            values = model_subset[metric].values
            summary_list.append({
                'Model': model, 'Metric': metric.upper(),
                'Mean': np.mean(values), 'Median': np.median(values),
                'Std': np.std(values), 'Min': np.min(values),
                'Max': np.max(values),
                'Q25': np.percentile(values, 25),
                'median': np.median(values),
                'Q75': np.percentile(values, 75),
                'IQR': np.percentile(values, 75) - np.percentile(values, 25)
            })
    
    # Add baseline
    first_model = model_data['model_name'].iloc[0]
    baseline_data = metrics_df[metrics_df['model_name'] == f"{first_model}_baseline"]
    
    for metric in ['cindex', 'dr', 'lr']:
        values = baseline_data[metric].values
        summary_list.append({
            'Model': 'Metabolomics', 'Metric': metric.upper(),
            'Mean': np.mean(values), 'Median': np.median(values),
            'Std': np.std(values), 'Min': np.min(values),
            'Max': np.max(values),
            'Q25': np.percentile(values, 25),
            'median': np.median(values),
            'Q75': np.percentile(values, 75),
            'IQR': np.percentile(values, 75) - np.percentile(values, 25)
        })
    
    return pd.DataFrame(summary_list)


def _build_pairwise_win_matrix(pairwise_results: pd.DataFrame) -> tuple:
    """
    Build an aggregated pairwise win matrix across all metrics (C-index, DR, LR).
    
    For each pair of models, counts the number of metrics where model A significantly
    beats model B (p < 0.05 and positive mean difference).
    
    Returns:
    --------
    win_matrix : np.ndarray
        Matrix where win_matrix[i, j] = number of metrics model i beats model j
    models : list
        Ordered list of model names
    """
    models = sorted(
        set(pairwise_results['Model_1'].unique()) | set(pairwise_results['Model_2'].unique())
    )
    n = len(models)
    win_matrix = np.zeros((n, n), dtype=int)
    
    for _, row in pairwise_results.iterrows():
        i = models.index(row['Model_1'])
        j = models.index(row['Model_2'])
        
        if row['Significant_t'] == 'Yes':
            if row['Mean_Difference'] > 0:
                win_matrix[i, j] += 1
            else:
                win_matrix[j, i] += 1
    
    return win_matrix, models



def create_heatmap_comparison(pairwise_results: pd.DataFrame, metric: str, output_path: str):
    """Create a heatmap showing pairwise p-values and mean differences."""
    metric_results = pairwise_results[pairwise_results['Metric'] == metric].copy()

    models = sorted(set(metric_results['Model_1'].unique()) | set(metric_results['Model_2'].unique()))

    pval_matrix = np.zeros((len(models), len(models)))
    diff_matrix = np.zeros((len(models), len(models)))

    for _, row in metric_results.iterrows():
        i = models.index(row['Model_1'])
        j = models.index(row['Model_2'])

        pval_matrix[i, j] = row['Wilcoxon_pvalue']
        pval_matrix[j, i] = row['Wilcoxon_pvalue']

        diff_matrix[i, j] = row['Mean_Difference']
        diff_matrix[j, i] = -row['Mean_Difference']

    fig, axes = plt.subplots(1, 2, figsize=(Style.width * 2, Style.height))

    sns.heatmap(pval_matrix, annot=True, fmt='.3f', cmap='RdYlGn_r',
                xticklabels=models, yticklabels=models, ax=axes[0],
                cbar_kws={'label': 'P-value'}, vmin=0, vmax=0.1,
                annot_kws={'size': 6})
    axes[0].set_title(f'{metric} - Pairwise P-values (Wilcoxon)\nGreen = Significant Difference')

    sns.heatmap(diff_matrix, annot=True, fmt='.3f', cmap='RdBu_r',
                xticklabels=models, yticklabels=models, ax=axes[1],
                cbar_kws={'label': 'Mean Difference'}, center=0,
                annot_kws={'size': 6})
    axes[1].set_title(f'{metric} - Mean Differences\nRed = Row > Column, Blue = Row < Column')

    plt.tight_layout()
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    print(f"Saved heatmap: {output_path}")
    plt.close()


def rank_models_by_metric(pairwise_results: pd.DataFrame, metric: str) -> pd.DataFrame:
    """Rank models based on pairwise comparison results."""
    metric_results = pairwise_results[pairwise_results['Metric'] == metric].copy()
    
    all_models = set(metric_results['Model_1'].unique()) | set(metric_results['Model_2'].unique())
    
    model_scores = {}
    for model in all_models:
        wins = 0
        total_comparisons = 0
        significant_wins = 0
        
        for _, row in metric_results.iterrows():
            if row['Model_1'] == model or row['Model_2'] == model:
                total_comparisons += 1
                if row['Winner'] == model:
                    wins += 1
                    if row['Significant_t'] == 'Yes':
                        significant_wins += 1
        
        model_scores[model] = {
            'Model': model, 'Wins': wins,
            'Total_Comparisons': total_comparisons,
            'Win_Rate': wins / total_comparisons if total_comparisons > 0 else 0,
            'Significant_Wins': significant_wins
        }
    
    ranking_df = pd.DataFrame(model_scores.values())
    ranking_df = ranking_df.sort_values('Win_Rate', ascending=False).reset_index(drop=True)
    ranking_df['Rank'] = ranking_df['Win_Rate'].rank(method='min', ascending=False).astype(int)
    
    return ranking_df[['Rank', 'Model', 'Wins', 'Total_Comparisons', 'Win_Rate', 'Significant_Wins']]


def visualize_and_test_models(model_dict: Dict[str, str], output_dir: str = ".",
                              model_selection_results_path: str = None):
    """
    Main function to visualize and statistically test all models.
    """
    print("="*80)
    print("MODEL VISUALIZATION AND STATISTICAL TESTING")
    print("="*80)
    
    # Load all data
    print("\n1. Loading data...")
    all_data = load_all_models(model_dict)
    
    # Extract metrics
    print("\n2. Extracting metrics...")
    metrics_df = extract_metrics(all_data)
    
    # Create summary statistics
    print("\n3. Computing summary statistics...")
    summary_stats = create_summary_statistics(metrics_df)
    summary_stats.to_csv(f"{output_dir}/summary_statistics.csv", index=False)
    print(f"Saved: {output_dir}/summary_statistics.csv")
    
    print("\n" + "="*80)
    print("SUMMARY STATISTICS")
    print("="*80)
    print(summary_stats.to_string(index=False))

    # Delta C-index (model - metabolomics/baseline) per model
    print("\n" + "="*80)
    print("DELTA C-INDEX (Model − Metabolomics)")
    print("="*80)
    model_rows = metrics_df[~metrics_df['model_name'].str.endswith('_baseline')]
    first_model = model_rows['model_name'].iloc[0]
    baseline_rows = metrics_df[metrics_df['model_name'] == f"{first_model}_baseline"][['disease', 'cindex']].rename(columns={'cindex': 'cindex_baseline'})
    for model in model_rows['model_name'].unique():
        model_subset = model_rows[model_rows['model_name'] == model][['disease', 'cindex']]
        merged = model_subset.merge(baseline_rows, on='disease')
        delta = merged['cindex'] - merged['cindex_baseline']
        print(f"  {model}: mean={delta.mean():.4f}, median={delta.median():.4f}, std={delta.std():.4f}, min={delta.min():.4f}, max={delta.max():.4f}")

    # Perform pairwise statistical tests
    print("\n5. Performing pairwise statistical tests...")
    all_pairwise_results = []
    
    for metric in ['cindex', 'dr', 'lr']:
        print(f"   - Testing {metric.upper()}...")
        pairwise_results = perform_pairwise_tests(metrics_df, metric)
        all_pairwise_results.append(pairwise_results)
    
    combined_pairwise = pd.concat(all_pairwise_results, ignore_index=True)
    combined_pairwise.to_csv(f"{output_dir}/pairwise_statistical_tests.csv", index=False)
    print(f"Saved: {output_dir}/pairwise_statistical_tests.csv")
    
    # Print pairwise results
    print("\n" + "="*80)
    print("PAIRWISE STATISTICAL TEST RESULTS")
    print("="*80)
    for metric in ['CINDEX', 'DR', 'LR']:
        print(f"\n{metric}:")
        print("-"*80)
        metric_results = combined_pairwise[combined_pairwise['Metric'] == metric]
        print(metric_results[['Model_1', 'Model_2', 'Mean_Difference',
                             'T_pvalue', 'Wilcoxon_pvalue', 'Significant_t']].to_string(index=False))

    # Create combined 3x2 figure
    print("\n6. Creating combined 3x2 figure...")
    model_selection_results = pd.read_csv(model_selection_results_path)
    # plot_combined_figure(
    #     metrics_df, combined_pairwise, model_selection_results,
    #     f"{output_dir}/fig1_combined_3x2.png"
    # )

    # Create fig1_row23 figure (panels b, c, d) and fig_1_row1 scatter
    print("\n6b. Creating fig1_row23 figure (panels b, c, d) and fig_1_row1 scatter...")
    plot_def_figure(all_data, model_selection_results, f"{output_dir}/fig1_cindex_row23.png",
                    scatter_output_path=f"{output_dir}/fig_1_scatter_row1.png",
                    pairwise_results=combined_pairwise)

    print("\n6c. Creating DR and LR figures...")
    plot_dr_lr_figures(all_data, combined_pairwise, output_dir=output_dir)

    print("\n" + "="*80)
    print("ANALYSIS COMPLETE!")
    print("="*80)

    return {
        'metrics_df': metrics_df,
        'summary_stats': summary_stats,
        'pairwise_results': combined_pairwise,
    }


def _plot_pairwise_heatmap_panel(ax, pairwise_results: pd.DataFrame, metric: str):
    """Render pairwise mean-difference heatmap onto ax.

    Parameters
    ----------
    metric : str
        Uppercase metric name matching the 'Metric' column, e.g. 'CINDEX', 'DR', 'LR'.
    """
    _display = {'CINDEX': 'C-Index', 'DR': 'DR', 'LR': 'LR'}
    metric_pw = pairwise_results[pairwise_results['Metric'] == metric].copy()
    all_pw_models = sorted(set(metric_pw['Model_1']) | set(metric_pw['Model_2']))
    n = len(all_pw_models)

    diff_mat = np.full((n, n), np.nan)
    pval_mat = np.full((n, n), np.nan)

    for _, row in metric_pw.iterrows():
        i = all_pw_models.index(row['Model_1'])
        j = all_pw_models.index(row['Model_2'])
        diff_mat[i, j] = row['Mean_Difference']
        diff_mat[j, i] = -row['Mean_Difference']
        pval_mat[i, j] = row['Wilcoxon_pvalue']
        pval_mat[j, i] = row['Wilcoxon_pvalue']

    pw_order = np.argsort(np.nanmean(diff_mat, axis=1))[::-1]
    diff_mat = diff_mat[np.ix_(pw_order, pw_order)]
    pval_mat = pval_mat[np.ix_(pw_order, pw_order)]
    pw_models = [all_pw_models[i] for i in pw_order]

    vmax_pw = np.nanmax(np.abs(np.where(np.eye(n, dtype=bool), 0, diff_mat)))
    cmap_pw = plt.cm.RdBu_r.copy()
    cmap_pw.set_bad(color='#f0f0f0')
    diff_masked = np.ma.masked_where(np.eye(n, dtype=bool), diff_mat)

    ax.imshow(diff_masked, aspect='equal', cmap=cmap_pw,
              vmin=-vmax_pw, vmax=vmax_pw, interpolation='nearest')

    def _sig_star(p):
        if np.isnan(p):
            return ''
        if p < 0.001:
            return '***'
        if p < 0.01:
            return '**'
        if p < 0.05:
            return '*'
        return ''

    for i in range(n):
        for j in range(n):
            if i == j:
                ax.text(j, i, '—', ha='center', va='center', fontsize=6, color='gray')
            else:
                val = diff_mat[i, j]
                star = _sig_star(pval_mat[i, j])
                if not np.isnan(val):
                    label = f'{val:.3f}{star}'
                    text_color = 'white' if abs(val) > vmax_pw * 0.65 else 'black'
                    ax.text(j, i, label, ha='center', va='center',
                            fontsize=6, color=text_color)

    ax.set_xticks(range(n))
    ax.set_xticklabels(pw_models, rotation=45, ha='right')
    ax.set_yticks(range(n))
    ax.set_yticklabels(pw_models)
    ax.set_title(f'Row − Column (Mean Delta {_display.get(metric, metric)})',
                 fontsize=8, fontstyle='italic', pad=3)


def _plot_delta_box_panel(ax, all_data: Dict[str, pd.DataFrame],
                          metric: str, model_color_map: dict):
    """Render delta (model − baseline) box plot onto ax.

    Parameters
    ----------
    metric : str
        Lowercase metric name, e.g. 'cindex', 'dr', 'lr'.
    """
    col_model = f'{metric}_1_mean'
    col_base  = f'{metric}_2_mean'
    _ylabel_map = {'cindex': 'Delta C-Index', 'dr': 'Delta DR', 'lr': 'Delta LR'}
    ylabel = f"{_ylabel_map.get(metric, f'Delta {metric.upper()}')} (Model − Metabolomics)"

    models = list(all_data.keys())
    first_model = models[0]
    diseases = all_data[first_model]['disease'].values

    delta_data = []
    for model_name, df in all_data.items():
        df_indexed = df.set_index('disease')
        delta = df_indexed.loc[diseases, col_model] - df_indexed.loc[diseases, col_base]
        delta_data.append(delta.dropna().values)

    sort_order = np.argsort([np.median(d) for d in delta_data])
    models_sorted = [models[i] for i in sort_order]
    delta_data_sorted = [delta_data[i] for i in sort_order]
    box_colors = [model_color_map[m] for m in models_sorted]

    bp = ax.boxplot(delta_data_sorted, positions=range(len(models_sorted)),
                    patch_artist=True, showmeans=True, widths=0.5,
                    meanprops=dict(marker='D', markerfacecolor='black', markersize=3),
                    medianprops=dict(color='black', linewidth=1),
                    whiskerprops=dict(linewidth=0.8), capprops=dict(linewidth=0.8),
                    flierprops=dict(marker='o', markersize=1.5, alpha=0.5))
    for patch, color in zip(bp['boxes'], box_colors):
        patch.set_facecolor(color)
        patch.set_alpha(0.7)

    ax.axhline(y=0, color='gray', linestyle='--', linewidth=1.25)
    ax.set_ylabel(ylabel)
    ax.set_xticks(range(len(models_sorted)))
    ax.set_xticklabels(wrap_labels(models_sorted, width=12), rotation=45, ha='right')


def plot_def_figure(all_data: Dict[str, pd.DataFrame],
                    model_selection_results: pd.DataFrame,
                    output_path: str,
                    scatter_output_path: str = None,
                    pairwise_results: pd.DataFrame = None,
                    prot_vs_met_path: str = PROT_VS_MET_PATH,
                    threshold: str = THRESHOLD):
    """
    Create figures for model comparison panels.

    Saves two files:
    - scatter_output_path (fig_1_row1.png): Scatter proteomics vs metabolomics C-index
    - output_path (fig1_row23.png): 1x3 figure with pairwise heatmap, delta box plot, bar chart
    """
    models = list(all_data.keys())
    model_color_map = build_color_map(models)

    # ---- Scatter figure: proteomics vs metabolomics (all diseases) ----
    if scatter_output_path is not None:
        full_prot_df = pd.read_csv(prot_vs_met_path)
        potential_codes, _ = get_disease_subsets(prot_vs_met_path, threshold)
        pot_set = set(potential_codes)

        x_all = full_prot_df['cindex_2_mean'].values   # metabolomics
        y_all = full_prot_df['cindex_1_mean'].values    # proteomics
        is_pot = np.array([c in pot_set for c in full_prot_df['code'].values])

        fig_scatter, ax_scatter = plt.subplots(1, 3, figsize=(Style.width * 1.5, Style.height * 0.55))
        ax_scatter[0].scatter(x_all[~is_pot], y_all[~is_pot], alpha=0.6, s=14,
                           color=Style.blue, edgecolors='none', label='Remaining set', zorder=2)
        ax_scatter[0].scatter(x_all[is_pot], y_all[is_pot], alpha=0.8, s=18,
                           color=Style.red, edgecolors='none', label='Proteomics-advantaged set', zorder=3)

        lo = min(x_all.min(), y_all.min()) - 0.03
        hi = max(x_all.max(), y_all.max()) + 0.03
        ax_scatter[0].plot([lo, hi], [lo, hi], color='gray', linestyle='--',
                        linewidth=0.8, alpha=0.7)
        ax_scatter[0].set_xlabel('Metabolomics C-Index')
        ax_scatter[0].set_ylabel('Proteomics C-Index')
        ax_scatter[0].set_xlim(lo, hi)
        ax_scatter[0].set_ylim(lo, hi)
        ax_scatter[0].set_aspect('equal', adjustable='box')
        ax_scatter[0].legend(loc='upper left', markerscale=0.8,
                          frameon=False, handletextpad=0.3)

        plt.tight_layout()
        plt.savefig(scatter_output_path, dpi=300, bbox_inches='tight')
        print(f"Saved: {scatter_output_path}")
        plt.close(fig_scatter)

    # ---- Main figure: 1x3 (heatmap, box, bar) ----
    fig, axes = plt.subplots(1, 3, figsize=(Style.width * 1.5, Style.height * 0.55))
    ax_heatmap = axes[0]
    ax_box    = axes[1]
    ax_bar    = axes[2]

    # ---- Panel b: Pairwise model comparison heatmap ----
    if pairwise_results is not None:
        cindex_pw = pairwise_results[pairwise_results['Metric'] == 'CINDEX'].copy()
        all_pw_models = sorted(set(cindex_pw['Model_1']) | set(cindex_pw['Model_2']))
        n = len(all_pw_models)

        diff_mat = np.full((n, n), np.nan)
        pval_mat = np.full((n, n), np.nan)

        for _, row in cindex_pw.iterrows():
            i = all_pw_models.index(row['Model_1'])
            j = all_pw_models.index(row['Model_2'])
            diff_mat[i, j] = row['Mean_Difference']
            diff_mat[j, i] = -row['Mean_Difference']
            pval_mat[i, j] = row['Wilcoxon_pvalue']
            pval_mat[j, i] = row['Wilcoxon_pvalue']

        # Sort: best mean C-Index difference on top/left
        pw_order = np.argsort(np.nanmean(diff_mat, axis=1))[::-1]
        diff_mat = diff_mat[np.ix_(pw_order, pw_order)]
        pval_mat = pval_mat[np.ix_(pw_order, pw_order)]
        pw_models = [all_pw_models[i] for i in pw_order]

        vmax_pw = np.nanmax(np.abs(np.where(np.eye(n, dtype=bool), 0, diff_mat)))
        cmap_pw = plt.cm.RdBu_r.copy()
        cmap_pw.set_bad(color='#f0f0f0')
        diff_masked = np.ma.masked_where(np.eye(n, dtype=bool), diff_mat)

        im = ax_heatmap.imshow(diff_masked, aspect='equal', cmap=cmap_pw,
                               vmin=-vmax_pw, vmax=vmax_pw, interpolation='nearest')

        def _sig_star(p):
            if np.isnan(p): return ''
            if p < 0.001: return '***'
            if p < 0.01: return '**'
            if p < 0.05: return '*'
            return ''

        for i in range(n):
            for j in range(n):
                if i == j:
                    ax_heatmap.text(j, i, '—', ha='center', va='center', fontsize=6,
                                    color='gray')
                else:
                    val = diff_mat[i, j]
                    star = _sig_star(pval_mat[i, j])
                    if not np.isnan(val):
                        label = f'{val:.3f}{star}'
                        text_color = 'white' if abs(val) > vmax_pw * 0.65 else 'black'
                        ax_heatmap.text(j, i, label, ha='center', va='center',
                                        fontsize=6, color=text_color)

        ax_heatmap.set_xticks(range(n))
        ax_heatmap.set_xticklabels(pw_models, rotation=45, ha='right')
        ax_heatmap.set_yticks(range(n))
        ax_heatmap.set_yticklabels(pw_models)
        ax_heatmap.set_title('Row − Column (Mean Delta C-Index)', fontsize=8, fontstyle='italic', pad=3)
        

    # ---- Panel c: Delta C-Index box plot ----
    first_model = models[0]
    diseases = all_data[first_model]['disease'].values

    delta_data = []
    for model_name, df in all_data.items():
        df_indexed = df.set_index('disease')
        delta = df_indexed.loc[diseases, 'cindex_1_mean'] - df_indexed.loc[diseases, 'cindex_2_mean']
        delta_data.append(delta.dropna().values)

    sort_order = np.argsort([np.median(d) for d in delta_data])
    models_sorted_box = [models[i] for i in sort_order]
    delta_data_sorted = [delta_data[i] for i in sort_order]
    box_colors = [model_color_map[m] for m in models_sorted_box]

    bp = ax_box.boxplot(delta_data_sorted, positions=range(len(models_sorted_box)),
                        patch_artist=True, showmeans=True, widths=0.5,
                        meanprops=dict(marker='D', markerfacecolor='black', markersize=3),
                        medianprops=dict(color='black', linewidth=1),
                        whiskerprops=dict(linewidth=0.8), capprops=dict(linewidth=0.8),
                        flierprops=dict(marker='o', markersize=1.5, alpha=0.5))
    for patch, color in zip(bp['boxes'], box_colors):
        patch.set_facecolor(color)
        patch.set_alpha(0.7)

    ax_box.axhline(y=0, color='gray', linestyle='--', linewidth=1.25)
    ax_box.set_ylabel('Delta C-Index (Model − Metabolomics)')
    ax_box.set_xticks(range(len(models_sorted_box)))
    ax_box.set_xticklabels(wrap_labels(models_sorted_box, width=12), rotation=45, ha='right')

    # ---- Panel d: Clinically predictable bar chart ----
    bar_labels = models + ['Metabolomics']
    bar_counts = []
    for m in models:
        row = model_selection_results[model_selection_results['Model'] == m]
        bar_counts.append(int(row['Clinically_Predictable_Count'].values[0]))
    metabolomics_count = int(model_selection_results.iloc[0]['Metabolomics_Clinically_Predictable_Count'])
    bar_counts.append(metabolomics_count)

    combined_bar = sorted(zip(bar_labels, bar_counts), key=lambda x: x[1])
    bar_labels, bar_counts = zip(*combined_bar)
    bar_labels, bar_counts = list(bar_labels), list(bar_counts)
    bar_colors = [model_color_map[m] for m in bar_labels]

    bars = ax_bar.bar(range(len(bar_labels)), bar_counts, color=bar_colors,
                      alpha=0.7, edgecolor='black', linewidth=1.0)
    for bar, count in zip(bars, bar_counts):
        ax_bar.text(bar.get_x() + bar.get_width() / 2., bar.get_height() + 0.3,
                    str(count), ha='center', va='bottom', fontsize=8, fontweight='bold')
    ax_bar.set_xticks(range(len(bar_labels)))
    ax_bar.set_xticklabels(wrap_labels(bar_labels, width=12), rotation=45, ha='right')
    ax_bar.set_ylabel('Diseases with C-index > 0.75')
    ax_bar.axhline(y=metabolomics_count, color='gray', linestyle='--', linewidth=1.25)
    ax_bar.set_ylim(bottom=0)

    # ---- Panel labels ----
    # for ax, label in zip(axes, ['b', 'c', 'd', 'e']):
    #     ax.text(-0.05, 1.02, label, transform=ax.transAxes,fontsize=8,
    #             fontweight='bold', va='bottom', ha='right')

    plt.tight_layout()
    plt.savefig(output_path, dpi=300, bbox_inches='tight')
    print(f"Saved: {output_path}")
    plt.close()


def plot_dr_lr_figures(all_data: Dict[str, pd.DataFrame],
                       pairwise_results: pd.DataFrame,
                       output_dir: str):
    """
    Create heatmap + delta box figures for DR and LR metrics.

    Saves two files:
    - {output_dir}/fig_dr_row23.png
    - {output_dir}/fig_lr_row23.png
    """
    model_color_map = build_color_map(all_data.keys())

    for metric_upper, metric_lower in [('DR', 'dr'), ('LR', 'lr')]:
        fig, axes = plt.subplots(1, 2, figsize=(Style.width, Style.height * 0.55))
        _plot_pairwise_heatmap_panel(axes[0], pairwise_results, metric_upper)
        _plot_delta_box_panel(axes[1], all_data, metric_lower, model_color_map)
        plt.tight_layout()
        out_path = os.path.join(output_dir, f'fig_{metric_lower}_row23.png')
        plt.savefig(out_path, dpi=300, bbox_inches='tight')
        print(f"Saved: {out_path}")
        plt.close()


class TeeOutput:
    """Redirect stdout to both console and a file."""
    def __init__(self, filepath):
        self.file = open(filepath, 'w')
        self.stdout = sys.stdout

    def write(self, text):
        self.stdout.write(text)
        self.file.write(text)

    def flush(self):
        self.stdout.flush()
        self.file.flush()

    def close(self):
        self.file.close()


def _run_subset(model_dict: Dict[str, str], output_dir: str,
                disease_codes, subset_label: str, summary_filename: str):
    """Run the full compare + visualize pipeline for a specific disease code subset."""
    os.makedirs(output_dir, exist_ok=True)
    vis_dir = os.path.join(output_dir, "visualization_and_stat_tests")
    os.makedirs(vis_dir, exist_ok=True)

    summary_path = os.path.join(output_dir, summary_filename)
    tee = TeeOutput(summary_path)
    sys.stdout = tee

    try:
        print(f"\n{'='*80}")
        print(f"SUBSET: {subset_label} ({len(disease_codes)} diseases)")
        print(f"{'='*80}\n")

        # Build filtered model_dict by loading with code filter
        # Write filtered CSVs to a temp location so existing functions work unchanged.
        tmp_dir = tempfile.mkdtemp()
        filtered_model_dict = {}
        try:
            for model_name, csv_path in model_dict.items():
                df = pd.read_csv(csv_path)
                df = df[df['code'].isin(disease_codes)].reset_index(drop=True)
                tmp_path = os.path.join(tmp_dir, f"{model_name}.csv")
                df.to_csv(tmp_path, index=False)
                filtered_model_dict[model_name] = tmp_path

            model_selection_results_path = os.path.join(output_dir, "model_selection_results.csv")
            ms_results = compare_models(filtered_model_dict, output_path=model_selection_results_path)
            print("\nClinically Predictable Diseases (C-index > 0.75):")
            print("-" * 50)
            for _, row in ms_results.sort_values('Clinically_Predictable_Count').iterrows():
                print(f"  {row['Model']}: {int(row['Clinically_Predictable_Count'])} / {int(row['Total_Diseases'])}")
            print(f"  Metabolomics: {int(ms_results.iloc[0]['Metabolomics_Clinically_Predictable_Count'])} / {int(ms_results.iloc[0]['Total_Diseases'])}")

            results = visualize_and_test_models(
                filtered_model_dict,
                output_dir=vis_dir,
                model_selection_results_path=model_selection_results_path
            )
        finally:
            shutil.rmtree(tmp_dir)

        print(f"\nSummary saved to: {summary_path}")
        return results

    finally:
        sys.stdout = tee.stdout
        tee.close()


def main(
    model_dict: Dict[str, str],
    output_dir: str,
    summary_filename: str = "summary.txt",
    threshold: str = THRESHOLD,
    prot_vs_met_path: str = PROT_VS_MET_PATH,
):
    """
    Main function to run model comparison and statistical analysis.

    Runs twice: once for diseases where proteomics > metabolomics ('potential'),
    and once for the remaining diseases ('rest').
    """
    os.makedirs(output_dir, exist_ok=True)

    potential_codes, _ = get_disease_subsets(prot_vs_met_path, threshold)

    # Save disease codes with names and proteomics_better flag
    prot_vs_met_df = pd.read_csv(prot_vs_met_path)[['code', 'disease']].drop_duplicates()
    rest_codes = np.setdiff1d(prot_vs_met_df['code'].values, potential_codes)
    disease_subset_df = prot_vs_met_df.copy()
    disease_subset_df['proteomics_better'] = disease_subset_df['code'].isin(potential_codes)
    disease_subset_df = disease_subset_df.rename(columns={'disease': 'disease_name', 'code': 'disease_code'})
    disease_subset_df = disease_subset_df[['disease_code', 'disease_name', 'proteomics_better']]
    disease_subset_path = os.path.join(output_dir, "disease_subsets.csv")
    disease_subset_df.to_csv(disease_subset_path, index=False)
    print(f"Disease subsets saved to: {disease_subset_path}")

    results = {}

    # --- Subset 1: diseases where proteomics is better (potential) ---
    potential_dir = os.path.join(output_dir, "prot_better")
    results['potential'] = _run_subset(
        model_dict, potential_dir, potential_codes,
        subset_label=f"prot_better (threshold='{threshold}')",
        summary_filename=summary_filename
    )

    # --- Subset 2: remaining diseases (rest) ---
    rest_dir = os.path.join(output_dir, "rest")
    results['rest'] = _run_subset(
        model_dict, rest_dir, rest_codes,
        subset_label=f"rest (not in prot_better, threshold='{threshold}')",
        summary_filename=summary_filename
    )

    return results


if __name__ == '__main__':
    ## model selection
    model_dict = {
            'CLIP': 'results/data/survival_metrics/paired_clip_vs_baseline_with_cov.csv',
            'FedCoder': 'results/data/survival_metrics/paired_fedcoder_vs_baseline_with_cov.csv',
            'Transformer': 'results/data/survival_metrics/paired_transformer_translation_vs_baseline_with_cov.csv',
            'MLP': 'results/data/survival_metrics/paired_mlp_vs_baseline_with_cov.csv',
            'Proteomics':'results/data/survival_metrics/paired_proteomics_vs_metabolomics_with_cov.csv',
            'Ridge':'results/data/survival_metrics/paired_ridge_vs_baseline_with_cov.csv'
        }
    output_dir = "results/main/model_selection"
    main(model_dict, output_dir)
    
    ## ablation
    model_dict_ablation = {
            'CLIP': 'results/data/survival_metrics/paired_clip_vs_baseline_no_cov.csv',
            'CLIP (τ=100)': 'results/data/ablation_tau/paired_clip_tau_100_vs_baseline_no_cov.csv',
    }
    output_dir_ablation = "results/main/ablation_contrastive"
    main(model_dict_ablation, output_dir_ablation)