multiple_task_analysis.py

"""
Post-training analysis of a multi-task MPN.

Loads a trained DeepMultiPlasticNet checkpoint and its recorded hidden states,
modulation matrices (M), and input activations, then characterizes how the
network organizes task-specific computation:

1. Weight structure visualization — heatmaps of W_initial_linear, W (recurrent),
   and W_output to inspect static connectivity.
2. Task-conditioned activity — averages hidden-layer and modulation responses
   per task to build neuron × task tuning profiles.
3. Hierarchical clustering — clusters input neurons (by their tuning across
   tasks) and hidden neurons (likewise), with automatic k selection via
   silhouette score. Also clusters modulation synapses (columns of M) by their
   task-tuning patterns.
4. Low-dimensional trajectory plots — PCA projections of hidden-state
   trajectories colored by task/stimulus/response.

Outputs:
  - Figures saved to ./multiple_tasks/{aname}/
  - cluster_info_{aname}.pkl — neuron cluster assignments for downstream
    lesion experiments
  - cluster_info_mod_{aname}.pkl — modulation synapse cluster assignments
"""
# %%
import os
import numpy as np
from numpy.linalg import norm
from pathlib import Path
import json
import psutil
import copy
import pickle

import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.ticker as ticker
from matplotlib.colors import LogNorm
from sklearn.metrics import adjusted_rand_score, normalized_mutual_info_score
from sklearn.metrics.cluster import contingency_matrix as sk_contingency_matrix
ticker.Locator.MAXTICKS = 10000 
import seaborn as sns 

mpl.rcParams.update({
    "font.family": "sans-serif",
    "font.sans-serif": ["Arial", "Helvetica", "DejaVu Sans"],  
    "font.size": 8,
    "axes.labelsize": 8,
    "axes.titlesize": 8,
    "xtick.labelsize": 7,
    "ytick.labelsize": 7,
    "pdf.fonttype": 42,  
    "ps.fonttype": 42,
})


from sklearn.metrics.pairwise import cosine_similarity
from scipy.spatial.distance import pdist, squareform
from sklearn.decomposition import PCA

import torch
from torch.serialization import add_safe_globals

import helper           
import clustering
import clustering_metric
import color_func
import mpn 
import mpn_tasks

c_vals = [
    "#e53e3e",  # red
    "#3182ce",  # blue
    "#38a169",  # green
    "#d69e2e",  # yellow-gold
    "#d53f8c",  # pink-magenta
    "#4c51bf",  # indigo
    "#dd6b20",  # orange
    "#0ea5e9",  # sky blue
    "#22c55e",  # bright green
    "#a855f7",  # purple
    "#f43f5e",  # red-pink
    "#0f766e",  # teal
    "#b83280",  # magenta-violet
    "#ca8a04",  # amber
    "#2b6cb0",  # deep blue
] * 10

c_vals_l = [
    "#feb2b2",  # light red
    "#90cdf4",  # light blue
    "#9ae6b4",  # light green
    "#faf089",  # light yellow-gold
    "#fbb6ce",  # light magenta
    "#c3dafe",  # light indigo
    "#fed7aa",  # light orange
    "#bae6fd",  # light sky blue
    "#bbf7d0",  # light bright green
    "#e9d5ff",  # light purple
    "#fecdd3",  # light red-pink
    "#a7f3d0",  # light teal
    "#f9a8d4",  # light violet-magenta
    "#fde68a",  # light amber
    "#bfdbfe",  # light deep blue
] * 10

c_vals_d = [
    "#9b2c2c",  # dark red
    "#2c5282",  # dark blue
    "#276749",  # dark green
    "#975a16",  # dark golden
    "#97266d",  # dark magenta
    "#4338ca",  # dark indigo
    "#7b341e",  # dark orange
    "#0369a1",  # dark sky blue
    "#15803d",  # dark bright green
    "#6b21a8",  # dark purple
    "#9f1239",  # dark red-pink
    "#0f4c3a",  # dark teal
    "#702459",  # dark violet-magenta
    "#854d0e",  # dark amber
    "#1e3a8a",  # dark deep blue
] * 10


l_vals = ['solid', 'dashed', 'dotted', 'dashdot', '-', '--', '-.', ':', (0, (3, 1, 1, 1)), (0, (5, 10))]
markers_vals = ['o', 'v', '*', '+', '>', '1', '2', '3', '4', 's', 'p', '*', 'h', 'H', '+', 'x', 'D', 'd', '|', '_']
linestyles = ["-", "--", "-."]
cs = "coolwarm"


