colocalization.py

###############################################################################################
'''
Process each condition separetly:
- create output path for each condition
- iterate over protein names and create pairs
- extract file path for each protein and its mask

Process each protein/pair:
- extract file names, locs and info, load mask/s
- filter locs if enabled
- uses additional mask if exists, otherwise uses locs outside of the mask
- optional: saves filtered locs and unspecific locs (must be uncommented)
- each protein: extracts all relevant data from original file as np.array -> performs self colocalization using query_ball_point from KDTree to find all neighbours within a threshold distance. 
OR query_ball_tree to find all neighbours.
- Calculates Euclidean distance (√[(x₁ - x₂)² + (y₁ - y₂)²]) between cluster center and each neighbour. Indx in locs_data and self_neighbours results are preserved!
- if dark_time_analysis=True -> extract Td values from unspecific locs -> calculate 1/Td -> filter extrime 1/Td values ->
performs gaussian fitting and extract xc (center), w(width) and number of Td values used for fitting. 
- adds calibration values to protein_df

Process protein pair:
-extract file names, locs and info, load mask/s
-uses query_ball_point from KDTree to find all neighbours within a threshold distance.  Calculates Euclidean distance (√[(x₁ - x₂)² + (y₁ - y₂)²]) between cluster center and each neighbour.
Indx in locs_data and self_neighbours results are preserved!
-saves all colocalzing clusters in hdf5 files
'''
###############################################################################################
import os
import numpy as np
from itertools import permutations
from scipy.spatial import KDTree
import gc
import pickle

from collectors import AnalysisCollector, add_calibration_to_protein_data, extract_protein_data, add_coloc_to_protein_data
from utils import load_mask, mask_locs, load_and_filter_locs, save_filtered_data, save_colocalizing_protein_data
from utils import get_patches, assign_locs_to_patches, simulate_locs_data, save_hdf5, load_info
from dark_time_analysis import calculate_calibration_values

###############################################################################################

  
def self_clocalization(locs_data, pxl_size, coloc_threshold):
    """
    For each cluster, count and get distances to all other clusters within threshold (excluding self).
    Calculate Euclidean Distances between colocalizing coordinates
    Returns a list of dicts: [{'n_self_neighbors': int, 'self_neighbors_distances': [float, ...]}, ...]
    The inx in locs_data is preserved in coords and later in enumerate(all_neighbors)
    """
    if locs_data is None:
        print("Warning: locs_data is None in self_clocalization")
        return []
    
    if len(locs_data) < 2:
        print("Warning: Not enough data points for self-neighbor analysis")
        return []
    
    coords = np.column_stack((locs_data['x'] * pxl_size, locs_data['y'] * pxl_size))
    tree = KDTree(coords)
    all_neighbors = tree.query_ball_tree(tree, r=coloc_threshold)
    results = []
    for idx, neighbors in enumerate(all_neighbors):
        # Exclude self
        neighbors = [n for n in neighbors if n != idx]
        distances = [np.linalg.norm(coords[idx] - coords[n]) for n in neighbors]
        
        self_n_locs_closest = 0
        if neighbors:
            closest_idx = neighbors[np.argmin(distances)]
            self_n_locs_closest = locs_data[closest_idx]['n']
            
        results.append({
            'n_self_neighbors': len(neighbors),
            'self_neighbors_distances': distances,
            'self_n_locs_closest': self_n_locs_closest
        })
    return results    

# def self_clocalization(locs_data, pxl_size, coloc_threshold):
#     """
#     For each cluster, find only the single nearest other cluster within threshold.
#     Uses query_ball_point to first find all neighbors within threshold, then picks the nearest one.
#     """
#     if locs_data is None or len(locs_data) < 2:
#         return []

#     coords = np.column_stack((locs_data['x'] * pxl_size, locs_data['y'] * pxl_size))
#     tree = KDTree(coords)

