{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Neural Data Clustering Analysis\n",
    "\n",
    "## Description\n",
    "This notebook performs unsupervised clustering analysis on the neuronal data.\n",
    "\n",
    "The analysis pipeline includes:\n",
    "- **PCA Analysis**: Identifies significant principal components using permutation testing\n",
    "- **Spectral Embedding**: Transforms data into spectral space for optimal cluster separation\n",
    "- **Gaussian Mixture Model Clustering**: Applies GMM with grid search for optimal parameters\n",
    "- **Cluster Validation**: Uses BIC criteria and hierarchical clustering for validation\n",
    "- **Results Export**: Saves clustering results in multiple formats for further analysis\n",
    "\n",
    "## Scientific Objectives\n",
    "- Identify functionally distinct neural populations in the recorded data\n",
    "- Determine optimal clustering parameters using statistical criteria\n",
    "- Analyze the relationship between brain areas, cell types, and cluster membership\n",
    "- Provide quantitative measures of cluster quality and stability\n",
    "\n",
    "## Methodology Overview\n",
    "1. **Dimensionality Reduction**: PCA identifies significant components via permutation testing\n",
    "2. **Spectral Embedding**: Transforms PCA space into spectral coordinates preserving local structure\n",
    "3. **Model Selection**: Grid search finds optimal GMM parameters using BIC\n",
    "4. **Validation**: Hierarchical clustering confirms cluster robustness\n",
    "5. **Export**: Results saved in multiple formats for downstream analysis\n",
    "\n",
    "## Requirements\n",
    "```bash\n",
    "pip install numpy matplotlib scipy scikit-learn pandas joblib\n",
    "```\n",
    "\n",
    "## Instructions\n",
    "1. Ensure the data files `Neuronal_data_100ms.npz` and `Neuronal_data_10ms_withroc.npz` are available in `processed_data/`\n",
    "2. Run all cells in order\n",
    "3. Results will be saved as `Data_Clustering.pkl` and `Data_Clustering.npz`\n",
    "\n",
    "## Expected Outputs\n",
    "- **Data_Clustering.pkl**: Python pickle file with complete clustering results\n",
    "- **Data_Clustering.npz**: NumPy-friendly clustering export for downstream notebooks\n",
    "- **GMM_100ms_Run0.pkl**: GMM model results and parameters\n",
    "- **Spectral_Data.npz**: Spectral embedding coordinates\n",
    "- **pca_analysis_data.npz**: PCA analysis results\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Import Libraries and Set Parameters\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Import required libraries\n",
    "import gc\n",
    "import os\n",
    "from pathlib import Path\n",
    "\n",
    "import joblib\n",
    "import matplotlib\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from scipy.cluster.hierarchy import dendrogram, linkage, fcluster\n",
    "from scipy.linalg import eigh\n",
    "from scipy.spatial.distance import pdist, squareform\n",
    "from scipy.spatial.distance import squareform as scipy_squareform\n",
    "from scipy.stats import ttest_1samp\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.metrics import pairwise_distances\n",
    "from sklearn.mixture import GaussianMixture\n",
    "\n",
    "# Set matplotlib backend and parameters for publication-quality figures\n",
    "# In notebooks, keep default inline backend to avoid Tk crashes.\n",
    "matplotlib.rcParams['pdf.fonttype'] = 42  # Embed fonts as editable text\n",
    "matplotlib.rcParams['font.size'] = 10\n",
    "matplotlib.rcParams['axes.linewidth'] = 0.5\n",
    "matplotlib.rcParams['xtick.major.width'] = 0.5\n",
    "matplotlib.rcParams['ytick.major.width'] = 0.5\n",
    "matplotlib.rcParams['figure.dpi'] = 300\n",
    "matplotlib.rcParams['savefig.dpi'] = 300\n",
    "\n",
    "# Set random seed for reproducibility\n",
    "np.random.seed(42)\n",
    "\n",
    "# Resolve repository and processed-data directory\n",
    "cwd = Path.cwd().resolve()\n",
    "REPO_ROOT = cwd\n",
    "for candidate in [cwd, *cwd.parents]:\n",
    "    if (candidate / 'functions').is_dir():\n",
    "        REPO_ROOT = candidate\n",
    "        break\n",
    "\n",
    "PROCESSED_DIR = REPO_ROOT / 'processed_data'\n",
    "PROCESSED_DIR.mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "print(\"All packages imported successfully!\")\n",
    "print(f\"NumPy version: {np.__version__}\")\n",
    "print(f\"Matplotlib version: {matplotlib.__version__}\")\n",
    "print(f\"REPO_ROOT: {REPO_ROOT}\")\n",
    "print(f\"PROCESSED_DIR: {PROCESSED_DIR}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Define Core Analysis Functions\n",
    "\n",
    "The following functions implement the main algorithms for PCA analysis, spectral embedding, and clustering. Each function is designed to be modular and reusable, with comprehensive documentation and error handling.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def pca_analysis(neural_data, plot_ind, path, seed):\n",
    "    \"\"\"\n",
    "    Perform PCA analysis on neural data with permutation testing to determine significant PCs.\n",
    "\n",
    "    Args:\n",
    "        neural_data (numpy.ndarray): Neural data matrix (neurons x time bins).\n",
    "        plot_ind (bool): If True, generate and display plots.\n",
    "        path (str): Path to save PCA analysis data.\n",
    "        seed (int): Random seed for reproducibility.\n",
    "\n",
    "    Returns:\n",
    "        int: Number of significant principal components.\n",
    "    \"\"\"\n",
    "    np.random.seed(0)\n",
    "    centered = False\n",
    "\n",
    "    # Mean-center the data (rows = neurons, columns = time bins)\n",
    "    temp_nd = neural_data - neural_data.mean(axis=1, keepdims=True)\n",
    "\n",
    "    # Apply PCA\n",
    "    pca = PCA()\n",
    "    pca.fit(temp_nd.T)\n",
    "    score = pca.transform(temp_nd.T)\n",
    "    latent = pca.explained_variance_  # Variances of principal components\n",
    "\n",
    "    # Plot explained variance ratios\n",
    "    if plot_ind:\n",
    "        plt.figure(figsize=(10, 5))\n",
    "        plt.subplot(1, 2, 1)\n",
    "        plt.plot(np.cumsum(latent) / np.sum(latent))\n",
    "        plt.grid()\n",
    "        plt.title(\"Cumulative Variance Ratio\")\n",
    "\n",
    "        plt.subplot(1, 2, 2)\n",
    "        plt.plot(latent / np.sum(latent))\n",
    "        plt.grid()\n",
    "        plt.title(\"Variance Ratio\")\n",
    "        plt.show(block=False)\n",
    "\n",
    "        # Plot first 80 principal components\n",
    "        num_kernels = min(80, score.shape[1])\n",
    "        num_rows = int(np.ceil(num_kernels / 6))\n",
    "        plt.figure(figsize=(15, num_rows * 2))\n",
    "        for i in range(num_kernels):\n",
    "            plt.subplot(num_rows, 6, i + 1)\n",
    "            plt.plot(score[:, i])\n",
    "            plt.title(f\"Kernel {i + 1}\")\n",
    "            plt.grid()\n",
    "            plt.xticks([])\n",
    "            plt.yticks([])\n",
    "        plt.tight_layout()\n",
    "        plt.show(block=False)\n",
    "\n",
    "    # Null distribution via permutation testing\n",
    "    null_sample = 1000\n",
    "    null_sigma = np.zeros((null_sample, len(latent)))\n",
    "    for j in range(null_sample):\n",
    "        temp = np.copy(temp_nd)\n",
    "        for i in range(temp.shape[0]):\n",
    "            np.random.shuffle(temp[i, :])\n",
    "        null_pca = PCA()\n",
    "        null_pca.fit(temp.T)\n",
    "        null_sigma[j, :] = null_pca.explained_variance_\n",
    "        if j % 100 == 0:\n",
    "            print(f\"Progress: {j}/{null_sample} iterations completed\")\n",
    "    # Save null distribution\n",
    "    np.savez(f\"{path}/pca_analysis_data.npz\", null_sigma=null_sigma, latent=latent)\n",
    "\n",
    "    # Compute mean and standard deviation of null distribution\n",
    "    y = np.mean(null_sigma, axis=0)\n",
    "    dy = np.std(null_sigma, axis=0)\n",
    "\n",
    "    # Plot comparison with null distribution\n",
    "    if plot_ind:\n",
    "        plt.figure()\n",
    "        plt.plot(latent, linewidth=2, color=\"#4B3F72\", label=\"Variance Ratio\")\n",
    "        plt.plot(y, linewidth=2, color=\"#0072BD\", label=\"Null Variance Ratio\")\n",
    "        plt.fill_between(range(len(y)), y - dy, y + dy, color=\"#0072BD\", alpha=0.3, linestyle=\"--\")\n",
    "        plt.grid()\n",
    "        plt.legend()\n",
    "        plt.title(\"Variance Ratio vs Null Distribution\")\n",
    "        plt.show(block=False)\n",
    "\n",
    "    # Significance testing with Bonferroni correction\n",
    "    p_vals = np.array([ttest_1samp(null_sigma[:, i], latent[i])[1] for i in range(len(latent))])\n",
    "    significant_pcs = np.where((latent > y) & (p_vals < 0.05 / len(latent)))[0]\n",
    "\n",
    "    return len(significant_pcs)\n",
    "\n",
    "def squareform(Y, dir=None):\n",
    "    \"\"\"\n",
    "    \n",
    "    Reformat a distance matrix between upper triangular and square form.\n",
    "\n",
    "    Args:\n",
    "        Y (ndarray): Input vector or matrix.\n",
    "        dir (str): 'tovector' or 'tomatrix' to force the format.\n",
    "\n",
    "    Returns:\n",
    "        ndarray: Reformatted distance matrix.\n",
    "    \"\"\"\n",
    "    np.random.seed(0)\n",
    "    if dir is None:\n",
    "        dir = 'tomatrix' if Y.ndim == 1 else 'tovector'\n",
    "\n",
    "    if dir == 'tomatrix':\n",
    "        return scipy_squareform(Y)\n",
    "    elif dir == 'tovector':\n",
    "        return scipy_squareform(Y)\n",
    "    else:\n",
    "        raise ValueError(\"Invalid direction specified.\")\n",
    "\n",
    "def spect_clust(S, K):\n",
    "    \"\"\"\n",
    "    Spectral clustering helper function.\n",
    "\n",
    "    Args:\n",
    "        S (ndarray): Similarity matrix.\n",
    "        K (int): Number of clusters (-1 for all eigenvalues).\n",
    "\n",
    "    Returns:\n",
    "        tuple: Eigenvalues, eigenvectors, and flag.\n",
    "    \"\"\"\n",
    "    np.random.seed(0)\n",
    "    DinvS = np.diag(1.0 / np.sqrt(S.sum(axis=1)))\n",
    "    L = np.eye(S.shape[0]) - DinvS @ S @ DinvS\n",
    "\n",
    "    if K == -1:\n",
    "        Lambda, V = eigh(L)\n",
    "        flag = \"All\"\n",
    "    else:\n",
    "        Lambda, V = eigh(L, subset_by_index=[0, K - 1])\n",
    "        flag = \"SmallestAbs\"\n",
    "\n",
    "    Lambda = np.sort(Lambda)\n",
    "\n",
    "    plt.figure()\n",
    "    plt.plot(np.diff(Lambda[1:40]), '*')\n",
    "    plt.title(\"Eigenvalue Differences\")\n",
    "    plt.show(block=False)\n",
    "\n",
    "    return Lambda, V, flag\n",
    "\n",
    "def A4_spectral_embedding(Neural_Data, N_PCA, plot_ind, Path, Seed):\n",
    "    \"\"\"\n",
    "    Perform spectral embedding on neural data.\n",
    "\n",
    "    Args:\n",
    "        Neural_Data_5T2P_preprocessed (dict): Preprocessed neural data.\n",
    "        N_PCA (int): Number of principal components.\n",
    "        plot_ind (bool): Indicator to enable/disable plots.\n",
    "        Path (str): Path to save results.\n",
    "        Seed (int): Random seed for reproducibility.\n",
    "\n",
    "    Returns:\n",
    "        dict: Spectral data and cluster information.\n",
    "    \"\"\"\n",
    "    np.random.seed(Seed)\n",
    "\n",
    "    temp_ND_pdf = Neural_Data['Neurons_activity_condition_area']\n",
    "    temp_ND_area = Neural_Data['Area']\n",
    "    sig = 0.07975 #  0.07975 for 100 ms             and 0.1065 for 10 ms\n",
    "    K_Spect = 17  # From elbow method 16 for 100 ms and 31 for 10 ms\n",
    "    print(\"Number of neurons:\", temp_ND_pdf.shape[0])           # Number of neurons: 10294\n",
    "    # PCA\n",
    "    temp_ND_pdf = temp_ND_pdf - np.mean(temp_ND_pdf, axis=1, keepdims=True)\n",
    "    PCA_Data= np.linalg.svd(temp_ND_pdf, full_matrices=False)[0][:, :N_PCA]\n",
    "    \n",
    "    print(\"Number of PCA_Data:\", PCA_Data.shape[0])\n",
    "    D_Both = squareform(pdist(PCA_Data, metric='euclidean'))\n",
    "    S_Both = np.exp(-D_Both / sig)\n",
    "    print(\"Average similarity:\", np.mean(S_Both))\n",
    "\n",
    "    Lambda_Both, V_Both, flag = spect_clust(S_Both, -1)\n",
    "    print(\"Spectral embedding flag:\", flag)\n",
    "\n",
    "    spectral_Data = V_Both[:, 1:K_Spect]  # Exclude the first eigenvector\n",
    "\n",
    "\n",
    "    Spectral_Data = {\n",
    "        'spectral_Data': spectral_Data,\n",
    "        'ClusterCounter': Neural_Data ['clustercounter']\n",
    "    }\n",
    "\n",
    "    np.savez(f\"{Path}/Spectral_Data\", Spectral_Data=Spectral_Data, sig=sig)\n",
    "\n",
    "    if plot_ind:\n",
    "        plt.figure()\n",
    "        plt.hist(S_Both.flatten(), bins=50, color='blue', alpha=0.7)\n",
    "        plt.grid()\n",
    "        plt.xlim([0, 1])\n",
    "        plt.title('Similarity Histogram')\n",
    "        plt.show(block=False)\n",
    "\n",
    "        N_Neurons = temp_ND_pdf.shape[0]\n",
    "        areas = np.unique(temp_ND_area)\n",
    "        Ind = np.arange(N_Neurons)\n",
    "        Ind2 = []\n",
    "\n",
    "        plt.figure()\n",
    "        plt.imshow(S_Both, aspect='auto', cmap='viridis')\n",
    "\n",
    "        for area in areas:\n",
    "            temp = np.max(Ind[temp_ND_area.flatten() == area])\n",
    "            Ind2.append(int(np.median(np.where(temp_ND_area == area)[0])))\n",
    "            plt.plot([0, N_Neurons], [temp, temp], 'k-', linewidth=0.5)\n",
    "            plt.plot([temp, temp], [0, N_Neurons], 'k-', linewidth=0.5)\n",
    "\n",
    "        plt.colorbar()\n",
    "        plt.title(\"Similarity Heatmap\")\n",
    "        plt.xticks(Ind2, areas, rotation=90)\n",
    "        plt.yticks(Ind2, areas)\n",
    "        plt.show(block=False)\n",
    "\n",
    "        plt.figure()\n",
    "        plt.plot(Lambda_Both[1:], '*')\n",
    "        plt.title(\"Eigenvalues\")\n",
    "        plt.axis([1, 50, 0.96, 1])\n",
    "        plt.grid()\n",
    "        plt.show(block=False)\n",
    "\n",
    "        plt.figure()\n",
    "\n",
    "        plt.scatter(V_Both[:, 0], V_Both[:, 1],V_Both[:, 2])\n",
    "        plt.title(\"Neurons in Spectral Space\")\n",
    "        plt.grid()\n",
    "        plt.show(block=False)\n",
    "\n",
    "\n",
    "    return Spectral_Data\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from sklearn.mixture import GaussianMixture\n",
    "import joblib\n",
    "\n",
    "def A6_GMM_grid_search(Spectral_Data, Path, Seed,Run):\n",
    "    \"\"\"\n",
    "    Perform a grid search to find the best Gaussian Mixture Model (GMM) based on BIC.\n",
    "    \n",
    "    Args:\n",
    "        Spectral_Data (dict): Dictionary with 'spectral_Data' containing the dataset.\n",
    "        Path (str): Path to save the results.\n",
    "        Seed (int): Random seed for reproducibility.\n",
    "    \n",
    "    Returns:\n",
    "        GaussianMixture: The best GMM model based on BIC.\n",
    "    \"\"\"\n",
    "    np.random.seed(Seed)\n",
    "    \n",
    "    X = Spectral_Data['spectral_Data']\n",
    "    \n",
    "    k = list(range(10, 25, 5)) + list(range(25, 45, 1)) + list(range(45, 56, 5))  # Number of components to test\n",
    "    Sigma = ['diag']  # Covariance types to test   diagonal or full\n",
    "    SharedCovariance = [False]  # Shared covariance (not directly in sklearn)\n",
    "    RegularizationValue = 0  # Regularization for GMM\n",
    "    MaxIter = 1000\n",
    "    Replicates =1000  # Number of initializations\n",
    "\n",
    "    nK = len(k)\n",
    "    nSigma = len(Sigma)\n",
    "    nSC = len(SharedCovariance)\n",
    "    \n",
    "    # Preallocate results\n",
    "    gm_models = []\n",
    "    aic = np.zeros((nK, nSigma, nSC))\n",
    "    bic = np.zeros((nK, nSigma, nSC))\n",
    "    converged = np.zeros((nK, nSigma, nSC), dtype=bool)\n",
    "    \n",
    "    results = {\n",
    "        \"Model\": np.empty((nK, nSigma, nSC), dtype=object),\n",
    "        \"AIC\": np.zeros((nK, nSigma, nSC)),\n",
    "        \"BIC\": np.zeros((nK, nSigma, nSC)),\n",
    "        \"SharedCov\": np.empty((nK, nSigma, nSC), dtype=object),\n",
    "        \"Rep\": np.full((nK, nSigma, nSC), Replicates),\n",
    "        \"Covtype\": np.empty((nK, nSigma, nSC), dtype=object),\n",
    "        \"MaxIte\": np.full((nK, nSigma, nSC), MaxIter),\n",
    "        \"Reg\": np.full((nK, nSigma, nSC), RegularizationValue),\n",
    "        \"k\": np.zeros((nK, nSigma, nSC)),\n",
    "    }\n",
    "    \n",
    "    for m, shared in enumerate(SharedCovariance):\n",
    "        for j, cov_type in enumerate(Sigma):\n",
    "            for i, n_components in enumerate(k):\n",
    "                print(f\"Training GMM with k={n_components}, CovType={cov_type}, SharedCov={shared}\")\n",
    "                gmm = GaussianMixture(\n",
    "                    n_components=n_components,\n",
    "                    covariance_type=cov_type,\n",
    "                    max_iter=MaxIter,\n",
    "                    reg_covar=RegularizationValue,\n",
    "                    random_state=Seed,\n",
    "                    n_init=Replicates\n",
    "                )\n",
    "                \n",
    "                Y=np.array(X)\n",
    "                gmm.fit(Y)\n",
    "                \n",
    "                # Save model and metrics\n",
    "                gm_models.append(gmm)\n",
    "                aic[i, j, m] = gmm.aic(X)\n",
    "                bic[i, j, m] = gmm.bic(X)\n",
    "                converged[i, j, m] = gmm.converged_\n",
    "                \n",
    "                results[\"Model\"][i, j, m] = gmm\n",
    "                results[\"AIC\"][i, j, m] = aic[i, j, m]\n",
    "                results[\"BIC\"][i, j, m] = bic[i, j, m]\n",
    "                results[\"SharedCov\"][i, j, m] = shared\n",
    "                results[\"Covtype\"][i, j, m] = cov_type\n",
    "                results[\"k\"][i, j, m] = n_components\n",
    "    \n",
    "    # Save results\n",
    "    joblib.dump(results, f\"{Path}/GMM_100ms_Run{str(Run)}.pkl\")\n",
    "    \n",
    "    all_converged = np.all(converged)\n",
    "    print(f\"All GMM models converged: {all_converged}\")\n",
    "    \n",
    "    # Plot BIC values\n",
    "    plt.figure()\n",
    "    for j in range(nSigma * nSC):\n",
    "        plt.plot(k, bic[:, j, 0], marker='o', label=f\"Model {j}\")\n",
    "    plt.title(\"BIC for Various k and Covariance Choices\")\n",
    "    plt.xlabel(\"Number of Components (k)\")\n",
    "    plt.ylabel(\"BIC\")\n",
    "    plt.legend()\n",
    "    plt.show(block=False)\n",
    "    \n",
    "    # Find the best model\n",
    "    best_index = np.unravel_index(np.argmin(bic), bic.shape)\n",
    "    best_gmm = results[\"Model\"][best_index]\n",
    "    \n",
    "    return best_gmm\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Load and Explore Data\n",
    "\n",
    "Before performing clustering analysis, we need to load the neural data and understand its structure. This section loads the data files and provides initial exploration to understand the dataset characteristics.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Load neural data and examine structure\n",
    "print(\"Loading neural data...\")\n",
    "\n",
    "data100_path = PROCESSED_DIR / 'Neuronal_data_100ms.npz'\n",
    "if not data100_path.exists():\n",
    "    raise FileNotFoundError(f\"Missing file: {data100_path}\")\n",
    "\n",
    "required_keys = ['Neurons_activity_condition_area', 'Area', 'Type', 'Depth', 'Layer', 'clustercounter']\n",
    "with np.load(data100_path, allow_pickle=True) as d:\n",
    "    Data = {k: d[k] for k in required_keys if k in d}\n",
    "\n",
    "missing = [k for k in required_keys if k not in Data]\n",
    "if missing:\n",
    "    raise KeyError(f\"Missing keys in {data100_path.name}: {missing}\")\n",
    "\n",
    "print(\"Loaded keys:\", list(Data.keys()))\n",
    "print(\"\\nData structure:\")\n",
    "print(\"=\" * 50)\n",
    "for key in Data.keys():\n",
    "    if isinstance(Data[key], np.ndarray):\n",
    "        print(f\"{key}: {Data[key].shape} - {Data[key].dtype}\")\n",
    "    else:\n",
    "        print(f\"{key}: {type(Data[key])}\")\n",
    "\n",
    "# Extract key variables\n",
    "neural_activity = Data['Neurons_activity_condition_area']\n",
    "brain_areas = Data['Area']\n",
    "neuron_types = Data['Type']\n",
    "depths = Data['Depth']\n",
    "layers = Data['Layer']\n",
    "cluster_counter = Data['clustercounter']\n",
    "\n",
    "print(f\"\\nDataset Summary:\")\n",
    "print(f\"Neural activity matrix: {neural_activity.shape}\")\n",
    "print(f\"Number of neurons: {neural_activity.shape[0]}\")\n",
    "print(f\"Number of time bins: {neural_activity.shape[1]}\")\n",
    "print(f\"Brain areas: {np.unique(brain_areas)}\")\n",
    "print(f\"Neuron types: {np.unique(neuron_types)}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3. PCA Analysis and Extraction of Significant Principal Components\n",
    "\n",
    "In this cell, we will perform PCA analysis on the neural data and extract the number of significant principal components (PCs). The PCA analysis involves mean-centering the data, applying PCA, and performing permutation testing to determine the significance of each PC. The number of significant PCs will be returned as the output."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The pca_analysis computes significant_pcs=17 on some computers / package versions. \n",
    "# However, that count is numerically unstable: the last ~0-variance PC is flagged 'significant'\n",
    "# only via ~1e-31 BLAS noise on (latent-null), flipping the count 16<->17 across environments. \n",
    "# The count sets how many leading PCs the spectral embedding uses. \n",
    "# In order to reproduce the clustering of Ghaderi et al. (2026) you may need to pin the value to 17.\n",
    "\n",
    "significant_pcs = pca_analysis(Data['Neurons_activity_condition_area'], plot_ind=False, path=str(PROCESSED_DIR), seed=0)\n",
    "# significant_pcs = 17\n",
    "print(\"Number of significant PCs:\", significant_pcs)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Perform Spectral Embedding Analysis\n",
    "\n",
    "Now we perform the spectral embedding analysis. This step transforms the neural data into a spectral space where similar neurons are close together, making clustering more effective.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Perform spectral embedding analysis\n",
    "print(\"Starting spectral embedding analysis...\")\n",
    "\n",
    "# Parameters for spectral embedding\n",
    "N_PCA = significant_pcs\n",
    "plot_ind = False  # Generate plots (set True only for inspection)\n",
    "Path = str(PROCESSED_DIR)  # Save path\n",
    "Seed = 0  # Random seed\n",
    "\n",
    "# Run spectral embedding\n",
    "spectral_data = A4_spectral_embedding(Data, N_PCA, plot_ind, Path, Seed)\n",
    "\n",
    "print(\"Spectral embedding analysis completed!\")\n",
    "print(f\"Spectral data shape: {spectral_data['spectral_Data'].shape}\")\n",
    "# Free large source matrix after embedding\n",
    "del Data\n",
    "_ = gc.collect()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Gaussian Mixture Model Clustering (Optional, then reused by Step 6)\n",
    "\n",
    "This step generates `GMM_100ms_Run0.pkl`. Step 6 loads that file.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Optional: run GMM grid search to generate GMM_100ms_Run0.pkl\n",
    "print(\"GMM training step\")\n",
    "\n",
    "RUN_GMM_GRID_SEARCH = True  # Set True to regenerate GMM_100ms_Run0.pkl\n",
    "FORCE_N_COMPONENTS = None    # Pin to Code_M's published cluster count. The auto min-BIC pick is numerically unstable (flat BIC across k=27-36), so force k=29 to reproduce Code_M exactly (needs N_PCA=17, pinned above).\n",
    "EXPORT_CLUSTERING_STEM = 'Data_Clustering'  # Canonical export stem; use another name only for comparison snapshots\n",
    "\n",
    "Path = str(PROCESSED_DIR)\n",
    "Seed = 0\n",
    "Run = 0\n",
    "\n",
    "if RUN_GMM_GRID_SEARCH:\n",
    "    print(\"Starting GMM clustering analysis...\")\n",
    "    GMModel = A6_GMM_grid_search(spectral_data, Path=Path, Seed=Seed, Run=Run)\n",
    "    print(\"GMM clustering analysis completed!\")\n",
    "    print(f\"Best model has {GMModel.n_components} components\")\n",
    "else:\n",
    "    print(\"Skipping GMM training. Step 6 will load existing GMM_100ms_Run0.pkl from PROCESSED_DIR.\")\n",
    "    print(f\"FORCE_N_COMPONENTS = {FORCE_N_COMPONENTS}\")\n",
    "    print(f\"EXPORT_CLUSTERING_STEM = {EXPORT_CLUSTERING_STEM}\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Cluster Analysis and Validation\n",
    "\n",
    "After determining the optimal number of clusters, we perform additional validation using hierarchical clustering and analyze the cluster assignments.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load GMM results and perform cluster validation\n",
    "print(\"Loading GMM results for cluster validation...\")\n",
    "\n",
    "gmm_path = PROCESSED_DIR / 'GMM_100ms_Run0.pkl'\n",
    "spectral_path = PROCESSED_DIR / 'Spectral_Data.npz'\n",
    "\n",
    "if not spectral_path.exists():\n",
    "    raise FileNotFoundError(\n",
    "        f\"Missing {spectral_path}. Run Step 4 (spectral embedding) first.\"\n",
    "    )\n",
    "\n",
    "with np.load(spectral_path, allow_pickle=True) as data:\n",
    "    spectral_data = data['Spectral_Data'].item()  # .item() converts back to dict\n",
    "\n",
    "if not gmm_path.exists():\n",
    "    raise FileNotFoundError(\n",
    "        f\"Missing {gmm_path}. Run Step 5 with RUN_GMM_GRID_SEARCH=True to generate it.\"\n",
    "    )\n",
    "\n",
    "results = joblib.load(gmm_path)\n",
    "print(f\"Loaded existing GMM results: {gmm_path}\")\n",
    "\n",
    "# Select the model explicitly if FORCE_N_COMPONENTS is set; otherwise fall back to min-BIC\n",
    "bic_flat = np.asarray(results['BIC']).ravel()\n",
    "k_flat = np.asarray(results['k']).ravel().astype(int)\n",
    "model_flat = np.asarray(results['Model'], dtype=object).ravel()\n",
    "\n",
    "if FORCE_N_COMPONENTS is None:\n",
    "    best_model_index = int(np.argmin(bic_flat))\n",
    "    selection_mode = 'min-BIC'\n",
    "else:\n",
    "    forced_k = int(FORCE_N_COMPONENTS)\n",
    "    candidate_idx = np.flatnonzero(k_flat == forced_k)\n",
    "    if candidate_idx.size == 0:\n",
    "        available_k = sorted(np.unique(k_flat).tolist())\n",
    "        raise ValueError(\n",
    "            f\"Requested FORCE_N_COMPONENTS={forced_k}, but available k values are {available_k}\"\n",
    "        )\n",
    "    local_best = int(candidate_idx[np.argmin(bic_flat[candidate_idx])])\n",
    "    best_model_index = local_best\n",
    "    selection_mode = f'forced-k={forced_k}'\n",
    "\n",
    "best_gmm = model_flat[best_model_index]\n",
    "best_k = int(best_gmm.n_components)\n",
    "best_bic = float(bic_flat[best_model_index])\n",
    "\n",
    "print(f\"Selected model: {best_k} components ({selection_mode})\")\n",
    "print(f\"BIC: {best_bic:.2f}\")\n",
    "\n",
    "# Get cluster means and perform hierarchical clustering\n",
    "cluster_means = best_gmm.means_\n",
    "print(f\"Cluster means shape: {cluster_means.shape}\")\n",
    "\n",
    "# Compute condensed pairwise distances; linkage expects observations or condensed distances, not a square matrix\n",
    "dist_vector = pdist(cluster_means, metric='euclidean')\n",
    "print(f\"Condensed distance vector length: {dist_vector.shape[0]}\")\n",
    "\n",
    "# Perform hierarchical clustering\n",
    "Z = linkage(dist_vector, method='ward')\n",
    "\n",
    "# Determine optimal number of clusters using dendrogram\n",
    "max_d = 0.032  # Maximum distance threshold\n",
    "clusters = fcluster(Z, max_d, criterion='distance')\n",
    "num_clusters = len(np.unique(clusters))\n",
    "\n",
    "print(f\"Optimal number of clusters (hierarchical): {num_clusters}\")\n",
    "\n",
    "# Plot dendrogram\n",
    "plt.figure(figsize=(3, 3))\n",
    "dendrogram(Z, labels=np.arange(1, len(cluster_means) + 1))\n",
    "plt.axhline(y=max_d, color='r', linestyle='--', linewidth=2, label=f'Cutoff at {max_d}')\n",
    "plt.title('Dendrogram for GMM Cluster Validation', fontsize=10, fontweight='bold')\n",
    "plt.xlabel('Cluster Index')\n",
    "plt.ylabel('Distance')\n",
    "plt.legend()\n",
    "plt.grid(True, alpha=0.3)\n",
    "plt.show(block=False)\n",
    "\n",
    "print(\"Hierarchical clustering validation completed!\")\n",
    "\n",
    "# Create a single figure and axes to overlay both runs\n",
    "plt.figure(figsize=(3, 3))\n",
    "\n",
    "# Extract and sort BIC values\n",
    "BIC_3D = results['BIC']\n",
    "k_3D = results['k']\n",
    "BIC_1D = BIC_3D[:, 0, 0]\n",
    "k_vals_1D = k_3D[:, 0, 0]\n",
    "indices_sorted = np.argsort(k_vals_1D)\n",
    "sorted_k = k_vals_1D[indices_sorted]\n",
    "sorted_BIC = BIC_1D[indices_sorted]\n",
    "\n",
    "# Find the minimum and plot\n",
    "min_bic_idx = np.argmin(sorted_BIC)\n",
    "min_bic_k = sorted_k[min_bic_idx]\n",
    "min_bic_val = sorted_BIC[min_bic_idx]\n",
    "\n",
    "plt.plot(sorted_k, sorted_BIC, '-o', color='b', label='Run 1')\n",
    "plt.plot(min_bic_k, min_bic_val, 's', color='b', ms=5)\n",
    "plt.annotate(\n",
    "    f\"Min BIC: {min_bic_val:.2f}\\nat k={int(min_bic_k)}\",\n",
    "    xy=(min_bic_k, min_bic_val),\n",
    "    xytext=(min_bic_k + 5, min_bic_val + 1000),\n",
    "    arrowprops=dict(arrowstyle='->', color='k'),\n",
    "    ha='left',\n",
    "    color='k'\n",
    ")\n",
    "\n",
    "if FORCE_N_COMPONENTS is not None:\n",
    "    forced_mask = sorted_k == int(FORCE_N_COMPONENTS)\n",
    "    if np.any(forced_mask):\n",
    "        forced_bic = sorted_BIC[forced_mask][0]\n",
    "        plt.plot(int(FORCE_N_COMPONENTS), forced_bic, 'o', color='crimson', ms=5, label=f'Forced k={int(FORCE_N_COMPONENTS)}')\n",
    "\n",
    "plt.xlabel('Number of Components (k)')\n",
    "plt.ylabel('BIC')\n",
    "plt.title('BIC vs. Number of GMM Components (Run0)')\n",
    "plt.legend(loc='best')\n",
    "plt.tick_params(direction='out', top=False, right=False)\n",
    "plt.box(False)\n",
    "plt.show(block=False)\n",
    "\n",
    "_ = gc.collect()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Export results in multiple formats\n",
    "print(\"Exporting clustering results...\")\n",
    "\n",
    "if 'best_model_index' not in globals():\n",
    "    raise RuntimeError('Missing best_model_index. Run the previous selection/validation step first.')\n",
    "\n",
    "bic_values = np.asarray(results['BIC']).ravel()\n",
    "model = np.asarray(results['Model'], dtype=object).ravel()[best_model_index]\n",
    "print(f\"Selected export model has {model.n_components} components\")\n",
    "\n",
    "X = spectral_data['spectral_Data']\n",
    "clusterLabels = model.predict(X)\n",
    "clusterLabels = clusterLabels + 1  # MATLAB-style 1-indexing\n",
    "nNeurons = len(clusterLabels)\n",
    "print(f\"Assigned cluster labels (first 10): {list(clusterLabels[:10])}\")\n",
    "print(f\"Number of neurons (data points): {nNeurons}\")\n",
    "\n",
    "# Reorder neurons by cluster\n",
    "nClusters = int(model.n_components)\n",
    "clusterLabels = clusterLabels.astype(int)\n",
    "newClusterLabels = clusterLabels.astype(np.float32)\n",
    "ind_temp = np.argsort(newClusterLabels)\n",
    "T_temp = newClusterLabels[ind_temp]\n",
    "newclusterTags = [str(i + 1) for i in range(nClusters)]\n",
    "\n",
    "# Load 10ms+ROC NPZ\n",
    "data10_path = PROCESSED_DIR / 'Neuronal_data_10ms_withroc.npz'\n",
    "if not data10_path.exists():\n",
    "    raise FileNotFoundError(f\"Missing file: {data10_path}\")\n",
    "\n",
    "required10 = [\n",
    "    'Neurons_activity_condition_area', 'clustercounter', 'Area', 'Type', 'Depth',\n",
    "    'Layer', 'CCF_location', 'CCF_xyz', 'Discrimination_index', 'Unit_ids'\n",
    "]\n",
    "with np.load(data10_path, allow_pickle=True) as d:\n",
    "    Data10 = {k: d[k] for k in required10 if k in d}\n",
    "\n",
    "missing10 = [k for k in required10 if k not in Data10]\n",
    "if missing10:\n",
    "    raise KeyError(f\"Missing keys in {data10_path.name}: {missing10}\")\n",
    "\n",
    "# Create comprehensive results dictionary\n",
    "Cl = {\n",
    "    'Cluster_Counter_Ordered': T_temp,\n",
    "    'Neural_Data_normalized_Ordered': np.asarray(Data10['Neurons_activity_condition_area'][ind_temp, :], dtype=np.float32),\n",
    "    'Neuron_Counter_Ordered': Data10['clustercounter'][ind_temp],\n",
    "    # MATLAB random-dropout uses Id_Ordered as per-unit global IDs\n",
    "    'Id_Ordered': Data10['Unit_ids'][ind_temp],\n",
    "    'Areas_Ordered': Data10['Area'][ind_temp],\n",
    "    'Type_Ordered': Data10['Type'][ind_temp],\n",
    "    'Depth_Ordered': Data10['Depth'][ind_temp],\n",
    "    'Layer_Ordered': Data10['Layer'][ind_temp],\n",
    "    'CCF_location_Ordered': Data10['CCF_location'][ind_temp],\n",
    "    'CCF_xyz_Ordered': Data10['CCF_xyz'][ind_temp],\n",
    "    'Cluster_Tags': newclusterTags,\n",
    "    'Selectivity_index': Data10['Discrimination_index'][ind_temp],\n",
    "}\n",
    "\n",
    "# Free large intermediates\n",
    "del Data10, X, model, bic_values\n",
    "_ = gc.collect()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save clustering results in Python-friendly formats\n",
    "import pickle\n",
    "\n",
    "export_pkl_path = PROCESSED_DIR / f'{EXPORT_CLUSTERING_STEM}.pkl'\n",
    "export_npz_path = PROCESSED_DIR / f'{EXPORT_CLUSTERING_STEM}.npz'\n",
    "\n",
    "# Save as Python pickle file\n",
    "with open(export_pkl_path, 'wb') as f:\n",
    "    pickle.dump(Cl, f)\n",
    "print(f'{export_pkl_path.name} saved successfully!')\n",
    "\n",
    "print('\\nClustering analysis completed successfully!')\n",
    "\n",
    "# Save Python-friendly NPZ\n",
    "np.savez_compressed(\n",
    "    export_npz_path,\n",
    "    **{k: (np.array(v, dtype=object) if isinstance(v, (list, tuple)) else np.asarray(v)) for k, v in Cl.items()}\n",
    ")\n",
    "print(f'{export_npz_path.name} saved successfully!')\n",
    "\n",
    "# Final cleanup\n",
    "del Cl\n",
    "_ = gc.collect()\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.14.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
