SPDM_5FU.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Protocol for Bone Analysis Workflow"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Import Libraries and Parameters Setting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Standard Libraries\n",
    "import os\n",
    "import math\n",
    "import pickle\n",
    "from collections import defaultdict\n",
    "from datetime import datetime\n",
    "import re\n",
    "from itertools import islice\n",
    "import joblib\n",
    "\n",
    "# Numerical and Data Manipulation\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "# Visualization\n",
    "from matplotlib import pyplot as plt\n",
    "from matplotlib.colors import Normalize\n",
    "from matplotlib.backends.backend_pdf import PdfPages\n",
    "from matplotlib.patches import Patch\n",
    "from matplotlib.path import Path\n",
    "import matplotlib.patches as mpatches\n",
    "import seaborn as sns\n",
    "from PIL import Image\n",
    "\n",
    "# Spatial Analysis and Geometry\n",
    "from scipy.spatial import KDTree, ConvexHull\n",
    "from scipy.spatial.distance import pdist, squareform, jensenshannon\n",
    "from shapely.geometry import Polygon, box\n",
    "import geopandas as gpd\n",
    "\n",
    "# Image Processing\n",
    "from skimage import measure, morphology\n",
    "from scipy.ndimage import gaussian_filter1d\n",
    "\n",
    "# Interpolation and Optimization\n",
    "from scipy.interpolate import Rbf, griddata, interp1d\n",
    "from scipy.optimize import minimize\n",
    "\n",
    "# Statistical Analysis\n",
    "from scipy.stats import gaussian_kde, entropy, mannwhitneyu\n",
    "from statsmodels.stats.multitest import multipletests\n",
    "\n",
    "# Clustering and Dimensionality Reduction\n",
    "from scipy.cluster.hierarchy import linkage, fcluster\n",
    "from sklearn.cluster import KMeans, SpectralClustering\n",
    "from sklearn.metrics import silhouette_score\n",
    "\n",
    "# Machine Learning\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from sklearn.model_selection import cross_val_score, StratifiedKFold, GroupKFold\n",
    "from sklearn.metrics import (\n",
    "    classification_report, roc_auc_score, precision_score,\n",
    "    recall_score, f1_score, accuracy_score, balanced_accuracy_score\n",
    ")\n",
    "\n",
    "# File and Data Formats\n",
    "from imaris_ims_file_reader.ims import ims"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.getcwd()\n",
    "# os.chdir('/Users/fm07-admin/Projects/DisVis/')\n",
    "# Increase the maximum image size limit\n",
    "Image.MAX_IMAGE_PIXELS = None\n",
    "data_dir_5fu = 'data_metabone_5fu'\n",
    "# 'data_metabone_5fu/reference_bone/200117KK st4 acat27a 5FU d5 mBone(1)  mGFP(2).ims'\n",
    "reference_bone_dir = '/' # replace it with the real path\n",
    "save_dir = 'results_metabone_5fu'\n",
    "os.makedirs(save_dir, exist_ok=True)\n",
    "\n",
    "mk_color_map = {\n",
    "    'MKs': \"#FF00FF\",  # Magenta \n",
    "    'Adipos': \"#C96500\",  # Brown \n",
    "    'Sinusoids': \"#00FFFF\",  # Cyan # Marker CD105 = Sinusoids\n",
    "    'MSCs': '#FFCC00',  # Yellow # Marker Cxcl12 = MSCs\n",
    "    'Neurons': \"#FF7F00\",  # Orange  \n",
    "    'Arteries': \"#33A02C\", # Green\n",
    "    'Bones': \"#A6CEE3\" # Light Blue\n",
    "}\n",
    "\n",
    "\n",
    "hscs_color_map = {\n",
    "    'HSCs': \"#A6CEE3\" ,  # Light Blue\n",
    "    'RDs' : \"#A9A9A9\" # Gray\n",
    "}\n",
    "\n",
    "condition_colors = {\n",
    "    'd0': '#E31A1C',  # Red\n",
    "    'd5': '#FF1493',  # Deep Pink\n",
    "    'd10': '#4B0082', # Indigo\n",
    "    'd15': '#20B2AA', # Light Sea Green\n",
    "    'd30': '#FFD700',  # Gold\n",
    "    'd15_Detex': '#008080'  # Teal '#FF4500'  # Orange Red '#808080'  # Grey \n",
    "}\n",
    "\n",
    "cluster_colors = {\n",
    "    1: \"#B15928\",  # Brown\n",
    "    2: \"#6A3D9A\",  # Purple\n",
    "    3: \"#FF7F00\",  # Orange\n",
    "    4: \"#FDBF6F\",  # Light Orange\n",
    "    5: \"#A6CEE3\",  # Light Blue\n",
    "    6: \"#FB9A99\",  # Pink\n",
    "    7: \"#CAB2D6\",  # Lavender\n",
    "    8: \"#1F78B4\",  # Blue\n",
    "    9: \"#B2DF8A\",  # Light Green\n",
    "    10: \"#33A02C\"  # Green\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Data Import\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.1 Load CSV Files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define a function to filter and get the Position.csv files by condition\n",
    "def get_position_files_by_conditions(data_dir_root='data_metabone_5fu'):\n",
    "    data_dir_5fu = data_dir_root\n",
    "    # Lists to store Position.csv files for each condition\n",
    "    position_files_d5 = []\n",
    "    position_files_d10 = []\n",
    "    position_files_d15 = []\n",
    "    position_files_d30 = []\n",
    "    position_files_ss = []\n",
    "    \n",
    "    # Find all directories in data_dir_5fu, excluding .ims files\n",
    "    all_dirs = [os.path.join(data_dir_5fu, d) for d in os.listdir(data_dir_5fu) if d.startswith('d')] # get all the directories starting with 2\n",
    "    \n",
    "    # Classify the directories into d5, d10, d15, d30, or \"ss\"\n",
    "    for data_dir in all_dirs:\n",
    "        \n",
    "        sub_dirs = [os.path.join(data_dir, f) for f in os.listdir(data_dir) if f.endswith('Statistics')]\n",
    "        for sub_dir in sub_dirs:\n",
    "            # Get the first file ending with 'Position.csv' in each sub-directory\n",
    "            # First check for Position_Reference_Frame.csv, if not found, fallback to Position.csv\n",
    "            position_file = [os.path.join(sub_dir, f) for f in os.listdir(sub_dir) if f.endswith('Position_Reference_Frame.csv')]\n",
    "            if not position_file:\n",
    "                # Only look for Position.csv if Position_Reference_Frame.csv is not found\n",
    "                position_file = [os.path.join(sub_dir, f) for f in os.listdir(sub_dir) if f.endswith('Position.csv')]\n",
    "        \n",
    "            if position_file:\n",
    "                if 'd5' in data_dir:\n",
    "                    position_files_d5.append(position_file[0])\n",
    "                elif 'd10' in data_dir:\n",
    "                    position_files_d10.append(position_file[0])\n",
    "                elif 'd15' in data_dir:\n",
    "                    position_files_d15.append(position_file[0])\n",
    "                elif 'd30' in data_dir:\n",
    "                    position_files_d30.append(position_file[0])\n",
    "                else:\n",
    "                    # Files that don't contain 'd5', 'd10', 'd15', or 'd30' are classified as \"ss\"\n",
    "                    position_files_ss.append(position_file[0])\n",
    "\n",
    "    # Return a dictionary with the results\n",
    "    return {\n",
    "        'd5': position_files_d5,\n",
    "        'd10': position_files_d10,\n",
    "        'd15': position_files_d15,\n",
    "        'd30': position_files_d30,\n",
    "        'ss': position_files_ss\n",
    "    }\n",
    "\n",
    "# Get position files for d5, d10, d15, d30, and the \"ss\" case\n",
    "position_files_by_conditions = get_position_files_by_conditions()\n",
    "# position_files_by_conditions_hsc = get_position_files_by_conditions_hsc()\n",
    "\n",
    "# Access individual lists for each condition\n",
    "position_dirs_d5 = position_files_by_conditions['d5']\n",
    "position_dirs_d10 = position_files_by_conditions['d10']\n",
    "position_dirs_d15 = position_files_by_conditions['d15']\n",
    "position_dirs_d30 = position_files_by_conditions['d30'] \n",
    "position_dirs_ss = position_files_by_conditions['ss']\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Read the position files and create the DataFrames with columns 'Position.x', 'Position.y', 'Position.z'\n",
    "# Regular expression to capture the word that starts with 'm' and extract the part after 'm'\n",
    "pattern = r'\\bm([A-Za-z0-9]+)'\n",
    "\n",
    "def process_position_dirs(position_dirs, scale_factors):\n",
    "    \"\"\"\n",
    "    Process a list of CSV files containing position data, scales the coordinates,\n",
    "    and stores the resulting DataFrame in the provided output dictionary.\n",
    "    \n",
    "    Parameters:\n",
    "    position_dirs (list): List of file paths to CSV files containing the position data.\n",
    "    scale_factors (dict): Dictionary with scaling factors for 'x', 'y', and 'z'.\n",
    "                        e.g., {'x': 0.7575, 'y': 0.7575, 'z': 2.5}\n",
    "    Returns:\n",
    "    None: The function modifies the provided output_dict in place.\n",
    "    \"\"\"\n",
    "    output_dict = {}\n",
    "    for position_dir in position_dirs:\n",
    "        # print(position_sub_dir.split('/')[-1])\n",
    "        df = pd.read_csv(position_dir, skiprows=3)\n",
    "        \n",
    "        # Take the first 3 columns as the position coordinates\n",
    "        df = df.iloc[:, :3]\n",
    "        df.columns = ['Position.X', 'Position.Y', 'Position.Z']\n",
    "        \n",
    "        # Scale the position coordinates\n",
    "        df['Position.X'] = df['Position.X'] / scale_factors['x']\n",
    "        df['Position.Y'] = df['Position.Y'] / scale_factors['y']\n",
    "        df['Position.Z'] = df['Position.Z'] / scale_factors['z']\n",
    "        position_sub_dir = position_dir.split('/')[-1]\n",
    "        # Extract a key name for the output dictionary from the file path\n",
    "        if 'CD105' in position_sub_dir:\n",
    "            df_name_prefix = '_'.join(position_sub_dir.split(' ')[:2])\n",
    "            df_name = '_'.join([df_name_prefix, 'Sinusoids'])\n",
    "            print(df_name)\n",
    "        elif 'Sinusoids' in position_sub_dir:\n",
    "            df_name_prefix = '_'.join(position_sub_dir.split(' ')[:2])\n",
    "            df_name = '_'.join([df_name_prefix, 'Sinusoids'])\n",
    "            print(df_name)\n",
    "        elif 'Cxcl12' in position_sub_dir and 'acatCxcl12' not in position_sub_dir:\n",
    "            df_name_prefix = '_'.join(position_sub_dir.split(' ')[:2])\n",
    "            df_name = '_'.join([df_name_prefix, 'MSCs'])\n",
    "            print(df_name)\n",
    "        elif 'MSCs' in position_sub_dir and 'acatCxcl12' not in position_sub_dir:\n",
    "            df_name_prefix = '_'.join(position_sub_dir.split(' ')[:2])\n",
    "            df_name = '_'.join([df_name_prefix, 'MSCs'])\n",
    "            print(df_name)\n",
    "        elif 'HSC' in position_sub_dir:\n",
    "            df_name_prefix = '_'.join(position_sub_dir.split(' ')[:2])\n",
    "            df_name = '_'.join([df_name_prefix, 'HSCs'])\n",
    "            print(df_name)\n",
    "            # df_name = position_sub_dir.split('/')[-1].split('.')[0]\n",
    "        elif 'RD' in position_sub_dir:\n",
    "            df_name_prefix = '_'.join(position_sub_dir.split(' ')[:2])\n",
    "            df_name = '_'.join([df_name_prefix, 'RDs'])\n",
    "            print(df_name)\n",
    "        elif 'Sca1' in position_sub_dir:\n",
    "            df_name_prefix = '_'.join(position_sub_dir.split(' ')[:2])\n",
    "            df_name = '_'.join([df_name_prefix, 'Arteries'])\n",
    "            print(df_name)\n",
    "        elif 'Arteries' in position_sub_dir:\n",
    "            df_name_prefix = '_'.join(position_sub_dir.split(' ')[:2])\n",
    "            df_name = '_'.join([df_name_prefix, 'Arteries'])\n",
    "            print(df_name)\n",
    "        elif 'GFAP' in position_sub_dir:\n",
    "            df_name_prefix = '_'.join(position_sub_dir.split(' ')[:2])\n",
    "            df_name = '_'.join([df_name_prefix, 'Neurons'])\n",
    "            print(df_name)\n",
    "        elif 'MK' in position_sub_dir:\n",
    "            df_name_prefix = '_'.join(position_sub_dir.split(' ')[:2])\n",
    "            df_name = '_'.join([df_name_prefix, 'MKs'])\n",
    "            print(df_name)\n",
    "        elif 'Adipo' in position_sub_dir:\n",
    "            df_name_prefix = '_'.join(position_sub_dir.split(' ')[:2])\n",
    "            df_name = '_'.join([df_name_prefix, 'Adipos'])\n",
    "            print(df_name)\n",
    "        elif 'Neurons' in position_sub_dir:\n",
    "            df_name_prefix = '_'.join(position_sub_dir.split(' ')[:2])\n",
    "            df_name = '_'.join([df_name_prefix, 'Neurons'])\n",
    "            print(df_name)\n",
    "        elif 'Bone' in position_sub_dir or 'bone' in position_sub_dir:\n",
    "\n",
    "            df_name_prefix = '_'.join(position_sub_dir.split(' ')[:2])\n",
    "            df_name = '_'.join([df_name_prefix, 'Bones'])\n",
    "            print(df_name)\n",
    "        else:\n",
    "            df_name = '_'.join(position_sub_dir.split(' ')[:2])\n",
    "            print('speical case')\n",
    "            print(position_sub_dir)\n",
    "            print(df_name)\n",
    "        \n",
    "        # Store the processed DataFrame in the output dictionary\n",
    "        output_dict[df_name] = df\n",
    "        \n",
    "        # Print summary statistics of the DataFrame\n",
    "        # print(df.describe())\n",
    "    return output_dict\n",
    "\n",
    "\n",
    "# Define the scale factors\n",
    "scale_factors = {'x': 0.7575, 'y': 0.7575, 'z': 2.5}\n",
    "\n",
    "# Call the function for position_dirs_200901KK_st1 and position_dirs_201002KK_st3\n",
    "positions_ss = process_position_dirs(position_dirs_ss, scale_factors)\n",
    "positions_d5 = process_position_dirs(position_dirs_d5, scale_factors)\n",
    "positions_d10 = process_position_dirs(position_dirs_d10, scale_factors)\n",
    "positions_d15 = process_position_dirs(position_dirs_d15, scale_factors)\n",
    "positions_d30 = process_position_dirs(position_dirs_d30, scale_factors)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.2 Load Image Files (ims)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get the GFP images of the bones\n",
    "def get_ims_files_by_conditions(data_dir_root='data_metabone_5fu'):\n",
    "    data_dir_5fu = data_dir_root\n",
    "    # Lists to store .ims files for each condition\n",
    "    ims_files_d5 = []\n",
    "    ims_files_d10 = []\n",
    "    ims_files_d15 = []\n",
    "    ims_files_d30 = []\n",
    "    ims_files_ss = []\n",
    "    \n",
    "    # Find all directories in data_dir_5fu, excluding .ims files\n",
    "    all_dirs = [os.path.join(data_dir_5fu, d) for d in os.listdir(data_dir_5fu) if d.startswith('d')] \n",
    "    \n",
    "    # Classify the directories into d5, d10, d15, d30, or \"ss\"\n",
    "    for data_dir in all_dirs:\n",
    "        sub_dirs = [os.path.join(data_dir, f) for f in os.listdir(data_dir) if f.endswith('ims') and 'GFP' in f]\n",
    "        if 'd5' in data_dir:\n",
    "            ims_files_d5.extend(sub_dirs)\n",
    "        elif 'd10' in data_dir:\n",
    "            ims_files_d10.extend(sub_dirs)\n",
    "        elif 'd15' in data_dir:\n",
    "            ims_files_d15.extend(sub_dirs)\n",
    "        elif 'd30' in data_dir:\n",
    "            ims_files_d30.extend(sub_dirs)\n",
    "        else:\n",
    "            # Files that don't contain 'd5', 'd10', 'd15', or 'd30' are classified as \"ss\"\n",
    "            ims_files_ss.extend(sub_dirs)\n",
    "\n",
    "    # Return a dictionary with the results\n",
    "    return {\n",
    "        'd5': ims_files_d5,\n",
    "        'd10': ims_files_d10,\n",
    "        'd15': ims_files_d15,\n",
    "        'd30': ims_files_d30,\n",
    "        'ss': ims_files_ss\n",
    "    }\n",
    "\n",
    "\n",
    "# Filter the ims files, keep the one with shorter name\n",
    "def filter_ims_files(ims_list):\n",
    "    # Dictionary to hold the shortest file for each unique prefix\n",
    "    grouped_files = defaultdict(list)\n",
    "\n",
    "    # Group filenames by the unique prefix after `d0/`\n",
    "    for filepath in ims_list:\n",
    "        # Extract the unique prefix (first two elements after 'd0/')\n",
    "        parts = filepath.split('/')\n",
    "        if len(parts) >= 3:\n",
    "            unique_prefix = parts[-1].split()[0] + ' ' + parts[-1].split()[1]  # Example: '200901KK st1'\n",
    "            grouped_files[unique_prefix].append(filepath)\n",
    "\n",
    "    # Filter to keep only the shortest name in each group\n",
    "    filtered_bone_ims = []\n",
    "    for prefix, files in grouped_files.items():\n",
    "        # Sort by filename length and take the shortest\n",
    "        shortest_file = min(files, key=len)\n",
    "        filtered_bone_ims.append(shortest_file)\n",
    "    return filtered_bone_ims\n",
    "        \n",
    "\n",
    "\n",
    "\n",
    "# Convert the bone images with non-zero values to dataframes with Position.X and Position.Y columns\n",
    "def image_to_df(image):\n",
    "    # Find the non-zero intensities and their coordinates\n",
    "    non_zero_coords = np.nonzero(image)  # Returns the y, x coordinates where intensity > 0\n",
    "    # intensity_values = image[non_zero_coords]  # Extract the intensity values\n",
    "\n",
    "    # Create a DataFrame from the non-zero coordinates and intensity values\n",
    "    df = pd.DataFrame({\n",
    "        'source': 'GFP',\n",
    "        'Position.Z': 0, #non_zero_coords[0],\n",
    "        'Position.Y': non_zero_coords[0],\n",
    "        'Position.X': non_zero_coords[1]\n",
    "    })\n",
    "    \n",
    "    return df\n",
    "\n",
    "\n",
    "\n",
    "\n",
    "# Function to process and store images for a given condition\n",
    "def process_bone_images(bone_ims_files):\n",
    "    bone_dict = {}\n",
    "    \n",
    "    for i in bone_ims_files:\n",
    "        # Extract bone name from the file path\n",
    "        bone_name = '_'.join(i.split('/')[-1].split(' ')[:2])\n",
    "        \n",
    "        # Load and process the image\n",
    "        bone_img = ims(i)  # You need to ensure that ims(i) correctly loads the image\n",
    "        # bone_img = bone_img[0, bone_img.shape[1] - 1, :, :, :].max(axis=0).copy()\n",
    "        bone_img = bone_img[0, 0, :, :, :].max(axis=0).copy()\n",
    "        # Morphological operations, labeling and remove small objects\n",
    "        # Threshold the image to create a binary mask\n",
    "        binary_img = bone_img > 0\n",
    "        # Label the binary image\n",
    "        labeled_img = measure.label(binary_img)\n",
    "        # Remove small objects\n",
    "        cleaned_img = morphology.remove_small_objects(labeled_img, min_size=10000)\n",
    "        # Convert the cleaned image back to binary\n",
    "        bone_img_df = image_to_df(cleaned_img)\n",
    "        # Store the processed image in the dictionary with the extracted name\n",
    "        bone_dict[bone_name] = bone_img_df\n",
    "\n",
    "    return bone_dict\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Access individual lists for each condition\n",
    "ims_files_by_conditions = get_ims_files_by_conditions()\n",
    "bone_ss_ims = ims_files_by_conditions['ss']\n",
    "bone_d5_ims = ims_files_by_conditions['d5']\n",
    "bone_d10_ims = ims_files_by_conditions['d10']\n",
    "bone_d15_ims = ims_files_by_conditions['d15']\n",
    "bone_d30_ims = ims_files_by_conditions['d30']\n",
    "\n",
    "\n",
    "# Filter the ims files, keep the one with shorter name\n",
    "filtered_bone_ss_ims = filter_ims_files(bone_ss_ims)\n",
    "filtered_bone_d5_ims = filter_ims_files(bone_d5_ims)\n",
    "filtered_bone_d10_ims = filter_ims_files(bone_d10_ims)\n",
    "filtered_bone_d15_ims = filter_ims_files(bone_d15_ims)\n",
    "filtered_bone_d30_ims = filter_ims_files(bone_d30_ims)\n",
    "\n",
    "# Process images for each condition\n",
    "bone_ss = process_bone_images(filtered_bone_ss_ims)\n",
    "bone_d5 = process_bone_images(filtered_bone_d5_ims)\n",
    "bone_d10 = process_bone_images(filtered_bone_d10_ims)\n",
    "bone_d15 = process_bone_images(filtered_bone_d15_ims)\n",
    "bone_d30 = process_bone_images(filtered_bone_d30_ims)\n",
    "\n",
    "\n",
    "del bone_ss_ims, bone_d5_ims, bone_d10_ims, bone_d15_ims, bone_d30_ims"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.3 Convert Image to Dataframe and Combine both Dataframes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# From the positions_ss, positions_d5, positions_d10, positions_d15, positions_d30\n",
    "def positions_concat(positions_dict, bone_dict):\n",
    "    positions_df_dict = dict()\n",
    "    keys = {}\n",
    "    \n",
    "    # Create a mapping of dataset and sources\n",
    "    for key in positions_dict.keys():\n",
    "        dataset = '_'.join(key.split('_')[:2])\n",
    "        source = key.split('_')[-1]\n",
    "        if dataset not in keys:\n",
    "            keys[dataset] = [source]\n",
    "        else:\n",
    "            keys[dataset].append(source)\n",
    "    \n",
    "    # Concatenate DataFrames, adding a 'source' column\n",
    "    for dataset, sources in keys.items():\n",
    "        df = pd.concat(\n",
    "            [positions_dict[f'{dataset}_{source}'].assign(source=source) for source in sources], \n",
    "            ignore_index=True\n",
    "        )\n",
    "        df = pd.concat([df, bone_dict[dataset]], ignore_index=True)\n",
    "        positions_df_dict[dataset] = df\n",
    "    return positions_df_dict\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "positions_ss_dfs = positions_concat(positions_ss, bone_ss)\n",
    "positions_d5_dfs = positions_concat(positions_d5, bone_d5)\n",
    "positions_d10_dfs = positions_concat(positions_d10, bone_d10)\n",
    "# delete bone_ss, bone_d5, bone_d10, bone_d15, bone_d30 to free up memory\n",
    "del bone_ss, bone_d5, bone_d10\n",
    "del filtered_bone_ss_ims, filtered_bone_d5_ims, filtered_bone_d10_ims"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "positions_d15_dfs = positions_concat(positions_d15, bone_d15)\n",
    "positions_d30_dfs = positions_concat(positions_d30, bone_d30)\n",
    "\n",
    "\n",
    "# delete bone_ss, bone_d5, bone_d10, bone_d15, bone_d30 to free up memory\n",
    "del bone_d15, bone_d30 \n",
    "del filtered_bone_d15_ims, filtered_bone_d30_ims"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1.4 Load Inhibitor Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_dir_inhibitor_Detex = 'data_metabone_inhibitors/metabone Detex inhibitor data'\n",
    "\n",
    "position_files_inhibitor_Detex = get_position_files_by_conditions(data_dir_root=data_dir_inhibitor_Detex)\n",
    "\n",
    "position_dirs_inhibitor_Detex_d15 = position_files_inhibitor_Detex['d15']\n",
    "\n",
    "ims_files_inhibitor_Detex = get_ims_files_by_conditions(data_dir_root=data_dir_inhibitor_Detex)\n",
    "\n",
    "bone_d15_inhibitor_Detex = process_bone_images(filter_ims_files(ims_files_inhibitor_Detex['d15']))\n",
    "\n",
    "# Call the function for position_dirs_200901KK_st1 and position_dirs_201002KK_st3\n",
    "positions_d15_inhibitor_Detex = process_position_dirs(position_dirs_inhibitor_Detex_d15, scale_factors)\n",
    "positions_d15_inhibitor_Detex_dfs = positions_concat(positions_d15_inhibitor_Detex, bone_d15_inhibitor_Detex)\n",
    "\n",
    "del bone_d15_inhibitor_Detex, ims_files_inhibitor_Detex\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "source": [
    "### 1.5 Load Reference Bone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "reference_bone = ims(reference_bone_dir)\n",
    "\n",
    "ref_bone_outer = reference_bone[0,0,:,:,:].max(axis=0).copy()\n",
    "ref_bone_inner = reference_bone[0,1,:,:,:].max(axis=0).copy()\n",
    "\n",
    "ref_bone_outer = image_to_df(ref_bone_outer)\n",
    "ref_bone_inner = image_to_df(ref_bone_inner)\n",
    "\n",
    "# Flip the y-axis to match the image coordinates\n",
    "y_max = np.max((ref_bone_outer['Position.Y'].max(), ref_bone_inner['Position.Y'].max()))\n",
    "# y_max = ref_bone_inner['Position.Y'].max()\n",
    "ref_bone_outer['Position.Y'] = y_max - ref_bone_outer['Position.Y']\n",
    "ref_bone_inner['Position.Y'] = y_max - ref_bone_inner['Position.Y']\n",
    "\n",
    "# delete reference bone to save memory\n",
    "del reference_bone"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Data Inspection and Adjustments"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Visual Inspection"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the x, y for each df \n",
    "# for test_df in [positions_ss_dfs, positions_d5_dfs, positions_d10_dfs, positions_d15_dfs, positions_d30_dfs]:\n",
    "for test_df in[positions_ss_dfs, positions_d5_dfs, positions_d10_dfs, positions_d15_dfs, positions_d30_dfs, positions_d15_inhibitor_Detex_dfs]:\n",
    "    fig, axs = plt.subplots(len(test_df.keys()), 9, figsize=(12*9, 4*len(test_df.keys())))\n",
    "\n",
    "    # Visualize the steady state positions row by key and the columns are source\n",
    "    for i, key in enumerate(test_df.keys()):\n",
    "        for j, source in enumerate(test_df[key]['source'].unique()):\n",
    "            if source == 'GFP':\n",
    "                sample_number = 10000\n",
    "                source_df = test_df[key].loc[test_df[key]['source'] == source, ['Position.X', 'Position.Y']].sample(sample_number)\n",
    "            else:\n",
    "                source_df = test_df[key].loc[test_df[key]['source'] == source, ['Position.X', 'Position.Y']]\n",
    "            if len(test_df.keys()) == 1:\n",
    "                ax = axs[j]\n",
    "            else:\n",
    "                ax = axs[i, j]\n",
    "            ax.scatter(source_df['Position.X'], source_df['Position.Y'], s=1)\n",
    "            ax.set_title(f\"{key} - {source}\")\n",
    "            ax.axis('equal')\n",
    "            ax.set_xlabel('Position.X')\n",
    "            ax.set_ylabel('Position.Y')\n",
    "            \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "vscode": {
     "languageId": "plaintext"
    }
   },
   "source": [
    "### 2.2 Data Adjustments (Flipping Coordinates, etc.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Based on the visual check above, we need to flip the y axis for dataset with key \n",
    "# The flip should be applied to all sources\n",
    "def y_flip(positions_df_dict, keys_to_flip):\n",
    "    for key in keys_to_flip:\n",
    "        # y_min = positions_df_dict[key]['Position.Y'].min()\n",
    "        y_max = positions_df_dict[key]['Position.Y'].max()\n",
    "        # Flip the y axis, we don't care about the source\n",
    "        positions_df_dict[key]['Position.Y'] = y_max - positions_df_dict[key]['Position.Y']\n",
    "    return positions_df_dict\n",
    "\n",
    "positions_ss_dfs = y_flip(positions_ss_dfs, ['201002KK_st5', '200901KK_st3', '210518KK_st12', '210518KK_st1', '200728KK_st2'])\n",
    "positions_d5_dfs = y_flip(positions_d5_dfs, ['200117KK_st3', '200117KK_st5', '200117KK_st4','200117KK_st1'])\n",
    "positions_d10_dfs = y_flip(positions_d10_dfs, ['200110KK_st3','200110KK_st1', '191202KK_st1', '200110KK_st4', '200110KK_st8'])\n",
    "positions_d15_dfs = y_flip(positions_d15_dfs, ['240819JH_st2', '201002KK_st7', '201020KK_st1', '200901KK_st5', '200728KK_st3', '200908KK_st4', '210108KK_st10'])\n",
    "positions_d30_dfs = y_flip(positions_d30_dfs, ['200901KK_st7', '201103KKI_st4', '201103KK_st2', '200908KK_st2', '240926JH_st2'])\n",
    "\n",
    "def x_flip(positions_df_dict, keys_to_flip):\n",
    "    for key in keys_to_flip:\n",
    "        # x_min = positions_df_dict[key]['Position.X'].min()\n",
    "        x_max = positions_df_dict[key]['Position.X'].max()\n",
    "        # Flip the x axis, we don't care about the source\n",
    "        positions_df_dict[key]['Position.X'] = x_max - positions_df_dict[key]['Position.X']\n",
    "    return positions_df_dict\n",
    "\n",
    "positions_d5_dfs = x_flip(positions_d5_dfs, ['200117KK_st5'])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "positions_d15_inhibitor_Detex_dfs = y_flip(positions_d15_inhibitor_Detex_dfs, ['201209KK_st4', '201121KK_st4'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the x, y for each df again after flipping the coordinates\n",
    "# for test_df in [positions_ss_dfs, positions_d5_dfs, positions_d10_dfs, positions_d15_dfs, positions_d30_dfs]:\n",
    "for test_df in [positions_ss_dfs, positions_d5_dfs, positions_d10_dfs, positions_d15_dfs, positions_d30_dfs, positions_d15_inhibitor_Detex_dfs]:\n",
    "    fig, axs = plt.subplots(len(test_df.keys()), 9, figsize=(12*9, 4*len(test_df.keys())))\n",
    "\n",
    "    # Visualize the steady state positions row by key and the columns are source\n",
    "    for i, key in enumerate(test_df.keys()):\n",
    "        for j, source in enumerate(test_df[key]['source'].unique()):\n",
    "            if source == 'GFP':\n",
    "                sample_number = 10000\n",
    "                source_df = test_df[key].loc[test_df[key]['source'] == source, ['Position.X', 'Position.Y']].sample(sample_number)\n",
    "            else:\n",
    "                source_df = test_df[key].loc[test_df[key]['source'] == source, ['Position.X', 'Position.Y']]\n",
    "            if len(test_df.keys()) == 1:\n",
    "                ax = axs[j]\n",
    "            else:\n",
    "                ax = axs[i, j]\n",
    "            ax.scatter(source_df['Position.X'], source_df['Position.Y'], s=1)\n",
    "            ax.set_title(f\"{key} - {source}\")\n",
    "            ax.axis('equal')\n",
    "            ax.set_xlabel('Position.X')\n",
    "            ax.set_ylabel('Position.Y')\n",
    "            \n",
    "    plt.tight_layout()\n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Bone Alignment and Transformation"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.1 Align Bones to the Reference Bone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Find the 2D outline of the bone in the x-y plane\n",
    "def find_outline(points, window_size=10):\n",
    "    \"\"\"\n",
    "    Find the outline of a set of points by finding the min and max y-values for each x-value within a window.\n",
    "    The outline is only in the x-y plane.\n",
    "    Parameters:\n",
    "    points : np.array of shape (n,2) points.\n",
    "    window_size : size of the window to smooth the outline.\n",
    "    Returns:\n",
    "    outline_points : np.array of outline points.\n",
    "    \"\"\"\n",
    "    df = pd.DataFrame(points, columns=['x', 'y'])\n",
    "    \n",
    "    min_y_points = []\n",
    "    max_y_points = []\n",
    "\n",
    "    # Sort points by x value\n",
    "    df_sorted = df.sort_values(by='x')\n",
    "    \n",
    "    # Slide over the x values with a window\n",
    "    for i in range(0, len(df_sorted), window_size):\n",
    "        window = df_sorted.iloc[i:i + window_size]\n",
    "        min_y = window.loc[window['y'].idxmin()]\n",
    "        max_y = window.loc[window['y'].idxmax()]\n",
    "        min_y_points.append(min_y)\n",
    "        max_y_points.append(max_y)\n",
    "    \n",
    "    # Ensure the outline is in order\n",
    "    min_y_points = pd.DataFrame(min_y_points).drop_duplicates().sort_values(by='x').values\n",
    "    max_y_points = pd.DataFrame(max_y_points).drop_duplicates().sort_values(by='x', ascending=False).values\n",
    "\n",
    "    # Combine min_y and max_y points and close the loop\n",
    "    outline_points = np.vstack([min_y_points, max_y_points, min_y_points[0]])\n",
    "\n",
    "    return outline_points\n",
    "\n",
    "# 0. Calculate the overlap\n",
    "# 1. Find the outline of each bone based on the x-y plane\n",
    "\n",
    "# 2. Find the center of each bone (outline) and put the center of the bones at the same position\n",
    "def calculate_centroid(outline_points):\n",
    "    \"\"\"\n",
    "    Calculate the centroid of the bone outline.\n",
    "    \n",
    "    Parameters:\n",
    "    outline_points : np.array of shape (n, 2)\n",
    "    \n",
    "    Returns:\n",
    "    centroid : tuple containing (centroid_x, centroid_y)\n",
    "    \"\"\"\n",
    "    # Use the weight centroid\n",
    "    centroid_x = np.mean(outline_points[:, 0])\n",
    "    centroid_y = np.mean(outline_points[:, 1])\n",
    "    \n",
    "    # Using the middle of the x and y values as the centroid\n",
    "    # x_min, x_max = outline_points[:, 0].min(), outline_points[:, 0].max()\n",
    "    # y_min, y_max = outline_points[:, 1].min(), outline_points[:, 1].max()\n",
    "    # centroid_x = (x_min + x_max) / 2\n",
    "    # centroid_y = (y_min + y_max) / 2\n",
    "    return centroid_x, centroid_y\n",
    "def translate_to_origin(outline_points, centroid):\n",
    "    \"\"\"\n",
    "    Translate the outline points so that the centroid is at the origin.\n",
    "    \n",
    "    Parameters:\n",
    "    outline_points : np.array of shape (n, 2)\n",
    "    centroid : tuple containing (centroid_x, centroid_y)\n",
    "    \n",
    "    Returns:\n",
    "    translated_points : np.array of shape (n, 2)\n",
    "    \"\"\"\n",
    "    translated_points = outline_points.astype(np.float64).copy()\n",
    "    \n",
    "    translated_points[:, 0] -= centroid[0]\n",
    "    translated_points[:, 1] -= centroid[1]\n",
    "    return translated_points\n",
    "# 3. Rescale the bones to the same size (bounding box)(optional)\n",
    "def get_max_dimensions(bone_dicts):\n",
    "    \"\"\"\n",
    "    Find the maximum width and height across all bone outlines in the given dictionaries.\n",
    "    \n",
    "    Parameters:\n",
    "    bone_dicts : list of dictionaries of bones (where each value is a DataFrame with 'Position.X' and 'Position.Y')\n",
    "    \n",
    "    Returns:\n",
    "    max_width : float, maximum width found across all bones\n",
    "    max_height : float, maximum height found across all bones\n",
    "    \"\"\"\n",
    "    max_width = 0\n",
    "    max_height = 0\n",
    "    \n",
    "    for bone_dict in bone_dicts:\n",
    "        for df in bone_dict.values():\n",
    "            # Extract bone points where 'source' == 'Bone'\n",
    "            bone_points = df[df['source'] == 'Bone'][['Position.X', 'Position.Y']].values\n",
    "            \n",
    "            # Find min and max values of x and y\n",
    "            min_x, max_x = bone_points[:, 0].min(), bone_points[:, 0].max()\n",
    "            min_y, max_y = bone_points[:, 1].min(), bone_points[:, 1].max()\n",
    "            \n",
    "            # Calculate width and height\n",
    "            width = max_x - min_x\n",
    "            height = max_y - min_y\n",
    "            \n",
    "            # Update maximum width and height if necessary\n",
    "            if width > max_width:\n",
    "                max_width = width\n",
    "            if height > max_height:\n",
    "                max_height = height\n",
    "                \n",
    "    return max_width, max_height\n",
    "def rescale_outline(outline_points, max_width, max_height):\n",
    "    \"\"\"\n",
    "    Rescale the bone outline to fit within the maximum width and height across all bones.\n",
    "    \n",
    "    Parameters:\n",
    "    outline_points : np.array of shape (n, 2)\n",
    "    max_width : float, the maximum width across all bones\n",
    "    max_height : float, the maximum height across all bones\n",
    "    \n",
    "    Returns:\n",
    "    scaled_points : np.array of shape (n, 2)\n",
    "    \"\"\"\n",
    "    min_x, max_x = outline_points[:, 0].min(), outline_points[:, 0].max()\n",
    "    min_y, max_y = outline_points[:, 1].min(), outline_points[:, 1].max()\n",
    "\n",
    "    current_width = max_x - min_x\n",
    "    current_height = max_y - min_y\n",
    "\n",
    "    scale_x = max_width / current_width\n",
    "    scale_y = max_height / current_height\n",
    "\n",
    "    scaled_points = outline_points.copy()\n",
    "    scaled_points[:, 0] *= scale_x\n",
    "    scaled_points[:, 1] *= scale_y\n",
    "\n",
    "    return scaled_points\n",
    "\n",
    "# Calculate the overlap area on the grid inside the outline\n",
    "def calculate_overlap_area(reference_outline, target_outline, resolution=200):\n",
    "    \"\"\"\n",
    "    Calculate the overlap area (in terms of pixels or points) between two outlines.\n",
    "    \n",
    "    Parameters:\n",
    "    reference_outline : np.array of shape (n, 2), outline of the reference bone\n",
    "    target_outline : np.array of shape (n, 2), outline of the target bone\n",
    "    resolution : int, the number of points or pixels to use for the area calculation.\n",
    "    \n",
    "    Returns:\n",
    "    overlap_area : float, the number of pixels or points where the areas overlap.\n",
    "    \"\"\"\n",
    "    # Ensure the outlines are 2D arrays of shape (N, 2)\n",
    "    reference_outline = np.asarray(reference_outline).reshape(-1, 2)\n",
    "    target_outline = np.asarray(target_outline).reshape(-1, 2)\n",
    "    \n",
    "    # Calculate centroids for both bones\n",
    "    reference_centroid = calculate_centroid(reference_outline)\n",
    "    target_centroid = calculate_centroid(target_outline)\n",
    "    \n",
    "    # Translate both bones to center them\n",
    "    reference_outline_centered = translate_to_origin(reference_outline, reference_centroid)\n",
    "    target_outline_centered = translate_to_origin(target_outline, target_centroid)\n",
    "    \n",
    "    # Get bounding box of the reference outline\n",
    "    min_x, max_x = reference_outline_centered[:, 0].min(), reference_outline_centered[:, 0].max()\n",
    "    min_y, max_y = reference_outline_centered[:, 1].min(), reference_outline_centered[:, 1].max()\n",
    "    \n",
    "    # Generate grid of points (pixels) covering the bounding box\n",
    "    x_grid = np.linspace(min_x, max_x, resolution)\n",
    "    y_grid = np.linspace(min_y, max_y, resolution)\n",
    "    xv, yv = np.meshgrid(x_grid, y_grid)\n",
    "    grid_points = np.vstack([xv.ravel(), yv.ravel()]).T\n",
    "\n",
    "    # Create Path objects for the reference and target outlines\n",
    "    reference_path = Path(reference_outline_centered)\n",
    "    target_path = Path(target_outline_centered)\n",
    "    \n",
    "    # Check which points of the grid are inside both outlines\n",
    "    points_in_reference = reference_path.contains_points(grid_points)\n",
    "    points_in_target = target_path.contains_points(grid_points)\n",
    "    \n",
    "    # Calculate the overlap area as the number of points (pixels) inside both outlines\n",
    "    overlap_area = np.sum(points_in_reference & points_in_target)\n",
    "    \n",
    "    return overlap_area\n",
    "\n",
    "def rotate_points(points, angle):\n",
    "    \"\"\"\n",
    "    Rotate a set of points by a given angle.\n",
    "    \n",
    "    Parameters:\n",
    "    points : np.array of shape (n, 2)\n",
    "    angle : float, angle to rotate by in radians\n",
    "    \n",
    "    Returns:\n",
    "    rotated_points : np.array of shape (n, 2)\n",
    "    \"\"\"\n",
    "    # Ensure points are a 2D array of shape (N, 2)\n",
    "    points = np.asarray(points).reshape(-1, 2)\n",
    "    \n",
    "    # Rotation matrix\n",
    "    rotation_matrix = np.array([[np.cos(angle), -np.sin(angle)], \n",
    "                                [np.sin(angle), np.cos(angle)]])\n",
    "    \n",
    "    # Rotate points\n",
    "    rotated_points = points.dot(rotation_matrix)\n",
    "    \n",
    "    return rotated_points\n",
    "\n",
    "# we are using the grid search\n",
    "def optimize_rotation(reference_outline, target_outline):\n",
    "    \"\"\"\n",
    "    Optimize the rotation of the target bone to maximize overlap with the reference bone,\n",
    "    constrained to -90 to 90 degrees.\n",
    "    \n",
    "    Parameters:\n",
    "    reference_outline : np.array of shape (n, 2), outline of the reference bone\n",
    "    target_outline : np.array of shape (n, 2), outline of the target bone\n",
    "    \n",
    "    Returns:\n",
    "    optimal_rotated_points : np.array of shape (n, 2), the rotated and recentered target bone outline with maximum overlap\n",
    "    optimal_angle : float, the optimal rotation angle in radians\n",
    "    \"\"\"\n",
    "    \n",
    "    def objective_function(angle):\n",
    "        # Rotate the target outline by the current angle\n",
    "        rotated_outline = rotate_points(target_outline, angle)\n",
    "        # Calculate the negative overlap (since we are minimizing)\n",
    "        return -calculate_overlap_area(reference_outline, rotated_outline)\n",
    "\n",
    "    # Try different initial guesses for the angle to improve convergence\n",
    "    initial_guesses = [-0.1, 0, 0.1]  # You can try different values here\n",
    "    \n",
    "    best_result = None\n",
    "    best_angle = None\n",
    "    best_rotated_points = None\n",
    "\n",
    "    for x0 in initial_guesses:\n",
    "        # Perform optimization with the initial guess x0\n",
    "        result = minimize(objective_function, x0=x0, bounds=[(-np.pi/4, np.pi/4)], method='L-BFGS-B', options={'maxiter': 50})\n",
    "        \n",
    "        # Check if this result is valid and better than the current best\n",
    "        if result.success:\n",
    "            if best_result is None or result.fun < best_result:\n",
    "                best_result = result.fun\n",
    "                best_angle = result.x[0]\n",
    "                best_rotated_points = rotate_points(target_outline, best_angle)\n",
    "\n",
    "    if best_rotated_points is None:\n",
    "        raise ValueError(\"Optimization failed to converge.\")\n",
    "\n",
    "    # Recenter the final rotated outline\n",
    "    final_centroid = calculate_centroid(best_rotated_points)\n",
    "    best_rotated_points = translate_to_origin(best_rotated_points, final_centroid)\n",
    "    \n",
    "    return best_rotated_points, best_angle, final_centroid\n",
    "\n",
    "def grid_search_rotation(reference_outline, target_outline, angle_step=np.pi/36):\n",
    "    \"\"\"\n",
    "    Perform a grid search over possible rotation angles to maximize overlap.\n",
    "    \n",
    "    Parameters:\n",
    "    reference_outline : np.array of shape (n, 2), outline of the reference bone\n",
    "    target_outline : np.array of shape (n, 2), outline of the target bone\n",
    "    angle_step : float, step size for angle search (in radians)\n",
    "    \n",
    "    Returns:\n",
    "    best_rotated_points : np.array of shape (n, 2), the rotated target bone outline with maximum overlap\n",
    "    best_angle : float, the optimal rotation angle in radians\n",
    "    \"\"\"\n",
    "    best_angle = None\n",
    "    max_overlap = -np.inf\n",
    "    best_rotated_points = None\n",
    "    \n",
    "    # Iterate over angles between -90 and 90 degrees (in radians)\n",
    "    for angle in np.arange(-np.pi/4, np.pi/4, angle_step):\n",
    "        rotated_outline = rotate_points(target_outline, angle)\n",
    "        overlap = calculate_overlap_area(reference_outline, rotated_outline)\n",
    "        \n",
    "        if overlap > max_overlap:\n",
    "            max_overlap = overlap\n",
    "            best_angle = angle\n",
    "            best_rotated_points = rotated_outline\n",
    "    \n",
    "    # Recenter the final rotated outline\n",
    "    final_centroid = calculate_centroid(best_rotated_points)\n",
    "    best_rotated_points = translate_to_origin(best_rotated_points, final_centroid)\n",
    "    \n",
    "    return best_rotated_points, best_angle, final_centroid\n",
    "\n",
    "# Modify the process_and_align_bones function to accept the max_width and max_height\n",
    "def process_and_align_bones_with_overlap(bone_dict, reference_bone, window_size=500, source_col = 'GFP'):\n",
    "    \"\"\"\n",
    "    Process and align all bones from the given dictionary to maximize overlap with a reference bone.\n",
    "    \n",
    "    Parameters:\n",
    "    bone_dict : dict of bones, where each value is a DataFrame with 'Position.X' and 'Position.Y'.\n",
    "    reference_bone_name : string, the name of the bone to use as the reference for alignment.\n",
    "    \n",
    "    Returns:\n",
    "    aligned_bones : dict of aligned bone outlines\n",
    "    \"\"\"\n",
    "    aligned_bones = {}\n",
    "    aligned_angles = {}\n",
    "    aligned_centroids = {}\n",
    "    # Check the type of the reference_bone_name\n",
    "    # If it is a dataframe, we can use the reference_bone_name directly\n",
    "    if isinstance(reference_bone, pd.DataFrame):\n",
    "        reference_df = reference_bone\n",
    "        reference_points = reference_df[reference_df['source'] == source_col][['Position.X', 'Position.Y']].values\n",
    "        reference_outline = find_outline(reference_points, window_size=window_size)\n",
    "        reference_centroid = calculate_centroid(reference_outline)\n",
    "        reference_outline = translate_to_origin(reference_outline, reference_centroid)\n",
    "        \n",
    "        for bone_name, df in bone_dict.items():\n",
    "\n",
    "            # Filter points where 'source' == 'Bone'\n",
    "            bone_points = df[df['source'] == source_col][['Position.X', 'Position.Y']].values\n",
    "        \n",
    "            # Find the outline of the target bone\n",
    "            target_outline = find_outline(bone_points, window_size=window_size)\n",
    "            # Optimize rotation to maximize overlap with the reference bone\n",
    "            aligned_outline, best_angle, final_centroid = grid_search_rotation(reference_outline, target_outline)\n",
    "            aligned_bones[bone_name] = aligned_outline\n",
    "            aligned_angles[bone_name] = best_angle\n",
    "            aligned_centroids[bone_name] = final_centroid\n",
    "    else:\n",
    "        raise ValueError(\"Invalid reference_bone_name. Must be a DataFrame.\")\n",
    "\n",
    "    return aligned_bones, aligned_angles, aligned_centroids\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Compare the outline found based on the ref_bone_outer and the ref_bone_inner\n",
    "# recenter and find the outline\n",
    "ref_bone_outer_points = ref_bone_outer[ref_bone_outer['source'] == 'GFP'][['Position.X', 'Position.Y']].values\n",
    "ref_bone_inner_points = ref_bone_inner[ref_bone_inner['source'] == 'GFP'][['Position.X', 'Position.Y']].values\n",
    "\n",
    "ref_bone_outer_outline = find_outline(ref_bone_outer_points, window_size=500)\n",
    "ref_bone_inner_outline = find_outline(ref_bone_inner_points, window_size=500)\n",
    "ref_bone_outer_centroid = calculate_centroid(ref_bone_outer_outline)\n",
    "\n",
    "# The two shared the same centroid\n",
    "ref_bone_outer_outline = translate_to_origin(ref_bone_outer_outline, ref_bone_outer_centroid)\n",
    "ref_bone_inner_outline = translate_to_origin(ref_bone_inner_outline, ref_bone_outer_centroid)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(15, 5))\n",
    "ax.plot(ref_bone_outer_outline[:, 0], ref_bone_outer_outline[:, 1], color='r', label='Outer')\n",
    "ax.plot(ref_bone_inner_outline[:, 0], ref_bone_inner_outline[:, 1], color='b', label='Inner')\n",
    "ax.set_title('Outer and Inner Bone Outlines')\n",
    "\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Apply the process to all bone dictionaries using the max width and height\n",
    "aligned_ss_bone_outlines, aligned_ss_angles, aligned_ss_centroids = process_and_align_bones_with_overlap(positions_ss_dfs, reference_bone=ref_bone_inner, window_size=500)\n",
    "aligned_d5_bone_outlines, aligned_d5_angles, aligned_d5_centroids= process_and_align_bones_with_overlap(positions_d5_dfs, reference_bone=ref_bone_inner, window_size=500)\n",
    "aligned_d10_bone_outlines, aligned_d10_angles, aligned_d10_centroids = process_and_align_bones_with_overlap(positions_d10_dfs, reference_bone=ref_bone_inner, window_size=500)\n",
    "\n",
    "# 49 mins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "aligned_d15_bone_outlines, aligned_d15_angles, aligned_d15_centroids = process_and_align_bones_with_overlap(positions_d15_dfs, reference_bone=ref_bone_inner, window_size=500)\n",
    "aligned_d30_bone_outlines, aligned_d30_angles, aligned_d30_centroids = process_and_align_bones_with_overlap(positions_d30_dfs, reference_bone=ref_bone_inner, window_size=500)\n",
    "# 37 mins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "aligned_bone_data = {\n",
    "    'aligned_ss_bone_outlines': aligned_ss_bone_outlines,\n",
    "    'aligned_ss_angles': aligned_ss_angles,\n",
    "    'aligned_ss_centroids': aligned_ss_centroids,\n",
    "    'aligned_d5_bone_outlines': aligned_d5_bone_outlines,\n",
    "    'aligned_d5_angles': aligned_d5_angles,\n",
    "    'aligned_d5_centroids': aligned_d5_centroids,\n",
    "    'aligned_d10_bone_outlines': aligned_d10_bone_outlines,\n",
    "    'aligned_d10_angles': aligned_d10_angles,\n",
    "    'aligned_d10_centroids': aligned_d10_centroids,\n",
    "    'aligned_d15_bone_outlines': aligned_d15_bone_outlines,\n",
    "    'aligned_d15_angles': aligned_d15_angles,\n",
    "    'aligned_d15_centroids': aligned_d15_centroids,\n",
    "    'aligned_d30_bone_outlines': aligned_d30_bone_outlines,\n",
    "    'aligned_d30_angles': aligned_d30_angles,\n",
    "    'aligned_d30_centroids': aligned_d30_centroids,\n",
    "}\n",
    "\n",
    "# Save the aligned bone data to a pickle file\n",
    "aligned_bone_data_filename = os.path.join(save_dir, f'{datetime.today().strftime(\"%y%m%d\")}_YL_aligned_bone_data.pkl')\n",
    "with open(aligned_bone_data_filename, 'wb') as file:\n",
    "    pickle.dump(aligned_bone_data, file)\n",
    "    \n",
    "del aligned_bone_data, aligned_bone_data_filename\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "aligned_bone_data_filename = f'{save_dir}/241219_YL_aligned_bone_data.pkl'\n",
    "aligned_bone_data = pickle.load(open(aligned_bone_data_filename, 'rb'))\n",
    "\n",
    "aligned_ss_bone_outlines = aligned_bone_data['aligned_ss_bone_outlines']\n",
    "aligned_ss_angles = aligned_bone_data['aligned_ss_angles']\n",
    "aligned_ss_centroids = aligned_bone_data['aligned_ss_centroids']\n",
    "\n",
    "aligned_d5_bone_outlines = aligned_bone_data['aligned_d5_bone_outlines']\n",
    "aligned_d5_angles = aligned_bone_data['aligned_d5_angles']\n",
    "aligned_d5_centroids = aligned_bone_data['aligned_d5_centroids']\n",
    "\n",
    "aligned_d10_bone_outlines = aligned_bone_data['aligned_d10_bone_outlines']\n",
    "aligned_d10_angles = aligned_bone_data['aligned_d10_angles']\n",
    "aligned_d10_centroids = aligned_bone_data['aligned_d10_centroids']\n",
    "\n",
    "aligned_d15_bone_outlines = aligned_bone_data['aligned_d15_bone_outlines']\n",
    "aligned_d15_angles = aligned_bone_data['aligned_d15_angles']\n",
    "aligned_d15_centroids = aligned_bone_data['aligned_d15_centroids']\n",
    "\n",
    "aligned_d30_bone_outlines = aligned_bone_data['aligned_d30_bone_outlines']\n",
    "aligned_d30_angles = aligned_bone_data['aligned_d30_angles']\n",
    "aligned_d30_centroids = aligned_bone_data['aligned_d30_centroids']\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Applied the rotation and recenter to all the populations and save the aligned bones\n",
    "\n",
    "def align_bones_with_centroids_angles(positions_df, aligned_centroids, aligned_angles):\n",
    "    aligned_bones = {}\n",
    "    \n",
    "    for bone_name, df in positions_df.items():\n",
    "        # Get the centroid of the bone\n",
    "        centroid = aligned_centroids[bone_name]\n",
    "        # Get the angle to rotate\n",
    "        angle = aligned_angles[bone_name]\n",
    "\n",
    "        # The order of the rotation and the translation matters!\n",
    "        # The centroids are based on the rotated and translated points\n",
    "        # Rotate the df\n",
    "        df_rotated = df.copy()\n",
    "        # Exclude the Bone points\n",
    "        # df_rotated = df_rotated[df_rotated['source'] != 'Bone']\n",
    "        df_rotated[['Position.X', 'Position.Y']] = rotate_points(df_rotated[['Position.X', 'Position.Y']].values, angle)\n",
    "        \n",
    "        # Translate the df\n",
    "        df_translated = df_rotated.copy()\n",
    "        df_translated['Position.X'] -= centroid[0]\n",
    "        df_translated['Position.Y'] -= centroid[1]\n",
    "        \n",
    "        aligned_bones[bone_name] = df_translated\n",
    "    return aligned_bones\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "aligned_ss_bones = align_bones_with_centroids_angles(positions_ss_dfs, aligned_ss_centroids, aligned_ss_angles)\n",
    "aligned_d5_bones = align_bones_with_centroids_angles(positions_d5_dfs, aligned_d5_centroids, aligned_d5_angles)\n",
    "aligned_d10_bones = align_bones_with_centroids_angles(positions_d10_dfs, aligned_d10_centroids, aligned_d10_angles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Delete positions_ss_dfs, positions_d5_dfs, positions_d10_dfs, positions_d15_dfs, positions_d30_dfs to save memory\n",
    "del positions_ss_dfs, positions_d5_dfs, position_dirs_d10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "aligned_d15_bones = align_bones_with_centroids_angles(positions_d15_dfs, aligned_d15_centroids, aligned_d15_angles)\n",
    "aligned_d30_bones = align_bones_with_centroids_angles(positions_d30_dfs, aligned_d30_centroids, aligned_d30_angles)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "del positions_d15_dfs, positions_d30_dfs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "aligned_d15_bone_outlines_Detex, aligned_d15_angles_Detex, aligned_d15_centroids_Detex = process_and_align_bones_with_overlap(positions_d15_inhibitor_Detex_dfs, reference_bone=ref_bone_inner, window_size=500)\n",
    "# 11 mins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Saved the aligned bones to a pickle file\n",
    "aligned_data_inhibitor_Detex = {\n",
    "    'aligned_d15_bone_outlines_Detex': aligned_d15_bone_outlines_Detex,\n",
    "    'aligned_d15_angles_Detex': aligned_d15_angles_Detex,\n",
    "    'aligned_d15_centroids_Detex': aligned_d15_centroids_Detex,\n",
    "}\n",
    "\n",
    "aligned_bone_data_filename_Detex = os.path.join(save_dir, f'{datetime.today().strftime(\"%y%m%d\")}_YL_aligned_bone_data_inhibitor_Detex.pkl')\n",
    "with open(aligned_bone_data_filename_Detex, 'wb') as file:\n",
    "    pickle.dump(aligned_data_inhibitor_Detex, file)\n",
    "    \n",
    "del aligned_data_inhibitor_Detex, aligned_bone_data_filename_Detex"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "aligned_data_inhibitor_Detex_filename = f'{save_dir}/241219_YL_aligned_bone_data_inhibitor_Detex.pkl'\n",
    "aligned_data_inhibitor_Detex = pickle.load(open(aligned_data_inhibitor_Detex_filename, 'rb'))\n",
    "\n",
    "aligned_d15_bone_outlines_Detex = aligned_data_inhibitor_Detex['aligned_d15_bone_outlines_Detex']\n",
    "aligned_d15_angles_Detex = aligned_data_inhibitor_Detex['aligned_d15_angles_Detex']\n",
    "aligned_d15_centroids_Detex = aligned_data_inhibitor_Detex['aligned_d15_centroids_Detex']\n",
    "'''\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "aligned_d15_bones_Detex = align_bones_with_centroids_angles(positions_d15_inhibitor_Detex_dfs, aligned_d15_centroids_Detex, aligned_d15_angles_Detex)\n",
    "del positions_d15_inhibitor_Detex_dfs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.2 Transform Bone Coordinates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Smooth the outlines\n",
    "def smooth_outline(outline, sigma=2):\n",
    "    if not np.array_equal(outline[0], outline[-1]):\n",
    "        outline = np.vstack([outline, outline[0]])\n",
    "    smoothed_x = gaussian_filter1d(outline[:, 0], sigma=sigma)\n",
    "    smoothed_y = gaussian_filter1d(outline[:, 1], sigma=sigma)\n",
    "    smoothed_outline = np.vstack((smoothed_x, smoothed_y)).T\n",
    "    if not np.array_equal(smoothed_outline[0], smoothed_outline[-1]):\n",
    "        smoothed_outline = np.vstack([smoothed_outline, smoothed_outline[0]])\n",
    "    return smoothed_outline\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Smooth the reference bone outlines\n",
    "ref_bone_inner_smoothed = smooth_outline(ref_bone_inner_outline, sigma=80)\n",
    "common_outline = ref_bone_inner_smoothed\n",
    "\n",
    "# Save the common outline to a pickle file\n",
    "common_outline_filename = os.path.join(save_dir, f'{datetime.today().strftime(\"%y%m%d\")}_YL_common_outline.pkl')\n",
    "with open(common_outline_filename, 'wb') as file:\n",
    "    pickle.dump(common_outline, file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Smoothed the alinged_bone_outlines\n",
    "aligned_ss_bone_outlines_smoothed = {}\n",
    "aligned_d5_bone_outlines_smoothed = {}\n",
    "aligned_d10_bone_outlines_smoothed = {}\n",
    "aligned_d15_bone_outlines_smoothed = {}\n",
    "aligned_d30_bone_outlines_smoothed = {}\n",
    "\n",
    "sigma = 50\n",
    "for bone_name, outline in aligned_ss_bone_outlines.items():\n",
    "    smoothed_outline = smooth_outline(outline, sigma=sigma)\n",
    "    aligned_ss_bone_outlines_smoothed[bone_name] = smoothed_outline\n",
    "for bone_name, outline in aligned_d5_bone_outlines.items():\n",
    "    smoothed_outline = smooth_outline(outline, sigma=sigma)\n",
    "    aligned_d5_bone_outlines_smoothed[bone_name] = smoothed_outline\n",
    "for bone_name, outline in aligned_d10_bone_outlines.items():\n",
    "    smoothed_outline = smooth_outline(outline, sigma=sigma)\n",
    "    aligned_d10_bone_outlines_smoothed[bone_name] = smoothed_outline\n",
    "for bone_name, outline in aligned_d15_bone_outlines.items():\n",
    "    smoothed_outline = smooth_outline(outline, sigma=sigma)\n",
    "    aligned_d15_bone_outlines_smoothed[bone_name] = smoothed_outline\n",
    "for bone_name, outline in aligned_d30_bone_outlines.items():\n",
    "    smoothed_outline = smooth_outline(outline, sigma=sigma)\n",
    "    aligned_d30_bone_outlines_smoothed[bone_name] = smoothed_outline\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Smooth the aligned bones with inhibitors\n",
    "aligned_d15_bone_outlines_Detex_smoothed = {}\n",
    "\n",
    "for bone_name, outline in aligned_d15_bone_outlines_Detex.items():\n",
    "    smoothed_outline = smooth_outline(outline, sigma=sigma)\n",
    "    aligned_d15_bone_outlines_Detex_smoothed[bone_name] = smoothed_outline\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_y_range_at_x(shape_points, x):\n",
    "    \"\"\"\n",
    "    Find the range of y-values where the vertical line at x intersects the shape.\n",
    "    \"\"\"\n",
    "    # Find all edges of the shape where x is between the x-coordinates of the endpoints\n",
    "    y_vals = []\n",
    "    for i in range(len(shape_points)):\n",
    "        p1 = shape_points[i]\n",
    "        p2 = shape_points[(i + 1) % len(shape_points)]  # wrap around the shape points\n",
    "        \n",
    "        # Check if the x value is between p1 and p2's x-coordinates\n",
    "        if (p1[0] <= x <= p2[0]) or (p2[0] <= x <= p1[0]):\n",
    "            # Linearly interpolate to find the corresponding y value at x\n",
    "            if p1[0] != p2[0]:  # Avoid division by zero\n",
    "                y = p1[1] + (p2[1] - p1[1]) * (x - p1[0]) / (p2[0] - p1[0])\n",
    "                y_vals.append(y)\n",
    "    \n",
    "    if y_vals:\n",
    "        return min(y_vals), max(y_vals)\n",
    "    else:\n",
    "        return None, None  # No intersection with the shape at this x\n",
    "\n",
    "def create_structured_grid(shape_points, x_num, y_num):\n",
    "    \"\"\"\n",
    "    Create a structured grid by dividing the bounding box of the shape into x_num vertical sections.\n",
    "    Then place y_num points along each vertical grid line where it intersects the shape.\n",
    "    \"\"\"\n",
    "    shape_points = np.array(shape_points)\n",
    "    \n",
    "    # Step 1: Compute the bounding box\n",
    "    min_x, max_x = np.min(shape_points[:, 0]), np.max(shape_points[:, 0])\n",
    "    \n",
    "    # Step 2: Divide the x-range into equal sections\n",
    "    x_vals = np.linspace(min_x, max_x, x_num + 1)  # x_num divisions create x_num + 1 grid lines\n",
    "    # Shift the x_vals to the left by half the grid spacing to center the grid\n",
    "    x_vals = x_vals[1:]  # Remove the first point (left edge of bounding box)\n",
    "    x_vals = x_vals - (x_vals[1] - x_vals[0]) / 2  # Shift left by half the grid spacing\n",
    "    \n",
    "    grid_points = []\n",
    "\n",
    "    # Step 3: For each x grid line, find the y range and then place points\n",
    "    for x in x_vals:  # Skip the first and last lines (already have the bounding box)\n",
    "        y_min, y_max = get_y_range_at_x(shape_points, x)\n",
    "        \n",
    "        if y_min is not None and y_max is not None:\n",
    "            # Get y points by placing y_num points between y_min and y_max\n",
    "            y_vals = np.linspace(y_min, y_max, y_num+1)\n",
    "            \n",
    "            # Shift the y_vals down by half the grid spacing to center the grid\n",
    "            y_vals = y_vals[1:]  # Remove the first point (bottom edge of bounding box)\n",
    "            y_vals = y_vals - (y_vals[1] - y_vals[0]) / 2  # Shift down by half the grid spacing\n",
    "            # Add grid points (x, y) for this vertical line\n",
    "            for y in y_vals: # Skip the first and last points (already have the y_min and y_max)\n",
    "                grid_points.append([x, y])\n",
    "\n",
    "    return np.array(grid_points)\n",
    "\n",
    "def is_point_inside_shape(point, shape_points):\n",
    "    \"\"\"\n",
    "    Determines if a point is inside an irregular shape using ray-casting.\n",
    "    \"\"\"\n",
    "    x, y = point\n",
    "    n = len(shape_points)\n",
    "    inside = False\n",
    "    p1x, p1y = shape_points[0]\n",
    "    for i in range(n + 1):\n",
    "        p2x, p2y = shape_points[i % n]\n",
    "        if y > min(p1y, p2y):\n",
    "            if y <= max(p1y, p2y):\n",
    "                if x <= max(p1x, p2x):\n",
    "                    if p1y != p2y:\n",
    "                        xinters = (y - p1y) * (p2x - p1x) / (p2y - p1y) + p1x\n",
    "                    if p1x == p2x or x <= xinters:\n",
    "                        inside = not inside\n",
    "        p1x, p1y = p2x, p2y\n",
    "    return inside\n",
    "\n",
    "\n",
    "def thin_plate_spline_transform(src_points, dst_points):\n",
    "    \"\"\"\n",
    "    Perform Thin Plate Spline (TPS) transformation from src_points to dst_points.\n",
    "    \"\"\"\n",
    "    # Create Radial Basis Function (RBF) interpolators for x and y coordinates\n",
    "    rbf_x = Rbf(src_points[:, 0], src_points[:, 1], dst_points[:, 0], function='thin_plate')\n",
    "    rbf_y = Rbf(src_points[:, 0], src_points[:, 1], dst_points[:, 1], function='thin_plate')\n",
    "    \n",
    "    def transform(points):\n",
    "        new_x = rbf_x(points[:, 0], points[:, 1])\n",
    "        new_y = rbf_y(points[:, 0], points[:, 1])\n",
    "        return np.vstack([new_x, new_y]).T\n",
    "    \n",
    "    return transform\n",
    "\n",
    "def transform_data(data_points, grid_shape_1, grid_shape_2):\n",
    "    \"\"\"\n",
    "    Apply the TPS transformation to the data points based on the grid transformation.\n",
    "    \"\"\"\n",
    "    # Perform Thin Plate Spline (TPS) transformation\n",
    "    tps_transform = thin_plate_spline_transform(grid_shape_2, grid_shape_1)\n",
    "    \n",
    "    # Apply the transformation to the data points\n",
    "    transformed_data_points = tps_transform(data_points)\n",
    "    \n",
    "    return transformed_data_points\n",
    "\n",
    "\n",
    "def transform_bone_positions(outline_dict, position_dict, common_outline, x_num=40, y_num=20, source_value=None):\n",
    "    \"\"\"\n",
    "    Transforms the bone positions from multiple datasets using Thin Plate Spline (TPS) based on the provided outlines and positions.\n",
    "    Parameters:\n",
    "        outline_dict: Dictionary containing outlines.\n",
    "        position_dict: Dictionary containing bone positions (DataFrames).\n",
    "        common_outline: The common outline (to which the other outlines will be aligned).\n",
    "        x_num: Number of vertical sections for structured grid.\n",
    "        y_num: Number of horizontal points along each vertical section.\n",
    "        source_value: If provided, filter the positions based on the specified source value.\n",
    "    Returns:\n",
    "        Dictionary containing the transformed bone positions with the 'source' column retained.\n",
    "    \"\"\"\n",
    "    transformed_dict = {}\n",
    "\n",
    "    # Create the structured grid for the common outline\n",
    "    grid_common_outline = create_structured_grid(common_outline, x_num=x_num, y_num=y_num)\n",
    "\n",
    "    # Loop through each dataset in the position_dict\n",
    "    for dataset_name, position_df in position_dict.items():\n",
    "        # Get the corresponding outline\n",
    "        outline_2 = outline_dict[dataset_name]\n",
    "        \n",
    "        # Create the structured grid for the specific dataset's outline\n",
    "        grid_outline_2 = create_structured_grid(outline_2, x_num=x_num, y_num=y_num)\n",
    "\n",
    "        # Filter the positions based on the source (if provided)\n",
    "        bone_positions_2, source_column = filter_bone_positions(position_df, source_value=source_value)\n",
    "\n",
    "        # Transform the filtered bone positions from the dataset outline to the common outline\n",
    "        transformed_bone_positions = transform_data(bone_positions_2, grid_common_outline, grid_outline_2)\n",
    "\n",
    "        # Convert the transformed positions to a DataFrame and include the source column\n",
    "        transformed_df = pd.DataFrame(transformed_bone_positions, columns=['Position.X', 'Position.Y'])\n",
    "        transformed_df['source'] = source_column  # Add the source column back\n",
    "        transformed_df['dataset'] = dataset_name  # Add the dataset name column\n",
    "        # Store the transformed DataFrame in the result dictionary\n",
    "        transformed_dict[dataset_name] = transformed_df\n",
    "\n",
    "    return transformed_dict\n",
    "\n",
    "\n",
    "def filter_bone_positions(df, source_value=None):\n",
    "    \"\"\"\n",
    "    Filters the DataFrame for a specific source if provided and returns the Position.X, Position.Y columns as a NumPy array.\n",
    "    If source_value is None, return all positions.\n",
    "    \"\"\"\n",
    "    if source_value:\n",
    "        filtered_df = df[df['source'] == source_value]\n",
    "    else:\n",
    "        # Exclude the data with source value 'GFP'\n",
    "        filtered_df = df[df['source'] != 'GFP']\n",
    "    positions = filtered_df[['Position.X', 'Position.Y']].to_numpy()\n",
    "    return positions, filtered_df['source'].to_numpy()  # Return positions and the source column\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_num = 200\n",
    "y_num = 35\n",
    "# 18 mins with 200*35\n",
    "# We are using the aligned_bone_date dict instead of using the dataframes, because for bones of the same day, they have different bone outline\n",
    "# Not giving the source value, because we will apply the transformation to all the data without the data with source value 'GFP'\n",
    "transformed_ss_bones = transform_bone_positions(aligned_ss_bone_outlines_smoothed, aligned_ss_bones, common_outline, x_num=x_num, y_num=y_num)\n",
    "transformed_d5_bones = transform_bone_positions(aligned_d5_bone_outlines_smoothed, aligned_d5_bones, common_outline, x_num=x_num, y_num=y_num)\n",
    "transformed_d10_bones = transform_bone_positions(aligned_d10_bone_outlines_smoothed, aligned_d10_bones, common_outline, x_num=x_num, y_num=y_num)\n",
    "\n",
    "# Visualize the original and transformed bone positions for all the data sources\n",
    "transformed_ss_bones_df = pd.concat(transformed_ss_bones)\n",
    "transformed_d5_bones_df = pd.concat(transformed_d5_bones)\n",
    "transformed_d10_bones_df = pd.concat(transformed_d10_bones)\n",
    "\n",
    "# 10 mins\n",
    "del transformed_ss_bones, transformed_d5_bones, transformed_d10_bones"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transformed_d15_bones = transform_bone_positions(aligned_d15_bone_outlines_smoothed, aligned_d15_bones, common_outline, x_num=x_num, y_num=y_num)\n",
    "transformed_d30_bones = transform_bone_positions(aligned_d30_bone_outlines_smoothed, aligned_d30_bones, common_outline, x_num=x_num, y_num=y_num)\n",
    "# Visualize the original and transformed bone positions for all the data sources\n",
    "\n",
    "transformed_d15_bones_df = pd.concat(transformed_d15_bones)\n",
    "transformed_d30_bones_df = pd.concat(transformed_d30_bones)\n",
    "\n",
    "# 8 mins\n",
    "del transformed_d15_bones, transformed_d30_bones"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the transformed bone positions to CSV files, with the data outside the common outline \n",
    "transformed_ss_bones_df.to_csv(f'{save_dir}/transformed_ss_bones_all_df.csv', index=False)\n",
    "transformed_d5_bones_df.to_csv(f'{save_dir}/transformed_d5_bones_all_df.csv', index=False)\n",
    "transformed_d10_bones_df.to_csv(f'{save_dir}/transformed_d10_bones_all_df.csv', index=False)\n",
    "transformed_d15_bones_df.to_csv(f'{save_dir}/transformed_d15_bones_all_df.csv', index=False)\n",
    "transformed_d30_bones_df.to_csv(f'{save_dir}/transformed_d30_bones_all_df.csv', index=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transformed_d15_bones_Detex = transform_bone_positions(aligned_d15_bone_outlines_Detex_smoothed, aligned_d15_bones_Detex, common_outline, x_num=x_num, y_num=y_num)\n",
    "transformed_d15_bones_Detex_df = pd.concat(transformed_d15_bones_Detex)\n",
    "transformed_d15_bones_Detex_df.to_csv(f'{save_dir}/transformed_d15_bones_Detex_all_df.csv', index=False)\n",
    "\n",
    "del transformed_d15_bones_Detex"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Spatial Probability Density Map Generation and Visualization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1 SPDM Generation (KDE Calculation)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the function to exclude the data points outside the bone outline\n",
    "def exclude_outside_bone_outline(df, bone_outline):\n",
    "    \"\"\"\n",
    "    Exclude the data points outside the bone outline using GeoPandas with bounding box filtering for more efficiency.\n",
    "\n",
    "    Parameters:\n",
    "    df : DataFrame containing 'Position.X', 'Position.Y', 'weights', and 'source'.\n",
    "    bone_outline : The outline of the bone to limit the KDE calculation within the bone.\n",
    "\n",
    "    Returns:\n",
    "    df_inside : DataFrame containing the data points inside the bone outline.\n",
    "    \"\"\"\n",
    "    # print(df.shape)\n",
    "    df = pd.DataFrame(df[df['source'] != 'GFP'])\n",
    "    # print(df.shape)\n",
    "    # Create a GeoDataFrame from the original DataFrame\n",
    "    gdf = gpd.GeoDataFrame(df, geometry=gpd.points_from_xy(df['Position.X'], df['Position.Y']))\n",
    "\n",
    "    # Convert the bone outline to a Shapely polygon\n",
    "    bone_polygon = Polygon(bone_outline)\n",
    "\n",
    "    # Create a bounding box polygon from the bounds of the bone_polygon\n",
    "    minx, miny, maxx, maxy = bone_polygon.bounds\n",
    "    bounding_box = box(minx, miny, maxx, maxy)\n",
    "\n",
    "    # First, filter by the bounding box of the polygon (faster operation)\n",
    "    gdf_in_bbox = gdf[gdf.geometry.within(bounding_box)]\n",
    "    \n",
    "    # Then, perform the more precise filtering with the actual polygon\n",
    "    gdf_inside = gdf_in_bbox[gdf_in_bbox.within(bone_polygon)]\n",
    "    \n",
    "    # Drop the 'geometry' column if you don't need it in the result\n",
    "    df_inside = gdf_inside.drop(columns='geometry')\n",
    "\n",
    "    return df_inside\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Use the bone outline to generate the KDE for each source\n",
    "#  Define the KDE function for each source, with weights based on z-aggregated points\n",
    "def kde_for_source(df, bw_method='scott', bone_outline = None, binsize = 20):\n",
    "    \"\"\"\n",
    "    Apply KDE to estimate the spatial distribution in the x-y plane for each source, \n",
    "    with weights aggregated for points with the same x, y but different z.\n",
    "\n",
    "    Parameters:\n",
    "    df : DataFrame containing 'Position.X', 'Position.Y', 'weights', and 'source'.\n",
    "    bw_method : Bandwidth method for KDE (default is 'scott').\n",
    "    bone_outline : The outline of the bone to limit the KDE calculation within the bone.\n",
    "    binsize : The size of the bins for the grid used in the KDE calculation.\n",
    "\n",
    "    Returns:\n",
    "    kde_result : A dictionary containing the grid and KDE values for each source.\n",
    "    \"\"\"\n",
    "    kde_results = {}\n",
    "    sources = df['source'].unique()\n",
    "    sources = sources[sources != 'GFP'] # Exclude the bone for the KDE calculation\n",
    "\n",
    "    for source in sources:\n",
    "        # Filter data for the current source\n",
    "        source_data = df[df['source'] == source]\n",
    "\n",
    "        # Group by Position.X and Position.Y, summing weights (or using counts as weights if no weights are given)\n",
    "        if 'weights' in source_data.columns:\n",
    "            source_data_agg = source_data.groupby(['Position.X', 'Position.Y'])['weights'].sum().reset_index()\n",
    "        else:\n",
    "            # If no weights are provided, use the count of occurrences as weights\n",
    "            source_data_agg = source_data.groupby(['Position.X', 'Position.Y']).size().reset_index(name='weights')\n",
    "\n",
    "        # Get the x and y values and aggregated weights\n",
    "        x_vals = source_data_agg['Position.X']\n",
    "        y_vals = source_data_agg['Position.Y']\n",
    "        weights = source_data_agg['weights']\n",
    "        if bone_outline is None:\n",
    "            x_min, x_max = x_vals.min(), x_vals.max()\n",
    "            y_min, y_max = y_vals.min(), y_vals.max()\n",
    "        else:\n",
    "            x_min, y_min = bone_outline.min(axis=0)\n",
    "            x_max, y_max = bone_outline.max(axis=0)\n",
    "\n",
    "        xi, yi = np.linspace(x_min, x_max, int((x_max - x_min)/binsize)+1), np.linspace(y_min, y_max, int((y_max - y_min)/binsize)+1)\n",
    "        xi, yi = np.meshgrid(xi, yi)\n",
    "        grid_points = np.vstack([xi.flatten(), yi.flatten()])\n",
    "        common_grid = (xi, yi, grid_points)\n",
    "\n",
    "        # Stack the x and y data for KDE input\n",
    "        xy = np.vstack([x_vals, y_vals])\n",
    "\n",
    "        # Perform the KDE with aggregated weights\n",
    "        kde = gaussian_kde(xy, weights=weights, bw_method=bw_method)\n",
    "        kde_values = kde(grid_points).reshape(xi.shape)\n",
    "\n",
    "        # Store the results for each source\n",
    "        kde_results[source] = kde_values\n",
    "\n",
    "    return kde_results, common_grid\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "# Read the transformed bone positions from the CSV files\n",
    "transformed_ss_bones_df = pd.read_csv(f'{save_dir}/transformed_ss_bones_all_df.csv')\n",
    "transformed_d5_bones_df = pd.read_csv(f'{save_dir}/transformed_d5_bones_all_df.csv')\n",
    "transformed_d10_bones_df = pd.read_csv(f'{save_dir}/transformed_d10_bones_all_df.csv')\n",
    "transformed_d15_bones_df = pd.read_csv(f'{save_dir}/transformed_d15_bones_all_df.csv')\n",
    "transformed_d30_bones_df = pd.read_csv(f'{save_dir}/transformed_d30_bones_all_df.csv')\n",
    "\n",
    "\n",
    "transformed_d15_bones_Detex_df = pd.read_csv(f'{save_dir}/transformed_d15_bones_Detex_all_df.csv')\n",
    "\n",
    "# Read the common outline from the pickle file\n",
    "common_outline_filename = os.path.join(save_dir, f'{datetime.today().strftime(\"%y%m%d\")}_YL_common_outline.pkl')\n",
    "with open(common_outline_filename, 'rb') as file:\n",
    "    common_outline = pickle.load(file)\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Exclude the data points outside the bone outline\n",
    "transformed_ss_bones_df = exclude_outside_bone_outline(transformed_ss_bones_df, common_outline)\n",
    "transformed_d5_bones_df = exclude_outside_bone_outline(transformed_d5_bones_df, common_outline)\n",
    "transformed_d10_bones_df = exclude_outside_bone_outline(transformed_d10_bones_df, common_outline)\n",
    "transformed_d15_bones_df = exclude_outside_bone_outline(transformed_d15_bones_df, common_outline)\n",
    "transformed_d30_bones_df = exclude_outside_bone_outline(transformed_d30_bones_df, common_outline)\n",
    "\n",
    "# Instead of using the common bone outline (outmost),  we use the overlap of the bone outlines for the KDE calculation\n",
    "kde_results_ss_overlap_bones, common_grid_overlap_bones = kde_for_source(transformed_ss_bones_df, binsize=20, bone_outline=common_outline)\n",
    "kde_results_d5_overlap_bones, _ = kde_for_source(transformed_d5_bones_df, binsize=20, bone_outline=common_outline)\n",
    "kde_results_d10_overlap_bones, _ = kde_for_source(transformed_d10_bones_df, binsize=20, bone_outline=common_outline)\n",
    "kde_results_d15_overlap_bones, _ = kde_for_source(transformed_d15_bones_df, binsize=20, bone_outline=common_outline)\n",
    "kde_results_d30_overlap_bones, _ = kde_for_source(transformed_d30_bones_df, binsize=20, bone_outline=common_outline)\n",
    "\n",
    "# Save the results to a file\n",
    "transformed_ss_bones_df.to_csv(f'{save_dir}/transformed_ss_bones_inside_df.csv', index=False)\n",
    "transformed_d5_bones_df.to_csv(f'{save_dir}/transformed_d5_bones_inside_df.csv', index=False)\n",
    "transformed_d10_bones_df.to_csv(f'{save_dir}/transformed_d10_bones_inside_df.csv', index=False)\n",
    "transformed_d15_bones_df.to_csv(f'{save_dir}/transformed_d15_bones_inside_df.csv', index=False)\n",
    "transformed_d30_bones_df.to_csv(f'{save_dir}/transformed_d30_bones_inside_df.csv', index=False)\n",
    "\n",
    "# 11 mins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(f'{save_dir}/kde_results_ss_overlap_bones.pkl', 'wb') as f:\n",
    "    pickle.dump(kde_results_ss_overlap_bones, f)\n",
    "with open(f'{save_dir}/common_grid_overlap_bones.pkl', 'wb') as f:\n",
    "    pickle.dump(common_grid_overlap_bones, f)\n",
    "    \n",
    "with open(f'{save_dir}/kde_results_d5_overlap_bones.pkl', 'wb') as f:\n",
    "    pickle.dump(kde_results_d5_overlap_bones, f)\n",
    "with open(f'{save_dir}/kde_results_d10_overlap_bones.pkl', 'wb') as f:\n",
    "    pickle.dump(kde_results_d10_overlap_bones, f)\n",
    "with open(f'{save_dir}/kde_results_d15_overlap_bones.pkl', 'wb') as f:\n",
    "    pickle.dump(kde_results_d15_overlap_bones, f)\n",
    "with open(f'{save_dir}/kde_results_d30_overlap_bones.pkl', 'wb') as f:\n",
    "    pickle.dump(kde_results_d30_overlap_bones, f)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For the Bone KDE, we will the strategy the placing date points on the outline and combined it with the inner bone\n",
    "# Step 1: Calculate the cumulative distance along the outline\n",
    "distances = np.sqrt(np.sum(np.diff(common_outline, axis=0)**2, axis=1))  # Euclidean distances\n",
    "cumulative_distance = np.cumsum(distances)\n",
    "cumulative_distance = np.insert(cumulative_distance, 0, 0)  # Start at 0\n",
    "\n",
    "# Step 2: Define the number of evenly spaced points\n",
    "n_points = 10000  # Adjust this to the desired number of points\n",
    "\n",
    "# For visualization, we set the n_points 100\n",
    "# n_points = 100\n",
    "even_distances = np.linspace(0, cumulative_distance[-1], n_points)\n",
    "\n",
    "# Step 3: Interpolate the outline for evenly spaced distances\n",
    "\n",
    "# Interpolation functions for x and y\n",
    "interp_x = interp1d(cumulative_distance, common_outline[:, 0], kind='linear')\n",
    "interp_y = interp1d(cumulative_distance, common_outline[:, 1], kind='linear')\n",
    "\n",
    "# Generate new points\n",
    "evenly_spaced_x = interp_x(even_distances)\n",
    "evenly_spaced_y = interp_y(even_distances)\n",
    "\n",
    "# Combine the new x and y into an array\n",
    "evenly_spaced_outline = np.vstack((evenly_spaced_x, evenly_spaced_y)).T\n",
    "\n",
    "# Step 4: Plot the original and evenly spaced outline\n",
    "plt.figure(figsize=(15, 5))\n",
    "\n",
    "# Original outline\n",
    "plt.plot(common_outline[:, 0], common_outline[:, 1], label='Original Outline', color='black', alpha=0.5)\n",
    "\n",
    "# Evenly spaced outline\n",
    "plt.scatter(evenly_spaced_outline[:, 0], evenly_spaced_outline[:, 1], label='Evenly Spaced Points', color='red', s=10)\n",
    "\n",
    "plt.title(\"Evenly Spaced Points Along Outline\")\n",
    "plt.xlabel(\"X\")\n",
    "plt.ylabel(\"Y\")\n",
    "plt.legend()\n",
    "plt.grid(False)\n",
    "# plt.savefig(f'{save_dir}/Evenly_Spaced_Points_Along_Outline.png')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate the KDE for the bone_outline assuming the weights are the same for each point\n",
    "# 1. transform the bone outline to the same format as the data\n",
    "# bone_outline_df = pd.DataFrame(common_outline, columns=['Position.X', 'Position.Y'])\n",
    "bone_outline_df = pd.DataFrame(evenly_spaced_outline, columns=['Position.X', 'Position.Y'])\n",
    "bone_outline_df['source'] = 'Bones_Outline'\n",
    "# 2. Use the same KDE function to calculate the KDE for the bone outline\n",
    "# bone_outline_df = exclude_outside_bone_outline(bone_outline_df, common_outline)\n",
    "kde_results_bone_outline, _ = kde_for_source(bone_outline_df, binsize=20, bone_outline=common_outline)\n",
    "\n",
    "    \n",
    "with open(f'{save_dir}/kde_results_bone_outline.pkl', 'wb') as f:\n",
    "    pickle.dump(kde_results_bone_outline, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Exclude the data points outside the bone outline\n",
    "transformed_d15_bones_Detex_df = exclude_outside_bone_outline(transformed_d15_bones_Detex_df, common_outline)\n",
    "\n",
    "# Instead of using the common bone outline (outmost),  we use the overlap of the bone outlines for the KDE calculation\n",
    "kde_results_d15_overlap_bones_Detex, _ = kde_for_source(transformed_d15_bones_Detex_df, binsize=20, bone_outline=common_outline)\n",
    "\n",
    "transformed_d15_bones_Detex_df.to_csv(f'{save_dir}/transformed_d15_bones_Detex_inside_df.csv', index=False)\n",
    "with open(f'{save_dir}/kde_results_d15_overlap_bones_Detex.pkl', 'wb') as f:\n",
    "    pickle.dump(kde_results_d15_overlap_bones_Detex, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2 SPDM Visualization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define function to create mask from outline \n",
    "def points_in_polygon(x_points, y_points, outline):\n",
    "    path = Path(outline)\n",
    "    points = np.vstack((x_points, y_points)).T\n",
    "    return path.contains_points(points)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the KDE with the bone outline and optionally with the data points\n",
    "def plot_kde_results_with_df(df, kde_results, common_grid, save_dir, color_map, bone_outline, filename=None, mesh = True, scatter = False):\n",
    "    raw_sources = df['source'].unique()\n",
    "    raw_sources = raw_sources[raw_sources != 'GFP']\n",
    "    \n",
    "\n",
    "    # Ensure that the sources are aligned with the keys of mk_color_map\n",
    "    # First, get the color map keys, excluding 'Bone'\n",
    "    color_map_keys = [key for key in color_map.keys() if key != 'GFP']\n",
    "    \n",
    "    # Reorder sources to match the order of color_map_keys\n",
    "    sources = [source for source in color_map_keys if source in raw_sources]\n",
    "    \n",
    "    \n",
    "    fig, axs = plt.subplots(len(sources), 1, figsize=(15 , 5 *len(sources)), sharex=True, sharey=True, constrained_layout=True)\n",
    "    if len(sources) == 1:\n",
    "        axs = [axs]  # since we only have one row, make it iterable\n",
    "\n",
    "    for i, source in enumerate(sources):\n",
    "        ax = axs[i]\n",
    "        source_positions = df[df['source'] == source]\n",
    "        \n",
    "        # Plot the x, y of each bone group as the background\n",
    "        if scatter:\n",
    "            ax.scatter(source_positions['Position.X'], source_positions['Position.Y'], s=1, c=color_map[source], alpha=1)\n",
    "        \n",
    "        xi, yi, _ = common_grid \n",
    "        \n",
    "        zi = kde_results[source]\n",
    "        # Masked the values outside the bone outline\n",
    "        mask = points_in_polygon(xi.flatten(), yi.flatten(), bone_outline).reshape(xi.shape)\n",
    "        zi[~mask] = 0\n",
    "\n",
    "        # Filter only the non-zero values of zi for percentile calculation (i.e., exclude outside the mask)\n",
    "        zi_inside_mask = zi[mask]\n",
    "\n",
    "        # Normalize the values of zi for the full plot\n",
    "        norm = Normalize(vmin=zi_inside_mask.min(), vmax=zi_inside_mask.max())\n",
    "        normed_z = norm(zi)\n",
    "        normed_z[~mask] = 0\n",
    "        # if source == 'Bones':\n",
    "        #     normed_z = kde_results[source]\n",
    "\n",
    "        # Set the alpha transparency to 1 for test only\n",
    "        # normed_z[mask] = 1    \n",
    "        colors = np.array(plt.cm.colors.hex2color(color_map[source]))\n",
    "        rgba_colors = np.zeros((*zi.shape, 4))\n",
    "        rgba_colors[..., :3] = colors[:3]  # RGB values\n",
    "        rgba_colors[..., -1] = normed_z  # Alpha transparency \n",
    "        \n",
    "        # rgba_colors[zi < np.percentile(zi, 80)] = 0\n",
    "        \n",
    "        ax.set_title(source)\n",
    "        xlim_min, ylim_min = bone_outline.min(axis=0)\n",
    "        xlim_max, ylim_max = bone_outline.max(axis=0)\n",
    "        # Convert them into integers\n",
    "        xlim_min, xlim_max = int(xlim_min)-500, int(xlim_max)+500\n",
    "        ylim_min, ylim_max = int(ylim_min)-300, int(ylim_max)+300\n",
    "        \n",
    "        ax.set_xlim(xlim_min, xlim_max)\n",
    "        ax.set_ylim(ylim_min, ylim_max)\n",
    "        ax.set_xlabel('Position.X')\n",
    "        ax.set_ylabel('Position.Y')\n",
    "        if mesh:\n",
    "            ax.pcolormesh(xi, yi, rgba_colors, shading='auto', rasterized=True)\n",
    "        \n",
    "        percentiles = np.arange(0, 81, 20)\n",
    "        contour_levels = np.unique(np.percentile(norm(zi_inside_mask), percentiles))\n",
    "        colors = color_map[source]\n",
    "        if len(contour_levels) > 1:\n",
    "            # print(contour_levels)\n",
    "            # With changing linewidths, from the largest to the smallest\n",
    "            # linewidths = np.linspace(0.5, 10, len(contour_levels))\n",
    "            contour = ax.contour(xi, yi, normed_z, levels=contour_levels, linewidths=1, colors='black', alpha=0.5)\n",
    "            \n",
    "\n",
    "        # fmt = {level: f'{perc}%' for level, perc in zip(contour_levels, percentiles)}\n",
    "        # ax.clabel(contour, contour_levels, inline=True, fmt=fmt, fontsize=8)\n",
    "\n",
    "        ax.plot(bone_outline[:,0], bone_outline[:, 1], color='black')\n",
    "        # TODO: optional\n",
    "        # Find the peak in the kde result\n",
    "        # max_idx = np.unravel_index(np.argmax(zi), zi.shape)  # Index of maximum value in zi\n",
    "        # peak_x, peak_y = xi[max_idx], yi[max_idx]  # Get the coordinates of the peak\n",
    "\n",
    "        # Plot the peak as a marker\n",
    "        # ax.plot(peak_x, peak_y, 's', markersize=5, label='Peak', color = 'black')  # Black square marker showing the global peak\n",
    "\n",
    "    # Save the figure for the current bone group in the corresponding directory\n",
    "    fig.suptitle(f'Spatial Distribution Estimation of {filename}', fontsize=16)\n",
    "    if save_dir:\n",
    "        if filename:\n",
    "            fig_filename = os.path.join(save_dir, f'{datetime.today().strftime(\"%y%m%d\")}_YL_KDE_{filename}.pdf')\n",
    "            # fig_filename_pdf = os.path.join(save_dir, f'{datetime.today().strftime(\"%y%m%d\")}_YL_KDE_{filename}.pdf')\n",
    "        else:\n",
    "            fig_filename = os.path.join(save_dir, f'{datetime.today().strftime(\"%y%m%d\")}_YL_KDE_{sources}.pdf') \n",
    "            # fig_filename_pdf = os.path.join(save_dir, f'{datetime.today().strftime(\"%y%m%d\")}_YL_KDE_{sources}.pdf')\n",
    "        # Save it as pdf file\n",
    "        fig.savefig(fig_filename, dpi=300, bbox_inches='tight')\n",
    "\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the KDE of MKs and Adipos with the bone outline\n",
    "def plot_kde_results_MK_Adipos(kde_results, common_grid, save_dir, color_map, bone_outline, filename=None, mesh = True):\n",
    "    sources = ['MKs','Adipos']\n",
    "\n",
    "    fig, ax = plt.subplots(1, 1, figsize=(15 , 5))\n",
    "    legend_patches = []\n",
    "    for i, source in enumerate(sources):\n",
    "\n",
    "        xi, yi, _ = common_grid \n",
    "        zi = kde_results[source]\n",
    "        # Masked the values outside the bone outline\n",
    "        mask = points_in_polygon(xi.flatten(), yi.flatten(), bone_outline).reshape(xi.shape)\n",
    "        zi[~mask] = 0\n",
    "        \n",
    "        # Make sure the values sum to 1\n",
    "        # zi = zi / zi.sum()\n",
    "        \n",
    "        # Filter only the non-zero values of zi for percentile calculation (i.e., exclude outside the mask)\n",
    "        zi_inside_mask = zi[mask]\n",
    "\n",
    "        # Normalize the values of zi for the full plot\n",
    "        norm = Normalize(vmin=zi_inside_mask.min(), vmax=zi_inside_mask.max())\n",
    "        normed_z = norm(zi)\n",
    "        normed_z[~mask] = 0\n",
    "        \n",
    "        # Set all no-zores values to 1\n",
    "        # zi[zi > np.percentile(zi, 80)] = 1\n",
    "\n",
    "        legend_patches.append(Patch(color=color_map[source], label=source, alpha = 0.5))\n",
    "\n",
    "        colors = np.array(plt.cm.colors.hex2color(color_map[source]))\n",
    "        rgba_colors = np.zeros((*zi.shape, 4))\n",
    "        rgba_colors[..., :3] = colors[:3]  # RGB values\n",
    "        \n",
    "        # set the transparency into 5 levels based on normed_z\n",
    "    \n",
    "        rgba_colors[..., -1] = normed_z  # Alpha transparency\n",
    "        # rgba_colors[normed_z < np.percentile(normed_z, 80)] = 0\n",
    "        \n",
    "\n",
    "        if mesh:\n",
    "            ax.pcolormesh(xi, yi, rgba_colors, shading='auto', rasterized=True)\n",
    "\n",
    "        \n",
    "        percentiles = np.arange(0, 81, 20) # for 5 levels and if we want 10 levels we need to use np.arange(10, 101, 10)\n",
    "        contour_levels = np.percentile(norm(zi_inside_mask), percentiles)\n",
    "        # contour_levels = np.unique(np.percentile(norm(zi), percentiles))\n",
    "        # Print the calculated contour levels\n",
    "        print(\"Calculated Contour Levels:\", contour_levels)\n",
    "        rgba_colors[normed_z < contour_levels[-1]] = 0\n",
    "        colors = color_map[source]\n",
    "        if len(contour_levels) > 1:\n",
    "            # With changing linewidths, from the largest to the smallest\n",
    "            # linewidths = np.linspace(1, 5, len(contour_levels))\n",
    "            # contour = ax.contour(xi, yi, normed_z, levels=contour_levels, linewidths=linewidths, colors=colors, alpha=0.5)\n",
    "\n",
    "            # Only show the last two levels\n",
    "            # linewidths = np.linspace(1, 5, len(contour_levels))\n",
    "            linewidths = [2,4]\n",
    "            contour = ax.contour(xi, yi, normed_z, levels=contour_levels[-2:], linewidths=linewidths, colors=colors, alpha=0.5)\n",
    "            \n",
    "            # fmt = {level: f'{100- perc}%' for level, perc in zip(contour_levels[-2:], percentiles[-2:])}\n",
    "            # ax.clabel(contour, contour_levels[-2:], inline=True, fmt=fmt, fontsize=8, colors=colors)\n",
    "\n",
    "    xlim_min, ylim_min = bone_outline.min(axis=0)\n",
    "    xlim_max, ylim_max = bone_outline.max(axis=0)\n",
    "    # Convert them into integers\n",
    "    xlim_min, xlim_max = int(xlim_min)-500, int(xlim_max)+500\n",
    "    ylim_min, ylim_max = int(ylim_min)-300, int(ylim_max)+300\n",
    "    ax.set_title(filename)\n",
    "    ax.set_xlim(xlim_min, xlim_max)\n",
    "    ax.set_ylim(ylim_min, ylim_max)\n",
    "    ax.set_xlabel('Position.X')\n",
    "    ax.set_ylabel('Position.Y')\n",
    "    ax.plot(bone_outline[:,0], bone_outline[:, 1], color='black')\n",
    "    ax.legend(handles = legend_patches, loc='upper right')\n",
    "    # Save the figure for the current bone group in the corresponding directory\n",
    "    fig.suptitle(f'Spatial Probability Density Map', fontsize=16)\n",
    "    if save_dir:\n",
    "        if filename:\n",
    "            fig_filename = os.path.join(save_dir, f'{datetime.today().strftime(\"%y%m%d\")}_YL_KDE_{filename}.pdf')\n",
    "        else:\n",
    "            fig_filename = os.path.join(save_dir, f'{datetime.today().strftime(\"%y%m%d\")}_YL_KDE_{sources}.pdf') \n",
    "\n",
    "        fig.savefig(fig_filename, dpi=300)\n",
    "\n",
    "    plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_1D_kde(kde_results, common_grid, save_dir, color_map, bone_outline, sources=['MKs', 'Adipos'], filename=None):\n",
    "\n",
    "    fig, ax = plt.subplots(1, 1, figsize=(15, 5))\n",
    "    # Take the common source based on the color_map.keys  and sources\n",
    "    sources = list(set(color_map.keys()).intersection(set(sources)))\n",
    "    # Take the common source based on the sources and the kde_results.keys\n",
    "    sources = list(set(kde_results.keys()).intersection(set(sources)))\n",
    "    # Keep the order of the sources based on the color_map.keys\n",
    "    sources = [source for source in color_map.keys() if source in sources]\n",
    "    \n",
    "    \n",
    "    for source in sources:\n",
    "        xi, yi, _ = common_grid  # Get the common grid (xi for X, yi for Y)\n",
    "        zi = kde_results[source]  # Get the KDE results for the current source\n",
    "        \n",
    "        # Mask values outside the bone outline\n",
    "        mask = points_in_polygon(xi.flatten(), yi.flatten(), bone_outline).reshape(xi.shape)\n",
    "        zi[~mask] = 0\n",
    "        \n",
    "        # Normalize the KDE results to ensure they sum to 1\n",
    "        zi = zi / zi.sum()\n",
    "        \n",
    "        # Sum along the y-axis (rows) to project the KDE along the x-axis\n",
    "        summed_zi_along_y = zi.sum(axis=0)  # Summing along rows gives us the KDE projection along X\n",
    "        # Smooth the summed_zi_along_y with Gaussian filter\n",
    "        smoothed_summed_zi = gaussian_filter1d(summed_zi_along_y, sigma=10)  # Adjust sigma as needed for smoothing\n",
    "    \n",
    "        # Plot the summed values as a 1D curve\n",
    "\n",
    "        name_label = source\n",
    "\n",
    "        # ax.plot(xi[0, :], summed_zi_along_y, label=f'{name_label} (original)', color=color_map[source], linewidth=2)\n",
    "    \n",
    "        # Plot the smoothed values as a dashed line in black\n",
    "        ax.plot(xi[0, :], smoothed_summed_zi, label=f'{name_label}', color=color_map[source], linewidth=2) # linestyle='--', \n",
    "    # Set plot title, labels, and legend\n",
    "    if filename:\n",
    "        name_end = filename.split('_')[-1]\n",
    "        ax.set_title(f'Marginal Distribution of {','.join(sources)} {name_end}', fontsize=16)\n",
    "    else:\n",
    "        ax.set_title('Marginal Distribution of MKs and Adipos', fontsize=16)\n",
    "    ax.set_xlabel('Position.X', fontsize=14)\n",
    "    ax.set_ylabel('Marginized Probability Density', fontsize=14)\n",
    "    ax.legend(loc='upper right')\n",
    "    # Set the min-max for y-axis\n",
    "    # Get the current y-axis limits\n",
    "    y_min, y_max = ax.get_ylim()\n",
    "    print(y_max)\n",
    "    y_min = 0\n",
    "    y_max = 0.005\n",
    "\n",
    "    # Set the y-ticks only at the min and max\n",
    "    ax.set_yticks([y_min, y_max])\n",
    "\n",
    "    # Set the y-tick labels to display min and max values\n",
    "    ax.set_yticklabels(['min', 'max'])\n",
    "    \n",
    "    ax.set_ylim(y_min, y_max)\n",
    "    \n",
    "    # Optionally, save the figure\n",
    "    if save_dir:\n",
    "        if filename:\n",
    "            fig_filename = os.path.join(save_dir, f'{datetime.today().strftime(\"%y%m%d\")}_YL_1D_KDE_min_max_{filename}.pdf')\n",
    "        else:\n",
    "            fig_filename = os.path.join(save_dir, f'{datetime.today().strftime(\"%y%m%d\")}_YL_1D_KDE_MKs_Adipos.pdf')\n",
    "        fig.savefig(fig_filename, dpi=300)\n",
    "    \n",
    "    plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 43,
   "metadata": {},
   "outputs": [],
   "source": [
    "'''\n",
    "# Read the saved files\n",
    "with open(f'{save_dir}/kde_results_ss_overlap_bones.pkl', 'rb') as f:\n",
    "    kde_results_ss_overlap_bones = pickle.load(f)\n",
    "with open(f'{save_dir}/common_grid_overlap_bones.pkl', 'rb') as f:\n",
    "    common_grid_overlap_bones = pickle.load(f)\n",
    "\n",
    "with open(f'{save_dir}/kde_results_d5_overlap_bones.pkl', 'rb') as f:\n",
    "    kde_results_d5_overlap_bones = pickle.load(f)\n",
    "with open(f'{save_dir}/kde_results_d10_overlap_bones.pkl', 'rb') as f:\n",
    "    kde_results_d10_overlap_bones = pickle.load(f)\n",
    "with open(f'{save_dir}/kde_results_d15_overlap_bones.pkl', 'rb') as f:\n",
    "    kde_results_d15_overlap_bones = pickle.load(f)\n",
    "with open(f'{save_dir}/kde_results_d30_overlap_bones.pkl', 'rb') as f:\n",
    "    kde_results_d30_overlap_bones = pickle.load(f)\n",
    "\n",
    "\n",
    "with open(f'{save_dir}/kde_results_bone_outline.pkl', 'rb') as f:\n",
    "    kde_results_bone_outline = pickle.load(f)\n",
    "    \n",
    "with open(f'{save_dir}/kde_results_d15_overlap_bones_Detex.pkl', 'rb') as f:\n",
    "    kde_results_d15_overlap_bones_Detex = pickle.load(f)\n",
    "    \n",
    "# Load the saved transformed bone positions\n",
    "transformed_ss_bones_df = pd.read_csv(f'{save_dir}/transformed_ss_bones_inside_df.csv')\n",
    "transformed_d5_bones_df = pd.read_csv(f'{save_dir}/transformed_d5_bones_inside_df.csv')\n",
    "transformed_d10_bones_df = pd.read_csv(f'{save_dir}/transformed_d10_bones_inside_df.csv')\n",
    "transformed_d15_bones_df = pd.read_csv(f'{save_dir}/transformed_d15_bones_inside_df.csv')\n",
    "transformed_d30_bones_df = pd.read_csv(f'{save_dir}/transformed_d30_bones_inside_df.csv')\n",
    "\n",
    "common_outline_filename = os.path.join(save_dir, '241219_YL_common_outline.pkl')\n",
    "with open(common_outline_filename, 'rb') as file:\n",
    "    common_outline = pickle.load(file)\n",
    "'''"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Combine the bone and bone-outline\n",
    "xi, yi, _ = common_grid_overlap_bones\n",
    "# Masked the values outside the bone outline\n",
    "mask = points_in_polygon(xi.flatten(), yi.flatten(), common_outline).reshape(xi.shape)\n",
    "\n",
    "for kde_results in [kde_results_ss_overlap_bones, kde_results_d5_overlap_bones, kde_results_d10_overlap_bones, kde_results_d15_overlap_bones, kde_results_d30_overlap_bones, kde_results_d15_overlap_bones_Detex]:\n",
    "    for source in kde_results.keys():\n",
    "        if source == 'Bones':\n",
    "            zi = kde_results[source]\n",
    "            zi[~mask] = 0\n",
    "        \n",
    "            # Normalize the bone KDE\n",
    "            zi_inside_mask = zi[mask]\n",
    "            norm_bone = Normalize(vmin=zi_inside_mask.min(), vmax=zi_inside_mask.max())\n",
    "            normed_z_bone = norm_bone(zi)\n",
    "            normed_z_bone[~mask] = 0\n",
    "            \n",
    "            # Normalize the bone outline KDE\n",
    "            zi_bone_outline = kde_results_bone_outline['Bones_Outline']\n",
    "            zi_bone_outline[~mask] = 0\n",
    "            zi_bone_outline_inside_mask = zi_bone_outline[mask]\n",
    "            normed_z_bone_outline = norm_bone(zi_bone_outline)\n",
    "            normed_z_bone_outline[~mask] = 0\n",
    "            \n",
    "            # Combine the normalized KDEs using the same weights\n",
    "            combined_normed_z = normed_z_bone * 0.5 + normed_z_bone_outline * 0.5\n",
    "            \n",
    "            # Optionally normalize the combined result again\n",
    "            combined_inside_mask = combined_normed_z[mask]\n",
    "            norm_combined = Normalize(vmin=combined_inside_mask.min(), \n",
    "            vmax=combined_inside_mask.max())\n",
    "            normed_combined_final = norm_combined(combined_normed_z)\n",
    "            normed_combined_final[~mask] = 0\n",
    "            \n",
    "            kde_results[source] = normed_combined_final  # Update with the combined result\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Plot the KDE results per bone with the corresponding outline\n",
    "plot_kde_results_with_df(transformed_ss_bones_df, kde_results_ss_overlap_bones, common_grid_overlap_bones, save_dir, mk_color_map, common_outline, 'd0')\n",
    "plot_kde_results_with_df(transformed_d5_bones_df, kde_results_d5_overlap_bones, common_grid_overlap_bones, save_dir, mk_color_map, common_outline, 'd5')\n",
    "plot_kde_results_with_df(transformed_d10_bones_df, kde_results_d10_overlap_bones, common_grid_overlap_bones, save_dir, mk_color_map, common_outline, 'd10')\n",
    "plot_kde_results_with_df(transformed_d15_bones_df, kde_results_d15_overlap_bones, common_grid_overlap_bones, save_dir, mk_color_map, common_outline, 'd15')\n",
    "plot_kde_results_with_df(transformed_d30_bones_df, kde_results_d30_overlap_bones, common_grid_overlap_bones, save_dir, mk_color_map, common_outline, 'd30')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the KDE results per bone with the corresponding outline\n",
    "plot_kde_results_with_df(transformed_ss_bones_df, kde_results_ss_overlap_bones, common_grid_overlap_bones, save_dir, hscs_color_map, common_outline, 'hsc_rd_d0', scatter=True)\n",
    "plot_kde_results_with_df(transformed_d5_bones_df, kde_results_d5_overlap_bones, common_grid_overlap_bones, save_dir, hscs_color_map, common_outline, 'hsc_rd_d5', scatter=True)\n",
    "plot_kde_results_with_df(transformed_d10_bones_df, kde_results_d10_overlap_bones, common_grid_overlap_bones, save_dir, hscs_color_map, common_outline, 'hsc_rd_d10', scatter=True)\n",
    "plot_kde_results_with_df(transformed_d15_bones_df, kde_results_d15_overlap_bones, common_grid_overlap_bones, save_dir, hscs_color_map, common_outline, 'hsc_rd_d15', scatter=True)\n",
    "plot_kde_results_with_df(transformed_d30_bones_df, kde_results_d30_overlap_bones, common_grid_overlap_bones, save_dir, hscs_color_map, common_outline, 'hsc_rd_d30', scatter=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# always the same y_min and y_max\n",
    "plot_1D_kde(kde_results_ss_overlap_bones, common_grid_overlap_bones, save_dir, mk_color_map, common_outline, filename='mk_adipos_d0')\n",
    "plot_1D_kde(kde_results_d5_overlap_bones, common_grid_overlap_bones, save_dir, mk_color_map, common_outline, filename='mk_adipos_d5')\n",
    "plot_1D_kde(kde_results_d10_overlap_bones, common_grid_overlap_bones, save_dir, mk_color_map, common_outline, filename='mk_adipos_d10')\n",
    "plot_1D_kde(kde_results_d15_overlap_bones, common_grid_overlap_bones, save_dir, mk_color_map, common_outline, filename='mk_adipos_d15')\n",
    "plot_1D_kde(kde_results_d30_overlap_bones, common_grid_overlap_bones, save_dir, mk_color_map, common_outline, filename='mk_adipos_d30')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# sources = ['MKs', 'Adipos', 'Sinusoids', 'MSCs', 'Neurons', 'Arteries', 'Bones']\n",
    "sources = ['MKs', 'Adipos', 'Sinusoids', 'MSCs', 'Neurons', 'Arteries']\n",
    "plot_1D_kde(kde_results_ss_overlap_bones, common_grid_overlap_bones, save_dir, mk_color_map, common_outline, sources = sources, filename='all_d0')\n",
    "plot_1D_kde(kde_results_d5_overlap_bones, common_grid_overlap_bones, save_dir, mk_color_map, common_outline, sources = sources, filename='all_d5')\n",
    "plot_1D_kde(kde_results_d10_overlap_bones, common_grid_overlap_bones, save_dir, mk_color_map, common_outline, sources = sources, filename='all_d10')\n",
    "plot_1D_kde(kde_results_d15_overlap_bones, common_grid_overlap_bones, save_dir, mk_color_map, common_outline, sources = sources, filename='all_d15')\n",
    "plot_1D_kde(kde_results_d30_overlap_bones, common_grid_overlap_bones, save_dir, mk_color_map, common_outline, sources = sources, filename='all_d30')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_kde_results_MK_Adipos(kde_results_ss_overlap_bones, common_grid_overlap_bones, save_dir, mk_color_map, common_outline, 'mk_adipos_d0')\n",
    "plot_kde_results_MK_Adipos(kde_results_d5_overlap_bones, common_grid_overlap_bones, save_dir, mk_color_map, common_outline, 'mk_adipos_d5')\n",
    "plot_kde_results_MK_Adipos(kde_results_d10_overlap_bones, common_grid_overlap_bones, save_dir, mk_color_map, common_outline, 'mk_adipos_d10')\n",
    "plot_kde_results_MK_Adipos(kde_results_d15_overlap_bones, common_grid_overlap_bones, save_dir, mk_color_map, common_outline, 'mk_adipos_d15')\n",
    "plot_kde_results_MK_Adipos(kde_results_d30_overlap_bones, common_grid_overlap_bones, save_dir, mk_color_map, common_outline, 'mk_adipos_d30')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Clustering, Prediciton and Visualization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.1 Extract Features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [],
   "source": [
    "def calculate_kde_features(conditions, density_features, interested_source, common_grid, bone_outline):\n",
    "    \"\"\"\n",
    "    Calculate KDE features for HSCs or RDs using nearest interpolation.\n",
    "\n",
    "    Parameters:\n",
    "    - conditions: List of tuples, where each tuple contains (condition_name, dataframe, kde_results_feature).\n",
    "    - density_features: List of feature names (e.g., ['MKs', 'Adipos', 'Sinusoids', ...]).\n",
    "    - interested_source: The source to calculate KDE for ('HSCs' or 'RDs').\n",
    "    - common_grid: Tuple containing xi, yi, and other grid parameters.\n",
    "    - bone_outline: Outline for masking KDE values outside the bone area.\n",
    "\n",
    "    Returns:\n",
    "    - DataFrame with KDE values for the specified source across all conditions.\n",
    "    \"\"\"\n",
    "    xi, yi, _ = common_grid\n",
    "    xi_flattened = xi.flatten()\n",
    "    yi_flattened = yi.flatten()\n",
    "\n",
    "    # Generate the bone mask once since it's the same for all conditions\n",
    "    mask = points_in_polygon(xi_flattened, yi_flattened, bone_outline).reshape(xi.shape)\n",
    "\n",
    "    # Prepare a list to collect DataFrames for each condition\n",
    "    kde_feature_list = []\n",
    "\n",
    "    # Loop over each condition and process KDE\n",
    "    for condition_name, df, kde_results_feature in conditions:\n",
    "        \n",
    "        # Filter the data based on the interested source (HSCs or RDs)\n",
    "        positions = df[df['source'] == interested_source][['Position.X', 'Position.Y']].values\n",
    "        dataset = df[df['source'] == interested_source]['dataset'].values\n",
    "        \n",
    "        # Prepare a dictionary to store KDE features for the current condition\n",
    "        kde_values_dict = {'Position.X': positions[:, 0], 'Position.Y': positions[:, 1], 'condition': condition_name, 'dataset': dataset, 'group': interested_source}\n",
    "        \n",
    "        # Loop over each density feature and interpolate KDE values\n",
    "        for source in density_features:\n",
    "            zi = kde_results_feature[source].copy()\n",
    "            # if source == 'Bones':\n",
    "                # we also take the KDE values using the bone outline\n",
    "            #     zi_bone_outline = kde_results_bone_outline['Bones_Outline']\n",
    "            #     zi =+ zi_bone_outline\n",
    "            zi[~mask] = 0  # Mask KDE values outside the bone outline\n",
    "            \n",
    "            # Normalize the KDE values within the mask to the range 0-1\n",
    "            zi_inside_mask = zi[mask]\n",
    "            norm = Normalize(vmin=zi_inside_mask.min(), vmax=zi_inside_mask.max())\n",
    "            normed_z = norm(zi)\n",
    "            normed_z[~mask] = 0\n",
    "            \n",
    "                \n",
    "            # Interpolate KDE values for the interested source positions\n",
    "            kde_values = griddata(\n",
    "                (xi_flattened, yi_flattened),\n",
    "                normed_z.flatten(),\n",
    "                (positions[:, 0], positions[:, 1]),\n",
    "                method='nearest',\n",
    "                fill_value=0\n",
    "            )\n",
    "            \n",
    "            # Add the interpolated KDE values to the dictionary\n",
    "            kde_values_dict[source] = kde_values\n",
    "        \n",
    "        # Append the dictionary as a DataFrame for the current condition\n",
    "        kde_feature_list.append(pd.DataFrame(kde_values_dict))\n",
    "    \n",
    "    # Concatenate all condition DataFrames into a single DataFrame\n",
    "    return pd.concat(kde_feature_list, ignore_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "metadata": {},
   "outputs": [],
   "source": [
    "conditions_pos_kde = [\n",
    "    (\"d0\", transformed_ss_bones_df, kde_results_ss_overlap_bones),\n",
    "    (\"d5\", transformed_d5_bones_df, kde_results_d5_overlap_bones),\n",
    "    (\"d10\", transformed_d10_bones_df, kde_results_d10_overlap_bones),\n",
    "    (\"d15\", transformed_d15_bones_df, kde_results_d15_overlap_bones),\n",
    "    (\"d30\", transformed_d30_bones_df, kde_results_d30_overlap_bones),\n",
    "]\n",
    "\n",
    "\n",
    "density_features = ['MKs', 'Adipos', 'Sinusoids', 'MSCs', 'Neurons', 'Arteries', 'Bones']\n",
    "\n",
    "kde_values_for_hscs_nearest = calculate_kde_features(\n",
    "    conditions=conditions_pos_kde,\n",
    "    density_features=density_features,\n",
    "    interested_source=\"HSCs\",\n",
    "    common_grid=common_grid_overlap_bones,\n",
    "    bone_outline=common_outline\n",
    ")\n",
    "kde_values_for_rds_nearest = calculate_kde_features(\n",
    "    conditions=conditions_pos_kde,\n",
    "    density_features=density_features,\n",
    "    interested_source=\"RDs\",\n",
    "    common_grid=common_grid_overlap_bones,\n",
    "    bone_outline=common_outline\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "metadata": {},
   "outputs": [],
   "source": [
    "condition_pos_kde_inhibitors = [(\"d15_Detex\", transformed_d15_bones_Detex_df, kde_results_d15_overlap_bones_Detex)]\n",
    "\n",
    "density_features_inhibitors = ['MKs', 'Adipos', 'Sinusoids', 'Arteries', 'Bones']\n",
    "\n",
    "kde_values_for_hscs_inhibitors_nearest = calculate_kde_features(\n",
    "    conditions=condition_pos_kde_inhibitors,\n",
    "    density_features=density_features_inhibitors,\n",
    "    interested_source=\"HSCs\",\n",
    "    common_grid=common_grid_overlap_bones,\n",
    "    bone_outline=common_outline\n",
    ")\n",
    "kde_values_for_rds_inhibitors_nearest = calculate_kde_features(\n",
    "    conditions=condition_pos_kde_inhibitors,\n",
    "    density_features=density_features_inhibitors,\n",
    "    interested_source=\"RDs\",\n",
    "    common_grid=common_grid_overlap_bones,\n",
    "    bone_outline=common_outline\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Concat the feature matrix and only keep the shared features\n",
    "kde_values_for_hscs_rds_nearest = pd.concat([kde_values_for_hscs_nearest, kde_values_for_rds_nearest], ignore_index=True)\n",
    "kde_values_for_hscs_rds_inhibitors_nearest = pd.concat([kde_values_for_hscs_inhibitors_nearest, kde_values_for_rds_inhibitors_nearest], ignore_index=True)\n",
    "\n",
    "shared_columns = list(set(kde_values_for_hscs_rds_nearest.columns) & set(kde_values_for_hscs_rds_inhibitors_nearest.columns))\n",
    "\n",
    "# concatenate the data with and without inhibitors\n",
    "kde_values_for_hscs_rds_all_nearest = pd.concat([kde_values_for_hscs_rds_nearest[shared_columns], kde_values_for_hscs_rds_inhibitors_nearest[shared_columns]], ignore_index=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.2 Consensus Clustering\n",
    "- Consensus clustering\n",
    "- Sihouette score"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. Generate multiple clustering results\n",
    "def generate_clusterings(features, methods):\n",
    "    clusterings = []\n",
    "    for method in methods:\n",
    "        if method['type'] == 'hierarchical':\n",
    "            linkage_matrix = linkage(features, method=method['linkage'])\n",
    "            for n_clusters in method['n_clusters']:\n",
    "                labels = fcluster(linkage_matrix, t=n_clusters, criterion='maxclust')\n",
    "                clusterings.append(labels)\n",
    "        elif method['type'] == 'kmeans':\n",
    "            for n_clusters in method['n_clusters']:\n",
    "                kmeans = KMeans(n_clusters=n_clusters, random_state=method['random_state']).fit(features)\n",
    "                clusterings.append(kmeans.labels_)\n",
    "        elif method['type'] == 'spectral':\n",
    "            for n_clusters in method['n_clusters']:\n",
    "                spectral = SpectralClustering(n_clusters=n_clusters, affinity='nearest_neighbors', random_state=method['random_state'])\n",
    "                labels = spectral.fit_predict(features)\n",
    "                clusterings.append(labels)\n",
    "    return clusterings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [],
   "source": [
    "def compute_distance_matrix(combined_kde_values, density_features, methods=None):\n",
    "    \"\"\"\n",
    "    Compute the distance matrix from a consensus clustering process.\n",
    "\n",
    "    Parameters:\n",
    "        combined_kde_values (pd.DataFrame): The DataFrame containing the KDE values and features.\n",
    "        density_features (list): A list of column names in the DataFrame to use as features for clustering.\n",
    "        save_dir (str): Directory path where necessary files are stored.\n",
    "\n",
    "    Returns:\n",
    "        np.ndarray: Distance matrix based on the consensus clustering.\n",
    "    \"\"\"\n",
    "    # Define clustering methods with different parameters\n",
    "    if methods is None:\n",
    "        methods = [\n",
    "            {'type': 'hierarchical', 'linkage': 'ward', 'n_clusters': range(2, 10)},\n",
    "            {'type': 'hierarchical', 'linkage': 'average', 'n_clusters': range(2, 10)},\n",
    "            {'type': 'kmeans', 'n_clusters': range(2, 10), 'random_state': 0},\n",
    "            {'type': 'spectral', 'n_clusters': range(2, 10), 'random_state': 0}\n",
    "        ]\n",
    "\n",
    "    # Load the combined KDE values and filter features\n",
    "    combined_features = combined_kde_values[density_features]\n",
    "\n",
    "    # Generate clusterings using a predefined function `generate_clusterings`\n",
    "    clusterings = generate_clusterings(combined_features, methods)\n",
    "\n",
    "    # Initialize the consensus matrix\n",
    "    n_samples = combined_features.shape[0]\n",
    "    consensus_matrix = np.zeros((n_samples, n_samples))\n",
    "\n",
    "    # Build the consensus matrix\n",
    "    for labels in clusterings:\n",
    "        for i in range(n_samples):\n",
    "            for j in range(n_samples):\n",
    "                if labels[i] == labels[j]:  # Increase count if two samples are in the same cluster\n",
    "                    consensus_matrix[i, j] += 1\n",
    "\n",
    "    # Normalize the consensus matrix by the number of clustering results\n",
    "    consensus_matrix /= len(clusterings)\n",
    "\n",
    "    # Convert consensus matrix to a distance matrix for clustering\n",
    "    distance_matrix = 1 - consensus_matrix  # Higher consensus becomes lower distance\n",
    "\n",
    "    return distance_matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "metadata": {},
   "outputs": [],
   "source": [
    "def find_optimal_clusters(distance_matrix, filename=None):\n",
    "    \"\"\"\n",
    "    Determine the optimal number of clusters using silhouette scores.\n",
    "\n",
    "    Parameters:\n",
    "        distance_matrix (np.ndarray): Precomputed distance matrix for clustering.\n",
    "        save_dir (str): Directory where the silhouette plot will be saved (if `plot_flag` is True).\n",
    "        plot_flag (bool): If True, displays and saves the silhouette plot.\n",
    "\n",
    "    Returns:\n",
    "        int: Optimal number of clusters based on the silhouette score.\n",
    "    \"\"\"\n",
    "    # Perform hierarchical clustering on the consensus matrix\n",
    "    linkage_matrix_consensus = linkage(squareform(distance_matrix), method='average')\n",
    "\n",
    "    # Range of cluster numbers to evaluate\n",
    "    cluster_range = range(2, 10)  # Adjust this range as needed\n",
    "    silhouette_scores = []\n",
    "\n",
    "    # Evaluate silhouette scores for each cluster count\n",
    "    for n_clusters in cluster_range:\n",
    "        # Assign clusters based on the number of clusters (cutting the dendrogram)\n",
    "        cluster_labels = fcluster(linkage_matrix_consensus, t=n_clusters, criterion='maxclust')\n",
    "\n",
    "        # Calculate silhouette score for the current number of clusters\n",
    "        score = silhouette_score(distance_matrix, cluster_labels, metric=\"precomputed\")\n",
    "        silhouette_scores.append(score)\n",
    "\n",
    "    # Find the optimal cluster number with the highest silhouette score\n",
    "    optimal_clusters = cluster_range[np.argmax(silhouette_scores)]\n",
    "    print(f\"Optimal number of clusters based on silhouette score: {optimal_clusters}\")\n",
    "\n",
    "    # Plot silhouette scores if plot_flag is True\n",
    "    if filename:\n",
    "        plt.figure(figsize=(8, 5))\n",
    "        plt.plot(cluster_range, silhouette_scores, marker='o')\n",
    "        plt.xlabel(\"Number of clusters\")\n",
    "        plt.ylabel(\"Silhouette Score\")\n",
    "        plt.title(\"Silhouette Score vs. Number of Clusters\")\n",
    "        fig_filename = os.path.join(save_dir, f'{datetime.today().strftime(\"%y%m%d\")}_YL_silhouette_score_{filename}.pdf')\n",
    "        plt.savefig(fig_filename, dpi=300)\n",
    "        plt.show()\n",
    "\n",
    "    return optimal_clusters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Distance matrix for combined KDE values with inhibitors 241205 \n",
    "distance_matrix_w_inhibitors = compute_distance_matrix(kde_values_for_hscs_rds_all_nearest, density_features_inhibitors)\n",
    "\n",
    "# Distance matrix for combined KDE values without inhibitors\n",
    "distance_matrix_wo_inhibitors = compute_distance_matrix(kde_values_for_hscs_rds_nearest, density_features)\n",
    "\n",
    "# Distance matrix for combined KDE values without inhibitors but only with 5 features\n",
    "# Compute the new distance matrix using the reduced features\n",
    "reduced_features = density_features_inhibitors  # Replace with your selected features\n",
    "reduced_distance_matrix = compute_distance_matrix(kde_values_for_hscs_rds_nearest, density_features_inhibitors)\n",
    "\n",
    "\n",
    "# Save the distance matrices\n",
    "# np.save(os.path.join(save_dir, 'distance_matrix_w_inhibitors_original_version.npy'), distance_matrix_w_inhibitors)\n",
    "# np.save(os.path.join(save_dir, 'distance_matrix_wo_inhibitors_original_version.npy'), distance_matrix_wo_inhibitors)\n",
    "# np.save(os.path.join(save_dir, 'distance_matrix_reduced_original_version.npy'), reduced_distance_matrix)\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate the optimal number of clusters for the distance matrices\n",
    "optimal_clusters = find_optimal_clusters(distance_matrix_wo_inhibitors, filename='wo_inhibitors')\n",
    "combined_kde_values = kde_values_for_hscs_rds_nearest.copy()\n",
    "\n",
    "# Perform hierarchical clustering on the consensus matrix\n",
    "linkage_matrix_consensus = linkage(squareform(distance_matrix_wo_inhibitors), method='average')\n",
    "final_cluster_labels = fcluster(linkage_matrix_consensus, t=optimal_clusters, criterion='maxclust')\n",
    "\n",
    "# Get the cluster label for each sample\n",
    "combined_kde_values['cluster'] = final_cluster_labels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.3 Random Forest for Cluster Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def perform_grouped_cross_validation_with_metrics(model, X, y, groups, n_splits=5):\n",
    "    gkf = GroupKFold(n_splits=n_splits)\n",
    "    accuracies = []\n",
    "    balanced_accuracies = []\n",
    "    precisions = []\n",
    "    recalls = []\n",
    "    f1_scores = []\n",
    "\n",
    "    for train_index, test_index in gkf.split(X, y, groups):\n",
    "        X_train, X_test = X.iloc[train_index], X.iloc[test_index]\n",
    "        y_train, y_test = y.iloc[train_index], y.iloc[test_index]\n",
    "\n",
    "        # Train the model\n",
    "        model.fit(X_train, y_train)\n",
    "\n",
    "        # Predict on the test set\n",
    "        y_pred = model.predict(X_test)\n",
    "        \n",
    "        # Handle cases where y_pred contains labels not in y_test\n",
    "        all_classes = np.unique(np.concatenate([y_test, y_pred]))\n",
    "        \n",
    "        # Calculate metrics for this fold\n",
    "        accuracies.append(accuracy_score(y_test, y_pred))\n",
    "        balanced_accuracies.append(balanced_accuracy_score(y_test, y_pred))\n",
    "        precisions.append(precision_score(y_test, y_pred, average='macro', zero_division=0, labels=all_classes))\n",
    "        recalls.append(recall_score(y_test, y_pred, average='macro', zero_division=0, labels=all_classes))\n",
    "        f1_scores.append(f1_score(y_test, y_pred, average='macro', zero_division=0, labels=all_classes))\n",
    "\n",
    "    # Print average metrics\n",
    "    print(f\"Grouped Cross-Validation Accuracy: {np.mean(accuracies):.4f} ± {np.std(accuracies):.4f}\")\n",
    "    print(f\"Grouped Cross-Validation Balanced Accuracy: {np.mean(balanced_accuracies):.4f} ± {np.std(balanced_accuracies):.4f}\")\n",
    "    print(f\"Macro-Averaged Precision: {np.mean(precisions):.4f} ± {np.std(precisions):.4f}\")\n",
    "    print(f\"Macro-Averaged Recall: {np.mean(recalls):.4f} ± {np.std(recalls):.4f}\")\n",
    "    print(f\"Macro-Averaged F1-Score: {np.mean(f1_scores):.4f} ± {np.std(f1_scores):.4f}\")\n",
    "    return accuracies, balanced_accuracies, precisions, recalls, f1_scores\n",
    "\n",
    "def perform_cross_validation_with_class_reports(model, X, y, cv=5):\n",
    "    skf = StratifiedKFold(n_splits=cv, shuffle=True, random_state=42)\n",
    "    for fold, (train_index, test_index) in enumerate(skf.split(X, y), 1):\n",
    "        X_train, X_test = X.iloc[train_index], X.iloc[test_index]\n",
    "        y_train, y_test = y.iloc[train_index], y.iloc[test_index]\n",
    "\n",
    "        # Train the model\n",
    "        model.fit(X_train, y_train)\n",
    "\n",
    "        # Predict on the test set\n",
    "        y_pred = model.predict(X_test)\n",
    "\n",
    "        # Print classification report for this fold\n",
    "        print(f\"\\nClassification Report for Fold {fold}:\")\n",
    "        print(classification_report(y_test, y_pred))\n",
    "\n",
    "def perform_cross_validation_with_metrics(model, X, y, cv=5):\n",
    "    skf = StratifiedKFold(n_splits=cv, shuffle=True, random_state=42)\n",
    "    accuracies = []\n",
    "    balanced_accuracies = []\n",
    "    precisions = []\n",
    "    recalls = []\n",
    "    f1_scores = []\n",
    "\n",
    "    for train_index, test_index in skf.split(X, y):\n",
    "        X_train, X_test = X.iloc[train_index], X.iloc[test_index]\n",
    "        y_train, y_test = y.iloc[train_index], y.iloc[test_index]\n",
    "\n",
    "        # Train the model\n",
    "        model.fit(X_train, y_train)\n",
    "\n",
    "        # Predict on the test set\n",
    "        y_pred = model.predict(X_test)\n",
    "        \n",
    "        # Handle cases where y_pred contains labels not in y_test\n",
    "        all_classes = np.unique(np.concatenate([y_test, y_pred]))\n",
    "        \n",
    "        # Calculate metrics for this fold\n",
    "        accuracies.append(accuracy_score(y_test, y_pred))\n",
    "        balanced_accuracies.append(balanced_accuracy_score(y_test, y_pred))\n",
    "        precisions.append(precision_score(y_test, y_pred, average='macro', zero_division=0, labels=all_classes))\n",
    "        recalls.append(recall_score(y_test, y_pred, average='macro', zero_division=0, labels=all_classes))\n",
    "        f1_scores.append(f1_score(y_test, y_pred, average='macro', zero_division=0, labels=all_classes))\n",
    "\n",
    "    # Print average metrics\n",
    "    print(f\"Cross-Validation Accuracy: {np.mean(accuracies):.4f} ± {np.std(accuracies):.4f}\")\n",
    "    print(f\"Cross-Validation Balanced Accuracy: {np.mean(balanced_accuracies):.4f} ± {np.std(balanced_accuracies):.4f}\")\n",
    "    print(f\"Macro-Averaged Precision: {np.mean(precisions):.4f} ± {np.std(precisions):.4f}\")\n",
    "    print(f\"Macro-Averaged Recall: {np.mean(recalls):.4f} ± {np.std(recalls):.4f}\")\n",
    "    print(f\"Macro-Averaged F1-Score: {np.mean(f1_scores):.4f} ± {np.std(f1_scores):.4f}\")\n",
    "    return accuracies, balanced_accuracies, precisions, recalls, f1_scores\n",
    "\n",
    "\n",
    "def perform_cross_validation(model, X, y, cv=5):\n",
    "    skf = StratifiedKFold(n_splits=cv, shuffle=True, random_state=42)\n",
    "    scores = []\n",
    "    for train_index, test_index in skf.split(X, y):\n",
    "        X_train, X_test = X.iloc[train_index], X.iloc[test_index]\n",
    "        y_train, y_test = y.iloc[train_index], y.iloc[test_index]\n",
    "        model.fit(X_train, y_train)\n",
    "        scores.append(model.score(X_test, y_test))\n",
    "    print(f\"Cross-Validation Scores: {scores}\")\n",
    "    print(f\"Mean Accuracy: {np.mean(scores):.4f} ± {np.std(scores):.4f}\")\n",
    "    return scores\n",
    "\n",
    "# Load and preprocess datasets\n",
    "ref_cluster_data = combined_kde_values  # Original dataset with labels\n",
    "new_data = kde_values_for_hscs_rds_inhibitors_nearest.copy()  # New dataset for prediction\n",
    "\n",
    "# Define feature sets\n",
    "all_features = ['MKs', 'Adipos', 'Sinusoids', 'MSCs', 'Neurons', 'Arteries', 'Bones']\n",
    "shared_features = ['MKs', 'Adipos', 'Sinusoids', 'Arteries', 'Bones']\n",
    "spatial_features = ['Position.X', 'Position.Y']\n",
    "meta_columns = ['condition', 'dataset', 'group', 'cluster']\n",
    "\n",
    "# Prepare training data\n",
    "X_all = ref_cluster_data[all_features]\n",
    "y = ref_cluster_data['cluster']  # Cluster labels from consensus clustering\n",
    "X_shared = ref_cluster_data[shared_features]\n",
    "X_all_with_position = ref_cluster_data[all_features + spatial_features]\n",
    "X_shared_with_position = ref_cluster_data[shared_features + spatial_features]\n",
    "\n",
    "\n",
    "\n",
    "# Random Forest models\n",
    "rf_all = RandomForestClassifier(n_estimators=100, random_state=42)\n",
    "rf_shared = RandomForestClassifier(n_estimators=100, random_state=42)\n",
    "rf_all_pos = RandomForestClassifier(n_estimators=100, random_state=42)\n",
    "rf_shared_pos = RandomForestClassifier(n_estimators=100, random_state=42)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def perform_and_store_metrics(model, X, y, groups=None, grouped=False, n_splits=5):\n",
    "    \"\"\"\n",
    "    Perform cross-validation (regular or grouped) and store metrics.\n",
    "    \"\"\"\n",
    "    if grouped:\n",
    "        results = perform_grouped_cross_validation_with_metrics(model, X, y, groups, n_splits)\n",
    "    else:\n",
    "        results = perform_cross_validation_with_metrics(model, X, y, n_splits)\n",
    "\n",
    "    accuracies, balanced_accuracies, precisions, recalls, f1_scores = results\n",
    "\n",
    "    means = [\n",
    "        np.mean(accuracies),\n",
    "        np.mean(balanced_accuracies),\n",
    "        np.mean(precisions),\n",
    "        np.mean(recalls),\n",
    "        np.mean(f1_scores)\n",
    "    ]\n",
    "    errors = [\n",
    "        np.std(accuracies),\n",
    "        np.std(balanced_accuracies),\n",
    "        np.std(precisions),\n",
    "        np.std(recalls),\n",
    "        np.std(f1_scores)\n",
    "    ]\n",
    "    return means, errors\n",
    "\n",
    "# Lists to store results\n",
    "regular_cv_means = []\n",
    "regular_cv_errors = []\n",
    "\n",
    "\n",
    "# Models and data to evaluate\n",
    "models = [\n",
    "    (\"7 features\", rf_all, X_all),\n",
    "    (\"5 features\", rf_shared, X_shared),\n",
    "    (\"5 features + spatial\", rf_shared_pos, X_shared_with_position)\n",
    "]\n",
    "\n",
    "# Evaluate each model for regular and grouped CV\n",
    "for model_name, model, X in models:\n",
    "    print(f\"\\nRegular CV ({model_name}):\")\n",
    "    means, errors = perform_and_store_metrics(model, X, y, grouped=False)\n",
    "    regular_cv_means.append(means)\n",
    "    regular_cv_errors.append(errors)\n",
    "\n",
    "\n",
    "# Print Results\n",
    "print(\"Regular CV Means:\", regular_cv_means)\n",
    "print(\"Regular CV Errors:\", regular_cv_errors)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Replace these with the values obtained from the `perform_cross_validation_with_metrics` and `perform_grouped_cross_validation_with_metrics` functions\n",
    "rf_models = ['7 features', '5 features', '5 features + spatial']\n",
    "rf_metrics = ['Accuracy', 'Balanced Accuracy', 'Precision', 'Recall', 'F1-Score']\n",
    "\n",
    "\n",
    "# Plot 1: Regular CV rf_metrics\n",
    "fig, ax1 = plt.subplots(figsize=(10, 6))\n",
    "x = np.arange(len(rf_metrics))\n",
    "bar_width = 0.2\n",
    "\n",
    "for i, model in enumerate(rf_models):\n",
    "    ax1.bar(x + i * bar_width, regular_cv_means[i], width=bar_width, label=model, yerr = regular_cv_errors[i], capsize=5)\n",
    "\n",
    "ax1.set_xlabel('Metrics')\n",
    "ax1.set_ylabel('Values')\n",
    "ax1.set_title('Regular Cross-Validation Metrics')\n",
    "ax1.set_xticks(x + bar_width)\n",
    "ax1.set_xticklabels(rf_metrics)\n",
    "ax1.legend(title='RF Models', bbox_to_anchor=(1.05, 0.5))\n",
    "fig.tight_layout()\n",
    "fig_filename = os.path.join(save_dir, f'{datetime.today().strftime(\"%y%m%d\")}_YL_regular_cv.pdf')\n",
    "fig.savefig(fig_filename, dpi=300)\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save the models for future use\n",
    "joblib.dump(rf_all, os.path.join(save_dir, 'rf_all.pkl'))\n",
    "joblib.dump(rf_shared, os.path.join(save_dir, 'rf_shared.pkl'))\n",
    "joblib.dump(rf_all_pos, os.path.join(save_dir, 'rf_all_pos.pkl'))\n",
    "joblib.dump(rf_shared_pos, os.path.join(save_dir, 'rf_shared_pos.pkl'))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Step 3: Predict Labels for the New Dataset\n",
    "# Prepare new data with and without spatial features\n",
    "new_data_features = new_data[shared_features]\n",
    "new_data_features_with_position = new_data[shared_features + spatial_features]\n",
    "\n",
    "# Predict\n",
    "new_predictions_no_position = rf_shared.predict(new_data_features)\n",
    "new_predictions_with_position = rf_shared_pos.predict(new_data_features_with_position)\n",
    "\n",
    "# Attach predictions to the new dataset\n",
    "new_data['predicted_cluster_no_position'] = new_predictions_no_position\n",
    "new_data['predicted_cluster_with_position'] = new_predictions_with_position\n",
    "\n",
    "\n",
    "# Save or inspect the results\n",
    "# print(\"\\nNew Data Predictions (Without Spatial):\")\n",
    "# print(new_data[['Position.X', 'Position.Y', 'predicted_cluster_no_position']].head())\n",
    "\n",
    "# print(\"\\nNew Data Predictions (With Spatial):\")\n",
    "# print(new_data[['Position.X', 'Position.Y', 'predicted_cluster_with_position']].head())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Calculate overlap between labels predicted with and without spatial information\n",
    "new_data['label_match'] = (\n",
    "    new_data['predicted_cluster_no_position'] == new_data['predicted_cluster_with_position']\n",
    ")\n",
    "\n",
    "# Compute the percentage of matching labels\n",
    "overlap_percentage = new_data['label_match'].mean() * 100\n",
    "\n",
    "# Print the results\n",
    "print(f\"Total number of data points: {len(new_data)}\")\n",
    "print(f\"Number of matching labels: {new_data['label_match'].sum()}\")\n",
    "print(f\"Percentage of overlap between labels: {overlap_percentage:.2f}%\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.4 Visualization\n",
    "- Heatmap\n",
    "- Stacked bar plot\n",
    "- Scatter plot (back projection)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_heatmap(combined_kde_values, linkage_matrix, features, condition_colors, cluster_colors, optimal_clusters, save_dir, filename, sort_col= 'cluster'):\n",
    "    \n",
    "    # Only sort when linkage_matrix is None\n",
    "    if linkage_matrix is None:\n",
    "        # Copy the DataFrame to avoid modifying the original\n",
    "        combined_kde_values = combined_kde_values.copy()\n",
    "\n",
    "        if 'condition' in sort_col:\n",
    "            # Define the predefined order for conditions\n",
    "            condition_order = ['d0', 'd5', 'd10', 'd15', 'd30', 'd15_Detex']\n",
    "\n",
    "            # Convert condition to a Categorical type with the predefined order\n",
    "            combined_kde_values['condition'] = pd.Categorical(\n",
    "                combined_kde_values['condition'],\n",
    "                categories=condition_order,\n",
    "                ordered=True\n",
    "            )\n",
    "\n",
    "        # Sort by both cluster and condition\n",
    "        combined_kde_values = combined_kde_values.sort_values(by=sort_col)\n",
    "        \n",
    "    # Select only the first `optimal_clusters` number of colors\n",
    "    cluster_colors = dict(islice(cluster_colors.items(), optimal_clusters))\n",
    "\n",
    "    # Define row colors based on condition, cluster, and group (only if they exist)\n",
    "    row_colors_combined_cluster = combined_kde_values['cluster'].map(cluster_colors)\n",
    "    row_colors_combined_condition = combined_kde_values['condition'].map(condition_colors)\n",
    "    if combined_kde_values['group'].nunique() > 1:\n",
    "        row_colors_combined_group = combined_kde_values['group'].map({'HSCs': 'gray', 'RDs': 'black'})\n",
    "    else:\n",
    "        row_colors_combined_group = None\n",
    "\n",
    "    # Combine row colors into a DataFrame with condition, cluster, and group colors\n",
    "    row_colors_df_combined = pd.DataFrame({\n",
    "        'Condition': row_colors_combined_condition,\n",
    "        'Cluster': combined_kde_values['cluster'].map(cluster_colors),\n",
    "    })\n",
    "    if row_colors_combined_group is not None:\n",
    "        row_colors_df_combined['Group'] = row_colors_combined_group\n",
    "\n",
    "    # Get the features for plotting\n",
    "    combined_features = combined_kde_values[features]\n",
    "\n",
    "    # Calculate the minimum and maximum value across all feature columns\n",
    "    min_value = combined_features.values.min()\n",
    "    max_value = combined_features.values.max()\n",
    "    \n",
    "    if linkage_matrix is None:\n",
    "        g = sns.clustermap(\n",
    "            combined_features,\n",
    "            row_cluster = False,\n",
    "            col_cluster=False,\n",
    "            cmap='coolwarm',\n",
    "            method='ward',\n",
    "            figsize=(9, 12),\n",
    "            row_colors=row_colors_df_combined,\n",
    "            xticklabels=True,\n",
    "            yticklabels=False,\n",
    "            cbar_pos=(1.2, 0.7, 0.07, 0.18),\n",
    "            colors_ratio=0.02,\n",
    "            vmin=min_value,  # Set the minimum color bar value\n",
    "            vmax=max_value   # Set the maximum color bar value to the max of the data\n",
    "        )\n",
    "        \n",
    "    else:\n",
    "        # Plot the clustermap with custom color bar limits\n",
    "        g = sns.clustermap(\n",
    "            combined_features,\n",
    "            row_linkage=linkage_matrix,\n",
    "            col_cluster=False,\n",
    "            cmap='coolwarm',\n",
    "            method='ward',\n",
    "            figsize=(9, 12),\n",
    "            row_colors=row_colors_df_combined,\n",
    "            xticklabels=True,\n",
    "            yticklabels=False,\n",
    "            cbar_pos=(1.2, 0.7, 0.07, 0.18),\n",
    "            colors_ratio=0.02,\n",
    "            vmin=min_value,  # Set the minimum color bar value\n",
    "            vmax=max_value   # Set the maximum color bar value to the max of the data\n",
    "        )\n",
    "\n",
    "    # Set the color bar ticks and labels\n",
    "    colorbar = g.cax  # Access the color bar\n",
    "    colorbar.set_yticks([min_value, 0.50, max_value])  # Set color bar ticks at min and max values\n",
    "    colorbar.set_yticklabels([f'{min_value:.2f}', f'{0.50:.2f}', f'{max_value:.2f}'])  # Annotate color bar with min and max\n",
    "\n",
    "    # Add legends dynamically based on available conditions, clusters, and groups\n",
    "    legend_handles = [\n",
    "        mpatches.Patch(color=color, label=label) for label, color in condition_colors.items() if label in combined_kde_values['condition'].unique()\n",
    "    ] + [\n",
    "        mpatches.Patch(color=color, label=f'Cluster {i+1}') for i, color in enumerate(cluster_colors.values())\n",
    "    ]\n",
    "    \n",
    "    if row_colors_combined_group is not None:\n",
    "        unique_groups = combined_kde_values['group'].unique()\n",
    "        group_colors = {'HSCs': 'gray', 'RDs': 'black'}\n",
    "        for group in unique_groups:\n",
    "            if group in group_colors:\n",
    "                legend_handles.append(mpatches.Patch(color=group_colors[group], label=group))\n",
    "\n",
    "    plt.legend(\n",
    "        handles=legend_handles,\n",
    "        title='Condition, Cluster, and Group',\n",
    "        bbox_to_anchor=(1.2, -1)\n",
    "    )\n",
    "\n",
    "    # Rasterization\n",
    "    heatmap_ax = g.ax_heatmap\n",
    "    \n",
    "    for artist in heatmap_ax.collections:\n",
    "        artist.set_rasterized(True)\n",
    "\n",
    "    plt.suptitle('Consensus Clustering Heatmap Based on Scaled KDE Values of Different Populations', fontsize=16, y=1)\n",
    "    if save_dir:\n",
    "        fig_filename = os.path.join(save_dir, f'{datetime.today().strftime(\"%y%m%d\")}_YL_heatmaps.pdf')\n",
    "        if filename:\n",
    "            fig_filename = os.path.join(save_dir, f'{datetime.today().strftime(\"%y%m%d\")}_YL_{filename}_heatmaps.pdf')\n",
    "        g.savefig(fig_filename, dpi=300)\n",
    "\n",
    "    plt.show()\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "combined_kde_values_detex = new_data.copy()\n",
    "# Drop the columns predictions with no spatial information and label match\n",
    "combined_kde_values_detex = combined_kde_values_detex.drop(columns=['predicted_cluster_no_position', 'label_match'])\n",
    "# Rename the column predicted cluster with position to cluster\n",
    "combined_kde_values_detex = combined_kde_values_detex.rename(columns={'predicted_cluster_with_position': 'cluster'})\n",
    "\n",
    "# Combined the data with and without inhibitors (drop the columns that do not match)\n",
    "combined_kde_values_detex = pd.concat([combined_kde_values.drop(columns=['MSCs', 'Neurons']), combined_kde_values_detex], ignore_index=True)\n",
    "combined_kde_values_detex.head()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate the heatmap with HSCs and RDs\n",
    "# plot_heatmap(combined_kde_values_detex, None, density_features_inhibitors, condition_colors, cluster_colors, optimal_clusters, save_dir, 'HSCs_RDs_5features_spatial', sort_col=['cluster', 'condition'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate the heatmap with HSCs only and RDs only\n",
    "combined_hsc_values = combined_kde_values[combined_kde_values['group'] == 'HSCs']\n",
    "plot_heatmap(combined_hsc_values, None, density_features, condition_colors, cluster_colors, optimal_clusters, save_dir, 'HSCs_7features_spatial', sort_col=['cluster', 'condition'])\n",
    "combined_rd_values = combined_kde_values[combined_kde_values['group'] == 'RDs']\n",
    "plot_heatmap(combined_rd_values, None, density_features, condition_colors, cluster_colors, optimal_clusters, save_dir, 'RDs_7features_spatial', sort_col=['cluster', 'condition'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate the heatmap with HSCs only and RDs only\n",
    "combined_hsc_values = combined_kde_values_detex[combined_kde_values_detex['group'] == 'HSCs']\n",
    "plot_heatmap(combined_hsc_values, None, density_features_inhibitors, condition_colors, cluster_colors, optimal_clusters, save_dir, 'HSCs_5features_spatial', sort_col=['cluster', 'condition'])\n",
    "combined_rd_values = combined_kde_values_detex[combined_kde_values_detex['group'] == 'RDs']\n",
    "plot_heatmap(combined_rd_values, None, density_features_inhibitors, condition_colors, cluster_colors, optimal_clusters, save_dir, 'RDs_5features_spatial', sort_col=['cluster', 'condition'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Stacked bar plot with error bars\n",
    "def plot_cluster_composition(\n",
    "    combined_kde_values,\n",
    "    group,\n",
    "    cluster_colors,\n",
    "    desired_order,\n",
    "    save_dir,\n",
    "    plot_type='percentage',  # Options: 'count', 'percentage'\n",
    "    linewidth=3  # Line width for error bars (applicable for percentage plot)\n",
    "):\n",
    "    \"\"\"\n",
    "    Plots the cluster composition by condition as a stacked bar chart, with optional error bars for percentage plots.\n",
    "\n",
    "    Parameters:\n",
    "    - combined_kde_values (pd.DataFrame): Data containing KDE values with 'group', 'condition', and 'cluster' columns.\n",
    "    - group (str): The group to filter by (e.g., 'HSCs').\n",
    "    - cluster_colors (dict): A dictionary mapping clusters to their respective colors.\n",
    "    - desired_order (list): The desired order for the x-axis categories.\n",
    "    - save_dir (str): Directory where the plot will be saved.\n",
    "    - plot_type (str): 'count' for absolute counts, 'percentage' for percentage plot with error bars.\n",
    "    - linewidth (int): Line width for error bars (only applicable for percentage plots).\n",
    "    \"\"\"\n",
    "    # Filter the data for the specified group\n",
    "    filtered_combined_kde_values = combined_kde_values[combined_kde_values['group'] == group]\n",
    "\n",
    "    # Group by condition, cluster, and dataset for detailed aggregation\n",
    "    grouped = filtered_combined_kde_values.groupby(['condition', 'dataset', 'cluster']).size().unstack(fill_value=0)\n",
    "\n",
    "    # Determine plot type\n",
    "    if plot_type == 'count':\n",
    "        # Aggregate across datasets for absolute counts\n",
    "        cluster_counts = grouped.groupby(level='condition').sum()\n",
    "\n",
    "        # Reindex to align with the desired order\n",
    "        cluster_counts = cluster_counts.reindex(desired_order)\n",
    "\n",
    "        # Extract custom colors for each cluster\n",
    "        color_list = [cluster_colors[cluster] for cluster in sorted(cluster_colors.keys())]\n",
    "\n",
    "        # Plot absolute counts\n",
    "        ax = cluster_counts.plot(kind='bar', stacked=True, color=color_list, figsize=(12, 8))\n",
    "        ax.set_title(f'Cluster Composition by Condition ({group})', fontsize=16)\n",
    "        ax.set_ylabel('Number of Cells')\n",
    "    elif plot_type == 'percentage':\n",
    "        # Normalize by row to get percentages\n",
    "        percentages = grouped.div(grouped.sum(axis=1), axis=0)\n",
    "\n",
    "        # Compute mean and SEM for each condition and cluster\n",
    "        mean_percentages = percentages.groupby(level='condition').mean()\n",
    "        sem_percentages = percentages.groupby(level='condition').sem()\n",
    "\n",
    "        # Reindex to align with the desired order\n",
    "        mean_percentages = mean_percentages.reindex(desired_order)\n",
    "        sem_percentages = sem_percentages.reindex(desired_order)\n",
    "\n",
    "        # Extract custom colors for each cluster\n",
    "        color_list = [cluster_colors[cluster] for cluster in sorted(cluster_colors.keys())]\n",
    "\n",
    "        # Plot percentages with error bars\n",
    "        fig, ax = plt.subplots(figsize=(12, 8))\n",
    "        x_positions = np.arange(len(desired_order))  # Positions for the bars\n",
    "        width = 1  # Bar width\n",
    "\n",
    "        # Initialize bottom positions for stacking bars\n",
    "        bottoms = np.zeros(len(desired_order))\n",
    "\n",
    "        # Plot each cluster with error bars\n",
    "        for cluster in mean_percentages.columns:\n",
    "            means = mean_percentages[cluster].fillna(0)\n",
    "            errors = sem_percentages[cluster].fillna(0)\n",
    "            color = cluster_colors[cluster]\n",
    "\n",
    "            ax.bar(\n",
    "                x_positions,\n",
    "                means,\n",
    "                width,\n",
    "                bottom=bottoms,\n",
    "                color=color,\n",
    "                label=cluster,\n",
    "                align='center'\n",
    "            )\n",
    "            # Add error bars\n",
    "            ax.errorbar(\n",
    "                x_positions,\n",
    "                bottoms + means,  # Position error bars at the top of each bar\n",
    "                yerr=[errors, np.zeros_like(errors)],  # Half error bars\n",
    "                fmt=\"none\",\n",
    "                ecolor=\"black\",\n",
    "                elinewidth=linewidth,  # Line width for error bars\n",
    "                capsize=0  # Remove caps\n",
    "            )\n",
    "            # Add bottom caps for error bars manually\n",
    "            for x, y, err in zip(x_positions, bottoms + means, errors):\n",
    "                if err > 0:  # Draw caps only if error exists\n",
    "                    ax.plot([x - 0.1 * width, x + 0.1 * width], [y - err, y - err], color='black', linewidth=linewidth)\n",
    "\n",
    "            bottoms += means  # Update bottoms for next cluster\n",
    "\n",
    "        # Labeling\n",
    "        ax.set_title(f'Cluster Composition by Condition ({group}, Percentage with Error Bars)', fontsize=16)\n",
    "        ax.set_ylabel(f'Percentage of {group}')\n",
    "        ax.set_ylim([0, 1])\n",
    "        ax.set_xlim([-0.5, len(desired_order) - 0.5])\n",
    "        ax.set_xticks(x_positions)\n",
    "        ax.set_xticklabels(desired_order)\n",
    "    else:\n",
    "        raise ValueError(\"Invalid plot_type. Choose 'count' or 'percentage'.\")\n",
    "\n",
    "    # Common labeling\n",
    "    ax.set_xlabel('Condition')\n",
    "    ax.legend(title='Cluster', bbox_to_anchor=(1.05, 1), loc='upper left')\n",
    "    plt.tight_layout()\n",
    "\n",
    "    # Save the plot\n",
    "    if save_dir:\n",
    "        fig_filename = os.path.join(\n",
    "            save_dir,\n",
    "            f\"{datetime.today().strftime('%y%m%d')}_YL_{group}_barplot_{plot_type}.pdf\"\n",
    "        )\n",
    "        plt.savefig(fig_filename, dpi=300)\n",
    "\n",
    "    # Show the plot\n",
    "    plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate the stacked bar plot for HSCs and RDs\n",
    "plot_cluster_composition(combined_kde_values_detex, 'HSCs', cluster_colors, ['d0', 'd5', 'd10', 'd15', 'd30', 'd15_Detex'], save_dir, plot_type='percentage')\n",
    "plot_cluster_composition(combined_kde_values_detex, 'RDs', cluster_colors, ['d0', 'd5', 'd10', 'd15', 'd30', 'd15_Detex'], save_dir, plot_type='percentage')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Scatter plot of the HSCs and RDs with the cluster information (Back projection)\n",
    "def plot_position_scatter(dataframe, hscs_rds_positions_df, cluster_colors, bone_outline, save_dir=None):\n",
    "    \"\"\"\n",
    "    Plot scatter plots for Position.X and Position.Y grouped by condition, dataset, and group.\n",
    "    \n",
    "    Parameters:\n",
    "        dataframe (pd.DataFrame): The input DataFrame with Position.X, Position.Y, condition, dataset, group, and cluster columns.\n",
    "        hscs_rds_positions_df (pd.DataFrame): The DataFrame containing the Position.X and Position.Y for all HSCs and RDs.\n",
    "        cluster_colors (dict): A dictionary mapping cluster IDs to colors.\n",
    "        bone_outline (np.ndarray): The bone outline coordinates.\n",
    "        save_dir (str): Directory to save the plots (optional).\n",
    "    \"\"\"\n",
    "    conditions = dataframe['condition'].unique()\n",
    "    groups = ['HSCs', 'RDs']  \n",
    "    \n",
    "    # Select only the first `optimal_clusters` number of colors\n",
    "    cluster_colors = dict(islice(cluster_colors.items(), len(dataframe['cluster'].unique())))\n",
    "\n",
    "    for condition in conditions:\n",
    "        # Filter data by condition\n",
    "        condition_data = dataframe[dataframe['condition'] == condition]\n",
    "        condition_hscs_rds = hscs_rds_positions_df[hscs_rds_positions_df['condition'] == condition]\n",
    "        datasets = condition_data['dataset'].unique()\n",
    "\n",
    "        # Create a figure with subplots\n",
    "        fig, axes = plt.subplots(len(datasets), len(groups), figsize=(30, 5 * 19.389 /16.486  * len(datasets)),sharex=True, sharey=True) # * 19.389 / 20.266\n",
    "\n",
    "        for i, dataset in enumerate(datasets):\n",
    "            for j, group in enumerate(groups):\n",
    "                # Filter data for current dataset and group\n",
    "                plot_data = condition_data[(condition_data['dataset'] == dataset) & (condition_data['group'] == group)]\n",
    "\n",
    "                outliers_data = condition_hscs_rds[(condition_hscs_rds['dataset'] == dataset) & (condition_hscs_rds['group'] == group)]\n",
    "                # We compare both dataframes to get the outliers\n",
    "                # Step 1: Extract relevant columns for comparison\n",
    "                plot_positions = plot_data[['Position.X', 'Position.Y']]\n",
    "                outlier_positions = outliers_data[['Position.X', 'Position.Y']].copy()\n",
    "\n",
    "                # Step 2: Preserve original indices for outliers_data\n",
    "                outlier_positions['original_index'] = outliers_data.index\n",
    "\n",
    "                # Step 3: Perform the merge using Position.X and Position.Y\n",
    "                matches = outlier_positions.merge(plot_positions, on=['Position.X', 'Position.Y'], how='inner')\n",
    "\n",
    "                # Step 4: Drop matched rows from outliers_data using original indices\n",
    "                outliers_data = outliers_data.drop(matches['original_index']).reset_index(drop=True)\n",
    "                                \n",
    "                # Get corresponding axis\n",
    "                \n",
    "                if len(datasets) == 1:\n",
    "                    ax = axes[j]\n",
    "                else:\n",
    "                    ax = axes[i, j]\n",
    "                \n",
    "                # Scatter plot with cluster coloring\n",
    "                scatter = ax.scatter(\n",
    "                    plot_data['Position.X'], \n",
    "                    plot_data['Position.Y'], \n",
    "                    c=plot_data['cluster'].map(cluster_colors),\n",
    "                    alpha=1, s=100\n",
    "                )\n",
    "                \n",
    "                scatter_outliers = ax.scatter(\n",
    "                    outliers_data['Position.X'], \n",
    "                    outliers_data['Position.Y'], \n",
    "                    c='black',\n",
    "                    alpha=0.5,\n",
    "                    marker='x', s=100\n",
    "                )\n",
    "                \n",
    "                ax.plot(bone_outline[:, 0], bone_outline[:, 1], color='black')\n",
    "                \n",
    "                # Set the xlim and ylim based on the bone_outline\n",
    "                x_min, x_max = bone_outline[:, 0].min(), bone_outline[:, 0].max()\n",
    "                y_min, y_max = bone_outline[:, 1].min(), bone_outline[:, 1].max()\n",
    "                ax.set_xlim(int(x_min)-500, int(x_max)+500)\n",
    "                ax.set_ylim(int(y_min)-300, int(y_max)+300)\n",
    "                \n",
    "                # Set titles and labels\n",
    "                ax.set_title(f\"{condition} - {dataset} - {group}\")\n",
    "                ax.set_xlabel(\"Position.X\")\n",
    "                ax.set_ylabel(\"Position.Y\")\n",
    "                ax.grid(False)\n",
    "                \n",
    "        # Create a shared legend\n",
    "        legend_patches = [\n",
    "            mpatches.Patch(color=color, label=f\"Cluster {cluster_id}\")\n",
    "            for cluster_id, color in cluster_colors.items()\n",
    "        ]\n",
    "        fig.legend(\n",
    "            handles=legend_patches, \n",
    "            title=\"Clusters\", \n",
    "            loc=\"upper center\", \n",
    "            bbox_to_anchor=(0.5, 1), \n",
    "            ncol=len(cluster_colors)\n",
    "        )\n",
    "\n",
    "        # Adjust layout with extra space for the legend\n",
    "        # plt.tight_layout(rect=[0, 0, 1, 0.95])  # Reserve space at the top\n",
    "\n",
    "        # Save or show plot\n",
    "        if save_dir:\n",
    "            plt.savefig(f\"{save_dir}/{condition}_scatter_plot.pdf\", dpi=300)\n",
    "        plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "transformed_ss_bones_df_all = pd.read_csv(f'{save_dir}/transformed_ss_bones_all_df.csv')\n",
    "transformed_d5_bones_df_all = pd.read_csv(f'{save_dir}/transformed_d5_bones_all_df.csv')\n",
    "transformed_d10_bones_df_all = pd.read_csv(f'{save_dir}/transformed_d10_bones_all_df.csv')\n",
    "transformed_d15_bones_df_all = pd.read_csv(f'{save_dir}/transformed_d15_bones_all_df.csv')\n",
    "transformed_d30_bones_df_all = pd.read_csv(f'{save_dir}/transformed_d30_bones_all_df.csv')\n",
    "transformed_d15_bones_Detex_df_all = pd.read_csv(f'{save_dir}/transformed_d15_bones_Detex_all_df.csv')\n",
    "\n",
    "# A function to take all the HSCs and RDs out of the dataframes and also the dataset\n",
    "# 1. keep the data with source HSCs and RDs\n",
    "# 2. keep the dataset information\n",
    "# 3. concatenate the dataframes using d0, d5, d10, d15, d30 by creating a new 'condition' coloumn\n",
    "# 4. return the combined dataframe\n",
    "\n",
    "def combine_dataframes(dataframes, conditions):\n",
    "    \"\"\"\n",
    "    Combine multiple dataframes into a single dataframe with a new condition column.\n",
    "\n",
    "    Parameters:\n",
    "        dataframes (list): A list of dataframes to combine.\n",
    "        conditions (list): A list of condition names to assign to each dataframe.\n",
    "\n",
    "    Returns:\n",
    "        pd.DataFrame: A combined dataframe with a new condition column.\n",
    "    \"\"\"\n",
    "    # Combine dataframes with a new condition column\n",
    "    combined_dataframes = []\n",
    "    for i, df in enumerate(dataframes):\n",
    "        df['condition'] = conditions[i]\n",
    "        combined_dataframes.append(df)\n",
    "    combined_data = pd.concat(combined_dataframes, ignore_index=True)\n",
    "    combined_data = combined_data[combined_data['source'].isin(['HSCs', 'RDs'])]\n",
    "    # change the source to group  \n",
    "    combined_data['group'] = combined_data['source']\n",
    "    return combined_data\n",
    "\n",
    "conditions = ['d0', 'd5', 'd10', 'd15', 'd30',  'd15_Detex'] # 'd15_GW',\n",
    "hscs_rds_positions_all_df = combine_dataframes(\n",
    "    [transformed_ss_bones_df_all, transformed_d5_bones_df_all, transformed_d10_bones_df_all, transformed_d15_bones_df_all, transformed_d30_bones_df_all,  transformed_d15_bones_Detex_df_all], #transformed_d15_bones_GW_df_all,\n",
    "    conditions\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 79,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Caculate the table with the data of each populations in used and all\n",
    "# transformed_ss_bones_df_all, transformed_ss_bones_df, transformed_d5_bones_df_all, etc.\n",
    "\n",
    "conditions = ['d0', 'd5', 'd10', 'd15', 'd30','d15_Detex'] \n",
    "population_order = ['MKs', 'Adipos', 'Sinusoids', 'MSCs', 'Neurons', 'Arteries', 'Bone', 'HSCs', 'RDs']\n",
    "second_level_order = ['Used', 'All', 'Percentage']\n",
    "dataframes_used = [\n",
    "    transformed_ss_bones_df,\n",
    "    transformed_d5_bones_df,\n",
    "    transformed_d10_bones_df,\n",
    "    transformed_d15_bones_df,\n",
    "    transformed_d30_bones_df,\n",
    "    transformed_d15_bones_Detex_df\n",
    "]\n",
    "dataframes_all = [\n",
    "    transformed_ss_bones_df_all,\n",
    "    transformed_d5_bones_df_all,\n",
    "    transformed_d10_bones_df_all,\n",
    "    transformed_d15_bones_df_all,\n",
    "    transformed_d30_bones_df_all,\n",
    "    transformed_d15_bones_Detex_df_all\n",
    "]\n",
    "\n",
    "\n",
    "# Initialize an empty dictionary to store results\n",
    "results = []\n",
    "\n",
    "# Process each condition\n",
    "for condition, df_used, df_all in zip(conditions, dataframes_used, dataframes_all):\n",
    "    # Get value counts for `source`\n",
    "    used_counts = df_used['source'].value_counts()\n",
    "    all_counts = df_all['source'].value_counts()\n",
    "\n",
    "    # Align the indices of the counts for consistent comparison\n",
    "    combined_sources = used_counts.index.union(all_counts.index)\n",
    "    used_counts = used_counts.reindex(combined_sources, fill_value=0)\n",
    "    all_counts = all_counts.reindex(combined_sources, fill_value=0)\n",
    "\n",
    "    # Calculate percentage\n",
    "    percentage = (used_counts / all_counts).fillna(0) * 100\n",
    "\n",
    "    # Add data to results\n",
    "    for source in combined_sources:\n",
    "        results.append({\n",
    "            'Condition': condition,\n",
    "            'Source': source,\n",
    "            'Used': used_counts[source],\n",
    "            'All': all_counts[source],\n",
    "            'Percentage': percentage[source]\n",
    "        })\n",
    "\n",
    "# Convert results to a DataFrame\n",
    "summary_table = pd.DataFrame(results)\n",
    "\n",
    "# Pivot to create the desired table format with sources as primary columns\n",
    "pivot_table = summary_table.pivot(index='Condition', columns='Source', values=['Used', 'All', 'Percentage'])\n",
    "\n",
    "# Reorder columns to have sources as primary headers\n",
    "pivot_table = pivot_table.swaplevel(axis=1).sort_index(axis=1, level=0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 80,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Reindex rows to follow conditions order\n",
    "pivot_table = pivot_table.reindex(conditions)\n",
    "\n",
    "# Reindex columns to follow population_order and second_level_order\n",
    "columns_sorted = sorted(\n",
    "    pivot_table.columns,\n",
    "    key=lambda x: (\n",
    "        population_order.index(x[0]) if x[0] in population_order else len(population_order),  # Primary level (sources)\n",
    "        second_level_order.index(x[1]) if x[1] in second_level_order else len(second_level_order)  # Second level\n",
    "    )\n",
    ")\n",
    "\n",
    "# Reconstruct the MultiIndex columns in the sorted order\n",
    "pivot_table = pivot_table.loc[:, pd.MultiIndex.from_tuples(columns_sorted)]\n",
    "pivot_table.to_csv(f'{save_dir}/data_used_summary_table.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Plot the HSCs and RDs with the bone outline\n",
    "plot_position_scatter(\n",
    "    dataframe=combined_kde_values_detex, \n",
    "    hscs_rds_positions_df=hscs_rds_positions_all_df,\n",
    "    cluster_colors=cluster_colors,\n",
    "    bone_outline=common_outline,\n",
    "    save_dir=save_dir  # Specify a directory or leave as None to just display\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5.5 Whitney-Mann U Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 82,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generate_cluster_compositions_with_proportions(df, cluster_column='cluster'):\n",
    "    # Initialize dictionaries for HSCs and RDs\n",
    "    hsc_data = {}\n",
    "    rd_data = {}\n",
    "    \n",
    "    # Extract unique conditions (dates) and clusters\n",
    "    conditions = df['condition'].unique()\n",
    "    clusters = df[cluster_column].unique()\n",
    "    \n",
    "    for condition in conditions:\n",
    "        # Filter data for the current condition\n",
    "        condition_data = df[df['condition'] == condition]\n",
    "        datasets_in_condition = condition_data['dataset'].unique()  # Datasets specific to this condition\n",
    "        \n",
    "        # Initialize dictionaries for HSCs and RDs for this condition\n",
    "        hsc_cluster_compositions = {}\n",
    "        rd_cluster_compositions = {}\n",
    "        \n",
    "        # Calculate total counts per dataset for HSCs and RDs\n",
    "        hsc_total_counts = condition_data[condition_data['group'] == 'HSCs'].groupby('dataset').size()\n",
    "        rd_total_counts = condition_data[condition_data['group'] == 'RDs'].groupby('dataset').size()\n",
    "        \n",
    "        for cluster in clusters:\n",
    "            # Filter data for HSCs and RDs for the current cluster\n",
    "            hsc_subset = condition_data[(condition_data['group'] == 'HSCs') & (condition_data[cluster_column] == cluster)]\n",
    "            rd_subset = condition_data[(condition_data['group'] == 'RDs') & (condition_data[cluster_column] == cluster)]\n",
    "            \n",
    "            # Calculate cluster-specific counts per dataset\n",
    "            hsc_cluster_counts = hsc_subset.groupby('dataset').size()\n",
    "            rd_cluster_counts = rd_subset.groupby('dataset').size()\n",
    "\n",
    "            \n",
    "            # Calculate proportions for HSCs\n",
    "            hsc_proportions = (hsc_cluster_counts / hsc_total_counts).reindex(datasets_in_condition, fill_value=0)\n",
    "            # Replace nan with 0\n",
    "            hsc_proportions = hsc_proportions.fillna(0)\n",
    "            hsc_cluster_compositions[f\"Cluster_{cluster}\"] = hsc_proportions.to_dict()\n",
    "            \n",
    "            # Calculate proportions for RDs\n",
    "            rd_proportions = (rd_cluster_counts / rd_total_counts).reindex(datasets_in_condition, fill_value=0)\n",
    "            rd_proportions = rd_proportions.fillna(0)\n",
    "            rd_cluster_compositions[f\"Cluster_{cluster}\"] = rd_proportions.to_dict()\n",
    "        \n",
    "        # Add compositions for this condition to the main dictionaries\n",
    "        hsc_data[condition] = hsc_cluster_compositions\n",
    "        rd_data[condition] = rd_cluster_compositions\n",
    "    \n",
    "    return hsc_data, rd_data\n",
    "\n",
    "\n",
    "\n",
    "# Perform Mann-Whitney U Test for each date and cluster with edge case handling\n",
    "def perform_whitney_u_test(hsc_compositions, rd_compositions):\n",
    "    results = []\n",
    "\n",
    "    # Iterate through each condition (date)\n",
    "    for condition in hsc_compositions.keys():\n",
    "        # Get cluster compositions for HSCs and RDs\n",
    "        hsc_clusters = hsc_compositions[condition]\n",
    "        rd_clusters = rd_compositions[condition]\n",
    "        # print(condition)\n",
    "\n",
    "        # Iterate through each cluster\n",
    "        for cluster in hsc_clusters.keys():\n",
    "            # Get proportions for this cluster in HSCs and RDs\n",
    "            hsc_values = list(hsc_clusters[cluster].values())\n",
    "            rd_values = list(rd_clusters[cluster].values())\n",
    "            # print(cluster)\n",
    "            # print(hsc_values)\n",
    "            # print(rd_values)\n",
    "\n",
    "            # Handle edge cases\n",
    "            if all(value == 0 for value in hsc_values + rd_values):\n",
    "                # All zeros in both groups\n",
    "                stat, p_value = None, 1.0  # No difference\n",
    "            # elif all(value == 0 for value in hsc_values) or all(value == 0 for value in rd_values):\n",
    "                # All zeros in one group\n",
    "            #     stat, p_value = None, None  # Test not meaningful\n",
    "            else:\n",
    "                # Perform Mann-Whitney U Test\n",
    "                stat, p_value = mannwhitneyu(hsc_values, rd_values, alternative='two-sided')\n",
    "\n",
    "            # Store results\n",
    "            results.append({\n",
    "                \"Date\": condition,\n",
    "                \"Cluster\": cluster,\n",
    "                \"U_statistic\": stat,\n",
    "                \"p_value\": p_value\n",
    "            })\n",
    "\n",
    "    # Convert results to DataFrame for easy handling\n",
    "    results_df = pd.DataFrame(results)\n",
    "    return results_df\n",
    "\n",
    "# Adjust p-values using Benjamini-Hochberg (FDR)\n",
    "def adjust_p_values(results_df, method='fdr_bh'):\n",
    "    # Extract p-values\n",
    "    p_values = results_df['p_value'].dropna()\n",
    "\n",
    "    # Apply multiple testing correction\n",
    "    _, p_adjusted, _, _ = multipletests(p_values, method=method)\n",
    "    \n",
    "    # Add adjusted p-values to the DataFrame\n",
    "    results_df.loc[results_df['p_value'].notna(), 'adjusted_p_value'] = p_adjusted\n",
    "    return results_df\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hsc_compositions, rd_compositions = generate_cluster_compositions_with_proportions(combined_kde_values_detex, cluster_column='cluster')\n",
    "# Print the compositions to verify\n",
    "print(\"HSC Data:\")\n",
    "print(hsc_compositions['d0']['Cluster_1'])\n",
    "\n",
    "print(\"\\nRD Data:\")\n",
    "print(rd_compositions['d0']['Cluster_1'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Run the Whitney U test\n",
    "utest_cluster_compositions = perform_whitney_u_test(hsc_compositions, rd_compositions)\n",
    "\n",
    "# Order the df by the Date first then Cluster\n",
    "# The Date should follow the order ['d0', 'd5', 'd10', 'd15', 'd30']\n",
    "\n",
    "utest_cluster_compositions['Date'] = pd.Categorical(utest_cluster_compositions['Date'], categories=['d0', 'd5', 'd10', 'd15', 'd30', 'd15_Detex'], ordered=True)\n",
    "\n",
    "utest_cluster_compositions = utest_cluster_compositions.sort_values(by=['Date', 'Cluster'])\n",
    "\n",
    "adjusted_utest_cluster_compositions = adjust_p_values(utest_cluster_compositions, method='fdr_bh')\n",
    "\n",
    "# Mark significant results\n",
    "adjusted_utest_cluster_compositions['significant'] = adjusted_utest_cluster_compositions['adjusted_p_value'] < 0.05\n",
    "\n",
    "\n",
    "# View significant results\n",
    "print(adjusted_utest_cluster_compositions[adjusted_utest_cluster_compositions['significant']])\n",
    "\n",
    "\n",
    "# Save the comparison results to a CSV file\n",
    "adjusted_utest_cluster_compositions.to_csv(os.path.join(save_dir, f'{datetime.today().strftime(\"%y%m%d\")}_YL_cluster_compositions_utest_feature_Detex.csv'), index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Perform the U test between HSCs of d15 and d15 Detex\n",
    "\n",
    "# Extract the data for d15 and d15 Detex\n",
    "# Create dictionaries to store the compositions for d15 and d15_Detex\n",
    "hsc_compositions_d15 = {'d15': hsc_compositions['d15']}\n",
    "hsc_compositions_d15_detex = {'d15': hsc_compositions['d15_Detex']}\n",
    "\n",
    "# Perform the U test\n",
    "utest_d15_d15_detex = perform_whitney_u_test(hsc_compositions_d15, hsc_compositions_d15_detex)\n",
    "\n",
    "# Order the df by the Cluster\n",
    "utest_d15_d15_detex = utest_d15_d15_detex.sort_values(by=['Cluster'])\n",
    "\n",
    "# Adjust the p-values\n",
    "adjusted_utest_d15_d15_detex = adjust_p_values(utest_d15_d15_detex, method='fdr_bh')\n",
    "\n",
    "# Mark significant results\n",
    "adjusted_utest_d15_d15_detex['significant'] = adjusted_utest_d15_d15_detex['adjusted_p_value'] < 0.05\n",
    "\n",
    "# View significant results\n",
    "print(adjusted_utest_d15_d15_detex[adjusted_utest_d15_d15_detex['significant']])\n",
    "\n",
    "# Save the comparison results to a CSV file\n",
    "adjusted_utest_d15_d15_detex.to_csv(os.path.join(save_dir, f'{datetime.today().strftime(\"%y%m%d\")}_YL_cluster_compositions_utest_d15_d15_Detex.csv'), index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "imgcap",
   "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.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}