#     results = []
#     for idx, coord in enumerate(coords):
#         # Find ALL neighbors within threshold first
#         neighbors = tree.query_ball_point(coord, r=coloc_threshold)
        
#         # Remove self from neighbors
#         neighbors = [n for n in neighbors if n != idx]
        
#         if neighbors:  # If any neighbors found within threshold
#             # Calculate distances to all neighbors within threshold
#             distances = [np.linalg.norm(coord - coords[n]) for n in neighbors]
            
#             # Find the nearest one among those within threshold
#             nearest_distance = min(distances)
            
#             # Return the nearest neighbor within threshold
#             self_neighbors_distances = [float(nearest_distance)]
#         else:
#             # No neighbors within threshold
#             self_neighbors_distances = []

#         results.append({
#             'n_self_neighbors': 1 if len(self_neighbors_distances) == 1 else 0,
#             'self_neighbors_distances': self_neighbors_distances
#         })
    
#     return results


def single_protein_data(protein, condition_data, idx, 
                    data_mask, calibration_mask,
                    condition_name, pxl_size, condition_out_path,
                    filter_params, filter_enabled, patch_size=None
                   ):
    """ 
    Processing individual prtein, extracting locs data, mask/s 
    Returns: masked data and unspecific signals data, file path
    """
    
        
    file_path = condition_data[protein][idx]    
    mask_path = file_path.replace('.hdf5', data_mask)
    calib_mask_path = file_path.replace('.hdf5', calibration_mask)
    
    locs_data, info = load_and_filter_locs(file_path, filter_params, filter_enabled)
    mask = load_mask(mask_path)
    calib_mask = load_mask(calib_mask_path)
    
    # Check if masks were loaded successfully
    if mask is None:
        print("Warning: Skipping processing due to missing mask files")
        return
    
    locs_data_masked, _ = mask_locs(locs_data, mask)
    
    # Use calibration mask for unspecific signal (if available)
    if calib_mask is not None:
        unspecific_locs, _ = mask_locs(locs_data, calib_mask)

    else:
        # Fallback to original method if calibration masks not available
        print("Calibration masks not found")
        _, unspecific_locs = mask_locs(locs_data, mask)
            
    patch_data = None
    if patch_size is not None:
        patches = get_patches(mask, pxl_size, patch_size)
        patch_ids = assign_locs_to_patches(locs_data_masked, patches)
        patch_data = {
            'patches': patches,
            'patch_ids': patch_ids
        }
        
    # Save filtered data for testing
    # save_filtered_data(file_path, locs_data_masked, unspecific_locs, info, condition_out_path)
    
    return locs_data_masked, unspecific_locs, mask, file_path, patch_data


def cross_colocalization(protein1, protein2, locs_data_masked_1, locs_data_masked_2, pxl_size, coloc_threshold):
    """
    For each cluster in protein1, count neighbors from protein2 within threshold
    Calculate Euclidean Distances between colocalizing coordinates
    Returns: [{'n_neighbors_protein2': int, 'distances_protein2': [float, ...]}, ...]
    The inx in locs_data is preserved in coords and later in neighbors
    """
    coords_1 = np.column_stack((locs_data_masked_1['x'] * pxl_size, locs_data_masked_1['y'] * pxl_size))
    coords_2 = np.column_stack((locs_data_masked_2['x'] * pxl_size, locs_data_masked_2['y'] * pxl_size))
    
    # tree_1 = KDTree(coords_1)
    tree_2 = KDTree(coords_2)  # Tree for protein2
    
    results = []
    for idx, coord_1 in enumerate(coords_1):  # For each protein1 cluster
        # Find protein2 neighbors within threshold of this protein1 cluster
        neighbors = tree_2.query_ball_point(coord_1, r=coloc_threshold)
        
        # Calculate Euclidean Distance (L2 norm) between protein1 cluster to all protein2 neighbours
        # Euclidean Distance = √[(x₁ - x₂)² + (y₁ - y₂)²]
        distances = [np.linalg.norm(coord_1 - coords_2[n]) for n in neighbors]
        indxs_protein2 = list(neighbors)
        
        n_locs_protein2 = [locs_data_masked_2[n]['n'] for n in neighbors]
        
        closest_n_locs_protein2 = 0
        if neighbors:
            closest_idx = neighbors[np.argmin(distances)]
            closest_n_locs_protein2 = locs_data_masked_2[closest_idx]['n']
        
        results.append({
            'Protein_pair': f'{protein1}->{protein2}',
            f'idx_{protein1}': idx,
            f'idxs_{protein2}': indxs_protein2,
            f'n_neighbors_{protein2}': len(neighbors),
            f'distances_{protein2}': distances,
            f'n_locs_{protein2}': n_locs_protein2,
            f'closest_n_locs_{protein2}': closest_n_locs_protein2
        })
        
    return results