def main(seed, feature): 
    """
    """
    mem = psutil.virtual_memory()
    print(f"Total: {mem.total / 1e9:.2f} GB")
    print(f"Available: {mem.available / 1e9:.2f} GB")
    print(f"Used: {mem.used / 1e9:.2f} GB")
    print(f"Percentage: {mem.percent}%")

    task = "everything"
    hidden = "300"
    batch = "128"
    accfeature = "+angle" 

    aname = f"{task}_seed{seed}_{feature}+hidden{hidden}+batch{batch}{accfeature}"
    out_path_name = "multiple_tasks/" + f"param_{aname}_result.npz"
    out_path = Path(out_path_name)

    size_bytes = out_path.stat().st_size  
    size_gb = size_bytes / 1024**3 
    print(f"{out_path} = {size_gb:.3f} GiB")

    with np.load(out_path_name, allow_pickle=True) as data:
        rules_epochs = data["rules_epochs"].item()
        hyp_dict = data["hyp_dict"].item()
        all_rules = data["all_rules"]
        test_task = data["test_task"]
        Ms_orig = data["Ms_orig"]
        hs = data["hs"]
        bs = data["bs"]
        xs = data["xs"]

    print(f"Ms_orig: {Ms_orig.shape}")
    print(f"hs: {hs.shape}")
    print(f"xs: {xs.shape}")
    print(f"bs: {bs.shape}")
    print(f"test_task: {test_task.shape}")
    print(f"all_rules: {all_rules}")

    out_param_path = "multiple_tasks/" + f"param_{aname}_param.json"
    out_param_path = Path(out_param_path)

    with out_param_path.open() as f: 
        raw_cfg_param = json.load(f)

    task_params, train_params, net_params = raw_cfg_param["task_params"], raw_cfg_param["train_params"], raw_cfg_param["net_params"]

    savefigure_name = f"{hyp_dict['ruleset']}_seed{seed}_{hyp_dict['addon_name']}"
    savefigure_name_base = copy.deepcopy(savefigure_name)

    # all outputs for this experiment go into their own subfolder
    save_dir = f"./multiple_tasks/{savefigure_name_base}"
    os.makedirs(save_dir, exist_ok=True)
    for _old_file in Path(save_dir).iterdir():
        if _old_file.is_file():
            _old_file.unlink()

    # %%
    # 2025-11-19: make sure the bias is only cell-dependent but not time- or trail-dependent
    ref = bs[0, 0, :]              
    same_per_k = np.all(bs == ref, axis=(0, 1))  
    all_k_constant = np.all(same_per_k)

    # %%
    add_safe_globals([np.core.multiarray._reconstruct])

    netpathname = "multiple_tasks/" + f"savednet_{aname}.pt"
    checkpoint = torch.load(netpathname, map_location="cpu")

    state_dict = checkpoint["state_dict"]
    print(state_dict.keys())
    load_net_params = checkpoint["net_params"]
    print(load_net_params)

    # reload the network
    model = mpn.DeepMultiPlasticNet(load_net_params, verbose=False, forzihan=True)

    missing, unexpected = model.load_state_dict(checkpoint["state_dict"], strict=True)
    print("missing:", missing)
    print("unexpected:", unexpected)
    del checkpoint  # free duplicated weights; state_dict already extracted above

    model.eval()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    print(f"Running on: {device}")

    task_params_c, train_params_c, net_params_c = mpn_tasks.convert_and_init_multitask_params(
        (task_params, train_params, net_params)
    )

    def shared_run(addtask): 
        """
        Trying to replicate Fig 3D-E in Driscoll etal 2024
        """
        assert addtask in ("dmcgo", "delaydm1", )
        
        task_params_dmcgo = copy.deepcopy(task_params_c)

        if addtask == "dmcgo":
            task_params_dmcgo["rules"] = ["dmcgo", "dmcnogo"]
        elif addtask == "delaydm1":
            task_params_dmcgo["rules"] = ["delaydm1", "delaydm2"]
            
        test_n_batch = 100

        task_params_dmcgo['hp']['batch_size_train'] = test_n_batch
        task_params_dmcgo["long_delay"] = "normal"
        
        test_data, test_trials_extra = mpn_tasks.generate_trials_wrap(task_params_dmcgo, 
                                                                    test_n_batch, 
                                                                    rules=task_params_dmcgo['rules'], 
                                                                    mode_input="random", 
                                                                    device="cpu",
                                                                    verbose=True)
        test_input, test_output, _ = test_data

        print(f"test_input.shape: {test_input.shape}")

        test_task = helper.find_task(task_params_dmcgo, test_input.detach().cpu().numpy(), 0)
        test_task = [int(test_task_ - min(test_task)) for test_task_ in test_task]

        _, test_trials, _ = test_trials_extra

        stim1_choices = np.concatenate([test_trials[0].meta["stim1"], test_trials[1].meta["stim1"]])
        print(f"stim1_choices: {stim1_choices}")

        epochs = test_trials[0].epochs
        
        if addtask in ("dmcgo", "delaydm1", ):
            delay_period = epochs["delay1"]
        
        # is delay period start inclusive and end exclusive?
        print(f"delay_period: {delay_period}")

        assert len(test_task) == len(stim1_choices)

        net_out, _, db_test = model.iterate_sequence_batch(test_input.to(device), run_mode='track_states')

        if addtask in ("dmcgo", "delaydm1", ):
            hidden_test = db_test["hidden1"][:, delay_period[0]:delay_period[1]-1, :]
            em_test = (db_test["M1"][:, delay_period[0]:delay_period[1]-1, :, :] * state_dict["mp_layer1.W"]).cpu().numpy()
            m_test = db_test["M1"][:, delay_period[0]:delay_period[1]-1, :, :].cpu().numpy()

        print(f"hidden_test: {hidden_test.shape}; em_test: {em_test.shape}; m_test: {m_test.shape}")

        fig, axs = plt.subplots(5,2,figsize=(5*2,5*2))
        for i in range(5):
            for inp in range(test_input.shape[2]):
                axs[i,0].plot(test_input[i,:,inp].detach().cpu().numpy(), color=c_vals[inp], alpha=0.5)
            for inp in range(net_out.shape[2]):
                axs[i,1].plot(net_out[i,:,inp].detach().cpu().numpy(), color=c_vals[inp], alpha=0.5)
            for outp in range(test_output.shape[2]):
                axs[i,1].plot(test_output[i,:,outp].detach().cpu().numpy(), color=c_vals[outp], 
                            alpha=0.5, linestyle="--")
        fig.tight_layout()  
        fig.savefig(f"{save_dir}/{addtask}_{savefigure_name}.png", dpi=300)

        def combo_indices_16(A, B):
            A = np.asarray(A)
            B = np.asarray(B)
            if A.shape != B.shape:
                raise ValueError(f"Length mismatch: A{A.shape} vs B{B.shape}")

            out = {}
            for a in (0, 1):
                for b in range(8):
                    out[(a, b)] = np.flatnonzero((A == a) & (B == b)).tolist()
            return out

        idx_map = combo_indices_16(test_task, stim1_choices)

        stacked_inputs = [[], []]
        stacked_inputs_labels = [[], []]
        for i in range(2):
            for b in range(8):
                # add multiple trials of the same kind of stimulus
                trial_num = 5
                for trial_idx in range(trial_num): 
                    # for trials that are considered to have the same stim1 and same task
                    # their input should be identical (quantified via sum) up to the end 
                    # of the delay period; only check for dmcgo since for delaydm1, the magnitude
                    # might be different across batches even the stim identity is the same 
                    if addtask == "dmcgo":
                        input_check = []
                        for key in idx_map[(i,b)]:
                            input_check.append(torch.sum(test_input[key,:delay_period[1],:]).item())
                        input_check = np.array(input_check, dtype=float)
                        ok = np.all(np.isclose(input_check, input_check[0]))
                        print(f"input_check: {input_check}")
                        assert ok 
                    
                    idxs = idx_map[(i, b)][trial_idx]
                    test_input_chosen = test_input[idxs]
                    stacked_inputs[i].append(test_input_chosen)
                    # add the stim1 label (for this trial) as well
                    stacked_inputs_labels[i].append(b)
                
            xs = stacked_inputs[i]  
            stacked_inputs[i] = torch.stack(
                [x if isinstance(x, torch.Tensor) else torch.as_tensor(x) for x in xs],
                dim=0
            )
            
        assert stacked_inputs_labels[0] == stacked_inputs_labels[1], "The two conditions should have the same set of stimuli for interpolation to make sense"

        # common delay period PCA
        as_hidden_wantperiod = hidden_test.reshape(-1, hidden_test.shape[-1])
        pca_delay_h = PCA(n_components=3, random_state=np.random.randint(0, 10000), svd_solver="auto")
        pca_delay_h.fit(as_hidden_wantperiod)

        as_em_wantperiod = em_test.reshape(-1, em_test.shape[-1] * em_test.shape[-2])
        pca_delay_em = PCA(n_components=3, random_state=np.random.randint(0, 10000), svd_solver="auto")
        pca_delay_em.fit(as_em_wantperiod)

        as_m_wantperiod = m_test.reshape(-1, m_test.shape[-1] * m_test.shape[-2])
        pca_delay_m = PCA(n_components=3, random_state=np.random.randint(0, 10000), svd_solver="auto")
        pca_delay_m.fit(as_m_wantperiod)

        plot_all = [["hidden", pca_delay_h], ["e_modulation", pca_delay_em], ["m_modulation", pca_delay_m]]

        for plot_name, pca_delay in plot_all:
            # interpolate between the two conditions (stim1 and stim2) in the input space, 
            # and track how the fixed points evolve in the PCA space of the delay period activity
            alpha_lst = np.linspace(0,1,20)
            interpolate_inputs = [(1-a) * stacked_inputs[0] + a * stacked_inputs[1] for a in alpha_lst]
            fixed_points_all = []
            traj_all = []

            for idx, interpolate_input in enumerate(interpolate_inputs):
                _, _, db = model.iterate_sequence_batch(interpolate_input.to(device), run_mode="track_states", save_to_cpu=True, detach_saved=True)

                if plot_name == "hidden":
                    data = db["hidden1"]
                    as_flat = data.reshape(-1, data.shape[-1])
                elif plot_name == "e_modulation":
                    data = (db["M1"] * state_dict["mp_layer1.W"]).cpu().numpy()
                    as_flat = data.reshape(-1, data.shape[-1] * data.shape[-2])
                elif plot_name == "m_modulation":
                    data = db["M1"]
                    as_flat = data.reshape(-1, data.shape[-1] * data.shape[-2])
                del db  # free activation dict before next iteration

                hidden_tf = pca_delay.transform(as_flat)
                projected_data = hidden_tf.reshape(data.shape[0], data.shape[1], -1)
                del data, as_flat  # intermediates no longer needed

                if addtask in ("dmcgo", "delaydm1", ):
                    fixed_points_all.append(projected_data[:,delay_period[1]-1,:])

                if idx in (0, len(interpolate_inputs) - 1, ):
                    if addtask in ("dmcgo", "delaydm1", ):
                        traj_all.append(projected_data[:, delay_period[0]:delay_period[1], :])
                del projected_data
                
            del interpolate_inputs  # free the 20 interpolated input tensors
            fixed_points_all_arr = np.stack(fixed_points_all, axis=0)  # (n_alpha, n_stim, n_pc)
            print(f"fixed_points_all_arr: {fixed_points_all_arr.shape}")
            n_alpha, n_stim, n_pc = fixed_points_all_arr.shape

            figmap, axsmap = plt.subplots(1,3,figsize=(5*3,5))
            pcs = [[0,1],[0,2],[1,2]]

            for pc_idx, (pc_x, pc_y) in enumerate(pcs):
                for stim in range(n_stim):
                    traj_fp = fixed_points_all_arr[:, stim, :] 

                    axsmap[pc_idx].plot(traj_fp[:, pc_x], traj_fp[:, pc_y], "-o", 
                                        color=c_vals[stacked_inputs_labels[0][stim]], linewidth=1.5, markersize=4, 
                                        alpha=0.5, 
                                        label=f"stim {(stim // trial_num + 1)}" if stim % trial_num == 0 else None)  

                    axsmap[pc_idx].scatter(traj_fp[0, pc_x],  traj_fp[0, pc_y],  color=c_vals[stacked_inputs_labels[0][stim]], 
                                        marker="s", s=50, zorder=3)
                    axsmap[pc_idx].scatter(traj_fp[-1, pc_x], traj_fp[-1, pc_y], color=c_vals[stacked_inputs_labels[0][stim]], 
                                        marker="^", s=50, zorder=3)

                # plot the trajectory during the delay period for the first and last alpha (i.e. the two original conditions)
                delay_traj = False
                if delay_traj: 
                    for be in range(len(traj_all)):
                        proj = traj_all[be]
                        for stim in range(proj.shape[0]):
                            xs = proj[stim, :, pc_x]
                            ys = proj[stim, :, pc_y]
                            axsmap[pc_idx].plot(xs, ys, color=c_vals[stacked_inputs_labels[0][stim]], 
                                                alpha=0.3, linewidth=1.0, linestyle=l_vals[be+1])
                            axsmap[pc_idx].scatter(xs[0], ys[0], color=c_vals[stacked_inputs_labels[0][stim]], 
                                                marker="o", s=35, zorder=4, alpha=0.9)

                axsmap[pc_idx].set_xlabel(f"Memory State PC{pc_x+1}", fontsize=12)
                axsmap[pc_idx].set_ylabel(f"Memory State PC{pc_y+1}", fontsize=12)
                axsmap[pc_idx].legend(frameon=True)
                # axsmap[pc_idx].set_xscale("symlog", linthresh=1e-3)
                # axsmap[pc_idx].set_yscale("symlog", linthresh=1e-3)
                
            figmap.tight_layout()
            figmap.savefig(f"{save_dir}/{addtask}_fixed_points_{plot_name}_{savefigure_name}.png", dpi=300)

    # shared_run("delaydm1") 
    # shared_run("dmcgo")

    # analyze the fitted weight matrices; we focus on the first layer of modulation and the output layer, since they are more interpretable than the hidden layer
    output_W = state_dict["W_output"].cpu().numpy()
    input_W = state_dict["W_initial_linear.weight"].cpu().numpy()
    modulation_W = state_dict["mp_layer1.W"].cpu().numpy()

    def heatmap_with_top_left_marginals(
        M,
        ax,
        cmap="coolwarm",
        center=0,
        vmin=None,
        vmax=None,
        xlabel="",
        ylabel="",
        label_fs=15,
        tick_fs=10,
        marginal_frac=0.18,   # thickness of top/left strips (in ax coords)
        marginal_pad=0.03,    # gap between heatmap and strips (in ax coords)
        marginal_lw=1.2,
        cbar=True,
        square=False
    ):
        """
        Heatmap on `ax` + marginals:
        - top inset:  col-wise sum(abs(M)) (length = n_cols)
        - left inset: row-wise sum(abs(M)) (length = n_rows), plotted vertically

        Layout: col-sum at top, row-sum at left. No child dirs, no extra axes elsewhere.
        """
        M = np.asarray(M)
        n_rows, n_cols = M.shape

        if vmin is None:
            vmin = np.nanmin(M)
        if vmax is None:
            vmax = np.nanmax(M)

        # ---- main heatmap
        sns.heatmap(
            M,
            ax=ax,
            cmap=cmap,
            vmin=vmin,
            vmax=vmax,
            center=center,
            cbar=cbar,
            square=square,
        )
        ax.set_xlabel(xlabel, fontsize=label_fs)
        ax.set_ylabel(ylabel, fontsize=label_fs)
        ax.tick_params(labelsize=tick_fs)

        # ---- marginals
        col_sum = np.nansum(np.abs(M), axis=0)  # (n_cols,)
        row_sum = np.nansum(np.abs(M), axis=1)  # (n_rows,)

        # inset axes positions are in the parent ax's coordinate system
        top_rect  = [0.0, 1.0 + marginal_pad, 1.0, marginal_frac]
        left_rect = [-(marginal_frac + marginal_pad), 0.0, marginal_frac, 1.0]

        ax_top  = ax.inset_axes(top_rect,  transform=ax.transAxes)
        ax_left = ax.inset_axes(left_rect, transform=ax.transAxes)

        # ---- top: col sums (aligned with heatmap columns)
        x = np.arange(n_cols)
        ax_top.plot(x, col_sum, lw=marginal_lw)
        ax_top.set_xlim(-0.5, n_cols - 0.5)
        ax_top.set_xticks([])
        ax_top.tick_params(axis="y", labelsize=tick_fs)
        ax_top.spines["top"].set_visible(False)
        ax_top.spines["right"].set_visible(False)

        # ---- left: row sums (vertical, aligned with heatmap rows)
        y = np.arange(n_rows)
        ax_left.plot(row_sum, y, lw=marginal_lw)  # x=row_sum, y=row index
        ax_left.set_ylim(n_rows - 0.5, -0.5)      # match heatmap's y-direction (top row at top)
        ax_left.set_yticks([])
        ax_left.tick_params(axis="x", labelsize=tick_fs)
        ax_left.spines["top"].set_visible(False)
        ax_left.spines["right"].set_visible(False)
        
        ax.set_xticks([]); ax.set_yticks([]) 

        return ax, ax_top, ax_left

    def plot_weight_triplet_with_top_left_marginals(
        output_W,
        input_W,
        modulation_W,
        aname,
        cs="coolwarm",
        save_dir="./multiple_tasks",
    ):
        figoutput, axsoutput = plt.subplots(1,1,figsize=(10, 12))

        heatmap_with_top_left_marginals(
            output_W,
            ax=axsoutput,
            cmap=cs,
            center=0,
            vmin=np.nanmin(output_W),
            vmax=np.nanmax(output_W),
            xlabel="Hidden 2 Index",
            ylabel="Output Index",
            square=False
        )
        
        figinput, axsinput = plt.subplots(1,1,figsize=(10, 12))
        heatmap_with_top_left_marginals(
            input_W.T,
            ax=axsinput,
            cmap=cs,
            center=0,
            vmin=np.nanmin(input_W),
            vmax=np.nanmax(input_W),
            xlabel="Hidden 1 Index",
            ylabel="Input Index",
            square=False
        )
        
        figmodulation, axsmodulation = plt.subplots(1,1,figsize=(10, 12))
        # modulation_W shape is (post, pre); .T gives (pre, post) so rows=pre, cols=post
        heatmap_with_top_left_marginals(
            modulation_W.T,
            ax=axsmodulation,
            cmap=cs,
            center=0,
            vmin=np.nanmin(modulation_W),
            vmax=np.nanmax(modulation_W),
            xlabel="Post (Hidden) Index",
            ylabel="Pre (Input) Index",
            square=False
        )

        figoutput.tight_layout()
        figoutput.savefig(f"{save_dir}/weight_matrices_{aname}.png", dpi=300)
        plt.close(figoutput)
        
        figinput.tight_layout()
        figinput.savefig(f"{save_dir}/weight_matrices_input_{aname}.png", dpi=300)
        plt.close(figinput)
        
        figmodulation.tight_layout()
        figmodulation.savefig(f"{save_dir}/weight_matrices_modulation_{aname}.png", dpi=300)
        plt.close(figmodulation)
        
        return

    def plot_static_pathway_analysis(
        output_W,
        input_W,
        modulation_W,
        input_labels,
        aname,
        save_dir="./multiple_tasks",
    ):
        """End-to-end pathway W_out @ W_mod @ W_in and its SVD spectrum."""
        W_path = output_W @ modulation_W @ input_W  # (n_output, n_input)
        _, s, _ = np.linalg.svd(W_path, full_matrices=False)

        n_output, n_input = W_path.shape
        fig, axs = plt.subplots(
            1, 2,
            figsize=(max(8, n_input * 0.35 + 4), max(4, n_output * 0.45 + 2)),
            gridspec_kw={"width_ratios": [3, 1]},
        )

        vabs = float(np.nanmax(np.abs(W_path)))
        sns.heatmap(
            W_path,
            ax=axs[0],
            cmap="coolwarm",
            center=0,
            vmin=-vabs,
            vmax=vabs,
            cbar=True,
            xticklabels=input_labels,
            yticklabels=np.arange(n_output),
        )
        axs[0].set_xticklabels(input_labels, rotation=90, fontsize=7)
        axs[0].set_xlabel("Input Feature", fontsize=10)
        axs[0].set_ylabel("Output Dim", fontsize=10)
        axs[0].set_title("End-to-end pathway  W_out @ W_mod @ W_in", fontsize=9)

        axs[1].bar(np.arange(len(s)), s, color="steelblue")
        axs[1].set_xlabel("Mode", fontsize=10)
        axs[1].set_ylabel("Singular value", fontsize=10)
        axs[1].set_title("SV spectrum", fontsize=9)
        axs[1].spines["top"].set_visible(False)
        axs[1].spines["right"].set_visible(False)

        fig.tight_layout()
        fig.savefig(f"{save_dir}/pathway_static_{aname}.png", dpi=300, bbox_inches="tight")
        plt.close(fig)

    plot_weight_triplet_with_top_left_marginals(
        output_W=output_W,
        input_W=input_W,
        modulation_W=modulation_W,
        aname=aname,
        cs=cs,
        save_dir=save_dir,
    )

    all_input = ["Fix On", "Fix Off", "Stim 1 Cos", "Stim 1 Sin", "Stim 2 Cos", "Stim 2 Sin"] + task_params_c["rules"]
    input_corr = np.corrcoef(input_W.T)
    mask = np.triu(np.ones_like(input_corr, dtype=bool), k=1)
    fig, ax = plt.subplots(1,1,figsize=(5,5))
    sns.heatmap(input_corr, ax=ax, cmap=cs, center=0, mask=mask, vmin=-1, vmax=1, square=True, cbar_kws={"shrink": 0.5})
    ax.set_xticks(np.arange(len(all_input)) + 0.5, labels=all_input, rotation=90)
    ax.set_yticks(np.arange(len(all_input)) + 0.5, labels=all_input, rotation=0)
    ax.set_xlabel("Input Index", fontsize=12)
    ax.set_ylabel("Input Index", fontsize=12)
    fig.tight_layout()
    fig.savefig(f"{save_dir}/input_weight_correlation_{aname}.png", dpi=300)
    plt.close(fig)

    plot_static_pathway_analysis(
        output_W=output_W,
        input_W=input_W,
        modulation_W=modulation_W,
        input_labels=all_input,
        aname=aname,
        save_dir=save_dir,
    )

    # should we re-evaluate the result? 
    reevaluate = True 
    if reevaluate:
        test_n_batch = 80 # number of batches for each task 
        task_params_c['hp']['batch_size_train'] = test_n_batch
            
        test_data, test_trials_extra = mpn_tasks.generate_trials_wrap(
            task_params_c, 
            test_n_batch, 
            rules=task_params_c['rules'],
            mode_input="random", 
            device="cpu", 
            verbose=False,
            long_all=True
        )
        
        test_input, test_output, test_mask = test_data
        print(f"test_input.shape: {test_input.shape}; test_output.shape: {test_output.shape}; test_mask.shape: {test_mask.shape}")
        _, test_trials, test_rule_idxs = test_trials_extra
        
        test_input = test_input.to(device)
        test_output = test_output.to(device)
        test_mask = test_mask.to(device)

        forward_chunk_size = 20  
        B_total = test_input.shape[0]
        chunk_net_outs, chunk_db = [], {}
        with torch.no_grad():
            for chunk_start in range(0, B_total, forward_chunk_size):
                chunk_end = min(chunk_start + forward_chunk_size, B_total)
                chunk_out, _, chunk_db_i = model.iterate_sequence_batch(
                    test_input[chunk_start:chunk_end], run_mode='track_states'
                )
                chunk_net_outs.append(chunk_out.cpu())
                for key, val in chunk_db_i.items():
                    if torch.is_tensor(val):
                        chunk_db.setdefault(key, []).append(val.cpu())
                    else:
                        chunk_db.setdefault(key, []).append(val)

        net_out = torch.cat(chunk_net_outs, dim=0)
        del chunk_net_outs  # free per-chunk output list now that net_out is assembled
        db_test = {
            key: torch.cat(chunks, dim=0) if torch.is_tensor(chunks[0]) else [x for c in chunks for x in c]
            for key, chunks in chunk_db.items()
        }
        del chunk_db  # free per-chunk state lists now that db_test is assembled
        acc, _ = model.compute_acc(net_out.to(device), test_output, test_mask, test_input, isvalid=True, mode=model.acc_measure)
        print(f"Accuracy: {acc}")
        del net_out, test_output, test_mask  # no longer needed after accuracy; test_input kept for find_task below

        Ms_orig = db_test["M1"].cpu().numpy()
        xs = db_test["input1"].cpu().numpy()
        hs = db_test["hidden1"].cpu().numpy()
        del db_test  # all needed arrays extracted; free the large activation tensors

        print(f"Ms_orig: {Ms_orig.shape}; xs: {xs.shape}; hs: {hs.shape}")
        print(f"modulation_W: {modulation_W.shape}")
        
        all_rules = np.array(task_params["rules"])
        
        def generate_response_stimulus(task_params, test_trials): 
            """
            """
            labels_resp, labels_stim1, labels_stim2 = [], [], []
            rules_epochs = {} 
            for rule_idx, rule in enumerate(task_params['rules']):
                rules_epochs[rule] = test_trials[rule_idx].epochs
                
                if rule in ('dmsgo','dmcgo','dmsnogo','dmcnogo',):
                    labels_resp.append(test_trials[rule_idx].meta['matches'])                
                else:
                    labels_resp.append(test_trials[rule_idx].meta['resp1'])           
                    
                labels_stim1.append(test_trials[rule_idx].meta['stim1']) 
                
                try: 
                    labels_stim2.append(test_trials[rule_idx].meta['stim2'])
                except KeyError:
                    labels_stim2.append(np.full_like(test_trials[rule_idx].meta['stim1'], np.nan))

            labels_resp = np.concatenate(labels_resp, axis=0).reshape(-1,1)
            labels_stim1 = np.concatenate(labels_stim1, axis=0).reshape(-1,1)
            labels_stim2 = np.concatenate(labels_stim2, axis=0).reshape(-1,1)

            return labels_resp, labels_stim1, labels_stim2, rules_epochs
        
        _, labels_stim1, _, rules_epochs = generate_response_stimulus(task_params, test_trials)
        
        test_task = helper.find_task(task_params, test_input.detach().cpu().numpy(), 0)
        test_task = np.array([int(c) for c in test_task]).flatten()

    # %%
    weighted_Ms_orig = Ms_orig * modulation_W
    print(f"weighted_Ms_orig: {weighted_Ms_orig.shape}")

    # 2026-03-29: for the modulation-weighted activity, we will not normalize across rules, 
    # otherwise per-neuron normalization
    # will wash out the multiplication effect
    # 2026-04-06: downstream analysis implicitly require the order of
    clustering_data_analysis = [xs, xs, hs, hs, Ms_orig, Ms_orig, weighted_Ms_orig, Ms_orig]
    clustering_data_analysis_names = ["input", "input", "hidden", "hidden", \
        "modulation_all", "modulation_all", "modulation_all_weighted", "modulation_all_var_weighted"]
    clustering_data_normalize = [True, False, True, False, True, False, False, False]
    c_metrics = ["euclidean", "euclidean", "euclidean", "euclidean", "euclidean", "euclidean", "euclidean", "euclidean"]
    c_methods = ["ward", "ward", "ward", "ward", "ward", "ward", "ward", "ward"]
    assert len(clustering_data_analysis) == len(clustering_data_analysis_names) \
        == len(clustering_data_normalize)

    # data registertion buffer
    clustering_data_hierarchy = {}
    clustering_corr_info = {}
    col_clusters_all, row_clusters_all = {}, {}
    row_cluster_breaker_all = {}
    input_hidden_comparison = {}
    base_data = {}
    metrics_all_all = {}
    rbreaks_all, cbreaks_all = {}, {}
    cluster_info_save = {}
    cluster_info_save_mod = {}

    selection_key = ["CH_blocks", "DB_blocks"]

    upper_cluster = 300
    lower_cluster = 5 # for input & hidden
    lower_cluster_mod = 10 # for modulation
    silhouette_tol = 0.05
    tol_mode = "relative"   # or "relative"
    tol_k_select = "gap"
    score_quantile = 0.50
    unresponsive_norm_frac = 1e-3

    assert tol_mode in ("absolute", "relative"), f"Invalid tol_mode: {tol_mode}"

    for clustering_index in range(len(clustering_data_analysis)): 
        print("======================================================")
        clustering_data = clustering_data_analysis[clustering_index]
        clustering_name = clustering_data_analysis_names[clustering_index]
        clustering_normalize = clustering_data_normalize[clustering_index]
        # 2026-04-09: only used for the modulation-related clustering analysis
        c_metric = c_metrics[clustering_index]
        c_method = c_methods[clustering_index]
        
        print(
            f"[clustering]\n"
            f"  name      : {clustering_name}\n"
            f"  data shape: {clustering_data.shape}\n"
            f"  normalize : {clustering_normalize}"
        )
        
        if hyp_dict['ruleset'] == "everything": 
            phase_to_indices = [
                ("stim1",  [0, 1, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]),
                ("stim2",  [6, 7, 8, 9, 10]),
                ("delay1", [4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14]),
                ("delay2", [6, 7, 8, 9, 10]),
                ("go1",    [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14])
            ]
        
        tb_break = [
            [idx, rules_epochs[all_rules[idx]][phase]]
            for phase, indices in phase_to_indices
            for idx in indices
        ]
        
        tb_break_name = [
            f"{all_rules[idx]}-{phase}"
            for phase, indices in phase_to_indices
            for idx in indices
        ]
        
        tb_break_name = np.array(tb_break_name)
        
        cell_vars_rules = [] 
        
        for el in range(len(tb_break)):
            n_rules = len(task_params['rules'])
            n_cells = clustering_data.shape[-1]
                
            rule_idx, period_time = tb_break[el][0], tb_break[el][1]
            
            # print('Rule {} (idx {}), {}'.format(all_rules[rule_idx], rule_idx, period_time))
            # print(f"tb_break_name[el]: {tb_break_name[el]}")
            
            # 2026-03-29: select the stimulus information
            labels_stim = labels_stim1
            # labels_stim = labels_stim2 if ("stim2" in tb_break_name[el] or "delay2" in tb_break_name[el]) else labels_stim1
            
            if len(clustering_data.shape) == 3: # input or hidden 
                # 2025-11-19: if period_time[1] is None, then go until the end along that axis
                rule_cluster = clustering_data[test_task == rule_idx, period_time[0]:period_time[1], :]
                varval = helper.task_variance_period_numpy(rule_cluster, labels_stim[test_task == rule_idx].flatten())
                cell_vars_rules.append(varval) 
            else:
                clustering_data_old = clustering_data                
                if "all" in clustering_name: # modulation 
                    clustering_data = clustering_data.reshape(clustering_data.shape[0], clustering_data.shape[1], -1)
                    rule_cluster = clustering_data[test_task == rule_idx, 
                                                   period_time[0]:period_time[1], :]
                    varval = helper.task_variance_period_numpy(rule_cluster, 
                                                               labels_stim[test_task == rule_idx].flatten())
                    cell_vars_rules.append(varval) 
                        
        cell_vars_rules = np.array(cell_vars_rules)    
        cell_vars_rules_norm = np.zeros_like(cell_vars_rules)

        print(f"cell_vars_rules.shape: {cell_vars_rules.shape}")
        
        # normalize for each neuron 
        if clustering_normalize: 
            cell_max_var = np.max(cell_vars_rules, axis=0) # Across rules
            print(f"cell_max_var.shape: {cell_max_var.shape}")

            for period_idx in range(len(tb_break)):
                cell_vars_rules_norm[period_idx] = np.where(
                    cell_max_var > 0., cell_vars_rules[period_idx] / cell_max_var, 0.
                )
                
            savefigure_name = savefigure_name_base + "_normalized"
            clustering_save_name = clustering_name + "_normalized"
            # with normalization, then when plotting the variance matrix, the entries 
            # are all between 0 and 1, so we can set a fixed color range for better comparability across periods
            vmins, vmaxs = 0, 1
        else:
            cell_vars_rules_norm = cell_vars_rules.copy()
            savefigure_name = savefigure_name_base + "_unnormalized"
            clustering_save_name = clustering_name + "_unnormalized"
            vmins, vmaxs = None, None

        # for var_weighted: compute variance on Ms_orig then weight by modulation_W
        if "var_weighted" in clustering_name:
            cell_vars_rules_norm = cell_vars_rules_norm * modulation_W.flatten()[np.newaxis, :]

        # modulation only, reshape to (N, post, pre) shape after calculating the variance
        # modulation_W shape is (post, pre): previously assumed (pre, post)
        # N here as the number of sessions after breakdown
        # 2026-04-06: the code implicitly requires the modulation to be a square matrix
        if "all" in clustering_name:
            N, MM = cell_vars_rules_norm.shape
            M = int(np.sqrt(MM))
            cell_vars_rules_norm_keepshape = cell_vars_rules_norm.reshape(N, M, M)
        
        # build rule-wise value lists and corresponding field names dynamically
        rule_vals  = [cell_vars_rules_norm[i].tolist() for i in range(n_rules)]
        # print(f"rule_vals: {rule_vals}")
        rule_names = [f"rule{i}" for i in range(n_rules)]
        
        # structured array whose fields are rule0, rule1, …, rule{n_rules-1}
        dtype = np.dtype([(name, float) for name in rule_names])
        rules_struct = np.array(list(zip(*rule_vals)), dtype=dtype)
        
        ## 2026-04-09: this part is important and should be treated delicately
        ## currently no sorting is applied, i.e. we keep the original order of the neurons; 
        # this is for better interpretability and also to avoid potential 
        # undesirable bugs in the clustering code that might be triggered by a non-identity sorting
        
        # sort_idxs = np.argsort(rules_struct, order=rule_names)[::-1] # descending lexicographic sort across all rule columns 
        sort_idxs = np.arange(rules_struct.shape[0], dtype=np.intp) # identity map for better interpretability

        # 2025-07-07: first sorting based on the normalized magnitude
        # all the following should be aligned with this change
        # 2025-08-22: sort based on the variance ordering OR using an identity map (i.e. do nothing)
        cell_vars_rules_sorted_norm = cell_vars_rules_norm[:, sort_idxs]
        base_data[clustering_save_name] = cell_vars_rules_sorted_norm
        print(f"cell_vars_rules_sorted_norm: {cell_vars_rules_sorted_norm.shape}")  

        # 2026-04-06: analyze input and hidden
        if not ("all" in clustering_name):
            # For unnormalized variance data (all non-negative), apply log1p to compress
            # the dynamic range before clustering so that cosine distance captures
            # selectivity profile shape rather than raw amplitude differences.
            # Normalized data is passed through unchanged.
            if clustering_normalize:
                V_for_clustering = cell_vars_rules_sorted_norm
            else:
                V_for_clustering = np.log1p(cell_vars_rules_sorted_norm)

            # clustering & grouping & re-ordering
            # first loop on input, second loop in hidden
            jitter_std = 0.02 
            resample_features_frac = 1.0 
            n_repeats = 100
            if jitter_std == 0 and resample_features_frac == 1.0:
                n_repeats = 1  # no need to repeat if no randomness is introduced
                
            result = clustering.cluster_variance_matrix_repeat(V_for_clustering,
                                                               k_min=lower_cluster,
                                                               k_max=upper_cluster,
                                                               metric=c_metric,
                                                               method=c_method,
                                                               n_repeats=n_repeats,
                                                               resample_features_frac=resample_features_frac,
                                                               jitter_std=jitter_std,
                                                               silhouette_tol=silhouette_tol,
                                                               tol_mode=tol_mode,
                                                               score_quantile=score_quantile,
                                                               unresponsive_norm_frac=unresponsive_norm_frac,
                                                               skip_unresponsive_detection=clustering_normalize,
                                                               tol_k_select=tol_k_select)
            
            # sanity check to make sure no undesirable bug happens with in the clustering code
            assert sorted(result["row_order"]) == list(range(cell_vars_rules_sorted_norm.shape[0]))
            assert sorted(result["col_order"]) == list(range(cell_vars_rules_sorted_norm.shape[1]))

            assert len(result["row_labels"]) == cell_vars_rules_sorted_norm.shape[0]
            assert len(result["col_labels"]) == cell_vars_rules_sorted_norm.shape[1]

            # Allow one extra label (k+1) for the unresponsive cluster that
            # clustering.py appends when unresponsive rows are detected.
            assert np.unique(result["row_labels"]).size in (result["row_k"], result["row_k"] + 1)
            assert np.unique(result["col_labels"]).size in (result["col_k"], result["col_k"] + 1)
            
            eval_res = clustering_metric.evaluate_bicluster_clustering(
                cell_vars_rules_sorted_norm, row_labels=result["row_tol_labels"], col_labels=result["col_tol_labels"]
            )
            
            eval_metrics = eval_res["metrics"]
            eval_blocks = eval_res["blocks"]
            eval_stdmean = np.nanmedian(eval_blocks["std"] / eval_blocks["means"])

            eval_random_metrics_all = []
            eval_random_blocks_all = []
            for _ in range(100):
                rng = np.random.default_rng(seed=np.random.randint(0, 10000))
                row_arr = rng.permutation(result["row_tol_labels"])
                col_arr = rng.permutation(result["col_tol_labels"])
                eval_res_random = clustering_metric.evaluate_bicluster_clustering(
                    cell_vars_rules_sorted_norm, row_labels=row_arr, col_labels=col_arr
                )
                eval_random_metrics_all.append(eval_res_random["metrics"])
                eval_random_blocks_all.append(eval_res_random["blocks"])

            eval_random_stdmean = [np.nanmedian(eval_random_blocks["std"] / eval_random_blocks["means"]) 
                                   for eval_random_blocks in eval_random_blocks_all]

            metrics_all = {}
            for metric_key in selection_key: 
                optimized_value = eval_metrics[metric_key]
                random_values = [eval_random_metrics[metric_key] 
                                for eval_random_metrics in eval_random_metrics_all]
                metrics_all[metric_key] = [optimized_value, np.mean(random_values), 
                                           np.std(random_values, ddof=1)/np.sqrt(len(random_values))]

            metrics_all["std/mean"] = [eval_stdmean, 
                                       np.mean(eval_random_stdmean), 
                                       np.std(eval_random_stdmean, ddof=1)/np.sqrt(len(eval_random_stdmean))]
            
            # registeration
            metrics_all_all[clustering_save_name] = metrics_all

            input_hidden_comparison[clustering_save_name] = [result, cell_vars_rules_sorted_norm]
            
            # reorder the original matrix based on the clustering result
            cell_vars_rules_sorted_norm_ordered = cell_vars_rules_sorted_norm[
                np.ix_(result["row_order"], result["col_order"])]
            
            rl = np.asarray(result["row_tol_labels"])[result["row_order"]]
            cl = np.asarray(result["col_tol_labels"])[result["col_order"]]
            rbreaks = clustering._breaks(rl)
            cbreaks = clustering._breaks(cl)
            rbreaks_all[clustering_save_name] = rbreaks
            cbreaks_all[clustering_save_name] = cbreaks

            # Print k decisions across all selection strategies for comparison.
            row_sil_mean = result["row_score_recording_mean"]
            col_sil_mean = result["col_score_recording_mean"]
            row_ari_mean = result["row_ari_recording_mean"]
            col_ari_mean = result["col_ari_recording_mean"]

            row_strict_k = result["row_strict_best_k"]
            col_strict_k = result["col_strict_best_k"]
            row_strict_sil = row_sil_mean.get(row_strict_k, float("nan"))
            col_strict_sil = col_sil_mean.get(col_strict_k, float("nan"))

            # Recompute tolerance band (mirrors clustering._score_threshold_from_best)
            def _tol_thresh(score):
                if tol_mode == "relative":
                    return score * (1.0 - silhouette_tol)
                return score - silhouette_tol

            row_band = sorted(
                k for k in result["row_k_values"]
                if row_sil_mean.get(k, -np.inf) >= _tol_thresh(row_strict_sil)
            )
            col_band = sorted(
                k for k in result["col_k_values"]
                if col_sil_mean.get(k, -np.inf) >= _tol_thresh(col_strict_sil)
            )

            def _band_select(band, strategy, fallback):
                if not band:
                    return fallback
                if strategy == "min":
                    return min(band)
                if strategy == "max":
                    return max(band)
                # "mean": k closest to the arithmetic mean; ties broken by smaller k
                m = float(np.mean(band))
                return min(band, key=lambda k: (abs(k - m), k))

            strat_ks = {
                "min":  (_band_select(row_band, "min",  row_strict_k), _band_select(col_band, "min",  col_strict_k)),
                "max":  (_band_select(row_band, "max",  row_strict_k), _band_select(col_band, "max",  col_strict_k)),
                "mean": (_band_select(row_band, "mean", row_strict_k), _band_select(col_band, "mean", col_strict_k)),
                "ari":  (result["row_ari_tol_k"],  result["col_ari_tol_k"]),
                "gap":  (result["row_gap_tol_k"],  result["col_gap_tol_k"]),
            }

            print(
                f"  strict argmax — row: {row_strict_k} (sil={row_strict_sil:.4f}); "
                f"col: {col_strict_k} (sil={col_strict_sil:.4f})"
            )
            for strat, (rk, ck) in strat_ks.items():
                marker = "  <-- selected" if strat == tol_k_select else ""
                r_sil = row_sil_mean.get(rk, float("nan"))
                c_sil = col_sil_mean.get(ck, float("nan"))
                r_ari = row_ari_mean.get(rk, float("nan"))
                c_ari = col_ari_mean.get(ck, float("nan"))
                print(
                    f"  {strat:4s}: row={rk:3d} (sil={r_sil:.4f}, ari={r_ari:.4f}); "
                    f"col={ck:3d} (sil={c_sil:.4f}, ari={c_ari:.4f}){marker}"
                )
            
            # extract the grouping information, i.e. which neuron belong to which cluster
            # instead of the view of dendrogram
            # 2025-10-20: we register the tolerant version of optimal cluster selection
            col_labels, col_k = result["col_tol_labels"], result["col_tol_k"]
            # group the neuron based on the labels
            col_clusters = {int(lab): np.where(col_labels == lab)[0] for lab in np.unique(col_labels)}
            # 2025-11-04: do similar things for row separation
            # we checked so that the cluster label (name) is monotonically increasing 
            # i.e. near by cluster (e.g. cluster 1 and cluster 2) should be more similiar in population 
            # as well, since increasing the cutoff threshold in dendrogram will "merge" these clusters
            row_labels, row_k = result["row_tol_labels"], result["row_tol_k"]
            row_clusters = {int(lab): np.where(row_labels == lab)[0] for lab in np.unique(row_labels)}
            print(f"col_clusters: {len(col_clusters)}; row_clusters: {len(row_clusters)}")
            
            # registeration
            col_clusters_all[clustering_save_name] = col_clusters
            row_clusters_all[clustering_save_name] = row_clusters
            
            cluster_info_save[clustering_save_name] = {
                "col_clusters": col_clusters,
                "row_clusters": row_clusters,
                "tb_break_name": tb_break_name,
                "cell_vars_rules_sorted_norm": cell_vars_rules_sorted_norm,
                "result": result,
            }

            # plot the optimization score as a function of number of clustering
            # also plot the indicator for the optimal number of cluster (and with tolerance version)
            # 2026-04-16: add more details in the plotting
            def _gap_curve(Z, k_vals):
                """Compute the Ward merge-height gap for each k from a linkage matrix."""
                n = Z.shape[0] + 1
                gaps = []
                for k in k_vals:
                    idx_above = n - k
                    idx_below = n - k - 1
                    h_above = float(Z[idx_above, 2]) if 0 <= idx_above < Z.shape[0] else np.inf
                    h_below = float(Z[idx_below, 2]) if 0 <= idx_below < Z.shape[0] else 0.0
                    gaps.append(h_above - h_below)
                return np.asarray(gaps, dtype=float)

            # solid vline = strict silhouette argmax; dashed vline = tol_k_select decision
            row_decision_k, col_decision_k = strat_ks[tol_k_select]

            figscore, (axscore, axari, axgap) = plt.subplots(3, 1, figsize=(6, 10), sharex=True)

            # ---------- silhouette: row ----------
            row_kvals   = np.asarray(result["row_k_values"], dtype=int)
            row_sil_all = np.asarray(result["row_score_recording_all"], dtype=float)  # (n_repeats, n_k)
            row_sil_arr = np.asarray([result["row_score_recording_mean"][k] for k in row_kvals], dtype=float)
            row_sil_std = np.asarray([result["row_score_recording_std"][k]  for k in row_kvals], dtype=float)

            for r in range(row_sil_all.shape[0]):
                axscore.plot(row_kvals, row_sil_all[r], color=c_vals[0], alpha=0.10, linewidth=0.8)
            axscore.fill_between(row_kvals, row_sil_arr - row_sil_std, row_sil_arr + row_sil_std,
                                 color=c_vals[0], alpha=0.20)
            axscore.plot(row_kvals, row_sil_arr, color=c_vals[0], linewidth=2.5, label="row")
            axscore.axvline(row_strict_k,   color=c_vals[0], linestyle="-",  linewidth=1.5)
            axscore.axvline(row_decision_k, color=c_vals[0], linestyle="--", linewidth=1.5)

            # ---------- silhouette: col ----------
            col_kvals   = np.asarray(result["col_k_values"], dtype=int)
            col_sil_all = np.asarray(result["col_score_recording_all"], dtype=float)  # (n_repeats, n_k)
            col_sil_arr = np.asarray([result["col_score_recording_mean"][k] for k in col_kvals], dtype=float)
            col_sil_std = np.asarray([result["col_score_recording_std"][k]  for k in col_kvals], dtype=float)

            for r in range(col_sil_all.shape[0]):
                axscore.plot(col_kvals, col_sil_all[r], color=c_vals[1], alpha=0.10, linewidth=0.8)
            axscore.fill_between(col_kvals, col_sil_arr - col_sil_std, col_sil_arr + col_sil_std,
                                 color=c_vals[1], alpha=0.20)
            axscore.plot(col_kvals, col_sil_arr, color=c_vals[1], linewidth=2.5, label="col")
            axscore.axvline(col_strict_k,   color=c_vals[1], linestyle="-",  linewidth=1.5)
            axscore.axvline(col_decision_k, color=c_vals[1], linestyle="--", linewidth=1.5)

            axscore.set_ylabel("Silhouette Score")
            axscore.legend(frameon=True)
            axscore.set_title(f"{clustering_name}", fontsize=15)

            # ---------- ARI stability: row ----------
            row_ari_arr = np.asarray([result["row_ari_recording_mean"][k] for k in row_kvals], dtype=float)
            row_ari_std = np.asarray([result["row_ari_recording_std"][k]  for k in row_kvals], dtype=float)
            axari.fill_between(row_kvals, row_ari_arr - row_ari_std, row_ari_arr + row_ari_std,
                               color=c_vals[0], alpha=0.20)
            axari.plot(row_kvals, row_ari_arr, color=c_vals[0], linewidth=2.5, label="row")
            axari.axvline(row_strict_k,   color=c_vals[0], linestyle="-",  linewidth=1.5)
            axari.axvline(row_decision_k, color=c_vals[0], linestyle="--", linewidth=1.5)

            # ---------- ARI stability: col ----------
            col_ari_arr = np.asarray([result["col_ari_recording_mean"][k] for k in col_kvals], dtype=float)
            col_ari_std = np.asarray([result["col_ari_recording_std"][k]  for k in col_kvals], dtype=float)
            axari.fill_between(col_kvals, col_ari_arr - col_ari_std, col_ari_arr + col_ari_std,
                               color=c_vals[1], alpha=0.20)
            axari.plot(col_kvals, col_ari_arr, color=c_vals[1], linewidth=2.5, label="col")
            axari.axvline(col_strict_k,   color=c_vals[1], linestyle="-",  linewidth=1.5)
            axari.axvline(col_decision_k, color=c_vals[1], linestyle="--", linewidth=1.5)

            axari.set_ylabel("Mean Pairwise ARI")
            axari.legend(frameon=True)

            # ---------- gap score: row ----------
            row_gap_arr = _gap_curve(result["row_linkage"], row_kvals)
            axgap.plot(row_kvals, row_gap_arr, color=c_vals[0], linewidth=2.5, label="row")
            axgap.axvline(row_strict_k,   color=c_vals[0], linestyle="-",  linewidth=1.5)
            axgap.axvline(row_decision_k, color=c_vals[0], linestyle="--", linewidth=1.5)

            # ---------- gap score: col ----------
            col_gap_arr = _gap_curve(result["col_linkage"], col_kvals)
            axgap.plot(col_kvals, col_gap_arr, color=c_vals[1], linewidth=2.5, label="col")
            axgap.axvline(col_strict_k,   color=c_vals[1], linestyle="-",  linewidth=1.5)
            axgap.axvline(col_decision_k, color=c_vals[1], linestyle="--", linewidth=1.5)

            axgap.set_xlabel("Number of Clusters")
            axgap.set_ylabel("Gap Score")
            axgap.set_xscale("log")
            axgap.legend(frameon=True)

            figscore.tight_layout()
            figscore.savefig(f"{save_dir}/{clustering_name}_variance_cluster_score_{savefigure_name}.png", dpi=300)
            plt.close(figscore)

            # 
            ordered_row_name = tb_break_name[result["row_order"]]
            row_breakers = [0]
            for row_group in row_clusters.values():
                row_breakers.append(row_breakers[-1] + len(row_group))
            row_breakers = row_breakers[1:]
            print(f"row_breakers: {row_breakers}")
            row_cluster_breaker_all[clustering_save_name] = row_breakers
            
            # pearson correlation matrix
            figcorr, axcorrs = plt.subplots(1,3,figsize=(8*3,8))
            metrics = ["Correlation", "Cosine Similarity", "L2 Distance"]
            cell_vars_rules_sorted_norm_ordered_measure = np.corrcoef(cell_vars_rules_sorted_norm_ordered, rowvar=True)
            cell_vars_rules_sorted_norm_ordered_measure_cos = cosine_similarity(cell_vars_rules_sorted_norm_ordered)
            cell_vars_rules_sorted_norm_ordered_measure_L2  = squareform(pdist(cell_vars_rules_sorted_norm_ordered, metric='euclidean'))

            # set uniform colorbar to cross-compare between analysis
            sns.heatmap(cell_vars_rules_sorted_norm_ordered_measure, cmap=cs, square=True, 
                        vmin=-0.5, vmax=1.0, ax=axcorrs[0])
            sns.heatmap(cell_vars_rules_sorted_norm_ordered_measure_cos, cmap=cs, square=True, 
                        vmin=-0.5, vmax=1.0, ax=axcorrs[1])
            sns.heatmap(cell_vars_rules_sorted_norm_ordered_measure_L2, cmap=cs, square=True, ax=axcorrs[2])

            for axcorr_index in range(len(axcorrs)): 
                axcorr = axcorrs[axcorr_index]
                # plot the group information (delimiter between different cluster)
                nn = cell_vars_rules_sorted_norm_ordered_measure.shape[0]
                boundaries = row_breakers
                for b in boundaries:
                    axcorr.axvline(b, 0, 1, color="k", linewidth=1.2)
                    axcorr.axhline(b, 0, 1, color="k", linewidth=1.2)
                
                axcorr.set_xticks(np.arange(len(tb_break_name)))
                axcorr.set_xticklabels(tb_break_name[result["row_order"]], rotation=45, ha='right', va='center', \
                                    rotation_mode='anchor', fontsize=9)    
                axcorr.set_yticks(np.arange(len(tb_break_name)))
                axcorr.set_yticklabels(tb_break_name[result["row_order"]], rotation=0, ha='right', va='center', \
                                    rotation_mode='anchor', fontsize=9) 
                axcorr.tick_params(axis="both", length=0)
                axcorr.set_title(f"{clustering_name}-{metrics[axcorr_index]}", fontsize=15)
            
            figcorr.tight_layout()
            figcorr.savefig(f"{save_dir}/{clustering_name}_variance_cluster_corr_{savefigure_name}.png", dpi=300)
            plt.close(figcorr)

            # register correlation information
            clustering_corr_info[clustering_save_name] = [
                cell_vars_rules_sorted_norm_ordered_measure, ordered_row_name, result["col_order"]
            ]

            # plot the effect of grouping & ordering through the feature axis
            # For unnormalized data use a log-scale colorbar spanning ~3 orders
            # of magnitude so that structure is visible despite large amplitude
            # differences across neurons. The data itself is NOT transformed.
            if clustering_normalize:
                heatmap_norm = None
                scale_label = ""
            else:
                data_max = np.max(cell_vars_rules_sorted_norm)
                data_min = max(np.min(cell_vars_rules_sorted_norm[cell_vars_rules_sorted_norm > 0]),
                               data_max / 1000.0)
                heatmap_norm = LogNorm(vmin=data_min, vmax=data_max)
                scale_label = " [log scale]"

            fig, ax = plt.subplots(2,1,figsize=(16,8*2))
            sns.heatmap(cell_vars_rules_sorted_norm, ax=ax[0], cmap=cs, cbar=True,
                        norm=heatmap_norm, vmin=vmins, vmax=vmaxs)
            sns.heatmap(cell_vars_rules_sorted_norm_ordered, ax=ax[1], cmap=cs, cbar=True,
                        norm=heatmap_norm, vmin=vmins, vmax=vmaxs)
            # Cluster boundaries: thick white line with a thin black outline for
            # visibility against any colormap value.
            for rb in rbreaks:
                ax[1].axhline(rb, color="w", lw=3.0, zorder=3)
                ax[1].axhline(rb, color="k", lw=0.8, zorder=4)
            for cb in cbreaks:
                ax[1].axvline(cb, color="w", lw=3.0, zorder=3)
                ax[1].axvline(cb, color="k", lw=0.8, zorder=4)
            ax[0].set_title(f"Before Clustering{scale_label}; best k row: {row_decision_k}; best k col: {col_decision_k}",
                            fontsize=15)
            ax[0].set_ylabel('Rule / Break-name', fontsize=12, labelpad=12)
            # Seaborn heatmap cells are centered at 0.5, 1.5, ...; use + 0.5 offset
            # so tick labels sit at the vertical center of each cell row.
            ax[0].set_yticks(np.arange(len(tb_break_name)) + 0.5)
            ax[0].set_yticklabels(tb_break_name, rotation=0, ha='right', va='center', fontsize=9)
            ax[1].set_title(f"After Clustering{scale_label}; best k row: {row_decision_k}; best k col: {col_decision_k}",
                            fontsize=15)
            ax[1].set_ylabel('Rule / Break-name', fontsize=12, labelpad=12)
            ax[1].set_yticks(np.arange(len(tb_break_name)) + 0.5)
            ax[1].set_yticklabels(tb_break_name[result["row_order"]], rotation=0, ha='right', va='center', fontsize=9)
            fig.savefig(f"{save_dir}/{clustering_name}_variance_cluster_{savefigure_name}.png", dpi=300)
            plt.close(fig)

            # 2025-11-04: for plotting purpose
            rbreaks_ = [0] + rbreaks + [cell_vars_rules_sorted_norm_ordered.shape[0]]
            cbreaks_ = [0] + cbreaks + [cell_vars_rules_sorted_norm_ordered.shape[1]]
            figvarmean, axsvarmean = plt.subplots(1,2,figsize=(4*2,4))
            varmean = np.zeros((len(rbreaks_) - 1, len(cbreaks_) - 1))
            for rr in range(len(rbreaks_) - 1):
                for cc in range(len(cbreaks_) - 1):
                    # 2025-11-04: mean covariance in this bicluster
                    varmean_ = np.mean(cell_vars_rules_sorted_norm_ordered[rbreaks_[rr]:rbreaks_[rr+1], 
                                                                           cbreaks_[cc]:cbreaks_[cc+1]])
                    varmean[rr,cc] = varmean_
            sns.heatmap(varmean, ax=axsvarmean[0], cmap=cs, cbar=True, vmin=vmins, vmax=vmaxs)
            axsvarmean[0].set_ylabel("Session Clusters", fontsize=15)
            axsvarmean[0].set_xlabel("Neuron Clusters", fontsize=15)
            axsvarmean[0].set_title("Mean Variance in Each Bicluster", fontsize=15)
            
            # 2025-11-04: for sanity check; since the neuron clusters are ordered so that the adjacent ones 
            # are more similar to each other than the ones that are further, therefore the correlation matrix
            # should have larger value near the diagonal 
            varmeanC = np.corrcoef(varmean, rowvar=False) 
            sns.heatmap(varmeanC, ax=axsvarmean[1], cmap=cs, cbar=True, vmin=0, vmax=1)
            axsvarmean[1].set_xlabel("Neuron Clusters", fontsize=15)
            axsvarmean[1].set_ylabel("Neuron Clusters", fontsize=15)
            axsvarmean[1].set_title("Correlation Between Neuron Clusters", fontsize=15)
            figvarmean.tight_layout()
            figvarmean.savefig(f"{save_dir}/{clustering_name}_variance_cluster_mean_{savefigure_name}.png", dpi=300)
            plt.close(figvarmean)

            # plot the norm of normalized variance of each session (across the neuron dimension)
            session_norm = cell_vars_rules_sorted_norm_ordered.sum(axis=1)
            norm_order = np.argsort(-session_norm)
            session_norm = session_norm[norm_order]
            session_norm_name = (tb_break_name[result["row_order"]])[norm_order]

            n_sess = len(session_norm)
            xs = np.arange(n_sess)
            norm_min, norm_max = session_norm.min(), session_norm.max()
            norm_range = norm_max - norm_min if norm_max > norm_min else 1.0

            # color each bar by its relative magnitude
            cmap_sn = plt.get_cmap("RdYlBu_r")
            bar_colors = [cmap_sn((v - norm_min) / norm_range) for v in session_norm]

            fig, ax = plt.subplots(1, 1, figsize=(max(12, n_sess * 0.55), 5))
            ax.bar(xs, session_norm, color=bar_colors, width=0.7, zorder=2)
            ax.plot(xs, session_norm, "-o", color="0.25", markersize=4,
                    linewidth=1.0, zorder=3)

            # horizontal reference line at the median
            median_val = np.median(session_norm)
            ax.axhline(median_val, color="0.5", linestyle="--", linewidth=1.0,
                       label=f"median = {median_val:.2f}", zorder=1)

            ax.set_xticks(xs)
            ax.set_xticklabels(session_norm_name, rotation=45,
                               ha="right", va="top", fontsize=9)
            ax.set_xlim(-0.7, n_sess - 0.3)
            ax.set_ylabel("Total Normalized Variance\n(sum across neurons)", fontsize=12)
            ax.set_xlabel("Task condition  (sorted by total variance)", fontsize=11)
            ax.set_title(f"{clustering_name} — per-condition total normalized variance",
                         fontsize=12, pad=8)
            ax.yaxis.grid(True, linestyle=":", linewidth=0.6, color="0.8", zorder=0)
            ax.set_axisbelow(True)
            ax.legend(fontsize=9, framealpha=0.7)
            ax.spines[["top", "right"]].set_visible(False)

            fig.tight_layout()
            fig.savefig(f"{save_dir}/{clustering_name}_variance_norm_{savefigure_name}.png", dpi=300)
            plt.close(fig)
            
            # register hierarchy clustering
            clustering_data_hierarchy[clustering_name] = result["col_linkage"]

        # align the correlation matrix for input and hidden based on an identical ordering 
        # this loop will be run during the hidden analysis iteration (not modulation iteration)
        # this IF condition will work when hidden is calculated but not yet move to the modulation iteration
        # 2026-04-10: == 4 instead, since we have 4 iterations before modulation
        if (len(clustering_corr_info) == 4) and ("all" not in clustering_name):
            corr_comparison_pairs = [
                ("input_normalized",   "hidden_normalized",   "normalized"),
                ("input_unnormalized", "hidden_unnormalized", "unnormalized"),
            ]
            for inp_key, hid_key, norm_desc in corr_comparison_pairs:
                input_order_row  = clustering_corr_info[inp_key][1]
                hidden_order_row = clustering_corr_info[hid_key][1]
                input_corr       = clustering_corr_info[inp_key][0]
                hidden_corr      = clustering_corr_info[hid_key][0]

                # reorder hidden rows/cols to match input's task-condition ordering
                shuffle_hidden_to_input = helper.permutation_indices_b_to_a(
                    input_order_row, hidden_order_row
                )
                hidden_corr_reordered = hidden_corr[
                    np.ix_(shuffle_hidden_to_input, shuffle_hidden_to_input)
                ]

                figinputhiddencorr, axinputhiddencorr = plt.subplots(1, 2, figsize=(8*2, 8))
                sns.heatmap(input_corr,           ax=axinputhiddencorr[0], cmap=cs, square=True, vmin=-0.5, vmax=1.0)
                sns.heatmap(hidden_corr_reordered, ax=axinputhiddencorr[1], cmap=cs, square=True, vmin=-0.5, vmax=1.0)

                for ax in axinputhiddencorr:
                    ax.set_xticks(np.arange(len(input_order_row)) + 0.5)
                    ax.set_xticklabels(input_order_row, rotation=45, ha='right', va='center',
                                       rotation_mode='anchor', fontsize=9)
                    ax.set_yticks(np.arange(len(input_order_row)) + 0.5)
                    ax.set_yticklabels(input_order_row, rotation=0, ha='right', va='center',
                                       rotation_mode='anchor', fontsize=9)
                    ax.tick_params(axis="both", length=0)

                axinputhiddencorr[0].set_title(f"Input Correlation ({norm_desc})", fontsize=15)
                axinputhiddencorr[1].set_title(
                    f"Hidden Correlation — reordered by input ({norm_desc})", fontsize=15
                )
                # draw cluster breaks using input's break positions
                for ax in axinputhiddencorr:
                    for b in row_cluster_breaker_all[inp_key]:
                        ax.axvline(b, 0, 1, color="k", linewidth=1.2)
                        ax.axhline(b, 0, 1, color="k", linewidth=1.2)

                figinputhiddencorr.suptitle(f"Input vs Hidden Correlation ({norm_desc})")
                figinputhiddencorr.savefig(
                    f"{save_dir}/input2hidden_variance_hierarchy_{norm_desc}_{savefigure_name_base}.png",
                    dpi=300
                )
                plt.close(figinputhiddencorr)
        
        # 2025-11-04: check the consistency of row clusters between input & hidden
        if (len(clustering_corr_info) == 4) and ("all" not in clustering_name):
            # save the clustering information for input & hidden
            with open(f"{save_dir}/cluster_info_{savefigure_name_base}.pkl", "wb") as f:
                pickle.dump(cluster_info_save, f)
                
            print(cluster_info_save.keys())

            # ----------------------------------------------------------------
            # 2026-04-12: Input vs Hidden clustering comparison
            # Row axis  — task conditions (always comparable: same tb_break).
            # Col axis  — neurons.  Comparable when input_dim == hidden_dim
            #             (square modulation layer); skipped otherwise.
            # Metrics: ARI, NMI (permutation-invariant, handle different k).
            # Figure layout: 2 rows × 2 cols.
            #   Row 0 — task-condition (row) clusters
            #   Row 1 — neuron (col) clusters  [omitted if dims differ]
            # Each column: left = raw counts, right = row-normalised fraction.
            # ----------------------------------------------------------------
            comparison_pairs = [
                ("input_normalized",   "hidden_normalized",   "normalized"),
                ("input_unnormalized", "hidden_unnormalized", "unnormalized"),
            ]

            # Fixed canvas size — content adapts via font/annotation scaling.
            _FIG_W       = 14.0   # total figure width  (2 panels side by side)
            _FIG_ROW_H   =  5.5   # height per subplot row
            _FIG_TOP_PAD =  0.9   # extra space for suptitle

            def _contingency_panels(ax_count, ax_frac, lbl_a, lbl_b,
                                    k_a, k_b, prefix_a, prefix_b,
                                    axis_label_a, axis_label_b):
                """Draw count and row-norm fraction heatmaps onto two axes.

                Font sizes adapt to k so panels stay readable at fixed canvas size.
                Annotations are suppressed when the matrix is too dense (>15×15).
                """
                cm      = sk_contingency_matrix(lbl_a, lbl_b)
                cm_frac = cm.astype(float) / cm.sum(axis=1, keepdims=True)
                n_a, n_b = cm.shape

                # cluster sizes from contingency matrix marginals
                sizes_a = cm.sum(axis=1)   # row sums → sizes of lbl_a clusters
                sizes_b = cm.sum(axis=0)   # col sums → sizes of lbl_b clusters
                xlbls = [f"{prefix_b}{i}\n| n={sizes_b[i]}" for i in range(n_b)]
                ylbls = [f"{prefix_a}{i}\n| n={sizes_a[i]}" for i in range(n_a)]

                # scale font sizes with matrix dimensions; clamp to readable range
                max_dim    = max(n_a, n_b)
                annot_fs   = max(6, min(10, int(130 / max(max_dim, 1))))
                tick_fs    = max(6, min(10, int(110 / max(max_dim, 1))))
                show_annot = (n_a <= 15 and n_b <= 15)

                _hm_kw = dict(
                    xticklabels=xlbls, yticklabels=ylbls,
                    annot=show_annot,
                    annot_kws={"size": annot_fs} if show_annot else {},
                )

                sns.heatmap(cm, ax=ax_count, fmt="d", cmap="Blues",
                            cbar=False, **_hm_kw)
                ax_count.set_xlabel(f"{axis_label_b}  (k={k_b})", fontsize=9)
                ax_count.set_ylabel(f"{axis_label_a}  (k={k_a})", fontsize=9)
                ax_count.tick_params(axis="both", labelsize=tick_fs)

                sns.heatmap(cm_frac, ax=ax_frac, fmt=".2f", cmap="Blues",
                            vmin=0, vmax=1, cbar=True, **_hm_kw)
                ax_frac.set_xlabel(f"{axis_label_b}  (k={k_b})", fontsize=9)
                ax_frac.set_ylabel(f"{axis_label_a}  (k={k_a})", fontsize=9)
                ax_frac.tick_params(axis="both", labelsize=tick_fs)

                ari = adjusted_rand_score(lbl_a, lbl_b)
                nmi = normalized_mutual_info_score(lbl_a, lbl_b)
                return ari, nmi, n_a, n_b

            def _make_comparison_figure(n_subplot_rows):
                """Create a fixed-size figure with n_subplot_rows × 2 axes."""
                fig_h = _FIG_ROW_H * n_subplot_rows + _FIG_TOP_PAD
                fig, axes = plt.subplots(
                    n_subplot_rows, 2,
                    figsize=(_FIG_W, fig_h),
                    gridspec_kw={"wspace": 0.45, "hspace": 0.6},
                )
                if n_subplot_rows == 1:
                    axes = axes[np.newaxis, :]
                return fig, axes

            # ----------------------------------------------------------------
            # Input vs Hidden clustering comparison
            # ----------------------------------------------------------------
            for input_key, hidden_key, norm_desc in comparison_pairs:
                if input_key not in cluster_info_save or hidden_key not in cluster_info_save:
                    continue

                r_input  = cluster_info_save[input_key]["result"]
                r_hidden = cluster_info_save[hidden_key]["result"]

                # ---- row (task-condition) labels ----
                lbl_row_inp = np.asarray(r_input["row_tol_labels"])
                lbl_row_hid = np.asarray(r_hidden["row_tol_labels"])
                k_row_inp   = r_input["row_tol_k"]
                k_row_hid   = r_hidden["row_tol_k"]
                assert len(lbl_row_inp) == len(lbl_row_hid), (
                    f"row_labels length mismatch: input={len(lbl_row_inp)}, "
                    f"hidden={len(lbl_row_hid)}"
                )

                # ---- col (neuron) labels — only when dims match ----
                lbl_col_inp = np.asarray(r_input["col_tol_labels"])
                lbl_col_hid = np.asarray(r_hidden["col_tol_labels"])
                k_col_inp   = r_input["col_tol_k"]
                k_col_hid   = r_hidden["col_tol_k"]
                col_comparable = (len(lbl_col_inp) == len(lbl_col_hid))
                if not col_comparable:
                    print(
                        f"[input vs hidden] {norm_desc} — col dims differ "
                        f"(input={len(lbl_col_inp)}, hidden={len(lbl_col_hid)}); "
                        f"skipping neuron-cluster comparison."
                    )

                n_subplot_rows = 2 if col_comparable else 1
                fig, axes = _make_comparison_figure(n_subplot_rows)

                # -- row 0: task-condition clusters --
                ari_row, nmi_row, _, _ = _contingency_panels(
                    axes[0, 0], axes[0, 1],
                    lbl_row_inp, lbl_row_hid,
                    k_row_inp, k_row_hid,
                    "I", "H",
                    "Input task-cond", "Hidden task-cond",
                )
                axes[0, 0].set_title(
                    f"Task conditions — Count\n(ARI={ari_row:.3f}  NMI={nmi_row:.3f})",
                    fontsize=10,
                )
                axes[0, 1].set_title(
                    f"Task conditions — Row-norm fraction\n(ARI={ari_row:.3f}  NMI={nmi_row:.3f})",
                    fontsize=10,
                )

                # -- row 1: neuron clusters (conditional) --
                if col_comparable:
                    ari_col, nmi_col, _, _ = _contingency_panels(
                        axes[1, 0], axes[1, 1],
                        lbl_col_inp, lbl_col_hid,
                        k_col_inp, k_col_hid,
                        "I", "H",
                        "Input neurons", "Hidden neurons",
                    )
                    axes[1, 0].set_title(
                        f"Neurons — Count\n(ARI={ari_col:.3f}  NMI={nmi_col:.3f})",
                        fontsize=10,
                    )
                    axes[1, 1].set_title(
                        f"Neurons — Row-norm fraction\n(ARI={ari_col:.3f}  NMI={nmi_col:.3f})",
                        fontsize=10,
                    )
                    col_info = (f"  |  neuron ARI={ari_col:.3f}  NMI={nmi_col:.3f}"
                                f"  (col k: input={k_col_inp}, hidden={k_col_hid})")
                else:
                    col_info = "  |  neuron clusters: dims differ — skipped"

                fig.suptitle(
                    f"Input vs Hidden  ({norm_desc})\n"
                    f"task-cond ARI={ari_row:.3f}  NMI={nmi_row:.3f}"
                    f"  (row k: input={k_row_inp}, hidden={k_row_hid})"
                    f"{col_info}",
                    fontsize=9, y=1.01,
                )
                out_path = (
                    f"{save_dir}/input_vs_hidden_"
                    f"{norm_desc}_{savefigure_name_base}.png"
                )
                fig.savefig(out_path, dpi=200, bbox_inches="tight")
                plt.close(fig)
                print(
                    f"[input vs hidden] {norm_desc} — task conditions: "
                    f"ARI={ari_row:.3f}  NMI={nmi_row:.3f}  "
                    f"(input k={k_row_inp}  hidden k={k_row_hid})"
                )
                if col_comparable:
                    print(
                        f"[input vs hidden] {norm_desc} — neurons: "
                        f"ARI={ari_col:.3f}  NMI={nmi_col:.3f}  "
                        f"(input k={k_col_inp}  hidden k={k_col_hid})"
                    )

            # ----------------------------------------------------------------
            # 2026-04-21: Normalized vs Unnormalized comparison (same modality)
            # Row axis  — task conditions (always comparable).
            # Col axis  — neurons (always comparable: same neurons, diff norm).
            # ----------------------------------------------------------------
            norm_vs_unnorm_pairs = [
                ("input_normalized",  "input_unnormalized",  "input"),
                ("hidden_normalized", "hidden_unnormalized", "hidden"),
            ]

            for norm_key, unnorm_key, modality in norm_vs_unnorm_pairs:
                if norm_key not in cluster_info_save or unnorm_key not in cluster_info_save:
                    continue

                r_norm   = cluster_info_save[norm_key]["result"]
                r_unnorm = cluster_info_save[unnorm_key]["result"]

                # ---- row (task-condition) labels ----
                lbl_row_norm   = np.asarray(r_norm["row_tol_labels"])
                lbl_row_unnorm = np.asarray(r_unnorm["row_tol_labels"])
                k_row_norm     = r_norm["row_tol_k"]
                k_row_unnorm   = r_unnorm["row_tol_k"]
                assert len(lbl_row_norm) == len(lbl_row_unnorm), (
                    f"row_labels length mismatch: norm={len(lbl_row_norm)}, "
                    f"unnorm={len(lbl_row_unnorm)}"
                )

                # ---- col (neuron) labels ----
                lbl_col_norm   = np.asarray(r_norm["col_tol_labels"])
                lbl_col_unnorm = np.asarray(r_unnorm["col_tol_labels"])
                k_col_norm     = r_norm["col_tol_k"]
                k_col_unnorm   = r_unnorm["col_tol_k"]
                assert len(lbl_col_norm) == len(lbl_col_unnorm), (
                    f"col_labels length mismatch: norm={len(lbl_col_norm)}, "
                    f"unnorm={len(lbl_col_unnorm)}"
                )

                fig, axes = _make_comparison_figure(2)

                # -- row 0: task-condition clusters --
                ari_row, nmi_row, _, _ = _contingency_panels(
                    axes[0, 0], axes[0, 1],
                    lbl_row_norm, lbl_row_unnorm,
                    k_row_norm, k_row_unnorm,
                    "N", "U",
                    f"{modality.capitalize()} norm task-cond",
                    f"{modality.capitalize()} unnorm task-cond",
                )
                axes[0, 0].set_title(
                    f"Task conditions — Count\n(ARI={ari_row:.3f}  NMI={nmi_row:.3f})",
                    fontsize=10,
                )
                axes[0, 1].set_title(
                    f"Task conditions — Row-norm fraction\n(ARI={ari_row:.3f}  NMI={nmi_row:.3f})",
                    fontsize=10,
                )

                # -- row 1: neuron clusters --
                ari_col, nmi_col, _, _ = _contingency_panels(
                    axes[1, 0], axes[1, 1],
                    lbl_col_norm, lbl_col_unnorm,
                    k_col_norm, k_col_unnorm,
                    "N", "U",
                    f"{modality.capitalize()} norm neurons",
                    f"{modality.capitalize()} unnorm neurons",
                )
                axes[1, 0].set_title(
                    f"Neurons — Count\n(ARI={ari_col:.3f}  NMI={nmi_col:.3f})",
                    fontsize=10,
                )
                axes[1, 1].set_title(
                    f"Neurons — Row-norm fraction\n(ARI={ari_col:.3f}  NMI={nmi_col:.3f})",
                    fontsize=10,
                )

                fig.suptitle(
                    f"{modality.capitalize()} — normalized vs unnormalized\n"
                    f"task-cond ARI={ari_row:.3f}  NMI={nmi_row:.3f}"
                    f"  (row k: norm={k_row_norm}, unnorm={k_row_unnorm})  |  "
                    f"neuron ARI={ari_col:.3f}  NMI={nmi_col:.3f}"
                    f"  (col k: norm={k_col_norm}, unnorm={k_col_unnorm})",
                    fontsize=9, y=1.01,
                )
                out_path = (
                    f"{save_dir}/norm_vs_unnorm_"
                    f"{modality}_{savefigure_name_base}.png"
                )
                fig.savefig(out_path, dpi=200, bbox_inches="tight")
                plt.close(fig)
                print(
                    f"[norm vs unnorm] {modality} — task conditions: "
                    f"ARI={ari_row:.3f}  NMI={nmi_row:.3f}  "
                    f"(norm k={k_row_norm}  unnorm k={k_row_unnorm})"
                )
                print(
                    f"[norm vs unnorm] {modality} — neurons: "
                    f"ARI={ari_col:.3f}  NMI={nmi_col:.3f}  "
                    f"(norm k={k_col_norm}  unnorm k={k_col_unnorm})"
                )

        # use the clustering result for input and hidden to order the modulation information
        # and/or cross-compare the clustering result from input & hidden 
        # trying to observe consistency in between
        # 2026-04-06: this part of analysis is specifically targeting to the modulation-related data
        ih_mod_run = False # 2026-04-15: control this part is running or not
        if (len(clustering_corr_info) == 4) and ("all" in clustering_name) and ih_mod_run:
            # Apply log1p to unnormalized modulation data before clustering,
            # matching the same transform used for unnormalized input/hidden data.
            # Variance values are non-negative, so log1p is valid and compresses
            # the dynamic range so Euclidean distance captures pattern shape
            # rather than raw amplitude. Normalized data is passed unchanged.
            if clustering_normalize:
                V_for_clustering_mod = cell_vars_rules_sorted_norm
            else:
                V_for_clustering_mod = np.log1p(cell_vars_rules_sorted_norm)

            input_order_col_ind  = clustering_corr_info["input_normalized"][2]
            hidden_order_col_ind = clustering_corr_info["hidden_normalized"][2]
            # cell_vars_rules_norm_keepshape: 3D array with shape (N, post, pre)
            # modulation_W shape is (post, pre): previously assumed (pre, post)
            # axis 1 = post (hidden), axis 2 = pre (input)
            cell_vars_rules_norm_keepshape_ih = cell_vars_rules_norm_keepshape[:, hidden_order_col_ind, :]
            cell_vars_rules_norm_keepshape_ih = cell_vars_rules_norm_keepshape_ih[:, :, input_order_col_ind]

            # shape is (N, post, pre): no transpose groups by post; transpose groups by pre
            flatten_by_post = cell_vars_rules_norm_keepshape_ih.reshape(N, M*M)
            flatten_by_pre = cell_vars_rules_norm_keepshape_ih.transpose(0,2,1).reshape(N, M*M)

            fig, axs = plt.subplots(2,1,figsize=(24,8*2))
            sns.heatmap(flatten_by_pre, ax=axs[0], cmap=cs, cbar=True, vmin=vmins, vmax=vmaxs)
            sns.heatmap(flatten_by_post, ax=axs[1], cmap=cs, cbar=True, vmin=vmins, vmax=vmaxs)
            axs[0].set_title("Flatten By Pre (Input)", fontsize=15)
            axs[1].set_title("Flatten By Post (Hidden)", fontsize=15)
            for ax in axs: 
                ax.set_yticks(np.arange(len(tb_break_name)))
                ax.set_yticklabels(tb_break_name, rotation=0, ha='right', va='center', fontsize=9)
            fig.tight_layout()
            fig.savefig(f"{save_dir}/modulation_input2hidden_variance_hierarchy_{savefigure_name}.png", dpi=300)
            plt.close(fig)

            #
            print(cell_vars_rules_sorted_norm.shape)
            if clustering_normalize:
                cluster_input  = col_clusters_all["input_normalized"]
                cluster_hidden = col_clusters_all["hidden_normalized"]
            else:
                cluster_input  = col_clusters_all["input_unnormalized"]
                cluster_hidden = col_clusters_all["hidden_unnormalized"]
                
            cluster_combine = {}
            newc = 0
            for c1 in cluster_input.keys(): # input cluster name
                for c2 in cluster_hidden.keys(): # hidden cluster name
                    newc += 1 # iterative counter of new cell type
                    # different newc stands for a unique cell type 
                    # (defined as a multiplication cluster of pre [input] and post [hidden])
                    nc1_input, nc2_hidden = cluster_input[c1], cluster_hidden[c2]
                    nc_combine = [(int(i),int(j)) for i in nc1_input for j in nc2_hidden]
                    # transform to reshaped matrix index
                    # modulation_W shape is (post, pre): previously assumed (pre, post)
                    # flat index in (post, pre) layout: post * M + pre
                    nc_combine = [cc[1] * M + cc[0] for cc in nc_combine]
                    cluster_combine[newc] = nc_combine

            # 2026-10-08: qualitatively, training MPN without regularization will cause more hidden cluster
            cluster_input_num, cluster_hidden_num = len(cluster_input), len(cluster_hidden)
            print(f"cluster_input_num: {cluster_input_num}; cluster_hidden_num: {cluster_hidden_num}")
            # by the setup, cluster_combine will be organized by shared pre
            # same pre neuron will be placed adjacently
            cluster_combine_pre = copy.deepcopy(cluster_combine)
            # create order of key to put post neuron together 
            post_order = []
            for c1 in range(cluster_input_num):
                for c2 in range(cluster_hidden_num):
                    post_order.append(c1 + c2 * cluster_input_num + 1) 

            pre_order = list(cluster_combine.keys())

            pre_order_group, post_order_group = [], [] 
            for ind in range(cluster_hidden_num):
                pre_order_group.append([i+1 for i in range(ind * cluster_input_num, (ind+1) * cluster_input_num)])
            for ind in range(cluster_input_num):
                block = []
                for ind2 in range(cluster_hidden_num):
                    block.append(ind + ind2 * cluster_input_num + 1)
                post_order_group.append(block)
            
            group_neurons_comb_pre, group_neurons_comb_post = [], [] 
            for preo in pre_order:
                group_neurons_comb_pre.extend(cluster_combine[preo])
            for posto in post_order:
                group_neurons_comb_post.extend(cluster_combine[posto])

            # 2025-09-02: we dont want random order of cell type
            # but based on certain pre or post neuron clustering 
            # random_celltype_orders = helper.concat_random_samples(cluster_combine, n_samples=1, seed=42)
            # random_celltype_orders = random_celltype_orders[0] # take the first sample
            # group_neurons_comb = random_celltype_orders[0] # concatenated list
            # assert len(group_neurons_comb) == len(set(group_neurons_comb)) # every element is unique
            # group_neurons = random_celltype_orders[1] # list of list, the order does not matter here

            group_neurons = [group for group in cluster_combine.values()]

            group_neurons_pre = []
            for pre_order_ in pre_order_group:
                block = []
                for ind in pre_order_:
                    block.extend(group_neurons[ind-1])
                group_neurons_pre.append(block)

            group_neurons_post = []
            for post_order_ in post_order_group:
                block = []
                for ind in post_order_:
                    block.extend(group_neurons[ind-1])
                group_neurons_post.append(block)

            # assert helper.all_leq(group_neurons_comb, cell_vars_rules_sorted_norm.shape[1])

            # cell_vars_rules_sorted_norm_r1 = cell_vars_rules_sorted_norm[:,group_neurons_comb]
            cell_vars_rules_sorted_norm_outer_bypre = cell_vars_rules_sorted_norm[:,group_neurons_comb_pre]
            cell_vars_rules_sorted_norm_outer_bypost = cell_vars_rules_sorted_norm[:,group_neurons_comb_post]

            fig, axs = plt.subplots(6,1,figsize=(24,8*6))
            sns.heatmap(cell_vars_rules_sorted_norm, ax=axs[0], cmap=cs, cbar=True, vmin=vmins, vmax=vmaxs)
            sns.heatmap(cell_vars_rules_sorted_norm_outer_bypre, ax=axs[1], cmap=cs, cbar=True, vmin=vmins, vmax=vmaxs)
            sns.heatmap(cell_vars_rules_sorted_norm_outer_bypost, ax=axs[2], cmap=cs, cbar=True, vmin=vmins, vmax=vmaxs)
            axs[0].set_title("Original", fontsize=15)
            axs[1].set_title(f"Ordering based on the Outer Product of Input & Hidden Clusters [Input Adjacent]", fontsize=15)
            axs[2].set_title(f"Ordering based on the Outer Product of Input & Hidden Clusters [Hidden Adjacent]", fontsize=15)
            for iii in range(3):
                axs[iii].set_yticks(np.arange(len(tb_break_name)))
                axs[iii].set_yticklabels(tb_break_name, rotation=0, ha='right', va='center', fontsize=9)

            group_all = [group_neurons, group_neurons_pre, group_neurons_post]
            group_all_names = ["Outer Product of Input & Hidden Clusters", "Shared Hidden", "Shared Input"]

            for group_neurons_index in range(len(group_all)):
                group_neurons_ = group_all[group_neurons_index]
                result_outer = clustering.cluster_variance_matrix_forgroup(V_for_clustering_mod,
                                                                           row_groups=None,
                                                                           col_groups_all_lst=[group_neurons_],
                                                                           k_min=lower_cluster,
                                                                           k_max=upper_cluster,
                                                                           silhouette_tol=silhouette_tol,
                                                                           tol_mode=tol_mode,
                                                                           metric=c_metric,
                                                                           score_quantile=score_quantile,
                                                                           unresponsive_norm_frac=unresponsive_norm_frac,
                                                                           tol_k_select=tol_k_select,
                                                                           skip_unresponsive_detection=clustering_normalize)
                prior_cluster_num = len(group_neurons_)
                cell_vars_rules_sorted_norm_inputhidden = cell_vars_rules_sorted_norm[np.ix_(result_outer["row_order"], 
                                                                                            result_outer["col_order"])]
        
                sns.heatmap(cell_vars_rules_sorted_norm_inputhidden, ax=axs[3+group_neurons_index], cmap=cs, cbar=True, 
                            vmin=vmins, vmax=vmaxs)
                axs[3+group_neurons_index].set_yticks(np.arange(len(tb_break_name)))
                axs[3+group_neurons_index].set_yticklabels(tb_break_name[result_outer["row_order"]], 
                                                           rotation=0, ha='right', va='center', fontsize=9)
                axs[3+group_neurons_index].set_title(f"Ordering based on the {group_all_names[group_neurons_index]}, #{prior_cluster_num} [Post-Clustering]", fontsize=15)
                
            fig.tight_layout()
            fig.savefig(f"{save_dir}/modulation_input2hiddentogether_variance_hierarchy_{savefigure_name}.png", dpi=300)
            plt.close(fig)
        
          
        if ("all" in clustering_name): 
            # ----------------------------------------------------------------
            # Modulation-specific clustering analysis
            # ----------------------------------------------------------------
            # The flattened modulation matrix (N_sessions x pre*post) is clustered
            # in three ways along the column (synapse) axis:
            #   1. "pre"  — columns grouped by shared presynaptic neuron
            #   2. "post" — columns grouped by shared postsynaptic neuron
            #   3. "all"  — columns grouped via MiniBatchKMeans with varying G
            # Row (session) clustering is always done without prior grouping.
            # For each scheme we obtain row/col orders, cluster labels, and
            # silhouette-based optimal k, then plot heatmaps and compute
            # pre/post membership statistics.
            # ----------------------------------------------------------------

            # Apply log1p to unnormalized modulation data before clustering,
            # matching the same transform used for unnormalized input/hidden data.
            # Normalized data is passed unchanged.
            if clustering_normalize:
                V_for_clustering_mod = cell_vars_rules_sorted_norm
            else:
                V_for_clustering_mod = np.log1p(cell_vars_rules_sorted_norm)

            # 2026-04-10: just a buffer to save the 4D shape of the modulation matrix
            assert len(clustering_data_old.shape) == 4
            # clustering_data_old shape is (N, T, post, pre) = (N, T, hidden, input)
            post_num, pre_num = clustering_data_old.shape[2], clustering_data_old.shape[3]
            # C-order flat: flat_idx = post_idx * pre_num + pre_idx
            # Group by shared pre (input) neuron: strided pattern {i, i+pre_num, ...}
            feature_group_pre = []
            for i in range(pre_num):
                feature_group_pre.append(helper.basic_sort([i + j * pre_num
                                                             for j in range(post_num)], sort_idxs))
            # Group by shared post (hidden) neuron: contiguous blocks of size pre_num
            feature_group_post = []
            for i in range(post_num):
                feature_group_post.append(helper.basic_sort([j for j in range(pre_num * i, pre_num * (i+1))], sort_idxs))

            result_pre = clustering.cluster_variance_matrix_forgroup(V_for_clustering_mod,
                                                                    row_groups=None,
                                                                    col_groups_all_lst=[feature_group_pre],
                                                                    k_min=lower_cluster,
                                                                    k_max=upper_cluster,
                                                                    silhouette_tol=silhouette_tol,
                                                                    tol_mode=tol_mode,
                                                                    metric=c_metric,
                                                                    score_quantile=score_quantile,
                                                                    unresponsive_norm_frac=unresponsive_norm_frac,
                                                                    tol_k_select=tol_k_select,
                                                                    skip_unresponsive_detection=clustering_normalize)

            result_post = clustering.cluster_variance_matrix_forgroup(V_for_clustering_mod,
                                                                    row_groups=None,
                                                                    col_groups_all_lst=[feature_group_post],
                                                                    k_min=lower_cluster,
                                                                    k_max=upper_cluster,
                                                                    silhouette_tol=silhouette_tol,
                                                                    tol_mode=tol_mode,
                                                                    metric=c_metric,
                                                                    score_quantile=score_quantile,
                                                                    unresponsive_norm_frac=unresponsive_norm_frac,
                                                                    tol_k_select=tol_k_select,
                                                                    skip_unresponsive_detection=clustering_normalize)

            for _lbl, _res in [("result_pre", result_pre), ("result_post", result_post)]:
                _row_k_strict = _res["row_strict_best_k"]
                _col_k_strict = _res["col_strict_best_k"]
                _row_k_tol    = _res["row_k"]
                _col_k_tol    = _res["col_k"]
                _row_sil_strict = _res["row_score_recording_mean"].get(_row_k_strict, float("nan"))
                _col_sil_strict = _res["col_score_recording_mean"].get(_col_k_strict, float("nan"))
                _row_sil_tol    = _res["row_score_recording_mean"].get(_row_k_tol, float("nan"))
                _col_sil_tol    = _res["col_score_recording_mean"].get(_col_k_tol, float("nan"))
                print(
                    f"[{_lbl}] row_k (strict): {_row_k_strict} "
                    f"(sil={_row_sil_strict:.4f}); "
                    f"col_k (strict): {_col_k_strict} "
                    f"(sil={_col_sil_strict:.4f})"
                )
                print(
                    f"[{_lbl}] row_k (tol):    {_row_k_tol} "
                    f"(sil={_row_sil_tol:.4f}); "
                    f"col_k (tol):    {_col_k_tol} "
                    f"(sil={_col_sil_tol:.4f})"
                )

            # 2025-10-06: do not give any prior grouping prior to the modulation information
            # simply grouping and considering each individual column as separate
            # having smaller G, e.g. 200, will make the following calculation in determining pre- and post-
            # belonging identity more time costly
            G_lst = [100, 300, 1000]
            figcol, axscol = plt.subplots(1,len(G_lst),figsize=(4*len(G_lst),4))
            figcolcluster, axscolcluster = plt.subplots(1,len(G_lst),figsize=(4*len(G_lst),4))
            figppshare, axsppshare = plt.subplots(2,len(G_lst),figsize=(4*len(G_lst),4*2))

            result_all_lst = []
            result_all_name_lst = []
            # Per-G breakdown of how many modulations survive each mask
            # (mask 1: unresponsive mod cluster; mask 2: endpoint in unresponsive input/hidden)
            mask_breakdown_lst = []
            prepost_belonging_results = []
            global_assignment_cache = None
            FIXED_K_OM = 20
            global_assignment_fixed_k_cache = None

            # 2025-11-04: input cluster & hidden cluster along the neuron dimension (N)
            if clustering_normalize:
                cluster_input  = col_clusters_all["input_normalized"]
                cluster_hidden = col_clusters_all["hidden_normalized"]
            else:
                cluster_input  = col_clusters_all["input_unnormalized"]
                cluster_hidden = col_clusters_all["hidden_unnormalized"]
                
            # sanity check: order it based on the key
            cluster_input = dict(sorted(cluster_input.items()))
            cluster_hidden = dict(sorted(cluster_hidden.items()))
            
            for G_idx in range(len(G_lst)): 
                G = G_lst[G_idx]
                assert G <= cell_vars_rules_sorted_norm.shape[0] * cell_vars_rules_sorted_norm.shape[1]
                
                col_groups_all_lst = []
                # 2025-10-06: since K-means has randomness, we run it multiple times to obtain 
                # the distribution of grouping statistics,
                krun = 20 
                print(f"Running K-means with G={G} for {krun} times to obtain the distribution of grouping statistics...")
                for _ in range(krun):
                    random_seed = np.random.randint(0, 10000)
                    col_groups_all = clustering.make_col_groups_with_kmeans(V_for_clustering_mod,
                                                                            n_groups=G,
                                                                            random_state=random_seed)[0]
                    col_groups_all_lst.append(col_groups_all)
        
                # plot the statistics of grouping on all modulations without considering their pre/post identity 
                for gi in range(5): 
                    group_lengths = [len(sublist) for sublist in col_groups_all_lst[gi]]
                    axscol[G_idx].hist(group_lengths, color=c_vals[gi], bins=20, edgecolor='black', alpha=0.7)
                # if uniform grouping, what is the expected length for each group 
                axscol[G_idx].axvline(MM / G, linestyle="--", color="black")
                axscol[G_idx].set_xlabel('Length of Formed Group', fontsize=15)
                axscol[G_idx].set_ylabel('Frequency', fontsize=15)

                # Modulation analysis: varying G changes the column-side grouped representation,
                # because columns are first averaged within K-means groups before hierarchical clustering.
                # Note that the row-side data are unchanged here; only the searched k-range also changes
                # because we currently pass k_max=G to both axes.
                result_all = clustering.cluster_variance_matrix_forgroup(V_for_clustering_mod,
                                                                         row_groups=None,
                                                                         col_groups_all_lst=col_groups_all_lst,
                                                                         k_min=lower_cluster_mod,
                                                                         k_max=G,
                                                                         silhouette_tol=silhouette_tol,
                                                                         tol_mode=tol_mode,
                                                                         metric=c_metric,
                                                                         score_quantile=score_quantile,
                                                                         unresponsive_norm_frac=unresponsive_norm_frac,
                                                                         tol_k_select=tol_k_select,
                                                                         skip_unresponsive_detection=clustering_normalize)
        
                _r_sil      = result_all["row_score_recording_mean"]
                _c_sil      = result_all["col_score_recording_mean"]
                _r_strict_k = result_all["row_strict_best_k"]
                _c_strict_k = result_all["col_strict_best_k"]
                # Tolerance band stored directly in the result — no need to recompute.
                _r_band     = result_all["row_primary_candidates"]
                _c_band     = result_all["col_primary_candidates"]

                def _bsel(band, strategy, fallback, Z=None):
                    if not band:
                        return fallback
                    if strategy == "min":
                        return min(band)
                    if strategy == "max":
                        return max(band)
                    if strategy == "gap":
                        return clustering._gap_k(Z, band) if Z is not None else fallback
                    m = float(np.mean(band))
                    return min(band, key=lambda kk: (abs(kk - m), kk))

                _strat_ks = {
                    "min":  (_bsel(_r_band, "min",  _r_strict_k),
                             _bsel(_c_band, "min",  _c_strict_k)),
                    "max":  (_bsel(_r_band, "max",  _r_strict_k),
                             _bsel(_c_band, "max",  _c_strict_k)),
                    "mean": (_bsel(_r_band, "mean", _r_strict_k),
                             _bsel(_c_band, "mean", _c_strict_k)),
                    "gap":  (_bsel(_r_band, "gap",  _r_strict_k, result_all["row_linkage"]),
                             _bsel(_c_band, "gap",  _c_strict_k, result_all["col_linkage"])),
                }

                print(
                    f"[G={G}]   strict argmax — "
                    f"row: {_r_strict_k} (sil={_r_sil.get(_r_strict_k, float('nan')):.4f}); "
                    f"col: {_c_strict_k} (sil={_c_sil.get(_c_strict_k, float('nan')):.4f})"
                )
                for _strat, (_rk, _ck) in _strat_ks.items():
                    _marker = "  <-- selected" if _strat == tol_k_select else ""
                    print(
                        f"[G={G}]   {_strat:4s}: row={_rk:3d} (sil={_r_sil.get(_rk, float('nan')):.4f}); "
                        f"col={_ck:3d} (sil={_c_sil.get(_ck, float('nan')):.4f}){_marker}"
                    )
                assert result_all["col_k"] < G
                axscol[G_idx].set_title(f"G={G}, Neuron Cluster={result_all['col_k']}", fontsize=15)

                # Plot the size distribution of the actual modulation clusters
                # after cluster_variance_matrix_forgroup expands labels back to
                # the original modulation columns.
                actual_cluster_ids, actual_cluster_sizes = np.unique(
                    np.asarray(result_all["col_labels"]), return_counts=True
                )
                # assert actual_cluster_ids.size in (result_all["col_k"], result_all["col_k"] + 1)
                axscolcluster[G_idx].hist(
                    actual_cluster_sizes,
                    color=c_vals[G_idx],
                    bins="auto",
                    edgecolor='black',
                    alpha=0.7
                )
                axscolcluster[G_idx].axvline(
                    MM / result_all["col_k"], linestyle="--", color="black"
                )
                axscolcluster[G_idx].set_xlabel('Length of Actual Cluster', fontsize=15)
                axscolcluster[G_idx].set_ylabel('Frequency', fontsize=15)
                axscolcluster[G_idx].set_title(
                    f"G={G}, Actual Cluster={result_all['col_k']}", fontsize=15
                )

                result_all_lst.append(result_all)
                result_all_name_lst.append(f"G={G}")
        
                # 2025-10-21: after grouping and clustering the modulation, 
                # we ask for two modulation within the same cluster, 
                # how likely they share the same presynaptic neuron (or neuron cluster), 
                # postsynaptic neuron (or neuron cluster), or neither, or both for neuron cluster 
                # it is not possible for different modulation sharing the same pre and post neuron,
                # but they may still share the same neuron cluster
                # here we are curious about their collective behavior and calculate the total summation of count
                row_all, col_all = result_all["row_labels"], result_all["col_labels"]

                print(f"col_all: {col_all}")

                # Two exclusion modes:
                #   "mod_only"        — exclude only the silent modulation cluster
                #   "mod_and_endpoint" — also exclude modulations whose pre/post neuron
                #                        belongs to the unresponsive input/hidden cluster
                # When clustering_normalize=True both modes are identical (no unresponsive
                # detection was run), so we only run "mod_only" to avoid redundant work.
                if clustering_normalize:
                    _exclusion_modes = ["mod_only"]
                else:
                    _exclusion_modes = ["mod_only"]

                for _excl_mode in _exclusion_modes:
                    print(f"  --- Exclusion mode: {_excl_mode} ---")

                    # Stage 1: exclude the unresponsive modulation cluster (label = col_k + 1).
                    # flat_idx preserves original positions so post=i//M and pre=i%M stay correct.
                    _unres_mod_label = result_all["col_k"] + 1
                    _active_mod_mask = col_all != _unres_mod_label
                    _col_all_active  = col_all[_active_mod_mask]
                    _flat_idx_active = np.where(_active_mod_mask)[0]
                    _n_drop_unres_mod = int((~_active_mod_mask).sum())
                    _n_drop_endpoint  = 0

                    # Sub-breakdown of unresponsive modulations by endpoint status:
                    # among modulations in the silent mod cluster, how many have
                    # unresponsive pre, post, both, or neither endpoint.
                    _n_unres_mod_pre_only  = 0
                    _n_unres_mod_post_only = 0
                    _n_unres_mod_both      = 0
                    _n_unres_mod_neither   = 0
                    _n_unres_input_neurons = 0
                    _n_unres_hidden_neurons = 0
                    if _n_drop_unres_mod > 0 and not clustering_normalize:
                        _input_col_k_um  = cluster_info_save["input_unnormalized"]["result"]["col_tol_k"]
                        _hidden_col_k_um = cluster_info_save["hidden_unnormalized"]["result"]["col_tol_k"]
                        _unres_input_um  = set(cluster_input.get(_input_col_k_um  + 1, []))
                        _unres_hidden_um = set(cluster_hidden.get(_hidden_col_k_um + 1, []))
                        _n_unres_input_neurons = len(_unres_input_um)
                        _n_unres_hidden_neurons = len(_unres_hidden_um)
                        _flat_unres_mod  = np.where(~_active_mod_mask)[0]
                        # modulation_W shape is (post, pre)
                        _um_post = _flat_unres_mod // M
                        _um_pre  = _flat_unres_mod % M
                        _um_is_pre  = np.isin(_um_pre,  list(_unres_input_um))
                        _um_is_post = np.isin(_um_post, list(_unres_hidden_um))
                        _n_unres_mod_pre_only  = int((_um_is_pre & ~_um_is_post).sum())
                        _n_unres_mod_post_only = int((~_um_is_pre & _um_is_post).sum())
                        _n_unres_mod_both      = int((_um_is_pre & _um_is_post).sum())
                        _n_unres_mod_neither   = int((~_um_is_pre & ~_um_is_post).sum())
                        print(f"    Unres mod endpoint breakdown: "
                              f"pre_only={_n_unres_mod_pre_only}, post_only={_n_unres_mod_post_only}, "
                              f"both={_n_unres_mod_both}, neither={_n_unres_mod_neither}")

                    # Stage 2 (mod_and_endpoint only): additionally exclude modulations
                    # whose pre-neuron (input) or post-neuron (hidden) is in the
                    # unresponsive neuron cluster.
                    _n_drop_pre_only  = 0
                    _n_drop_post_only = 0
                    _n_drop_both      = 0
                    if _excl_mode == "mod_and_endpoint" and not clustering_normalize:
                        _input_col_k  = cluster_info_save["input_unnormalized"]["result"]["col_tol_k"]
                        _hidden_col_k = cluster_info_save["hidden_unnormalized"]["result"]["col_tol_k"]
                        _unres_input_neurons  = set(cluster_input.get(_input_col_k  + 1, []))
                        _unres_hidden_neurons = set(cluster_hidden.get(_hidden_col_k + 1, []))

                        if _unres_input_neurons or _unres_hidden_neurons:
                            # modulation_W shape is (post, pre): previously assumed (pre, post)
                            _post_idx = _flat_idx_active // M
                            _pre_idx  = _flat_idx_active % M
                            _is_unres_pre  = np.isin(_pre_idx,  list(_unres_input_neurons))
                            _is_unres_post = np.isin(_post_idx, list(_unres_hidden_neurons))
                            _endpoint_mask = ~(_is_unres_pre | _is_unres_post)

                            # Sub-breakdown: pre-only, post-only, both silent
                            _n_drop_pre_only  = int((_is_unres_pre & ~_is_unres_post).sum())
                            _n_drop_post_only = int((~_is_unres_pre & _is_unres_post).sum())
                            _n_drop_both      = int((_is_unres_pre & _is_unres_post).sum())

                            _col_all_active  = _col_all_active[_endpoint_mask]
                            _flat_idx_active = _flat_idx_active[_endpoint_mask]
                            _n_drop_endpoint = _n_drop_pre_only + _n_drop_post_only + _n_drop_both
                            print(f"    Endpoint mask: dropped {_n_drop_endpoint} modulations "
                                  f"(pre_only={_n_drop_pre_only}, post_only={_n_drop_post_only}, "
                                  f"both={_n_drop_both}; "
                                  f"{len(_unres_input_neurons)} unresponsive input, "
                                  f"{len(_unres_hidden_neurons)} unresponsive hidden neurons)")

                    print(f"    kept={_col_all_active.size} / {col_all.size} "
                          f"(drop_mod={_n_drop_unres_mod}, drop_endpoint={_n_drop_endpoint})")

                    # Record per-G mask breakdown.
                    mask_breakdown_lst.append({
                        "G": G,
                        "exclusion_mode":    _excl_mode,
                        "n_total":           int(col_all.size),
                        "n_drop_unres_mod":          _n_drop_unres_mod,
                        "n_unres_mod_pre_only":      _n_unres_mod_pre_only,
                        "n_unres_mod_post_only":     _n_unres_mod_post_only,
                        "n_unres_mod_both":          _n_unres_mod_both,
                        "n_unres_mod_neither":       _n_unres_mod_neither,
                        "n_drop_endpoint":           _n_drop_endpoint,
                        "n_drop_pre_only":           _n_drop_pre_only,
                        "n_drop_post_only":          _n_drop_post_only,
                        "n_drop_both":               _n_drop_both,
                        "n_kept":                    int(_col_all_active.size),
                        "n_unres_input_neurons":     _n_unres_input_neurons,
                        "n_unres_hidden_neurons":    _n_unres_hidden_neurons,
                    })

                    # Membership using the actual clustering information
                    same_pre_all, same_post_all, no_same_pre_post_all, same_pre_cluster_all, same_post_cluster_all, \
                        same_pre_post_cluster_all, no_same_pre_post_cluster_all = clustering_metric.count_pairs_with_clusters(
                            _col_all_active, M, cluster_input, cluster_hidden, flat_idx=_flat_idx_active)
                    # (control) membership using random shuffled clustering information
                    same_pre_all_c, same_post_all_c, no_same_pre_post_all_c, same_pre_cluster_all_c, same_post_cluster_all_c, \
                        same_pre_post_cluster_all_c, no_same_pre_post_cluster_all_c = clustering_metric.count_pairs_with_clusters_control(
                            _col_all_active, M, cluster_input, cluster_hidden, repeat=10000, flat_idx=_flat_idx_active)

                    print(f"    same_pre_all: {same_pre_all}; same_post_all: {same_post_all}; no_same_pre_post_all: {no_same_pre_post_all}")
                    print(f"    same_pre_cluster_all: {same_pre_cluster_all}; same_post_cluster_all: {same_post_cluster_all}")
                    print(f"    same_pre_post_cluster_all: {same_pre_post_cluster_all}; no_same_pre_post_cluster_all: {no_same_pre_post_cluster_all}")

                    print(f"    same_pre_all_c: {same_pre_all_c}; same_post_all_c: {same_post_all_c}; no_same_pre_post_all_c: {no_same_pre_post_all_c}")
                    print(f"    same_pre_cluster_all_c: {same_pre_cluster_all_c}; same_post_cluster_all_c: {same_post_cluster_all_c}")
                    print(f"    same_pre_post_cluster_all_c: {same_pre_post_cluster_all_c}; no_same_pre_post_cluster_all_c: {no_same_pre_post_cluster_all_c}")

                    bar_all_lst = [[same_pre_all, same_post_all, no_same_pre_post_all],
                                   [same_pre_cluster_all, same_post_cluster_all,
                                    same_pre_post_cluster_all, no_same_pre_post_cluster_all]]
                    bar_all_ctrl_lst = [[same_pre_all_c, same_post_all_c, no_same_pre_post_all_c],
                                        [same_pre_cluster_all_c, same_post_cluster_all_c,
                                         same_pre_post_cluster_all_c, no_same_pre_post_cluster_all_c]]
                    bar_name_lst = [["Share-Pre", "Share-Post", "Neither"],
                                    ["Share-Pre-Cluster", "Share-Post-Cluster", "Share-Both-Cluster", "Neither"]]

                    N_cluster = np.max(col_all)

                    ppshare_row_names = ["Same Neuron Check", "Same Neuron Cluster Check"]
                    n_pre_clusters  = len(cluster_input)
                    n_post_clusters = len(cluster_hidden)
                    for idx in range(len(bar_all_lst)):
                        bar_all = np.array(bar_all_lst[idx])
                        bar_all_c = np.array(bar_all_ctrl_lst[idx])

                        over_membership = (bar_all - bar_all_c) / bar_all_c

                        axsppshare[idx,G_idx].bar([i for i in range(len(over_membership))], over_membership,
                                                  alpha=0.7 if _excl_mode == "mod_only" else 1.0,
                                                  label=_excl_mode)
                        axsppshare[idx,G_idx].set_xticks([i for i in range(len(over_membership))])
                        axsppshare[idx,G_idx].set_xticklabels(bar_name_lst[idx], rotation=45, ha="right")
                        axsppshare[idx,G_idx].set_ylabel("Over-membership", fontsize=15)
                        if idx == 0:
                            title = f"{ppshare_row_names[idx]} | G={G}; #Cluster={N_cluster}"
                        else:
                            title = (f"{ppshare_row_names[idx]} | G={G}; #Cluster={N_cluster}\n"
                                     f"#PreCluster={n_pre_clusters}; #PostCluster={n_post_clusters}")
                        axsppshare[idx,G_idx].set_title(title)
                        axsppshare[idx,G_idx].legend(fontsize=7)

                    if _excl_mode == "mod_only":
                        prepost_belonging_results.append({
                            "G": G,
                            "N_cluster": int(N_cluster),
                            "n_pre_clusters": n_pre_clusters,
                            "n_post_clusters": n_post_clusters,
                            "bar_all_lst": [np.array(b) for b in bar_all_lst],
                            "bar_all_ctrl_lst": [np.array(b) for b in bar_all_ctrl_lst],
                            "bar_name_lst": bar_name_lst,
                        })

                # ── Fixed-k prepost_belonging (k=FIXED_K_OM for input, hidden, modulation) ──
                from scipy.cluster.hierarchy import fcluster as _fcluster_ppb

                def _fk_col_clusters_ppb(ci_entry, fk):
                    """Cut dendrogram at fk, return {1-based label: neuron_indices}."""
                    res = ci_entry["result"]
                    lnk = res["col_linkage"]
                    tol_k = res["col_tol_k"]
                    tol_labels = res["col_tol_labels"]
                    unres_mask = tol_labels == (tol_k + 1)
                    active_labels = _fcluster_ppb(lnk, fk, criterion="maxclust")
                    full = np.zeros(len(tol_labels), dtype=int)
                    full[~unres_mask] = active_labels
                    if unres_mask.any():
                        full[unres_mask] = fk + 1
                    return {int(lab): np.where(full == lab)[0] for lab in np.unique(full) if lab > 0}

                _ppb_input_key = "input_normalized" if clustering_normalize else "input_unnormalized"
                _ppb_hidden_key = "hidden_normalized" if clustering_normalize else "hidden_unnormalized"
                _ppb_fk_input = _fk_col_clusters_ppb(cluster_info_save[_ppb_input_key], FIXED_K_OM)
                _ppb_fk_hidden = _fk_col_clusters_ppb(cluster_info_save[_ppb_hidden_key], FIXED_K_OM)
                _ppb_fk_input = dict(sorted(_ppb_fk_input.items()))
                _ppb_fk_hidden = dict(sorted(_ppb_fk_hidden.items()))

                # Fixed-k modulation cluster labels
                if "col_labels_by_k" in result_all and FIXED_K_OM in result_all["col_labels_by_k"]:
                    _ppb_fk_mod_labels = result_all["col_labels_by_k"][FIXED_K_OM]
                elif "col_labels_by_k" in result_all and result_all["col_labels_by_k"]:
                    _ppb_fk_avail = sorted(result_all["col_labels_by_k"].keys())
                    _ppb_fk_mod_labels = result_all["col_labels_by_k"][_ppb_fk_avail[0]]
                    print(f"    [fixed-k ppb] mod k={FIXED_K_OM} not available, using k={_ppb_fk_avail[0]}")
                else:
                    _ppb_fk_mod_labels = None

                if _ppb_fk_mod_labels is not None:
                    # Exclude unresponsive modulation cluster (mod_only mode)
                    _ppb_unres_label = result_all["col_k"] + 1
                    _ppb_active_mask = col_all != _ppb_unres_label
                    _ppb_fk_col_active = _ppb_fk_mod_labels[_ppb_active_mask]
                    _ppb_fk_flat_idx = np.where(_ppb_active_mask)[0]

                    _ppb_fk_same_pre, _ppb_fk_same_post, _ppb_fk_no_same, \
                        _ppb_fk_same_preC, _ppb_fk_same_postC, \
                        _ppb_fk_same_bothC, _ppb_fk_no_bothC = \
                        clustering_metric.count_pairs_with_clusters(
                            _ppb_fk_col_active, M, _ppb_fk_input, _ppb_fk_hidden,
                            flat_idx=_ppb_fk_flat_idx)

                    _ppb_fk_same_pre_c, _ppb_fk_same_post_c, _ppb_fk_no_same_c, \
                        _ppb_fk_same_preC_c, _ppb_fk_same_postC_c, \
                        _ppb_fk_same_bothC_c, _ppb_fk_no_bothC_c = \
                        clustering_metric.count_pairs_with_clusters_control(
                            _ppb_fk_col_active, M, _ppb_fk_input, _ppb_fk_hidden,
                            repeat=10000, flat_idx=_ppb_fk_flat_idx)

                    _ppb_fk_bar_all = [
                        [_ppb_fk_same_pre, _ppb_fk_same_post, _ppb_fk_no_same],
                        [_ppb_fk_same_preC, _ppb_fk_same_postC,
                         _ppb_fk_same_bothC, _ppb_fk_no_bothC],
                    ]
                    _ppb_fk_bar_ctrl = [
                        [_ppb_fk_same_pre_c, _ppb_fk_same_post_c, _ppb_fk_no_same_c],
                        [_ppb_fk_same_preC_c, _ppb_fk_same_postC_c,
                         _ppb_fk_same_bothC_c, _ppb_fk_no_bothC_c],
                    ]
                    _ppb_fk_bar_names = [
                        ["Share-Pre", "Share-Post", "Neither"],
                        ["Share-Pre-Cluster", "Share-Post-Cluster",
                         "Share-Both-Cluster", "Neither"],
                    ]

                    _ppb_fk_N_cluster = int(np.max(_ppb_fk_col_active))
                    prepost_belonging_results.append({
                        "G": G,
                        "fixed_k": FIXED_K_OM,
                        "N_cluster": _ppb_fk_N_cluster,
                        "n_pre_clusters": len(_ppb_fk_input),
                        "n_post_clusters": len(_ppb_fk_hidden),
                        "bar_all_lst": [np.array(b) for b in _ppb_fk_bar_all],
                        "bar_all_ctrl_lst": [np.array(b) for b in _ppb_fk_bar_ctrl],
                        "bar_name_lst": _ppb_fk_bar_names,
                    })

                    print(f"    [fixed-k ppb] G={G}, mod_clusters={_ppb_fk_N_cluster}, "
                          f"input_clusters={len(_ppb_fk_input)}, hidden_clusters={len(_ppb_fk_hidden)}")

                # 2025-10-21: next analyze for each individual modulation cluster,
                # the belonging to different individual
                # pre and post cluster; only test in the minimal modulation cluster selection case
                # 2025-11-17: revise to plot for G=300 case
                if G_idx == 1:
                    print(f"Plot for G={G_lst[G_idx]} Case")
                    # Use active (responsive) modulations only, consistent with count_pairs_with_clusters.
                    # _col_all_active / _flat_idx_active are all-inclusive no-ops when
                    # clustering_normalize=True (unresponsive detection is skipped).
                    all_choice_order_dict = helper.value_counts_desc(_col_all_active)
                    all_choice_order = list(all_choice_order_dict.keys())

                    # cluster sizes come directly from value_counts_desc (no need to re-scan col_all)
                    cluster_size_all = np.array([all_choice_order_dict[k] for k in all_choice_order])
                    cluster_size_percent = cluster_size_all / cluster_size_all.sum()
                    assert np.isclose(cluster_size_percent.sum(), 1.0)

                    # all_num is the outer product of input and hidden cluster sizes
                    in_num = np.array([len(cluster_input[k]) for k in cluster_input])
                    hid_num = np.array([len(cluster_hidden[k]) for k in cluster_hidden])
                    all_num = np.outer(in_num, hid_num).astype(float)
                    flat_all_num = all_num.flatten()
                    n_in, n_hid = len(cluster_input), len(cluster_hidden)

                    # Precompute neuron -> cluster index (0-based) lookup arrays once.
                    # Keys in cluster_input/cluster_hidden are 1-indexed, so subtract 1.
                    max_pre = max(max(v) for v in cluster_input.values())
                    max_post = max(max(v) for v in cluster_hidden.values())
                    pre_lookup = np.empty(max_pre + 1, dtype=int)
                    for key, neurons in cluster_input.items():
                        pre_lookup[np.asarray(neurons)] = key - 1
                    post_lookup = np.empty(max_post + 1, dtype=int)
                    for key, neurons in cluster_hidden.items():
                        post_lookup[np.asarray(neurons)] = key - 1

                    # Precompute cluster assignment using original flat positions (_flat_idx_active).
                    # modulation_W shape is (post, pre): previously assumed (pre, post)
                    # so flat_idx // M = post (hidden), flat_idx % M = pre (input).
                    pre_clusters  = pre_lookup[_flat_idx_active % M]    # shape (N_active,), 0-based
                    post_clusters = post_lookup[_flat_idx_active // M]  # shape (N_active,), 0-based

                    over_membership_lst = []
                    z_count_lst = []

                    for cidx, cluster_num in enumerate(all_choice_order):
                        idx = np.where(_col_all_active == cluster_num)[0]

                        flat_idx = pre_clusters[idx] * n_hid + post_clusters[idx]
                        Z_count = np.bincount(flat_idx, minlength=n_in * n_hid).reshape(n_in, n_hid).astype(float)

                        all_num_avg = all_num * cluster_size_percent[cidx]
                        over_membership = Z_count / all_num_avg
                        over_membership_lst.append(over_membership)
                        z_count_lst.append(Z_count)

                    # Directional asymmetry analysis:
                    # Each modulation cluster groups synapses (pre, post) with similar Hebbian dynamics.
                    # The question is whether co-clustered synapses share a common pre-synaptic (input)
                    # neuron cluster, a common post-synaptic (hidden) neuron cluster, or neither.
                    # Over-membership (OM) = observed count / expected count under random assignment.
                    # OM > 1 means the modulation cluster is enriched for that (input or hidden) cluster.
                    # The log2 asymmetry ratio log2(peak_input_OM / peak_hidden_OM) > 0 means
                    # a modulation cluster is better explained by shared input identity than shared
                    # hidden identity, suggesting pre-synaptic organisation of plasticity.
                    om_stack = np.stack(over_membership_lst)           # (N_cls, n_in, n_hid)
                    N_cls = len(over_membership_lst)

                    global_assignment_cache = {
                        "om_stack": om_stack,
                        "all_choice_order": all_choice_order,
                        "cluster_size_percent": cluster_size_percent,
                        "n_in": n_in,
                        "n_hid": n_hid,
                    }

                    # ── Fixed-k overmembership (for leison_plot.py) ──
                    # Recompute overmembership using FIXED_K_OM for input, hidden,
                    # and modulation clusters so it aligns with the fixed-k lesion.
                    from scipy.cluster.hierarchy import fcluster as _fcluster_om

                    def _fk_col_clusters_om(ci_entry, fk):
                        """Cut dendrogram at fk, return {1-based label: neuron_indices}."""
                        res = ci_entry["result"]
                        lnk = res["col_linkage"]
                        tol_k = res["col_tol_k"]
                        tol_labels = res["col_tol_labels"]
                        unres_mask = tol_labels == (tol_k + 1)
                        active_labels = _fcluster_om(lnk, fk, criterion="maxclust")
                        full = np.zeros(len(tol_labels), dtype=int)
                        full[~unres_mask] = active_labels
                        if unres_mask.any():
                            full[unres_mask] = fk + 1
                        return {int(lab): np.where(full == lab)[0] for lab in np.unique(full) if lab > 0}

                    _input_base_om = "input_normalized" if clustering_normalize else "input_unnormalized"
                    _hidden_base_om = "hidden_normalized" if clustering_normalize else "hidden_unnormalized"
                    _fk_input = _fk_col_clusters_om(cluster_info_save[_input_base_om], FIXED_K_OM)
                    _fk_hidden = _fk_col_clusters_om(cluster_info_save[_hidden_base_om], FIXED_K_OM)
                    _fk_n_in = len(_fk_input)
                    _fk_n_hid = len(_fk_hidden)

                    # Build neuron -> fixed-k cluster lookup (0-based)
                    _fk_max_pre = max(max(v) for v in _fk_input.values())
                    _fk_max_post = max(max(v) for v in _fk_hidden.values())
                    _fk_pre_lookup = np.empty(_fk_max_pre + 1, dtype=int)
                    for key, neurons in _fk_input.items():
                        _fk_pre_lookup[np.asarray(neurons)] = key - 1
                    _fk_post_lookup = np.empty(_fk_max_post + 1, dtype=int)
                    for key, neurons in _fk_hidden.items():
                        _fk_post_lookup[np.asarray(neurons)] = key - 1

                    _fk_pre_clusters = _fk_pre_lookup[_flat_idx_active % M]
                    _fk_post_clusters = _fk_post_lookup[_flat_idx_active // M]

                    # Fixed-k modulation cluster labels
                    if "col_labels_by_k" in result_all and FIXED_K_OM in result_all["col_labels_by_k"]:
                        _fk_mod_labels_full = result_all["col_labels_by_k"][FIXED_K_OM]
                    elif "col_labels_by_k" in result_all and result_all["col_labels_by_k"]:
                        _fk_avail = sorted(result_all["col_labels_by_k"].keys())
                        _fk_mod_labels_full = result_all["col_labels_by_k"][_fk_avail[0]]
                        print(f"    [fixed-k OM] mod k={FIXED_K_OM} not available, using k={_fk_avail[0]}")
                    else:
                        _fk_mod_labels_full = None

                    if _fk_mod_labels_full is not None:
                        # Get mod labels for active synapses only
                        _fk_mod_active = _fk_mod_labels_full[_flat_idx_active]
                        _fk_mod_order_dict = helper.value_counts_desc(_fk_mod_active)
                        _fk_mod_order = list(_fk_mod_order_dict.keys())
                        _fk_mod_sizes = np.array([_fk_mod_order_dict[k] for k in _fk_mod_order])
                        _fk_mod_pct = _fk_mod_sizes / _fk_mod_sizes.sum()

                        _fk_in_num = np.array([len(_fk_input[k]) for k in sorted(_fk_input.keys())])
                        _fk_hid_num = np.array([len(_fk_hidden[k]) for k in sorted(_fk_hidden.keys())])
                        _fk_all_num = np.outer(_fk_in_num, _fk_hid_num).astype(float)

                        _fk_om_lst = []
                        for cidx, cnum in enumerate(_fk_mod_order):
                            idx = np.where(_fk_mod_active == cnum)[0]
                            flat_idx_fk = _fk_pre_clusters[idx] * _fk_n_hid + _fk_post_clusters[idx]
                            Z_count_fk = np.bincount(flat_idx_fk, minlength=_fk_n_in * _fk_n_hid).reshape(_fk_n_in, _fk_n_hid).astype(float)
                            all_num_avg_fk = _fk_all_num * _fk_mod_pct[cidx]
                            _fk_om_lst.append(Z_count_fk / all_num_avg_fk)

                        global_assignment_fixed_k_cache = {
                            "om_stack": np.stack(_fk_om_lst),
                            "all_choice_order": _fk_mod_order,
                            "cluster_size_percent": _fk_mod_pct,
                            "n_in": _fk_n_in,
                            "n_hid": _fk_n_hid,
                            "fixed_k": FIXED_K_OM,
                        }
                        print(f"    [fixed-k OM] computed: {len(_fk_mod_order)} mod clusters × "
                              f"{_fk_n_in} input × {_fk_n_hid} hidden (k={FIXED_K_OM})")

                    # Mean |OM| collapsed over the other dimension
                    over_membership_array_input  = np.mean(np.abs(om_stack), axis=2)  # (N_cls, n_in)
                    over_membership_array_hidden = np.mean(np.abs(om_stack), axis=1)  # (N_cls, n_hid)

                    # Per-mod-cluster peak alignment strength in each dimension
                    max_input_om  = over_membership_array_input.max(axis=1)   # (N_cls,)
                    max_hidden_om = over_membership_array_hidden.max(axis=1)  # (N_cls,)

                    # Log2 asymmetry ratio: >0 = input-dominant, <0 = hidden-dominant
                    log_asym = np.log2(
                        np.clip(max_input_om, 1e-6, None) /
                        np.clip(max_hidden_om, 1e-6, None)
                    )
                    input_dominant_frac = float(np.mean(log_asym > 0))
                    mean_log_asym = float(np.mean(log_asym))

                    figomcluster, axsomcluster = plt.subplots(2, 2, figsize=(14, 10))
                    x_cls = np.arange(N_cls)

                    # Panel [0,0]: scatter — max input OM vs max hidden OM, sized by cluster size
                    ax_sc = axsomcluster[0, 0]
                    sizes = cluster_size_percent / cluster_size_percent.max() * 200 + 20
                    sc = ax_sc.scatter(
                        max_input_om, max_hidden_om,
                        s=sizes, c=x_cls, cmap="viridis",
                        alpha=0.8, edgecolors="k", linewidths=0.5,
                    )
                    lim = max(float(max_input_om.max()), float(max_hidden_om.max())) * 1.15
                    ax_sc.plot([0, lim], [0, lim], "--", color="gray", lw=1.2, label="y = x")
                    ax_sc.set_xlim(0, lim); ax_sc.set_ylim(0, lim)
                    ax_sc.set_xlabel("Peak OM — input cluster", fontsize=12)
                    ax_sc.set_ylabel("Peak OM — hidden cluster", fontsize=12)
                    ax_sc.set_title(
                        f"Per-cluster peak alignment\n"
                        f"{input_dominant_frac:.0%} input-dominant  |  mean log₂ ratio = {mean_log_asym:+.2f}",
                        fontsize=10,
                    )
                    ax_sc.legend(fontsize=8)
                    plt.colorbar(sc, ax=ax_sc, label="Mod cluster order (0=largest)")

                    # Panel [0,1]: log asymmetry ratio bar chart, sorted by cluster size
                    ax_bar = axsomcluster[0, 1]
                    bar_colors = ["#d73027" if v > 0 else "#4575b4" for v in log_asym]
                    ax_bar.bar(x_cls, log_asym, color=bar_colors, edgecolor="k", linewidth=0.4)
                    ax_bar.axhline(0, color="black", lw=1)
                    ax_bar.set_xlabel("Modulation cluster (largest → smallest)", fontsize=12)
                    ax_bar.set_ylabel("log₂(peak input OM / peak hidden OM)", fontsize=11)
                    ax_bar.set_title(
                        "Directional asymmetry per mod cluster\n"
                        "red = input-dominant   blue = hidden-dominant",
                        fontsize=10,
                    )

                    # Panel [1,0]: input-cluster OM profiles across mod clusters
                    ax_in = axsomcluster[1, 0]
                    in_colors = color_func.rainbow_generate(n_in)
                    for ci in range(n_in):
                        ax_in.plot(
                            x_cls, over_membership_array_input[:, ci],
                            color=in_colors[ci], alpha=0.8, lw=1.5, label=f"In C{ci}",
                        )
                    ax_in.axhline(1.0, color="gray", lw=1, linestyle="--", label="Expected = 1")
                    ax_in.set_xlabel("Modulation cluster (largest → smallest)", fontsize=12)
                    ax_in.set_ylabel("Mean |OM| — input cluster", fontsize=12)
                    ax_in.set_title("Input-cluster alignment profiles", fontsize=10)
                    ax_in.legend(fontsize=8, loc="upper right")

                    # Panel [1,1]: hidden-cluster OM profiles across mod clusters
                    ax_hid = axsomcluster[1, 1]
                    hid_colors = color_func.rainbow_generate(n_hid)
                    for ci in range(n_hid):
                        ax_hid.plot(
                            x_cls, over_membership_array_hidden[:, ci],
                            color=hid_colors[ci], alpha=0.8, lw=1.5, label=f"Hid C{ci}",
                        )
                    ax_hid.axhline(1.0, color="gray", lw=1, linestyle="--", label="Expected = 1")
                    ax_hid.set_xlabel("Modulation cluster (largest → smallest)", fontsize=12)
                    ax_hid.set_ylabel("Mean |OM| — hidden cluster", fontsize=12)
                    ax_hid.set_title("Hidden-cluster alignment profiles", fontsize=10)
                    ax_hid.legend(fontsize=8, loc="upper right")

                    figomcluster.suptitle(
                        f"Modulation cluster ↔ Input / Hidden cluster alignment  "
                        f"| {n_in} input clusters × {n_hid} hidden clusters × {N_cls} mod clusters",
                        fontsize=11,
                    )
                    figomcluster.tight_layout()
                    figomcluster.savefig(f"{save_dir}/{clustering_name}_om_across_cluster_{savefigure_name}.png", dpi=300)
                    plt.close(figomcluster)

                    # Global assignment heatmap: all modulation clusters × flattened
                    # (input-cluster, hidden-cluster) pairs.  Each row is one modulation
                    # cluster (largest → smallest); each column is one neuron-cluster pair
                    # grouped by input cluster.  Color = over-membership (1 = chance level).
                    om_global = om_stack.reshape(N_cls, n_in * n_hid)  # (N_cls, n_in*n_hid)

                    y_labels = [
                        f"MC{cid} ({pct:.1%})"
                        for cid, pct in zip(all_choice_order, cluster_size_percent)
                    ]
                    x_labels = [
                        f"({i},{j})" for i in range(n_in) for j in range(n_hid)
                    ]

                    vmin_gah = 0.0
                    vmax_gah = 2.0

                    # Second ordering: same hidden cluster adjacent — (0,0),(1,0),...,(n_in-1,0),(0,1),...
                    col_order_hid = [i * n_hid + j for j in range(n_hid) for i in range(n_in)]
                    om_global_hid = om_global[:, col_order_hid]
                    x_labels_hid = [f"({i},{j})" for j in range(n_hid) for i in range(n_in)]

                    _gah_xfs = max(5, min(9, int(110 / max(n_in * n_hid, 1))))
                    _gah_yfs = max(5, min(9, int(110 / max(N_cls, 1))))

                    fig_gah, (ax_gah, ax_gah2) = plt.subplots(2, 1, figsize=(14, 14))

                    # --- Top subplot: grouped by input cluster ---
                    sns.heatmap(
                        om_global,
                        ax=ax_gah,
                        cmap="RdBu_r",
                        vmin=vmin_gah, vmax=vmax_gah,
                        center=1.0,
                        xticklabels=x_labels,
                        yticklabels=y_labels,
                        linewidths=0.3,
                        linecolor="lightgray",
                        cbar_kws={"label": "Over-membership (1 = expected)"},
                    )
                    for _i in range(1, n_in):
                        ax_gah.axvline(_i * n_hid, color="black", lw=1.5)
                    ax_gah.set_xticklabels(x_labels, rotation=45, ha="right", fontsize=_gah_xfs)
                    ax_gah.set_yticklabels(y_labels, fontsize=_gah_yfs)
                    ax_gah.set_xlabel("(Input cluster, Hidden cluster) pair  [grouped by input cluster]", fontsize=12)
                    ax_gah.set_ylabel("Modulation cluster (largest → smallest)", fontsize=12)
                    ax_gah.set_title(
                        f"Global assignment heatmap — grouped by input cluster\n"
                        f"{N_cls} mod clusters  ×  {n_in} input clusters  ×  {n_hid} hidden clusters",
                        fontsize=11,
                    )

                    # --- Bottom subplot: grouped by hidden cluster ---
                    sns.heatmap(
                        om_global_hid,
                        ax=ax_gah2,
                        cmap="RdBu_r",
                        vmin=vmin_gah, vmax=vmax_gah,
                        center=1.0,
                        xticklabels=x_labels_hid,
                        yticklabels=y_labels,
                        linewidths=0.3,
                        linecolor="lightgray",
                        cbar_kws={"label": "Over-membership (1 = expected)"},
                    )
                    for _j in range(1, n_hid):
                        ax_gah2.axvline(_j * n_in, color="black", lw=1.5)
                    ax_gah2.set_xticklabels(x_labels_hid, rotation=45, ha="right", fontsize=_gah_xfs)
                    ax_gah2.set_yticklabels(y_labels, fontsize=_gah_yfs)
                    ax_gah2.set_xlabel("(Input cluster, Hidden cluster) pair  [grouped by hidden cluster]", fontsize=12)
                    ax_gah2.set_ylabel("Modulation cluster (largest → smallest)", fontsize=12)
                    ax_gah2.set_title(
                        f"Global assignment heatmap — grouped by hidden cluster\n"
                        f"{N_cls} mod clusters  ×  {n_in} input clusters  ×  {n_hid} hidden clusters",
                        fontsize=11,
                    )

                    fig_gah.tight_layout()
                    fig_gah.savefig(
                        f"{save_dir}/{clustering_name}_global_assignment_heatmap_{savefigure_name}.png",
                        dpi=300,
                    )
                    plt.close(fig_gah)

                    # -------------------------------------------------------
                    # Analysis 2: Block purity / entropy
                    #
                    # Direction A — mod-cluster entropy over neuron-cluster blocks:
                    #   For each modulation cluster c, treat the distribution of its
                    #   synapses over (input-cluster, hidden-cluster) pairs as a
                    #   probability vector and compute its Shannon entropy.
                    #   Low entropy  → cluster is concentrated in one block (specialised).
                    #   High entropy → cluster spans many blocks (diffuse).
                    #
                    # Direction B — block entropy over mod clusters:
                    #   For each (input-cluster, hidden-cluster) pair, treat the
                    #   distribution of its synapses across modulation clusters as a
                    #   probability vector.
                    #   Low entropy  → block is "owned" by one mod cluster (clean partition).
                    #   High entropy → block is shared across many mod clusters.
                    # -------------------------------------------------------
                    z_count_stack = np.stack(z_count_lst)  # (N_cls, n_in, n_hid)

                    # Direction A: entropy of each mod cluster over (n_in*n_hid) blocks
                    z_flat = z_count_stack.reshape(N_cls, -1).astype(float)     # (N_cls, n_in*n_hid)
                    row_sum = z_flat.sum(axis=1, keepdims=True)
                    p_mod = z_flat / np.maximum(row_sum, 1e-12)                 # (N_cls, n_in*n_hid)
                    with np.errstate(divide="ignore", invalid="ignore"):
                        h_mod = -np.sum(np.where(p_mod > 0, p_mod * np.log2(p_mod), 0.0), axis=1)
                    h_mod_norm = h_mod / np.log2(n_in * n_hid)                  # (N_cls,) in [0,1]

                    # Direction B: entropy of each (i,j) block over N_cls mod clusters
                    col_sum = z_count_stack.sum(axis=0, keepdims=True)          # (1, n_in, n_hid)
                    p_block = z_count_stack / np.maximum(col_sum, 1e-12)        # (N_cls, n_in, n_hid)
                    with np.errstate(divide="ignore", invalid="ignore"):
                        h_block = -np.sum(np.where(p_block > 0, p_block * np.log2(p_block), 0.0), axis=0)
                    h_block_norm = h_block / np.log2(max(N_cls, 2))             # (n_in, n_hid) in [0,1]

                    fig_ent, axs_ent = plt.subplots(2, 3, figsize=(16, 9))

                    # [0,0] Bar chart: mod-cluster normalised entropy, sorted by cluster size
                    ax = axs_ent[0, 0]
                    bar_colors = plt.cm.RdYlGn_r(h_mod_norm)
                    ax.bar(np.arange(N_cls), h_mod_norm, color=bar_colors, edgecolor="k", linewidth=0.4)
                    ax.axhline(1.0, color="gray", lw=1, linestyle="--", label="Max entropy (uniform)")
                    ax.set_xticks(np.arange(N_cls))
                    ax.set_xticklabels(
                        [f"MC{cid}" for cid in all_choice_order],
                        rotation=45, ha="right", fontsize=7,
                    )
                    ax.set_ylim(0, 1.1)
                    ax.set_xlabel("Modulation cluster (largest → smallest)", fontsize=11)
                    ax.set_ylabel("Normalised entropy", fontsize=11)
                    ax.set_title(
                        "Mod-cluster entropy over neuron-cluster blocks\n"
                        "0 = one block owns it all,  1 = uniform across all blocks",
                        fontsize=10,
                    )
                    ax.legend(fontsize=8)

                    # [0,1] Scatter: cluster size vs mod-cluster entropy
                    ax = axs_ent[0, 1]
                    sc = ax.scatter(
                        cluster_size_percent * 100, h_mod_norm,
                        c=h_mod_norm, cmap="RdYlGn_r", vmin=0, vmax=1,
                        s=60, edgecolors="k", linewidths=0.5, alpha=0.85,
                    )
                    plt.colorbar(sc, ax=ax, label="Normalised entropy")
                    ax.set_xlabel("Modulation cluster size (%)", fontsize=11)
                    ax.set_ylabel("Normalised entropy", fontsize=11)
                    ax.set_title("Cluster size vs specialisation", fontsize=10)

                    # [0,2] Histogram of mod-cluster normalised entropy
                    ax = axs_ent[0, 2]
                    ax.hist(h_mod_norm, bins="auto", color="#4575b4", edgecolor="k", alpha=0.8)
                    ax.axvline(np.mean(h_mod_norm), color="red", lw=1.5,
                               label=f"Mean = {np.mean(h_mod_norm):.2f}")
                    ax.set_xlabel("Normalised entropy", fontsize=11)
                    ax.set_ylabel("Count", fontsize=11)
                    ax.set_title("Distribution of mod-cluster entropy", fontsize=10)
                    ax.legend(fontsize=8)

                    # [1,0] Heatmap: block entropy over mod clusters (n_in × n_hid)
                    ax = axs_ent[1, 0]
                    sns.heatmap(
                        h_block_norm,
                        ax=ax,
                        cmap="RdYlGn_r",
                        vmin=0, vmax=1,
                        annot=False,
                        cbar_kws={"label": "Normalised entropy"},
                        xticklabels=[f"Hid C{j}" for j in range(n_hid)],
                        yticklabels=[f"In C{i}" for i in range(n_in)],
                    )
                    ax.set_xlabel("Hidden cluster", fontsize=11)
                    ax.set_ylabel("Input cluster", fontsize=11)
                    ax.set_title(
                        "Block entropy over mod clusters\n"
                        "0 = owned by one mod cluster,  1 = shared equally",
                        fontsize=10,
                    )

                    # [1,1] Histogram of block normalised entropy
                    ax = axs_ent[1, 1]
                    ax.hist(h_block_norm.flatten(), bins="auto", color="#d73027", edgecolor="k", alpha=0.8)
                    ax.axvline(np.mean(h_block_norm), color="navy", lw=1.5,
                               label=f"Mean = {np.mean(h_block_norm):.2f}")
                    ax.set_xlabel("Normalised entropy", fontsize=11)
                    ax.set_ylabel("Count", fontsize=11)
                    ax.set_title("Distribution of block entropy", fontsize=10)
                    ax.legend(fontsize=8)

                    # [1,2] Scatter: mod-cluster entropy vs block entropy of the peak block
                    # (the (i,j) block where that mod cluster has the highest count)
                    peak_block_flat = z_flat.argmax(axis=1)                     # (N_cls,)
                    peak_block_entropy = h_block_norm.flatten()[peak_block_flat] # (N_cls,)
                    ax = axs_ent[1, 2]
                    sc2 = ax.scatter(
                        h_mod_norm, peak_block_entropy,
                        c=cluster_size_percent * 100, cmap="viridis",
                        s=60, edgecolors="k", linewidths=0.5, alpha=0.85,
                    )
                    plt.colorbar(sc2, ax=ax, label="Cluster size (%)")
                    ax.set_xlabel("Mod-cluster entropy (specialisation)", fontsize=11)
                    ax.set_ylabel("Entropy of its peak block", fontsize=11)
                    ax.set_title(
                        "Specialised clusters → do they own exclusive blocks?\n"
                        "Bottom-left = both specific",
                        fontsize=10,
                    )

                    fig_ent.suptitle(
                        f"Block purity / entropy analysis  |  "
                        f"{N_cls} mod clusters  ×  {n_in} input clusters  ×  {n_hid} hidden clusters",
                        fontsize=12,
                    )
                    fig_ent.tight_layout()
                    fig_ent.savefig(
                        f"{save_dir}/{clustering_name}_block_entropy_{savefigure_name}.png",
                        dpi=300,
                    )
                    plt.close(fig_ent)

            figcol.tight_layout()
            figcol.savefig(f"{save_dir}/{clustering_name}_allneuron_grouplength_{savefigure_name}.png", dpi=300)  
            plt.close(figcol)
            figcolcluster.tight_layout()
            figcolcluster.savefig(f"{save_dir}/{clustering_name}_actualcluster_length_{savefigure_name}.png", dpi=300)
            plt.close(figcolcluster)
            figppshare.tight_layout()
            figppshare.savefig(f"{save_dir}/{clustering_name}_prepost_belonging_{savefigure_name}.png", dpi=300)
            plt.close(figppshare)

            with open(f"{save_dir}/{clustering_name}_prepost_belonging_{savefigure_name}.pkl", "wb") as f:
                pickle.dump({
                    "prepost_belonging_results": prepost_belonging_results,
                    "G_lst": G_lst,
                    "clustering_name": clustering_name,
                    "clustering_normalize": clustering_normalize,
                    "n_input_clusters": len(cluster_input),
                    "n_hidden_clusters": len(cluster_hidden),
                    "FIXED_K": FIXED_K_OM,
                }, f)

            # Fixed-k prepost_belonging figure
            _fk_entries = [r for r in prepost_belonging_results if "fixed_k" in r]
            if _fk_entries:
                figppfk, axsppfk = plt.subplots(2, len(_fk_entries), figsize=(4 * len(_fk_entries), 4 * 2))
                if len(_fk_entries) == 1:
                    axsppfk = axsppfk[:, np.newaxis]
                ppfk_row_names = ["Same Neuron Check", "Same Neuron Cluster Check"]
                for gi, entry in enumerate(_fk_entries):
                    for idx in range(len(entry["bar_all_lst"])):
                        bar_all = np.array(entry["bar_all_lst"][idx])
                        bar_all_c = np.array(entry["bar_all_ctrl_lst"][idx])
                        over_membership = (bar_all - bar_all_c) / bar_all_c
                        axsppfk[idx, gi].bar(range(len(over_membership)), over_membership, alpha=0.7)
                        axsppfk[idx, gi].set_xticks(range(len(over_membership)))
                        axsppfk[idx, gi].set_xticklabels(entry["bar_name_lst"][idx], rotation=45, ha="right")
                        axsppfk[idx, gi].set_ylabel("Over-membership", fontsize=12)
                        if idx == 0:
                            title = (f"{ppfk_row_names[idx]} | G={entry['G']}; "
                                     f"#ModCluster={entry['N_cluster']} (fixed k={entry['fixed_k']})")
                        else:
                            title = (f"{ppfk_row_names[idx]} | G={entry['G']}; "
                                     f"#ModCluster={entry['N_cluster']} (fixed k={entry['fixed_k']})\n"
                                     f"#PreCluster={entry['n_pre_clusters']}; "
                                     f"#PostCluster={entry['n_post_clusters']}")
                        axsppfk[idx, gi].set_title(title, fontsize=9)
                figppfk.suptitle(
                    f"Fixed-k={FIXED_K_OM} prepost belonging "
                    f"({'normalized' if clustering_normalize else 'unnormalized'})",
                    fontsize=12,
                )
                figppfk.tight_layout()
                figppfk.savefig(
                    f"{save_dir}/{clustering_name}_prepost_belonging_fixedk{FIXED_K_OM}_{savefigure_name}.png",
                    dpi=300,
                )
                plt.close(figppfk)

            # Modulation masking breakdown: 5 stacked categories per bar.
            # 1) Kept (active), 2) Drop: unresponsive mod cluster,
            # 3) Drop: pre (input) silent only, 4) Drop: post (hidden) silent only,
            # 5) Drop: both pre & post silent.
            # For clustering_normalize=True or mod_only mode the endpoint sub-counts
            # are all zero, so only the first two segments are visible.
            figmask, axmask = plt.subplots(1, 1, figsize=(max(4, 1.4 * len(mask_breakdown_lst) + 2), 5))
            _x = np.arange(len(mask_breakdown_lst))
            _m1       = np.array([d["n_drop_unres_mod"] for d in mask_breakdown_lst], dtype=float)
            _ep_pre   = np.array([d["n_drop_pre_only"]  for d in mask_breakdown_lst], dtype=float)
            _ep_post  = np.array([d["n_drop_post_only"] for d in mask_breakdown_lst], dtype=float)
            _ep_both  = np.array([d["n_drop_both"]      for d in mask_breakdown_lst], dtype=float)
            _kept     = np.array([d["n_kept"]           for d in mask_breakdown_lst], dtype=float)
            _totals   = np.array([d["n_total"]          for d in mask_breakdown_lst], dtype=float)
            _bar_w = 0.6
            _seg_specs = [
                (_kept,    "Kept (active)",              c_vals[2], np.zeros_like(_kept)),
                (_m1,      "Drop: unresponsive mod",     c_vals[0], _kept),
                (_ep_pre,  "Drop: pre (input) silent",   c_vals[1], _kept + _m1),
                (_ep_post, "Drop: post (hidden) silent", c_vals[3], _kept + _m1 + _ep_pre),
                (_ep_both, "Drop: both pre & post silent", c_vals[4], _kept + _m1 + _ep_pre + _ep_post),
            ]
            for _vals, _lbl, _col, _bot in _seg_specs:
                axmask.bar(_x, _vals, width=_bar_w, bottom=_bot, color=_col,
                           edgecolor="black", linewidth=0.5, label=_lbl)
                for _xi, (_v, _b, _tot) in enumerate(zip(_vals, _bot, _totals)):
                    if _v <= 0 or _tot <= 0:
                        continue
                    axmask.text(_xi, _b + _v / 2.0,
                                f"{int(_v)}\n({100.0 * _v / _tot:.1f}%)",
                                ha="center", va="center", fontsize=8,
                                color="white" if _lbl.startswith("Kept") else "black")
            axmask.set_xticks(_x)
            axmask.set_xticklabels(
                [f"G={d['G']}\n{d['exclusion_mode']}\n(MM={d['n_total']})" for d in mask_breakdown_lst]
            )
            axmask.set_ylabel("Number of modulations", fontsize=13)
            axmask.set_title(
                f"Modulation mask breakdown ({'normalized' if clustering_normalize else 'unnormalized'})",
                fontsize=13,
            )
            axmask.legend(loc="upper center", bbox_to_anchor=(0.5, -0.12), ncol=2, frameon=False, fontsize=8)
            figmask.savefig(f"{save_dir}/{clustering_name}_mask_breakdown_{savefigure_name}.png",
                            dpi=300, bbox_inches="tight")
            plt.close(figmask)

            # Companion figure: grouped bars on log scale (mod_and_endpoint only).
            # 8 categories: Kept, then 4 sub-categories of unresponsive mod (by endpoint
            # status: pre-only, post-only, both, neither silent endpoint), then 3 active-
            # endpoint drop categories (pre-only, post-only, both).
            _mbd_log = [d for d in mask_breakdown_lst if d["exclusion_mode"] == "mod_and_endpoint"]
            if not _mbd_log:
                _mbd_log = mask_breakdown_lst  # fallback for normalized (mod_only only)
            _n_groups_log = len(_mbd_log)
            _x_log = np.arange(_n_groups_log)
            _kept_log    = np.array([d["n_kept"]               for d in _mbd_log], dtype=float)
            _m1_log      = np.array([d["n_drop_unres_mod"]     for d in _mbd_log], dtype=float)
            _um_pre_log  = np.array([d["n_unres_mod_pre_only"] for d in _mbd_log], dtype=float)
            _um_post_log = np.array([d["n_unres_mod_post_only"]for d in _mbd_log], dtype=float)
            _um_both_log = np.array([d["n_unres_mod_both"]     for d in _mbd_log], dtype=float)
            _um_neit_log = np.array([d["n_unres_mod_neither"]  for d in _mbd_log], dtype=float)
            _ep_pre_log  = np.array([d["n_drop_pre_only"]      for d in _mbd_log], dtype=float)
            _ep_post_log = np.array([d["n_drop_post_only"]     for d in _mbd_log], dtype=float)
            _ep_both_log = np.array([d["n_drop_both"]          for d in _mbd_log], dtype=float)

            # Order follows the filtering pipeline:
            # 1) Stage 1: unresponsive modulation cluster (broken down by endpoint)
            # 2) Stage 2: responsive mod cluster but unresponsive pre/post endpoint
            # 3) Kept: survived both stages
            _cat_labels_log = [
                "Unres mod: pre silent",  "Unres mod: post silent",
                "Unres mod: both silent", "Unres mod: neither silent",
                "Drop: pre silent",  "Drop: post silent", "Drop: both silent",
                "Kept",
            ]
            _cat_colors_log = [
                c_vals_l[0], c_vals_l[1], c_vals_l[3], c_vals_l[6],
                c_vals[1], c_vals[3], c_vals[4],
                c_vals[2],
            ]
            _cat_vals_log = np.array([
                _um_pre_log, _um_post_log, _um_both_log, _um_neit_log,
                _ep_pre_log, _ep_post_log, _ep_both_log,
                _kept_log,
            ])  # (8, n_groups_log)
            _totals_log = np.array([d["n_total"] for d in _mbd_log], dtype=float)
            assert np.allclose(_cat_vals_log.sum(axis=0), _totals_log), (
                f"Mask breakdown categories do not sum to total: "
                f"{_cat_vals_log.sum(axis=0)} vs {_totals_log}"
            )
            _n_cats_log = len(_cat_labels_log)
            _grp_w = 0.85
            _bar_w_log = _grp_w / _n_cats_log
            _offsets_log = np.arange(_n_cats_log) * _bar_w_log - (_n_cats_log - 1) * _bar_w_log / 2

            figmask_log, axmask_log = plt.subplots(
                1, 1, figsize=(max(8, 2.5 * _n_groups_log + 2), 6)
            )
            for ci in range(_n_cats_log):
                vals = _cat_vals_log[ci]
                axmask_log.bar(
                    _x_log + _offsets_log[ci], vals, width=_bar_w_log,
                    color=_cat_colors_log[ci], edgecolor="black", linewidth=0.4,
                    label=_cat_labels_log[ci],
                )
                for xi, v in enumerate(vals):
                    if v > 0:
                        axmask_log.text(
                            xi + _offsets_log[ci], v * 1.15,
                            f"{int(v)}",
                            ha="center", va="bottom", fontsize=6,
                        )
            axmask_log.set_yscale("log")
            axmask_log.set_xticks(_x_log)
            axmask_log.set_xticklabels(
                [f"G={d['G']}  (MM={d['n_total']})\n"
                 f"unres pre={d['n_unres_input_neurons']}, "
                 f"unres post={d['n_unres_hidden_neurons']}"
                 for d in _mbd_log],
                fontsize=8,
            )
            axmask_log.set_ylabel("Number of modulations (log scale)", fontsize=13)
            axmask_log.set_title(
                f"Modulation mask breakdown — log scale, mod_and_endpoint "
                f"({'normalized' if clustering_normalize else 'unnormalized'})",
                fontsize=12,
            )
            axmask_log.legend(loc="upper center", bbox_to_anchor=(0.5, -0.15),
                              ncol=3, frameon=False, fontsize=7)
            figmask_log.savefig(
                f"{save_dir}/{clustering_name}_mask_breakdown_log_{savefigure_name}.png",
                dpi=300, bbox_inches="tight",
            )
            plt.close(figmask_log)

            # Decision-tree diagram of the filtering hierarchy (one per G).
            # Skip for normalized data: no unresponsive detection, so the tree is trivial
            # (all synapses are kept, all endpoint drops are zero).
            _mbd_tree = [] if clustering_normalize else _mbd_log
            for _di, _d in enumerate(_mbd_tree):
                _total = _d["n_total"]
                _n_unres = _d["n_drop_unres_mod"]
                _n_resp  = _total - _n_unres
                _pct_of = lambda v, base: f"{100.0 * v / base:.1f}%" if base > 0 else "—"

                fig_tree, ax_tree = plt.subplots(1, 1, figsize=(20, 9))
                ax_tree.set_xlim(0, 20)
                ax_tree.set_ylim(0, 10)
                ax_tree.axis("off")

                # Node style
                _box = dict(boxstyle="round,pad=0.35", facecolor="white", edgecolor="black", linewidth=1.2)
                _arrow = dict(arrowstyle="->, head_width=0.15", color="black", lw=1.2)

                # Level 0: root
                ax_tree.annotate(
                    f"Total\n{_total}", xy=(10, 9.5), fontsize=10, ha="center", va="center", bbox=_box)

                # Level 1: two branches
                _l1_left, _l1_right = 5.0, 15.0
                ax_tree.annotate("", xy=(_l1_left, 8.3), xytext=(10, 9.1), arrowprops=_arrow)
                ax_tree.annotate("", xy=(_l1_right, 8.3), xytext=(10, 9.1), arrowprops=_arrow)

                ax_tree.annotate(
                    f"Unres mod cluster\n{_n_unres} ({_pct_of(_n_unres, _total)} of total)",
                    xy=(_l1_left, 8.0), fontsize=9, ha="center", va="center",
                    bbox={**_box, "facecolor": "#fee2e2"})
                ax_tree.annotate(
                    f"Responsive mod cluster\n{_n_resp} ({_pct_of(_n_resp, _total)} of total)",
                    xy=(_l1_right, 8.0), fontsize=9, ha="center", va="center",
                    bbox={**_box, "facecolor": "#dbeafe"})

                # Level 2 left: unres mod sub-breakdown
                _um_vals = [
                    ("pre silent",  _d["n_unres_mod_pre_only"],  c_vals_l[0]),
                    ("post silent", _d["n_unres_mod_post_only"], c_vals_l[1]),
                    ("both silent", _d["n_unres_mod_both"],      c_vals_l[3]),
                    ("neither",     _d["n_unres_mod_neither"],   c_vals_l[6]),
                ]
                _left_xs = [1.5, 3.5, 6.5, 8.5]
                _left_ys = [5.5, 4.5, 5.5, 4.5]
                for (_lbl, _v, _fc), _lx, _ly in zip(_um_vals, _left_xs, _left_ys):
                    ax_tree.annotate("", xy=(_lx, _ly + 0.8), xytext=(_l1_left, 7.6), arrowprops=_arrow)
                    ax_tree.annotate(
                        f"{_lbl}\n{_v}\n({_pct_of(_v, _total)} of total"
                        f", {_pct_of(_v, _n_unres)} of parent)",
                        xy=(_lx, _ly), fontsize=7, ha="center", va="center",
                        bbox={**_box, "facecolor": _fc})

                # Level 2 right: responsive mod sub-breakdown
                _ep_vals = [
                    ("pre silent",  _d["n_drop_pre_only"],  c_vals[1]),
                    ("post silent", _d["n_drop_post_only"], c_vals[3]),
                    ("both silent", _d["n_drop_both"],      c_vals[4]),
                    ("Kept",        _d["n_kept"],           c_vals[2]),
                ]
                _right_xs = [11.5, 13.5, 16.5, 18.5]
                _right_ys = [5.5, 4.5, 5.5, 4.5]
                for (_lbl, _v, _fc), _rx, _ry in zip(_ep_vals, _right_xs, _right_ys):
                    ax_tree.annotate("", xy=(_rx, _ry + 0.8), xytext=(_l1_right, 7.6), arrowprops=_arrow)
                    ax_tree.annotate(
                        f"{_lbl}\n{_v}\n({_pct_of(_v, _total)} of total"
                        f", {_pct_of(_v, _n_resp)} of parent)",
                        xy=(_rx, _ry), fontsize=7, ha="center", va="center",
                        bbox={**_box, "facecolor": _fc},
                        fontweight="bold" if _lbl == "Kept" else "normal")

                # Subtitle with G and unresponsive neuron counts
                ax_tree.text(
                    10, 2.5,
                    f"G={_d['G']}  |  MM={_total}  |  "
                    f"unres input neurons={_d['n_unres_input_neurons']}, "
                    f"unres hidden neurons={_d['n_unres_hidden_neurons']}",
                    fontsize=9, ha="center", va="center", style="italic",
                )

                fig_tree.suptitle(
                    f"Modulation filtering hierarchy "
                    f"({'normalized' if clustering_normalize else 'unnormalized'})",
                    fontsize=12,
                )
                fig_tree.savefig(
                    f"{save_dir}/{clustering_name}_mask_tree_G{_d['G']}_{savefigure_name}.png",
                    dpi=300, bbox_inches="tight",
                )
                plt.close(fig_tree)

            # all following analysis based on the maximal G (G_lst[-1])
            cell_vars_rules_sorted_norm_pre = cell_vars_rules_sorted_norm[np.ix_(result_pre["row_order"], result_pre["col_order"])]
            cell_vars_rules_sorted_norm_post = cell_vars_rules_sorted_norm[np.ix_(result_post["row_order"], result_post["col_order"])]

            # 2025-11-12: for all G results (after reordering)
            # up to here, the covariance matrix is determined. 
            cell_vars_rules_sorted_norm_all_lst = []
            for result_all_ in result_all_lst:
                cell_vars_rules_sorted_norm_all = cell_vars_rules_sorted_norm[np.ix_(result_all_["row_order"], result_all_["col_order"])]
                cell_vars_rules_sorted_norm_all_lst.append(cell_vars_rules_sorted_norm_all)

            clustering_data_hierarchy["modulation_all_pre"] = result_pre["col_linkage"]
            clustering_data_hierarchy["modulation_all_post"] = result_post["col_linkage"]

            # check if length is consistent
            assert len(G_lst) == len(result_all_lst)

            # plot original, pre, post, plus all results under different G
            tf = 3 + len(result_all_lst)
            figprepost, axprepost = plt.subplots(tf,1,figsize=(24,8*tf))

            # For unnormalized data use a log colorbar bounded below at 1e-4;
            # compute a shared vmax across all subplots for consistent color scale.
            if not clustering_normalize:
                _all_mats = (
                    [cell_vars_rules_sorted_norm,
                     cell_vars_rules_sorted_norm_pre,
                     cell_vars_rules_sorted_norm_post]
                    + list(cell_vars_rules_sorted_norm_all_lst)
                )
                _vmax_log = float(max(m.max() for m in _all_mats))
                _heatmap_kw = dict(norm=LogNorm(vmin=1e-4, vmax=_vmax_log))
            else:
                _heatmap_kw = dict(vmin=vmins, vmax=vmaxs)

            sns.heatmap(cell_vars_rules_sorted_norm, ax=axprepost[0], cmap=cs, cbar=True, **_heatmap_kw)
            sns.heatmap(cell_vars_rules_sorted_norm_pre, ax=axprepost[1], cmap=cs, cbar=True, **_heatmap_kw)
            sns.heatmap(cell_vars_rules_sorted_norm_post, ax=axprepost[2], cmap=cs, cbar=True, **_heatmap_kw)

            # plot the headmap for all G result
            for k in range(len(cell_vars_rules_sorted_norm_all_lst)):
                sns.heatmap(cell_vars_rules_sorted_norm_all_lst[k], ax=axprepost[3+k], cmap=cs, cbar=True, **_heatmap_kw)

            for ax in axprepost:
                ax.set_yticks(np.arange(len(tb_break_name)))
                
            axprepost[0].set_yticklabels(tb_break_name, rotation=0, ha='right', va='center', fontsize=9)
            axprepost[1].set_yticklabels(tb_break_name[result_pre["row_order"]], rotation=0, ha='right', va='center', fontsize=9)
            axprepost[2].set_yticklabels(tb_break_name[result_post["row_order"]], rotation=0, ha='right', va='center', fontsize=9)
            
            # save for external plotting 
            cell_vars_rules_sorted_norm_all_save = {}
            cell_vars_rules_sorted_norm_all_save["cell_vars_rules_sorted_norm_all_lst"] = cell_vars_rules_sorted_norm_all_lst

            # 2025-10-30: modulation, clustering by G groups from MinibatchKmeans, "border" between clusters
            modulation_cluster_norm = []
            modulation_cluster_boundary = []
            for k in range(len(cell_vars_rules_sorted_norm_all_lst)):
                axprepost[3+k].set_yticklabels(tb_break_name[result_all_lst[k]["row_order"]], 
                                               rotation=0, ha='right', va='center', fontsize=9)
                result_all = result_all_lst[k]
                rl_all = np.asarray(result_all["row_labels"])[result_all["row_order"]]
                cl_all = np.asarray(result_all["col_labels"])[result_all["col_order"]]
                rbreaks_all_ = clustering._breaks(rl_all)
                cbreaks_all_ = clustering._breaks(cl_all)
                for rb in rbreaks_all_:
                    axprepost[3+k].axhline(rb, color="w", lw=3.0, zorder=3)
                    axprepost[3+k].axhline(rb, color="k", lw=0.8, zorder=4)
                for cb in cbreaks_all_:
                    axprepost[3+k].axvline(cb, color="w", lw=3.0, zorder=3)
                    axprepost[3+k].axvline(cb, color="k", lw=0.8, zorder=4)
                    
                modulation_cluster_boundary.append([rbreaks_all_, cbreaks_all_])

                # 2025-11-12: calculate the mean of the reordered covaraince matrix under different G value
                # notice G indicates the group size, not the actual 
                rbreaks_all_end = [0] + rbreaks_all_ + [cell_vars_rules_sorted_norm_all_lst[0].shape[0]]
                cbreaks_all_end = [0] + cbreaks_all_ + [cell_vars_rules_sorted_norm_all_lst[0].shape[1]]

                varmeanmod = np.zeros((len(rbreaks_all_end) - 1, len(cbreaks_all_end) - 1))
                for rr in range(len(rbreaks_all_end) - 1):
                    for cc in range(len(cbreaks_all_end) - 1):
                        varmeanmod[rr,cc] = np.mean(cell_vars_rules_sorted_norm_all_lst[k][rbreaks_all_end[rr]:rbreaks_all_end[rr+1], 
                                                                                        cbreaks_all_end[cc]:cbreaks_all_end[cc+1]])

                modulation_cluster_norm.append(varmeanmod)    
                
            cell_vars_rules_sorted_norm_all_save["modulation_cluster_boundary"] = modulation_cluster_boundary      
                
            axprepost[0].set_title("Original", fontsize=15)
            axprepost[1].set_title(f"Group by Pre (row_k={result_pre['row_k']}, col_k={result_pre['col_k']})", fontsize=15)
            axprepost[2].set_title(f"Group by Post (row_k={result_post['row_k']}, col_k={result_post['col_k']})", fontsize=15)

            # 2025-10-30: modulation, clustering by pre/post, "border" between clusters
            rl_pre = np.asarray(result_pre["row_labels"])[result_pre["row_order"]]
            cl_pre = np.asarray(result_pre["col_labels"])[result_pre["col_order"]]
            rbreaks_pre = clustering._breaks(rl_pre)
            cbreaks_pre = clustering._breaks(cl_pre)
            for rb in rbreaks_pre:
                axprepost[1].axhline(rb, color="w", lw=3.0, zorder=3)
                axprepost[1].axhline(rb, color="k", lw=0.8, zorder=4)
            for cb in cbreaks_pre:
                axprepost[1].axvline(cb, color="w", lw=3.0, zorder=3)
                axprepost[1].axvline(cb, color="k", lw=0.8, zorder=4)

            rl_post = np.asarray(result_post["row_labels"])[result_post["row_order"]]
            cl_post = np.asarray(result_post["col_labels"])[result_post["col_order"]]
            rbreaks_post = clustering._breaks(rl_post)
            cbreaks_post = clustering._breaks(cl_post)
            for rb in rbreaks_post:
                axprepost[2].axhline(rb, color="w", lw=3.0, zorder=3)
                axprepost[2].axhline(rb, color="k", lw=0.8, zorder=4)
            for cb in cbreaks_post:
                axprepost[2].axvline(cb, color="w", lw=3.0, zorder=3)
                axprepost[2].axvline(cb, color="k", lw=0.8, zorder=4)
            
            for k in range(len(cell_vars_rules_sorted_norm_all_lst)):
                axprepost[3+k].set_title(f"Group by All, G={G_lst[k]} (row_k={result_all_lst[k]['row_k']}, col_k={result_all_lst[k]['col_k']})", fontsize=15)
                
            figprepost.tight_layout()
            figprepost.savefig(f"{save_dir}/{clustering_name}_bygroup_variance_{savefigure_name}.png", dpi=300)
            plt.close(figprepost)
            
            # save the fitting result under different number of predefined modulation cluster
            with open(f"{save_dir}/{clustering_name}_clustering_result_all_{savefigure_name}.pkl", "wb") as f:
                pickle.dump(cell_vars_rules_sorted_norm_all_save, f)

            # plot the norm 
            figmodnorm, axsmodnorm = plt.subplots(len(modulation_cluster_norm),1,figsize=(10,4*len(modulation_cluster_norm)))
            for idx in range(len(modulation_cluster_norm)):
                sns.heatmap(modulation_cluster_norm[idx], ax=axsmodnorm[idx], cmap=cs)
                axsmodnorm[idx].set_xlabel("Modulation Index", fontsize=15)
                axsmodnorm[idx].set_ylabel("Session Index", fontsize=15)

            figmodnorm.tight_layout()
            figmodnorm.savefig(f"{save_dir}/{clustering_name}_modulation_clusternorm_{savefigure_name}.png", dpi=300)
            plt.close(figmodnorm)

            # plot the score
            def _gap_curve_mod(Z, k_vals):
                """Compute Ward merge-height gap for each k from a linkage matrix."""
                if Z is None:
                    return np.full(len(k_vals), np.nan)
                n = Z.shape[0] + 1
                gaps = []
                for k in k_vals:
                    idx_above = n - k
                    idx_below = n - k - 1
                    h_above = float(Z[idx_above, 2]) if 0 <= idx_above < Z.shape[0] else np.inf
                    h_below = float(Z[idx_below, 2]) if 0 <= idx_below < Z.shape[0] else 0.0
                    gaps.append(h_above - h_below)
                return np.asarray(gaps, dtype=float)

            figscore, (axscore, axgap) = plt.subplots(2, 1, figsize=(10, 6), sharex=True)

            # result_pre / result_post: single grouping → col_score_recording_mean has no meaningful std
            pre_ks  = np.asarray(result_pre["col_k_values"],  dtype=int)
            post_ks = np.asarray(result_post["col_k_values"], dtype=int)

            # ---------- silhouette: pre ----------
            axscore.plot(pre_ks, [result_pre["col_score_recording_mean"][k] for k in pre_ks],
                         label="col pre", color=c_vals[1])
            axscore.scatter(result_pre["col_k"],
                            result_pre["col_score_recording_mean"][result_pre["col_k"]],
                            color=c_vals[1], zorder=5, s=60, marker="*")
            axscore.axvline(result_pre["col_strict_best_k"], color=c_vals[1], linestyle="-",  linewidth=1.0)
            axscore.axvline(result_pre["col_k"],             color=c_vals[1], linestyle="--", linewidth=1.0)

            # ---------- silhouette: post ----------
            axscore.plot(post_ks, [result_post["col_score_recording_mean"][k] for k in post_ks],
                         label="col post", color=c_vals[3])
            axscore.scatter(result_post["col_k"],
                            result_post["col_score_recording_mean"][result_post["col_k"]],
                            color=c_vals[3], zorder=5, s=60, marker="*")
            axscore.axvline(result_post["col_strict_best_k"], color=c_vals[3], linestyle="-",  linewidth=1.0)
            axscore.axvline(result_post["col_k"],             color=c_vals[3], linestyle="--", linewidth=1.0)

            # ---------- silhouette: result_all_lst (mean ± std across groupings) ----------
            for k_idx in range(len(result_all_lst)):
                r     = result_all_lst[k_idx]
                ks    = np.asarray(r["col_k_values"], dtype=int)
                mean  = np.asarray([r["col_score_recording_mean"][ki] for ki in ks], dtype=float)
                std   = np.asarray([r["col_score_recording_std"][ki]  for ki in ks], dtype=float)
                color = c_vals[5 + k_idx]
                axscore.plot(ks, mean, label=f"col all, G={G_lst[k_idx]}", color=color, linestyle="--")
                axscore.fill_between(ks, mean - std, mean + std, color=color, alpha=0.2)
                axscore.scatter(r["col_k"], r["col_score_recording_mean"][r["col_k"]],
                                color=color, zorder=5, s=60, marker="*")
                axscore.axvline(r["col_strict_best_k"], color=color, linestyle="-",  linewidth=1.0)
                axscore.axvline(r["col_k"],             color=color, linestyle="--", linewidth=1.0)

            axscore.set_ylabel("Silhouette Score", fontsize=15)
            axscore.legend()

            # ---------- gap: pre ----------
            axgap.plot(pre_ks, _gap_curve_mod(result_pre["col_linkage"], pre_ks),
                       label="col pre", color=c_vals[1])
            axgap.axvline(result_pre["col_strict_best_k"], color=c_vals[1], linestyle="-",  linewidth=1.0)
            axgap.axvline(result_pre["col_k"],             color=c_vals[1], linestyle="--", linewidth=1.0)

            # ---------- gap: post ----------
            axgap.plot(post_ks, _gap_curve_mod(result_post["col_linkage"], post_ks),
                       label="col post", color=c_vals[3])
            axgap.axvline(result_post["col_strict_best_k"], color=c_vals[3], linestyle="-",  linewidth=1.0)
            axgap.axvline(result_post["col_k"],             color=c_vals[3], linestyle="--", linewidth=1.0)

            # ---------- gap: result_all_lst ----------
            for k_idx in range(len(result_all_lst)):
                r     = result_all_lst[k_idx]
                ks    = np.asarray(r["col_k_values"], dtype=int)
                color = c_vals[5 + k_idx]
                axgap.plot(ks, _gap_curve_mod(r["col_linkage"], ks),
                           label=f"col all, G={G_lst[k_idx]}", color=color, linestyle="--")
                axgap.axvline(r["col_strict_best_k"], color=color, linestyle="-",  linewidth=1.0)
                axgap.axvline(r["col_k"],             color=color, linestyle="--", linewidth=1.0)

            axgap.set_xlabel("Number of Clusters", fontsize=15)
            axgap.set_ylabel("Gap Score", fontsize=15)
            axgap.set_xscale("log")
            axgap.legend()

            figscore.tight_layout()
            figscore.savefig(f"{save_dir}/{clustering_name}_variance_cluster_score_{savefigure_name}.png", dpi=300)
            plt.close(figscore)


            # compare modulation grouping result
            # either using pre, post, or different level of K-means pre-clusters
            all_mod = [result_pre, result_post] + result_all_lst
            all_mod_name = ["Pre", "Post"] + result_all_name_lst

            all_mod_metric_allk = []
            # loop through downsampled clusters along the neuron dimension
            select_col_allk = np.arange(lower_cluster_mod, 30, 3)
            valid_select_col_allk = []
            for select_col_k_ in select_col_allk:
                if all(
                    (select_col_k_ in mod_["row_labels_by_k"]) and
                    (select_col_k_ in mod_["col_labels_by_k"])
                    for mod_ in all_mod
                ):
                    valid_select_col_allk.append(select_col_k_)

            if len(valid_select_col_allk) == 0:
                print("No shared modulation clustering keys found for comparison; skipping metric plot.")
                all_mod_metric_allk = np.empty((0, len(all_mod), 3))
            else:
                for select_col_k_ in valid_select_col_allk:
                    all_mod_metric = []
                    for mod_ in all_mod: 
                        eval_res_modulation = clustering_metric.evaluate_bicluster_clustering(
                            cell_vars_rules_sorted_norm, 
                            row_labels=mod_["row_labels_by_k"][select_col_k_], 
                            col_labels=mod_["col_labels_by_k"][select_col_k_]
                        )
                        
                        all_mod_metric.append([
                            eval_res_modulation["metrics"]["CH_blocks"],
                            eval_res_modulation["metrics"]["DB_blocks"],
                            eval_res_modulation["metrics"]["XB_blocks"], 
                        ])

                    all_mod_metric_allk.append(all_mod_metric)

                all_mod_metric_allk = np.array(all_mod_metric_allk)
            print(all_mod_metric_allk.shape)
            
            # compare the effect by clustering based on presynaptic neuron, postsynaptic neuron, or all
            # hypothesis: all >≈ postsynaptic > presynaptic
            metric_labels = [
                ("Calinski-Harabasz (↑)", True),   # (label, higher_is_better)
                ("Davies-Bouldin (↓)", False),
                ("Xie-Beni (↓)", False),
            ]
            if all_mod_metric_allk.shape[0] > 0:
                n_metrics = all_mod_metric_allk.shape[2]
                fig, ax = plt.subplots(1, n_metrics, figsize=(5 * n_metrics, 4.5))
                for metric_index in range(n_metrics):
                    for model_index in range(all_mod_metric_allk.shape[1]):
                        ax[metric_index].plot(
                            valid_select_col_allk,
                            all_mod_metric_allk[:, model_index, metric_index],
                            "-o", markersize=4, linewidth=1.5,
                            label=all_mod_name[model_index],
                            color=c_vals[model_index],
                        )
                    ax[metric_index].set_ylabel(metric_labels[metric_index][0], fontsize=13)
                    ax[metric_index].set_xlabel("# Cluster (col k)", fontsize=13)
                    ax[metric_index].set_yscale("log")
                    ax[metric_index].legend(fontsize=8, frameon=True, loc="best")
                    ax[metric_index].tick_params(labelsize=10)
                    ax[metric_index].grid(True, which="both", ls="--", alpha=0.3)
                fig.tight_layout()
                fig.savefig(f"{save_dir}/{clustering_name}_between_modulation_{savefigure_name}.png", dpi=300)
                plt.close(fig)

            cluster_info_save_mod[clustering_save_name] = {
                "result_pre": result_pre,
                "result_post": result_post,
                "result_all_lst": result_all_lst,
                "result_all_name_lst": result_all_name_lst,
                "tb_break_name": tb_break_name,
                "cell_vars_rules_sorted_norm": cell_vars_rules_sorted_norm,
                "global_assignment": global_assignment_cache,
                f"global_assignment_fixed_k{FIXED_K_OM}": global_assignment_fixed_k_cache,
            }
    
    # save this only at the end     
    with open(f"{save_dir}/cluster_info_mod_{savefigure_name_base}.pkl", "wb") as f:
        pickle.dump(cluster_info_save_mod, f)
        

if __name__ == "__main__":
    # Clean up old output files
    clean = True 
    if clean: 
        for f in Path("multiple_tasks").glob("*.png"):
            f.unlink()
        for f in Path("multiple_tasks").glob("*.pkl"):
            f.unlink()
        
    import re
    saved_nets = sorted(Path("multiple_tasks").glob("savednet_everything_seed*+angle.pt"))
    param_lst = []
    for p in saved_nets:
        m = re.match(r"savednet_everything_seed(\d+)_(\w+)\+hidden\d+\+batch\d+\+angle\.pt", p.name)
        if m:
            param_lst.append((int(m.group(1)), m.group(2)))
            
    print(f"Found {len(param_lst)} saved models: {param_lst}")
    
    param_lst = [[749, "L21e4"]]
    
    for seed, feature in param_lst:
        main(seed, feature)