# def cross_colocalization(protein1, protein2, locs_data_masked_1, locs_data_masked_2, pxl_size, coloc_threshold):
#     """
#     For each cluster in protein1, find only the single nearest neighbor from protein2 within threshold.
#     Uses query_ball_point to first find all neighbors within threshold, then picks the nearest one.
#     """
#     coords_1 = np.column_stack((locs_data_masked_1['x'] * pxl_size, locs_data_masked_1['y'] * pxl_size))
#     coords_2 = np.column_stack((locs_data_masked_2['x'] * pxl_size, locs_data_masked_2['y'] * pxl_size))
    
#     tree_2 = KDTree(coords_2)
    
#     results = []
#     for idx, coord_1 in enumerate(coords_1):
#         # Find ALL neighbors within threshold first
#         neighbors = tree_2.query_ball_point(coord_1, r=coloc_threshold)
        
#         if neighbors:  # If any neighbors found within threshold
#             # Calculate distances to all neighbors within threshold
#             distances = [np.linalg.norm(coord_1 - coords_2[n]) for n in neighbors]
            
#             # Find the nearest one among those within threshold
#             nearest_idx = neighbors[np.argmin(distances)]
#             nearest_distance = min(distances)
            
#             # Return the nearest neighbor within threshold
#             indxs_protein2 = [int(nearest_idx)]
#             distances_list = [float(nearest_distance)]
#             n_locs_protein2 = [locs_data_masked_2[int(nearest_idx)]['n']]
#         else:
#             # No neighbors within threshold
#             indxs_protein2 = []
#             distances_list = []
#             n_locs_protein2 = []
        
#         results.append({
#             'Protein_pair': f'{protein1}->{protein2}',
#             f'idx_{protein1}': idx,
#             f'idxs_{protein2}': indxs_protein2,
#             f'n_neighbors_{protein2}': len(indxs_protein2),
#             f'distances_{protein2}': distances_list,
#             f'n_locs_{protein2}': n_locs_protein2
#         })

#     return results

def process_protein_pair(protein1, protein2, condition_data, idx, 
                         data_mask, calibration_mask, 
                         condition_name, 
                         pxl_size, coloc_threshold, 
                         condition_out_path,
                         filter_params, filter_enabled, 
                         patching, patch_size, patch_id
                         ):
    
    """Extracting locs data from two proteins and find colocalization"""
    # Load both proteins
    locs_data_masked_1, _, _, file_path_1, patch_data_1 = single_protein_data(protein1, condition_data, idx, 
                                                    data_mask, calibration_mask,
                                                    condition_name, pxl_size, condition_out_path,
                                                    filter_params, filter_enabled, patch_size if patching else None
                                                    )
    locs_data_masked_2, _, _, file_path_2, patch_data_2 = single_protein_data(protein2, condition_data, idx, 
                                                    data_mask, calibration_mask,
                                                    condition_name, pxl_size, condition_out_path,
                                                    filter_params, filter_enabled, patch_size if patching else None
                                                    )
    
    # Cross-colocalization
    coloc_results = cross_colocalization(protein1, protein2, locs_data_masked_1, locs_data_masked_2, pxl_size, coloc_threshold)
    
    # Save colocalizing clusters    
    coloc_indices_protein1 = [result[f'idx_{protein1}'] for result in coloc_results if result[f'n_neighbors_{protein2}'] > 0]
    if len(coloc_indices_protein1) > 0:
        save_colocalizing_protein_data(locs_data_masked_1, coloc_indices_protein1, protein1, protein2, file_path_1, condition_out_path, patch_id)
    
    coloc_indices_protein2 = list(set([idx for result in coloc_results if result[f'n_neighbors_{protein2}'] > 0 for idx in result[f'idxs_{protein2}']]))
    if len(coloc_indices_protein1) > 0:    
        save_colocalizing_protein_data(locs_data_masked_2, coloc_indices_protein2, protein2, protein1, file_path_2, condition_out_path, patch_id)
    
    return locs_data_masked_1, locs_data_masked_2, coloc_results


def process_condition(
    filter_params: dict,
    filter_enabled: bool,
    condition_data: dict,
    condition_out_path: str,
    condition_name: str,
    collector: AnalysisCollector,
    protein_names,
    coloc_threshold,
    pxl_size,
    data_mask,
    calibration_mask,
    unspecific_dark_list,
    dark_time_analysis=False,
    cross_coloc=False,
    patching=False,
    patch_size=None,
    patch_id=-1, 
    simulation=False
):
    """
    Process all proteins per current condition - one file at a time for memory efficiency
    """
    
    os.makedirs(condition_out_path, exist_ok=True)
    
    # Get the number of files (assuming all proteins must have same number of files)
    num_files = len(condition_data[protein_names[0]])
    
    # Process each file completely 
    for file_idx in range(num_files):
        print(f"\nProcessing file {file_idx + 1}/{num_files}")
        
        # data for this file only
        file_protein_data = []
        file_coloc_results = {}
        
        # === Process single protein data ===
        print("Processing single protein data extraction and self colocalization")
        for protein in protein_names:
            try:
                locs_data_masked, unspecific_locs, mask, file_path, patch_data = single_protein_data(
                    protein, condition_data, file_idx, 
                    data_mask, calibration_mask, 
                    condition_name, pxl_size, condition_out_path,
                    filter_params, filter_enabled, 
                    patch_size 
                )
                
                self_neighbors_results = self_clocalization(locs_data_masked, pxl_size, coloc_threshold)
                
                # Store unspecific data for dark time analysis
                if dark_time_analysis and unspecific_locs is not None and len(unspecific_locs) > 0:
                    unspecific_dark_list.append({
                        'condition': condition_name,
                        'day': os.path.basename(file_path).split('_')[0],
                        'protein': protein,
                        'file_name': os.path.basename(file_path),
                        'unspecific_locs': unspecific_locs,
                        'pxl_size': pxl_size
                    })
                
                # Extract data for this file (handles all patches at once)
                protein_data_list = extract_protein_data(collector, 
                                                      locs_data_masked, mask, 
                                                      condition_name, protein, file_path, 
                                                      pxl_size, 
                                                      self_neighbors_results,
                                                      patch_data, patch_id)
                
                file_protein_data.extend(protein_data_list)
                
                # Clear large data structures immediately
                del locs_data_masked, self_neighbors_results
                if patch_data is not None:
                    del patch_data
                
            except Exception as e:
                print(f"Error loading {protein}-{file_idx}: {str(e)}")
        
        # === Process cross-colocalization for this file ===
        if cross_coloc:
            print(f"  Processing cross colocalization between protein pairs")
            protein_pairs = list(permutations(protein_names, 2))
            print(f"  Protein pairs: {protein_pairs}")
            
            for protein1, protein2 in protein_pairs:
                try:
                    locs_data_masked_1, locs_data_masked_2, coloc_results = process_protein_pair(
                        protein1, protein2, condition_data, file_idx, 
                        data_mask, calibration_mask, 
                        condition_name, 
                        pxl_size, coloc_threshold, 
                        condition_out_path,
                        filter_params, filter_enabled, 
                        patching, patch_size, patch_id
                    )
                    
                    pair_key = f"{protein1}->{protein2}"
                    if pair_key not in file_coloc_results:
                        file_coloc_results[pair_key] = []
                    file_coloc_results[pair_key].extend(coloc_results)
                    
                    # Clear large data structures immediately
                    del locs_data_masked_1, locs_data_masked_2, coloc_results
                    
                except Exception as e:
                    print(f"Error processing {protein1}-{protein2} pair {file_idx}: {str(e)}")
            
            # Add colocalization data for this file only
            add_coloc_to_protein_data(file_protein_data, file_coloc_results, protein_names)
        
        # === Add everything to collector for this file ===  
        for protein_data in file_protein_data:
            collector.add_protein_data(protein_data)
        
        # Clear data for this file from memory
        del file_protein_data, file_coloc_results
        gc.collect()  # Force garbage collection
        
        print(f"  Completed file {file_idx + 1}/{num_files}")
        
        # === Process simulated data for this file ===
        if simulation:
            print(f"  Processing simulated data")
            process_simulated_data(
                condition_data, file_idx, condition_name, collector, protein_names,
                filter_params, filter_enabled, data_mask, calibration_mask,
                pxl_size, coloc_threshold, condition_out_path,
                dark_time_analysis, cross_coloc, unspecific_dark_list,
                patching, patch_size
            )



def run_analysis(output_dir, 
                 structure,
                filter_params,
                filter_enabled,
                condition_names, 
                collector, 
                protein_names,
                coloc_threshold,
                pxl_size,
                data_mask,
                calibration_mask,
                dark_time_analysis,
                cross_coloc,
                patching, patch_size,
                simulation):
    
    unspecific_dark_list = []
    # coloc_results = {}
    for condition in structure:
        print(f"\n=== Processing condition: {condition} ===")
        condition_out_path = os.path.join(output_dir, condition)
        os.makedirs(condition_out_path, exist_ok=True)
        
        # Process all protein combinations and self-analysis for this condition
        process_condition(
            filter_params,
            filter_enabled,
            structure[condition],   # = condition_data 
            condition_out_path, 
            condition, 
            collector, 
            protein_names,
            coloc_threshold=coloc_threshold,
            pxl_size=pxl_size,
            data_mask=data_mask,
            calibration_mask=calibration_mask,
            unspecific_dark_list=unspecific_dark_list,
            dark_time_analysis=dark_time_analysis,
            cross_coloc=cross_coloc,
            patching=patching,
            patch_size=patch_size,
            patch_id=-1, 
            simulation=simulation
        )

    calibration_results = None
    if dark_time_analysis and unspecific_dark_list:
        print("=== Performing Dark Time Analysis ===")
        print("\nCalibration values:")
        
        calibration_results = calculate_calibration_values(unspecific_dark_list, protein_names)
        # Add calibration values to existing protein data
        add_calibration_to_protein_data(collector, calibration_results)
        
        save_path = os.path.join(output_dir, "unspecific_dark_list.pkl")
        with open(save_path, "wb") as fp:
            pickle.dump(unspecific_dark_list, fp)
        

            
                
def process_simulated_data(condition_data, file_idx, condition_name, collector, protein_names,
                          filter_params, filter_enabled, data_mask, calibration_mask,
                          pxl_size, coloc_threshold, condition_out_path,
                          dark_time_analysis, cross_coloc, unspecific_dark_list,
                          patching, patch_size):
    """
    Process simulated data for a single file
    """
    
    # Store data for this file only
    file_protein_data = []
    file_coloc_results = {}
    
    # === Process single protein data for this file ===
    for protein in protein_names:
        try:
            # Load real data first
            locs_data_masked, unspecific_locs, mask, file_path, patch_data = single_protein_data(
                protein, condition_data, file_idx, 
                data_mask, calibration_mask, 
                condition_name, pxl_size, condition_out_path,
                filter_params, filter_enabled, 
                patch_size if patching else None
            )
            
            # Simulate the data
            simulated_locs_data = simulate_locs_data(locs_data_masked, n_simulations=5, random_seed=None)
            
            if patch_data is not None:
                simulated_patch_ids = assign_locs_to_patches(simulated_locs_data, patch_data['patches'])
                # Add patch_id to simulated data
                new_dtype = simulated_locs_data.dtype.descr + [('patch_id', 'i4')]
                new_data = np.zeros(len(simulated_locs_data), dtype=new_dtype)
                for field in simulated_locs_data.dtype.names:
                    new_data[field] = simulated_locs_data[field]
                new_data['patch_id'] = simulated_patch_ids
                simulated_locs_data = new_data
    
            save_hdf5(os.path.join(condition_out_path, f"simulated_{os.path.basename(file_path)}"), simulated_locs_data, load_info(file_path))
            
            # Process simulated data
            self_neighbors_results = self_clocalization(simulated_locs_data, pxl_size, coloc_threshold)
            
            # Extract data for this file with data_type='simulated'
            protein_data_list = extract_protein_data(collector, 
                                                  simulated_locs_data, mask, 
                                                  condition_name, protein, file_path, 
                                                  pxl_size, 
                                                  self_neighbors_results,
                                                  patch_data, -1, 'simulated')  # ADD 'simulated' HERE
            
            file_protein_data.extend(protein_data_list)
            
            # Clear large data structures immediately
            del locs_data_masked, simulated_locs_data, self_neighbors_results
            if patch_data is not None:
                del patch_data
            
        except Exception as e:
            print(f"Error loading {protein}-{file_idx} (simulated): {str(e)}")
    
    # === Process cross-colocalization for this file ===
    if cross_coloc:
        protein_pairs = list(permutations(protein_names, 2))
        
        for protein1, protein2 in protein_pairs:
            try:
                # Load real data for both proteins
                locs_data_masked_1, _, mask1, file_path_1, patch_data_1 = single_protein_data(
                    protein1, condition_data, file_idx, 
                    data_mask, calibration_mask, 
                    condition_name, pxl_size, condition_out_path,
                    filter_params, filter_enabled, 
                    patch_size if patching else None
                )
                locs_data_masked_2, _, mask2, file_path_2, patch_data_2 = single_protein_data(
                    protein2, condition_data, file_idx, 
                    data_mask, calibration_mask, 
                    condition_name, pxl_size, condition_out_path,
                    filter_params, filter_enabled, 
                    patch_size if patching else None
                )
                    
                # Simulate both datasets
                simulated_locs_1 = simulate_locs_data(locs_data_masked_1, mask1)
                simulated_locs_2 = simulate_locs_data(locs_data_masked_2, mask2)
                
                # Cross-colocalization
                coloc_results = cross_colocalization(protein1, protein2, simulated_locs_1, simulated_locs_2, pxl_size, coloc_threshold)
                
                pair_key = f"{protein1}->{protein2}"
                if pair_key not in file_coloc_results:
                    file_coloc_results[pair_key] = []
                file_coloc_results[pair_key].extend(coloc_results)
                
                # Clear large data structures immediately
                del locs_data_masked_1, locs_data_masked_2, simulated_locs_1, simulated_locs_2, coloc_results
                
            except Exception as e:
                print(f"Error processing {protein1}-{protein2} pair {file_idx} (simulated): {str(e)}")
        
        # Add colocalization data for this file only
        add_coloc_to_protein_data(file_protein_data, file_coloc_results, protein_names)
    
    # === Add everything to collector for this file ===  
    for protein_data in file_protein_data:
        collector.add_protein_data(protein_data)
    
    # Clear data for this file from memory
    del file_protein_data, file_coloc_results
    gc.collect()