library.py

import numpy as np
import scanpy as sc
from collections import Counter
import random
import matplotlib.pyplot as plt
import os
import shutil
from sklearn.metrics import silhouette_samples, calinski_harabasz_score, pairwise_distances, silhouette_score
import umap 
from scipy.spatial.distance import cosine as cosine_distance, euclidean as euclidean_distance
import argparse
from anndata import AnnData
import pickle
import pandas as pd
import warnings 
import gseapy
import csv
import traceback
from sklearn.metrics.pairwise import cosine_distances
from sklearn.preprocessing import StandardScaler, normalize as sklearn_normalize
import socket
from sklearn.cluster import HDBSCAN, KMeans, AgglomerativeClustering
from scipy.cluster.hierarchy import dendrogram
import itertools
import time
from scipy.stats import spearmanr, pearsonr
import psutil
from adjustText import adjust_text
import multiprocessing as mp

def get_filtered_perturbations(adata, control_strings, exclude_perturbations):
    """
    return list of unique perturbations in adata minus control strings and exclude_perturbations
    """
    pert_candidates = set(adata.obs["gene_target"])
    ##remove control strings and unique perturbations from candidates 
    for control in control_strings:
        pert_candidates.remove(control)
    for u_p in exclude_perturbations:
        if u_p in pert_candidates:
            pert_candidates.remove(u_p)
    return list(pert_candidates)

def get_umap(X):
    """
    return UMAP reduction of X
    """
    reducer = umap.UMAP()
    embedding = reducer.fit_transform(X)
    return embedding 

def rank_perturbations_by_similarity(adata, anchor_state="Untreated_SAFE_TARGET", normalize=False, distance_metric="cosine", embedding_type="method_embedding"):
    """
    Given adata with perturbation labels in .obs["gene_target"] and embeddings within adata.X, 
    ranks label classes by their distance (distance type determined by distance_metric: either cosine or eulicean) from average of label embedding to average anchor state A.
    Anchor state is {condition}_{perturbation}, e.g. Untreated_SAFE_TARGET
    embedding_type will dictate what embeddings to use for calculating the distances:
        "method_embedding" will be the raw embedding of the method
        "umap_embedding" will be the UMAP of the raw embedding of the method
    returns list of tuples: (perturbation, distance) 
    """
    if embedding_type == "method_embedding":
        matrix = adata.X
    elif embedding_type == "umap_embedding":
        matrix = adata.obsm["umap"]
    else:
        raise Exception()
    if normalize: 
        matrix = preprocess_matrix(matrix, standardize=False, normalize=normalize)
    labels = list(adata.obs["gene_target"])
    ##add condition to target so that final labels are {condition}_{gene_target}
    labels = [adata.obs["condition"].iloc[i] + "_" + labels[i] for i in range(0, len(labels))]
    all_labels = set(labels)
    label_to_indices = {label: [i for i in range(0, len(labels)) if labels[i] == label] for label in all_labels}
    label_to_centroid = {label: np.mean(matrix[label_to_indices[label]], axis=0) for label in all_labels}
    if distance_metric == "cosine":
        anchor_to_labels_distance = {label: cosine_distance(label_to_centroid[label], label_to_centroid[anchor_state]) for label in all_labels}
    elif distance_metric == "euclidean":
        anchor_to_labels_distance = {label: euclidean_distance(label_to_centroid[label], label_to_centroid[anchor_state]) for label in all_labels}  
    else:
        raise Exception()
    sorted_ = sorted(anchor_to_labels_distance.items(), key=lambda x: x[1], reverse=False) ##sorted from lowest distance to highest 
    return sorted_

def is_control(string, control_strings):
    """
    given a string and a set of control strings, will return True if string is a superset of any of the control strings else False, case insensitive 
    """
    for control_string in control_strings:
        if control_string.upper() in string.upper():
            return True 
    return False

def get_stats(l):
    """
    from a list of values, return basic stats
    """
    mean = np.mean(l)
    median = np.median(l)
    std = np.std(l)
    mini = np.min(l)
    maxi = np.max(l)
    return f"mean: {mean:.2E}, median: {median:.2E}, std: {std:.2E}, min: {mini:.2E}, max: {maxi:.2E}"

def perturbation_vs_control_umap(adata, control_strings, unique_perturbations, plot_perts=None, n=50, method=None, direction=None, dataset=None, save_dir=None):
    """
    if plot_perts is NOT set:
        plot UMAP of 4 control populations:
            (1) treated control
            (2) untreated control
            (3) treated perturbation
            (4) untreated perturbation
        will choose n perturbations at random
    else if plot_perts is set:
            must also specify direction:
                direction is either "to_healthy" or "to_disease" and will be used to decide which perturbed treatment type to plot for reversion maps
            plot_perts specifies a list of specific perturbations to plot
                if direction == "to_healthy" will remove set (4) because we want to visualize just treated perturbed
                if direction == "to_disease" will remove set (3) because we want to visualize just untreated perturbed
            will also plot each pert in plot_perts individually 
    """
    if not os.path.isdir(save_dir):
        os.makedirs(save_dir)
    if plot_perts != None:
        assert(direction != None)
        selected_perturbations = plot_perts
        optional_save_prefix = "_" + direction
    else:
        pert_candidates = get_filtered_perturbations(adata, control_strings, unique_perturbations)
        selected_perturbations = set(random.sample(pert_candidates, n))
        optional_save_prefix = ""
    control_and_selected_perts = control_strings.union(selected_perturbations)
    sub_adata = adata[(adata.obs["gene_target"].isin(control_and_selected_perts))]
    coordinates = sub_adata.obsm["umap"]
    kds = sub_adata.obs["knockdown_efficiencies"]
    adata_treated_control_indices = np.array([i for i in range(0, len(sub_adata)) if sub_adata.obs["condition"].iloc[i] == "Treated" and sub_adata.obs["gene_target"].iloc[i] in control_strings])
    adata_untreated_control_indices = np.array([i for i in range(0, len(sub_adata)) if sub_adata.obs["condition"].iloc[i] == "Untreated" and sub_adata.obs["gene_target"].iloc[i] in control_strings])
    adata_treated_perturbed_indices = np.array([i for i in range(0, len(sub_adata)) if sub_adata.obs["condition"].iloc[i] == "Treated" and sub_adata.obs["gene_target"].iloc[i] in selected_perturbations])
    adata_untreated_perturbed_indices = np.array([i for i in range(0, len(sub_adata)) if sub_adata.obs["condition"].iloc[i] == "Untreated" and sub_adata.obs["gene_target"].iloc[i] in selected_perturbations])
    mapp = {"untreated control": adata_untreated_control_indices, "untreated perturbed": adata_untreated_perturbed_indices, "treated control": adata_treated_control_indices, "treated perturbed": adata_treated_perturbed_indices}
    color_map = {"untreated control": "lime", "untreated perturbed": "violet", "treated control": "deepskyblue", "treated perturbed": "midnightblue"}
    if plot_perts != None:
        if direction == "to_healthy":
            del mapp["untreated perturbed"] ##for reversion to healthy, want to visualize just treated perturbed
        if direction == "to_disease":
            del mapp["treated perturbed"] ##for reversion to disease, want to visualize just untreated perturbed
    ##make scatter plot of categories
    fig, ax = plt.subplots()
    for key in mapp:
        indices = mapp[key] 
        sub_coordinates = coordinates[indices]
        ax.scatter(sub_coordinates[:,0], sub_coordinates[:,1], label=key, s=1.0, color=color_map[key])
    ax.set_xlabel("UMAP 1")
    ax.set_ylabel("UMAP 2")
    box = ax.get_position()
    ax.set_position([box.x0, box.y0, box.width, box.height * 0.85])
    ax.legend(loc='upper right', prop={"size":7}, bbox_to_anchor=(1, 1.32))
    plt.gcf().subplots_adjust(top=.76)
    if dataset == "full":
        if method in ["umap", "raw"]:
            prefix1 = f"UMAP"
        else:
            prefix1 = f"UMAP of {get_method_title(method)} Embeddings\n"
        prefix2 = f"UMAP of {get_method_title(method)} Embeddings\n"
    else:
        prefix1 = f"{get_method_title(method)} on {get_dataset_title(dataset)}:\nUMAP"
        prefix2 = f"{get_method_title(method)} on {get_dataset_title(dataset)}:\nUMAP"
    if plot_perts == None:
        plt.title(f"{prefix1} for {n} Randomly Chosen Perturbations + Controls")
    else:
        plt.title(f"{prefix2} of Top {len(plot_perts)} Perturbations Inducing Shift {direction.replace('_',' ').title()}\n{plot_perts}", fontsize=8)
    plt.savefig(os.path.join(save_dir, f"{method}{optional_save_prefix}_controls_vs_perturbations_plot_perts={plot_perts}.png"), dpi=300)
    ##also plot each perturbation individually if plot_perts specifies a list of genes 
    if plot_perts != None:
        for pert in plot_perts:
            fig, ax = plt.subplots()
            p_indices_treated = np.array([i for i in range(0, len(sub_adata)) if sub_adata.obs["condition"].iloc[i] == "Treated" and sub_adata.obs["gene_target"].iloc[i] == pert])
            p_indices_untreated = np.array([i for i in range(0, len(sub_adata)) if sub_adata.obs["condition"].iloc[i] == "Untreated" and sub_adata.obs["gene_target"].iloc[i] == pert])
            if direction == "to_healthy":
                mapp = {"untreated control": adata_untreated_control_indices, "treated control": adata_treated_control_indices, f"treated {pert}": p_indices_treated} 
            if direction == "to_disease": 
                mapp = {"untreated control": adata_untreated_control_indices, "treated control": adata_treated_control_indices, f"untreated {pert}": p_indices_untreated} 
            color_map = {"untreated control": ("lime", 0), "treated control": ("deepskyblue", 1), f"untreated {pert}": ("violet", 2), f"treated {pert}": ("midnightblue", 3)} ##colors and z-order
            for key in mapp:
                indices = mapp[key] 
                if len(indices) == 0:
                    print(f"WARNING: {method} {pert} {key} has no indices, skipping plot")
                    continue 
                sub_coordinates = coordinates[indices]
                ax.scatter(sub_coordinates[:,0], sub_coordinates[:,1], label=f"{key}, n={len(indices)}", s=1.0, color=color_map[key][0], zorder=color_map[key][1])
                ##annotate reversion points with knockdown ratio
                if (direction == "to_healthy" and key == f"treated {pert}") or (direction == "to_disease" and key == f"untreated {pert}"):
                    sub_kds = kds.iloc[indices]
                    sub_kds = [round(k, 2) for k in sub_kds]
                    x = sub_coordinates[:,0]
                    y = sub_coordinates[:,1]
                    for i, txt in enumerate(sub_kds):
                        plt.annotate(
                            txt,  # The text to display
                            (x[i], y[i]),  # The coordinates of the point to annotate
                            textcoords="offset points",  # How to position the text relative to the point
                            xytext=(0, 3),  # Distance from text to point (x, y)
                            ha='center',  # Horizontal alignment of the text
                            # arrowprops=dict(arrowstyle="->", color='gray'), # Optional: add an arrow
                            fontsize=3)
            box = ax.get_position()
            ax.set_position([box.x0, box.y0, box.width, box.height * 0.85])
            ax.legend(loc='upper right', prop={"size":7}, bbox_to_anchor=(1, 1.32))
            plt.gcf().subplots_adjust(top=.76)
            plt.title(f"{get_method_title(method)} on {get_dataset_title(dataset)}:\nUMAP of Untreated Control vs {pert}")
            plt.savefig(os.path.join(save_dir, f"{method}{optional_save_prefix}_controls_vs_pert_{pert}.png"), dpi=300)

def untreated_with_treated_control_umap(adata, control_strings, method, dataset, normalize, just_inflammatory_genes=False):
    """
    plots UMAP of 4 control populations: untreated safe target, untreated no target, treated safe target, treated no target
    plots UMAP of treated vs untreated controls
    """
    adata = adata[adata.obs["gene_target"].isin(control_strings)]
    
    if just_inflammatory_genes:
        if method != "raw": ##this will only work on raw data, scFMs do not preserve raw gene counts for (+) controls
            print(f"skipping untreated_with_treated_control_umap() with method={method}")
            return
        full_matrix = preprocess_matrix(adata.X, standardize=False, normalize=True)
        library_sets = get_relevant_pathways(just_this_study=True, include_imru_internal=True, key_type="tuple")
        inflammatory_genes = set()
        for pathway in library_sets:
            inflammatory_genes = inflammatory_genes.union(set(library_sets[pathway]))
        print("inflammatory genes: ", sorted(inflammatory_genes))
        feature_names = list(adata.var.index)
        feature_coordinates = np.array([i for i in range(0, len(feature_names)) if feature_names[i] in inflammatory_genes])
        coordinates = full_matrix[:, feature_coordinates] ##select just the columns of the positive controls
        coordinates = get_umap(coordinates) ##perform UMAP on just these positive control features as final coordinates to plot 
    else: 
        coordinates = adata.obsm["umap"] ##take existing UMAP coordinates pre-computed over all feature space 

    ##scatter plot of 4 categories - treated safe, treated no, untreated safe, untreated no
    treated_safe = np.array([i for i in range(0, len(adata)) if adata.obs["condition"].iloc[i] == "Treated" and adata.obs["gene_target"].iloc[i] == "SAFE_TARGET"])
    treated_no = np.array([i for i in range(0, len(adata)) if adata.obs["condition"].iloc[i] == "Treated" and adata.obs["gene_target"].iloc[i] == "NO-TARGET"])
    untreated_safe = np.array([i for i in range(0, len(adata)) if adata.obs["condition"].iloc[i] == "Untreated" and adata.obs["gene_target"].iloc[i] == "SAFE_TARGET"])
    untreated_no = np.array([i for i in range(0, len(adata)) if adata.obs["condition"].iloc[i] == "Untreated" and adata.obs["gene_target"].iloc[i] == "NO-TARGET"])
    mapp = {"untreated safe target": untreated_safe, "untreated no target": untreated_no, "treated safe target": treated_safe, "treated no target": treated_no}
    fig, ax = plt.subplots()
    for key in mapp:
        indices = mapp[key]
        sub_coordinates = coordinates[indices]
        ax.scatter(sub_coordinates[:,0], sub_coordinates[:,1], label=key, s=1.0)
    box = ax.get_position()
    ax.set_position([box.x0, box.y0, box.width, box.height * 0.85])
    ax.legend(loc='upper right', prop={"size":7}, bbox_to_anchor=(1, 1.32))
    plt.gcf().subplots_adjust(top=.76)
    plt.title(f"UMAP of {get_method_title(method)} Embeddings\nfor Control Cells")
    plt.savefig(f"outputs/{dataset}/{method}_controls.png", dpi=300)
    ##scatter plot of treated vs untreated
    treated = np.concatenate((treated_safe, treated_no))
    untreated = np.concatenate((untreated_safe, untreated_no))
    mapp = {"untreated": (untreated, "lime", 0), "treated": (treated, "deepskyblue", 1)}
    fig, ax = plt.subplots()
    for key in mapp:
        sub_coordinates = coordinates[mapp[key][0]]
        ax.scatter(sub_coordinates[:,0], sub_coordinates[:,1], label=key.capitalize(), s=1.0, color=mapp[key][1], zorder=mapp[key][2])
    ax.set_xlabel("UMAP 1")
    ax.set_ylabel("UMAP 2")
    box = ax.get_position()
    ax.set_position([box.x0, box.y0, box.width, box.height * 0.85])
    ax.legend(loc='upper right', prop={"size":7}, bbox_to_anchor=(1, 1.15))
    plt.gcf().subplots_adjust(top=.76)
    if method == "umap":
        plt.title(f"UMAP of Raw Transcriptome\nfor Control Cells")
    else:
        plt.title(f"UMAP of {get_method_title(method)} Embeddings\nfor Control Cells")

    save_string = f"outputs/{dataset}/{method}_controls_treatment_type_only" 
    if method == "raw" and just_inflammatory_genes:
        plt.title("UMAP of Just Inflammatory Genes\nfor Control Cells")
        save_string += "_only_inflammatory_gene_features"
    plt.savefig(f"{save_string}.png", dpi=300)

def get_knockdown_efficiencies_Replogle(control_strings):
    """
    calculate KD stats by method of Reploge 2020 paper (https://pmc.ncbi.nlm.nih.gov/articles/PMC7416462/):
    do a sum normalization to the median UMI and ignore genes with an average of <1 UMI
    """
    adata = sc.read_h5ad("data/raw_imru_full_processed.h5ad")
    ##set up basic data structures 
    gene_columns = list(adata.var.index)
    target_set = list(set(adata.obs["gene_target"])) 
    target_set = [x for x in target_set if x not in control_strings] ##remove controls from the gene target set 
    target_to_col_index = {t: gene_columns.index(t) for t in target_set} ##index of each target in the columns of X
    target_list = list(adata.obs["gene_target"])
    median_UMI = np.median(adata.obs["nUMI"])
    print(f"median UMI over all cells: {median_UMI}")
    X = adata.X
    X = sum_normalize_matrix(X, total=median_UMI, log1p=False) ##sum normalize matrix to median_UMI
    ##get control average (average after sum normalization)
    control_row_indices = np.array([i for i in range(0, len(target_list)) if target_list[i] in control_strings])
    control_col_means = np.mean(X[control_row_indices], axis=0)
    ##calculate KD % for all cells 
    kds = [] ##list of kds for each cell
    for i in range(0, len(target_list)):
        print(i, len(target_list))
        if target_list[i] in control_strings:
            kds.append(-1.0) ##controls will get -1.0 (actual value won't matter when we calculate KD on a per target basis - these values will be ignored)
        else:
            target_col_index = target_to_col_index[target_list[i]]
            kds.append(1.0 - (X[i][target_col_index] / control_col_means[target_col_index]))
    kds = np.array(kds)
    ##calculate avg KD% for all targets
    target_to_row_indices = {t: [i for i in range(0, len(target_list)) if target_list[i] == t] for t in target_set}
    target_to_avg_kd = {t: np.mean(kds[target_to_row_indices[t]]) for t in target_set}
    print("all targets: ", np.mean(list(target_to_avg_kd.values())), np.median(list(target_to_avg_kd.values())))
    ##calculate avg KD% for just high UMI targets
    avg_columns = np.mean(adata.X, axis=0) ##use raw matrix for filtering out gene targets (not normalized) for computing gene UMI averages
    gene_indices_with_high_UMI = [i for i in range(0, len(avg_columns)) if avg_columns[i] >= 1.0] ##ignore genes with an average of <1 UMI
    targets_with_high_UMI = set([t for t in target_to_col_index if target_to_col_index[t] in gene_indices_with_high_UMI]) ##set of targets with UMI >= 1.0 in matrix
    print(f"{len(targets_with_high_UMI)} targets with UMI >= 1.0, keeping for KD calculation")
    target_to_avg_kd = {t: np.mean(kds[target_to_row_indices[t]]) for t in targets_with_high_UMI}
    print("just targets with high UMI: ", np.mean(list(target_to_avg_kd.values())), np.median(list(target_to_avg_kd.values())))
    ##save results
    pickle.dump((kds, target_list, target_to_row_indices, targets_with_high_UMI), open("pickles/kd_calculation.pkl", "wb"))

def plot_knockdown_efficiencies():
    """
    plot results of KD efficiency calculations according to the Replogle paper 
    will use this for main results figures because more standard
    """
    kds, target_list, target_to_row_indices, targets_with_high_UMI = pickle.load(open("pickles/kd_calculation.pkl", "rb"))
    target_to_avg_kd = {t: np.mean(kds[target_to_row_indices[t]]) for t in target_to_row_indices}
    target_kds = list(target_to_avg_kd.values())
    print(f"KD % stats: mean: {np.mean(target_kds)}, median: {np.median(target_kds)}, min: {np.min(target_kds)}, max: {np.max(target_kds)}, std: {np.std(target_kds)}")
    ##high UMI targets only
    high_target_to_avg_kd = {t: np.mean(kds[target_to_row_indices[t]]) for t in targets_with_high_UMI}
    high_target_kds = list(high_target_to_avg_kd.values())
    print(f"KD % stats (high UMI): mean: {np.mean(high_target_kds)}, median: {np.median(high_target_kds)}, min: {np.min(high_target_kds)}, max: {np.max(high_target_kds)}, std: {np.std(high_target_kds)}")
    ##make plots off all targets
    less_than_eq_0 = sum([1 for x in target_kds if x <= 0])
    less_than_eq_neg_10 =  sum([1 for x in target_kds if x <= -10])
    greater_than_0 = sum([1 for x in target_kds if x > 0])
    print(f"frequency of cells > 0: {float(greater_than_0 / (greater_than_0 + less_than_eq_0))}")
    print(greater_than_0, less_than_eq_0, less_than_eq_neg_10)
    ##bar chart
    fig, ax = plt.subplots()
    x = np.array([0, 1, 2])
    y = [greater_than_0, less_than_eq_0, less_than_eq_neg_10]
    bars = ax.bar(x, y)
    ax.bar_label(bars, fmt="%.2e", fontsize=7)
    ax.set_xticks(x)
    ax.set_xticklabels(["KD > 0.0", "KD <= 0.0", "KD <= -10.0"])
    ax.set_ylabel("Cell Count")
    plt.title("Cell Count by Knockdown Percentage")
    plt.ticklabel_format(axis='y', style='scientific', scilimits=(0,0))
    plt.savefig("outputs/knockdown_percentage_bar.png", dpi=300)
    ##histogram 
    fig, ax = plt.subplots()
    ax.hist(target_kds, bins=100)
    ax.set_xlim((-1.0, 1.0))
    ax.set_ylabel("Target Count")
    ax.set_xlabel("Knockdown Efficiency")
    plt.title("Knockdown Efficiency Histogram")
    ax.set_xticklabels(['{:,.0%}'.format(x_) for x_ in ax.get_xticks()])
    plt.savefig("outputs/knockdown_percentage_histogram.png", dpi=300)

def data_stats(adata_path):
    """
    basic stats about data written to outputs/
    """
    adata = sc.read_h5ad(adata_path)
    X = adata.X.toarray()
    ##get basic stats about data
    print(f"read count stats: median {np.median(X)}, fraction zero: {1.0 - (np.count_nonzero(X) / (X.shape[0] * X.shape[1]))}")    
    ##counts per cell 
    axis1_sum = np.sum(X, axis=1).flatten() ##sum for each cell
    fig, ax = plt.subplots()
    ax.violinplot(axis1_sum, showmedians=True)
    ax.set_ylabel("Transcript Count")
    plt.title(f"Violin Plot of Transcripts per Cell\n {get_stats(axis1_sum)}", fontsize=9)
    plt.savefig(f"outputs/violin_transcripts_per_cell.png", dpi=300)
    fig, ax = plt.subplots()
    n, bins, _ = plt.hist(axis1_sum, bins=100)
    ax.set_xlabel("Transcript Count")
    ax.set_ylabel("Cell Count")
    plt.ticklabel_format(axis='both', style='sci', scilimits=(0,0))
    plt.title(f"Histogram of Transcripts per Cell")
    plt.savefig(f"outputs/histogram_transcripts_per_cell.png", dpi=300)
    ##counts per gene
    axis0_sum = np.sum(X, axis=0).flatten() ##sum for each gene 
    print("genes that have 0 sum: ", len(X[0]) - np.count_nonzero(axis0_sum))
    fig, ax = plt.subplots()
    ax.violinplot(axis0_sum, showmedians=True)
    ax.set_ylabel("# of transcripts")
    plt.title(f"Violin Plot of Transcripts per Gene\n{get_stats(axis0_sum)}", fontsize=9)
    plt.savefig(f"outputs/violin_transcripts_per_gene.png", dpi=300)
    fig, ax = plt.subplots()
    n, bins, _ = plt.hist(axis0_sum, bins=100)
    ax.set_xlabel("Number of Transcripts")
    ax.set_ylabel("# of Genes")
    plt.title(f"Histogram of Transcripts per Gene\n {get_stats(axis1_sum)}", fontsize=9)
    plt.savefig(f"outputs/histogram_transcripts_per_gene.png", dpi=300)
    ##histogram of ratio of zeros for each gene 
    zeros = (len(X) - np.count_nonzero(X, axis=0)) / len(X)
    fig, ax = plt.subplots()
    n, bins, _ = plt.hist(zeros, bins=100)
    ax.set_xlabel(f"% of Cells with Zero Expression")
    ax.set_ylabel("# of Genes")
    plt.title(f"Histogram of Cells \n with Zero Frequencies per Gene Median={round(np.median(zeros), 3)}")
    plt.savefig(f"outputs/histogram_zeros_per_gene.png", dpi=300)
    ##plot x-axis: sorted genes by zero frequency, y-axis: 0 frequency
    fig, ax = plt.subplots()
    sorted_zeros = np.sort(zeros)
    ax.plot(np.array(range(0, len(sorted_zeros))), sorted_zeros)
    plt.ticklabel_format(axis='y', style='sci', scilimits=(0,0))
    plt.title(f"% Cells with Zero Read Count for each Gene")
    ax.set_xlabel("Gene Indices Sorted by % Cells with Zero Read Count for G")
    ax.set_ylabel("% Cells with Zero Read Count for G")
    plt.savefig(f"outputs/line_graph_zeros_per_gene.png", dpi=300)

def preprocess_matrix(X, standardize=False, normalize=False):
    """
    Given matrix X, will perform preprocessing operations and return a copy
    """
    X_new = X.copy()
    if standardize: ##standardize features (zero mean, unit variance)
        scaler = StandardScaler()
        X_new = scaler.fit_transform(X_new)
        print("standardized matrix")
    if normalize: # normalize each sample to unit norm (L2)
        X_new = sklearn_normalize(X_new, norm='l2')
        print("(l2) normalized matrix")
    return X_new

def subsample_matrix(X, y, sample_pct=0.5, n=10, random_state=None, constraint="keep_ratio"):
    """
    if constraint == "keep_ratio"
        for each class in y, sample `sample_pct` percent of them to keep and return a subsampled version of X, y.
        n must equal None
     if constraint == "equal_count"
        for each class in y, sample n at random and return a subsampled version of X, y.
        sample_pct must equal None
        if n is greater than total size of class y_i, will keep all of y_i
    Parameters:
    - X: np.ndarray or similar, shape (n_samples, n_features)
    - y: list of strings
    - sample_pct: float, percentage of samples to keep per class (between 0 and 1) when constraint == "keep_ratio"
    _ sample_n: n to use when constraint == "equal_count"
    - random_state: int or None, for reproducibility
    Returns:
    - X_sub: subsampled version of X
    - y_sub: subsampled version of y
    """
    if constraint == "keep_ratio":
        assert n == None
    if constraint == "equal_count":
        assert sample_pct == None
    if random_state is not None:
        np.random.seed(random_state)
    unique_classes = set(y)
    indices_to_keep = []
    for cls_ in unique_classes:
        cls_indices = np.array([i for i in range(0, len(y)) if y[i] == cls_])
        if constraint == "keep_ratio":
            n_samples = int(len(cls_indices) * sample_pct)
        elif constraint == "equal_count":
            n_samples = min(n, len(cls_indices))
        sampled_indices = np.random.choice(cls_indices, n_samples, replace=False)
        indices_to_keep.extend(sampled_indices)
    indices_to_keep = np.array(indices_to_keep)
    return X[indices_to_keep], [y[i] for i in indices_to_keep]

def subsample_adata_to_min_replicate(adata, n=10):
    """
    given adata and a minimum replicate count will return an adata object with just samples that have at least n replicates for gene_target for each treatment condition (needs to have min count for treated and min count for untreated)
    """
    treated_adata = adata[adata.obs["condition"] == "Treated"]
    treated_counts = Counter(treated_adata.obs["gene_target"])
    untreated_adata = adata[adata.obs["condition"] == "Untreated"]
    untreated_counts = Counter(untreated_adata.obs["gene_target"])
    key_intersection = [key for key in treated_counts if key in untreated_counts] ##if a key is only present in either Treated OR Untreated (but not both), then remove it 
    keep_pert = set([key for key in key_intersection if treated_counts[key] > 10 and untreated_counts[key] > 10])
    indices = [i for i in range(0, len(adata.obs["gene_target"])) if adata.obs["gene_target"].iloc[i] in keep_pert]
    return adata[indices]

def get_true_and_random_clustering_scores(matrix, labels, clustering_metric):
    """
    given a matrix of embeddings, a list of labels, and a clustering method to use
    returns a list of clustering scores for true labels, and a list of clustering scores for random labels (for each trial)
    if clustering_metric == "Silhouette"
        will subsample matrix for each unique label (such that each label has 10 unique samples), and will do 30 trials of random labels (Silhouette is computationally expensive)
    if clusering_metric == "Calinski-Harabasz"
        will not do any subsampling (will use entire matrix), and will run 30 trials of randomized labels
    """
    true_scores = []
    random_scores = []
    if clustering_metric == "Calinski-Harabasz":
        true_scores.append(calinski_harabasz_score(matrix, labels))
        for trial in range(0, 30):
            random_labels = random.sample(labels, len(labels))
            r_s = calinski_harabasz_score(matrix, random_labels)
            random_scores.append(r_s)
    elif clustering_metric == "Silhouette": ##silhouette is computationally expensive, parallelize this
        for trial in range(0, 30):
            init_time = time.time() 
            ##pairwise distance matrix of all points for this dataset is infeasible and will require more memory than we can possibly have (about 2.5 TB at 16 point precision to just store the resulting distance matrix) --> subsample the matrix
            sub_matrix, sub_labels = subsample_matrix(matrix, labels, sample_pct=None, n=10, random_state=None, constraint="equal_count")
            sub_matrix = sub_matrix.astype(np.float16) ##huge memory requirements to compute pairwise distance, cast as 16 bit
            sub_random_labels = random.sample(sub_labels, len(sub_labels)) ##need to get random labels again if subsampled labels
            distance_matrix = pairwise_distances(sub_matrix, metric='cosine', n_jobs=-1)
            # print(f"completed pairwise calculation of matrix with shape: {sub_matrix.shape} in {round(time.time() - init_time, 2)}")
            if np.sum(distance_matrix.diagonal()) != 0:
                print(f"WARNING: diagonal of distance_matrix (with shape {distance_matrix.shape}) was NOT all zeros, filling with zeros\ndiagonal of distance_matrix max: {np.max(distance_matrix.diagonal())}, avg: {np.mean(distance_matrix.diagonal())}, sum: {np.sum(distance_matrix.diagonal())}")
                np.fill_diagonal(distance_matrix, 0)
            true_scores.append(silhouette_score(distance_matrix, sub_labels, metric='precomputed'))
            random_scores.append(silhouette_score(distance_matrix, sub_random_labels, metric='precomputed'))
    else:
        raise Exception("clustering metric must be in [Calinski-Harabasz, Silhouette]")
    return true_scores, random_scores

def score_perturbation_clusters(adata, control_strings, method="umap", dataset=None, knockdown_threshold=500.0, normalize=False):
    """
    score clustering based on true labels vs random labels
    for two setup types: 
        (1) labels are treatment condition - will do separately for Control and Perturbed cells
            (e.g. 4 clustering scores using treatment condition as cluster label: true control, random control, true perturbed, random perturbed)
        (2) labels are gene targets (with same treatment condition) - will remove the control cells
            and use all perturbations from adata, and do separately for Treated and Untreated
            (e.g. 4 clustering scores using gene target as cluster label: true treated, random treated, true untreated, random untreated
    will filter adata and select cells where .obs["knockdown_efficiencies"] (in this case ratio of perturbed / control expression) is < knockdown_threshold
    writes figure to outputs/
    writes dictionary to pickle file with key: label_type, key: treatment, key: random or true, key: clustering_metric, value: clustering score
    """
    adata = adata[adata.obs["knockdown_efficiencies"] < knockdown_threshold]
    clustering_metrics = ["Calinski-Harabasz"]
    label_types = ["Gene Target", "Treatment Condition"] ## label points either by treatment condition or by perturbation (where perturbations all have the same treatment condition)
    shuffle_states = ["true", "random"]
    embedding_types = ["method_embedding", "umap_embedding"] #embedding_type will dictate what embeddings to use for labeling and scoring: "method_embedding" will be the raw embedding of the method, "umap_embedding" will be the UMAP of the raw embedding of the method
    results_map = {embedding_type: {label_type: {condition: {shuffle_state: {clustering_metric: "" for clustering_metric in clustering_metrics} for shuffle_state in shuffle_states} for condition in ["Treated", "Untreated", "Control", "Perturbed"]} for label_type in label_types} for embedding_type in embedding_types}
    for embedding_type in embedding_types:
        for clustering_metric in clustering_metrics:        
            for label_type in label_types:
                if label_type == "Gene Target":
                    split_bys = ["Treated", "Untreated"] ##split by treated and untreated
                elif label_type == "Treatment Condition":
                    split_bys = ["Control", "Perturbed"] ##split by control or perturbed
                else: 
                    raise Exception()
                ##for label_type == Gene Target: ensure that all perturbations are either treated or untreated, and we don't mix them together
                ##for label_type == Treatment Condition: split_bys will be either control cells or perturbed cells
                for s_b in split_bys: 
                    ##get matrix and labels for clustering based on label_type
                    if label_type == "Gene Target": ##for given condition, use all genetic perturbations in analysis
                        print(f"for label_type={label_type} removing controls, old shape: {adata.X.shape}")
                        sub_adata = adata[~adata.obs["gene_target"].isin(control_strings)]
                        print(f"new shape: {sub_adata.X.shape}")
                        sub_adata = sub_adata[sub_adata.obs["condition"] == s_b] 
                        labels = list(sub_adata.obs["gene_target"])
                    elif label_type == "Treatment Condition":
                        if s_b == "Control":
                            print(f"for label_type={label_type}, split={s_b} keeping just controls, old shape: {adata.X.shape}")
                            sub_adata = adata[adata.obs["gene_target"].isin(control_strings)]
                            print(f"new shape: {sub_adata.X.shape}")
                        elif s_b == "Perturbed":
                            print(f"for label_type={label_type}, split={s_b} keeping just perturbed cells, old shape: {adata.X.shape}")
                            sub_adata = adata[~adata.obs["gene_target"].isin(control_strings)]
                        else:
                            raise Exception()
                        labels = list(sub_adata.obs["condition"])
                    else:
                        raise Exception()
                    if embedding_type == "method_embedding":
                        matrix = sub_adata.X
                    else:
                        matrix = sub_adata.obsm["umap"]
                    if normalize: 
                        matrix = preprocess_matrix(matrix, standardize=False, normalize=normalize)
                    ##now that we have the matrix and labels, get clustering score of true and score of random, write to dictionary 
                    true_scores, random_scores = get_true_and_random_clustering_scores(matrix, labels, clustering_metric)
                    results_map[embedding_type][label_type][s_b]["true"][clustering_metric] = true_scores
                    results_map[embedding_type][label_type][s_b]["random"][clustering_metric] = random_scores
                ##print results of true vs random for each s_b
                for s_b in split_bys:
                    print(f"{dataset} {method}")
                    print(f"    {clustering_metric} {label_type} {s_b} delta: true - random: {np.mean(results_map[embedding_type][label_type][s_b]['true'][clustering_metric]) - np.mean(results_map[embedding_type][label_type][s_b]['random'][clustering_metric])}")
                    if clustering_metric != "Silhouette": ##ratio of silhouette scores can be misleading (if both are negative ratio < 1 indicates true is better than random, but if both are positive a ratio < 1 means true is worse than random
                        print(f"    {clustering_metric} {label_type} {s_b} ratio: {s_b} true / random: {np.mean(results_map[embedding_type][label_type][s_b]['true'][clustering_metric]) / np.mean(results_map[embedding_type][label_type][s_b]['random'][clustering_metric])}")
                ##bar plot of clustering metric scores
                fig, ax = plt.subplots()
                x_labels = split_bys
                width = 0.20
                x = np.array(range(0, len(x_labels)))
                ax.set_xticks(x)
                ax.set_xticklabels(x_labels)
                for shuffle_state in shuffle_states:
                    y = [np.mean(results_map[embedding_type][label_type][x_label][shuffle_state][clustering_metric]) for x_label in x_labels]
                    bars = ax.bar(x, y, width=width, label=f"{shuffle_state} target labels")        
                    x = x + width
                    if max(y) > 100:
                        plt.ticklabel_format(axis='y', style='scientific', scilimits=(0,0))
                        fmt = '%.2E'
                    else: 
                        fmt = '%.2f'
                    ax.bar_label(bars, fmt=fmt, fontsize=7)
                box = ax.get_position()
                ax.set_position([box.x0, box.y0, box.width, box.height * 0.85])
                ax.legend(loc='upper right', prop={"size":7}, bbox_to_anchor=(1, 1.32))
                plt.gcf().subplots_adjust(top=.76)    
                ax.set_ylabel(f"{clustering_metric} Score")
                plt.title(f"{clustering_metric} Scores for\nTrue vs Randomized {label_type} in {get_method_title(method)} Space")
                plt.savefig(f"outputs/{dataset}/cluster_{embedding_type}_{method}_{label_type}_{clustering_metric}_vs_random_mean_{knockdown_threshold}.png", dpi=300, bbox_inches="tight")
    ##save pickle of results
    pickle.dump(results_map, open(f"pickles/clustering_scores/{dataset}_{method}_{knockdown_threshold}.pkl", "wb"))

def compare_method_replicate_similarity(): 
    """
    Compare each method's clustering of true vs random using two different clustering metrics: 
    (1) Calinski-Harabasz
        will plot the ratio of true score / random score
    (2) Silhouette
        will pot the difference (true score - random score)
    will look at label = Gene Target or label = Treatment Condition
        for Gene Target, will separate Treated vs Untreated
        for Treatment Conditon, will just plot control cells for comparison (eliminates confounder of perturbation)
    will plot both using the method's embeddings as the representations vs using the UMAP of the method's embeddings as the representations (embedding_type)
    saves plot to outputs/
    """
    knockdown_threshold = 500.0
    width = 0.35
    color_map = {"Treated": "midnightblue", "Untreated": "violet", "Perturbed": "gold", "Control": "blue"}
    for embedding_type in ["method_embedding"]:
        if embedding_type == "method_embedding":
            methods = ["random", "raw", "umap", "scimilarity", "scgpt", "state"] 
            x_labels = [get_method_title(method) for method in methods]
        if embedding_type == "umap_embedding":
            methods = ["random", "umap", "scimilarity", "scgpt", "state"] ##embedding_type == "umap_embedding" will take the .obsm["umap"] of the adata object, the raw adata .obsm["umap"] and umap data .obsm["umap"] will be different (two independent UMAP reductions of raw space with different random seeds - see preprocessing.py) --> just show one of them
            x_labels = [f"UMAP({get_method_title(method)})" if method != "umap" else f"UMAP(Raw)" for method in methods]
        x = np.arange(0, len(x_labels))
        method_map = {m: pickle.load(open(f"pickles/clustering_scores/full_{m}_{knockdown_threshold}.pkl", "rb")) for m in methods}
        for label_type in ["Gene Target", "Treatment Condition"]:
            if label_type == "Gene Target":
                split_bys = ["Treated", "Untreated"] 
            if label_type == "Treatment Condition":
                split_bys = ["Control"] ##let's just look at the control cells to reduce variability and the added complexity of also introducing a perturbation (split_by="Perturbation")
            for clustering_metric in method_map["raw"][embedding_type][label_type][split_bys[0]]["true"]:
                fig, ax = plt.subplots()
                ax.set_xticks(x)
                ax.set_xticklabels(x_labels, rotation=40, ha="right", rotation_mode='anchor')
                if clustering_metric == "Calinski-Harabasz": ##ratio between true labels and random 
                    operator = "Ratio"
                    ax.set_ylabel(f"{clustering_metric} Ratio (True / Random)", fontsize=11)
                if clustering_metric == "Silhouette": ##difference between true labels and random 
                    operator = "Difference"
                    ax.set_ylabel(f"{clustering_metric} Difference (True - Random)", fontsize=11)
                ax.set_xlabel("Embedding Type", fontsize=11)
                for split_by in split_bys:
                    true_list = [method_map[method][embedding_type][label_type][split_by]["true"][clustering_metric] for method in methods] ##list of lists ordered by method
                    random_list = [method_map[method][embedding_type][label_type][split_by]["random"][clustering_metric] for method in methods]
                    if clustering_metric == "Calinski-Harabasz": ##ratio between true labels and random 
                        y = [np.mean(true_list[i]) / np.mean(random_list[i]) for i in range(0, len(methods))]
                        ##yerr will be the std of a list of each true / random from each trial
                        yerr = [np.std([true_list[j][i] / random_list[j][i] for i in range(0, len(true_list[j]))]) for j in range(0, len(methods))] 
                    if clustering_metric == "Silhouette": ##difference between true labels and random 
                        y = [np.mean(true_list[i]) - np.mean(random_list[i]) for i in range(0, len(methods))]
                        yerr = [np.std([true_list[j][i] - random_list[j][i] for i in range(0, len(true_list[j]))]) for j in range(0, len(methods))] 
                    bars = ax.bar(x, y,  width=width,yerr=yerr, capsize=5, label=split_by, color=color_map[split_by])
                    if max(y) > 100 : 
                        fmt = '%.1E' 
                        plt.ticklabel_format(axis='y', style='scientific', scilimits=(0,0))
                    else:
                        fmt = '%.1f'
                    ax.bar_label(bars, fmt=fmt, fontsize=6)
                    x = x + width
                if label_type == "Gene Target":
                    plt.title(f"{clustering_metric} {operator} Scores\nwith Cluster Labels = {label_type}")
                if label_type == "Treatment Condition":
                    if split_bys == ["Control"]: ##if just looking at control cells and cluster labels = Treatment Condition
                        plt.title(f"{clustering_metric} {operator} Scores for Control Cells\nwith Cluster Labels = {label_type}")
                    else:
                        plt.title(f"{clustering_metric} {operator} Scores\nwith Cluster Labels = {label_type}")
                if len(split_bys) > 1: ##if more than 1 split_bys add a legend
                    box = ax.get_position()
                    ax.set_position([box.x0, box.y0, box.width, box.height * 0.90])
                    ax.legend(loc='upper right', prop={"size":7}, bbox_to_anchor=(1, 1.25))
                    plt.gcf().subplots_adjust(top=.80)
                plt.savefig(f"outputs/method_clustering_{knockdown_threshold}_{embedding_type}_{label_type}_{clustering_metric}.png", dpi=300, bbox_inches="tight")

def get_perturbation_counts_per_treatment(adata):
    """
    returns dictionary with key treatment, value: Counter map for gene targets
    """
    treatment_types = ["Treated", "Untreated"]
    return_map = {"Treated": {}, "Untreated": {}}
    for treatment_type in treatment_types:
        sub_adata = adata[adata.obs["condition"] == treatment_type]
        counts = Counter(sub_adata.obs["gene_target"])
        return_map[treatment_type] = counts
    return return_map

def get_opposite_anchor(anchor_state):
    """
    return the opposite anchor string given anchor_state string
    """
    if "Treated" in anchor_state:
        return anchor_state.replace("Treated", "Untreated")
    elif "Untreated" in anchor_state:
        return anchor_state.replace("Untreated", "Treated")
    else:
        raise Exception()

def get_top_perturbations(adata, method="raw", top_n=5, anchor_state_candidate_treatment_types=[], min_sample=0, dataset=None, knockdown_threshold=500.0, normalize=False, distance_metric="cosine", embedding_type="method_embedding"):
    """
    anchor_state_candidate_treatment_types is a list of tuples (anchor_state, treatment_type) to each use as an independent 'anchor' to compute cosine distances towards in the embedding space specified by method 
    treatment_type is the condition of candidates that we are interested in (e.g. just find the treatment_type = "Treated" candidates closest to my anchor_state)
    will call rank_perturbations_by_similarity
    will also call perturbation_vs_control_umap to visualize the top n of these perturbations
    min_sample is the minimum number of cells for a condition+perturbation to be included in the list, will filter out any condition+perturbation that does not meet this minimal size
    returns:
        (1) dictionary with key: (anchor_state, treatment_type), value: ranked list of tuples(gene target, cosine distance) (ascending by cosine distance)
        (2) count map with key: treatment_type, key: perturbation, value: # of cells
    """
    ##filter cells out that are >= knockdown threshold
    subadata = adata[(adata.obs["knockdown_efficiencies"] < knockdown_threshold)] ##will keep control cells because these have a knockdown efficiency of -1.0
    print(f"knockdown_threshold of {knockdown_threshold} filtered {len(adata) - len(subadata)} cells out")
    anchor_map = {a_t: [] for a_t in anchor_state_candidate_treatment_types} ##key: (anchor_state, treatment_type) tuple, value: ranked list of (gene target, cosine) tuples
    count_map = get_perturbation_counts_per_treatment(subadata) ##key: condition, key: perturbation name, value: sample size / number of cells 
    for (anchor_state, treatment_type) in anchor_state_candidate_treatment_types:
        assert "Treated" in anchor_state or "Untreated" in anchor_state
        sorted_perts = rank_perturbations_by_similarity(subadata, anchor_state=anchor_state, normalize=normalize, distance_metric=distance_metric, embedding_type=embedding_type) ##list of tuples: (perturbation, cosine)
        condition_gene_target_cosines = [x for x in sorted_perts if treatment_type in x[0]] ##from the sorted_perts, get just the treated perts for to_healthy, or untreated perts for to_disease
        gene_target_cosines = [(x[0].replace(f"{treatment_type}_", ""), x[1]) for x in condition_gene_target_cosines] ##get rid of treatment condition in candidate strings
        gene_target_cosines = [(x[0], x[1]) for x in gene_target_cosines if count_map[treatment_type][x[0]] >= min_sample] ##filter by sample size
        anchor_map[(anchor_state, treatment_type)] = gene_target_cosines
        plot_perts = [x[0] for x in gene_target_cosines[0:top_n]] ##top n perts - just perturbation target names  
        print(f"{dataset} {method} knockdown_thresh: {knockdown_threshold}, top {top_n} perts most similar to {anchor_state}: {plot_perts}")
    return anchor_map, count_map

def create_state_precursor_h5ads():
    """
    state requires a column called gene_name in .var dataframe to use for gene names, 
    writes a new AnnData object data/state_imru_precursor.h5ad
    after method completes, can process written AnnData object further via:
        uv run state emb transform \
        --model-folder /home/wongd26/mlcs/mlhub/state_SE-600M \
        --checkpoint /home/wongd26/mlcs/mlhub/state_SE-600M/se600m_epoch16.ckpt \
        --input /hpfs/userws/wongd26/imru-reversion/data/state_imru_precursor.h5ad \
        --output /hpfs/userws/wongd26/imru-reversion/data/state_imru.h5ad
    """
    ##full dataset
    adata = sc.read_h5ad("data/imru_full.h5ad")
    adata.var["gene_name"] = adata.var.index
    adata.write("data/state_imru_precursor_full.h5ad")
    ##indication specific subsets
    dataset_to_gene_names = extract_indication_genes()
    for dataset in dataset_to_gene_names:
        adata = sc.read_h5ad(f"data/imru_{dataset}.h5ad")
        adata.var["gene_name"] = adata.var.index
        adata.write(f"data/state_imru_precursor_{dataset}.h5ad")

def sum_normalize_matrix(X, total=1e6, log1p=True):
    """
    given 2D matrix X, will have it such that each row of X sums to total,
    if log1p is true, will apply log1p (ln(1 + x))
    returns new matrix
    if a row is all zeros, will keep it as is
    """
    row_sums = X.sum(axis=1, keepdims=True)
    row_sums[row_sums == 0] = 1.0
    normalized_X = (X / row_sums) * total
    if log1p:
        normalized_X = np.log1p(normalized_X)
    return normalized_X

def get_knockdown_efficiency(adata, control_strings):
    """
    given adata, will compute the knockdown efficiency for each cell:
        count of cell-normalized perturbed gene G / cell-normalized avg expression of G in control population in same treatment condition 
    return adata with knockdown efficiencies added to .obs["knockdown_efficiencies"]
    Note that this is not the traditional KD efficiency as presented in Replogle, but a ratio of raw count of target gene relative to control
    See get_knockdown_efficiencies_Replogle(control_strings) for the more standard approach
    """
    treated_control_indices = np.array([i for i in range(0, len(adata)) if adata.obs["condition"].iloc[i] == "Treated" and adata.obs["gene_target"].iloc[i] in control_strings])
    untreated_control_indices = np.array([i for i in range(0, len(adata)) if adata.obs["condition"].iloc[i] == "Untreated" and adata.obs["gene_target"].iloc[i] in control_strings])
    X = sum_normalize_matrix(adata.X.toarray(), log1p=True)
    condition_to_avg_control_vector = {"Treated": np.mean(X[treated_control_indices], axis=0), "Untreated": np.mean(X[untreated_control_indices], axis=0)}
    columns = list(adata.var.index)
    efficiencies = []
    for i in range(0, len(adata.obs["gene_target"])): ##i will be cell index (row)
        gene_pert = adata.obs["gene_target"].iloc[i]
        if gene_pert in ["NO-TARGET", "SAFE_TARGET"]: ##if any of the control strings, use -1.0
            efficiencies.append(-1.0)
        else:
            condition = adata.obs["condition"].iloc[i]
            column_index = columns.index(gene_pert)
            perturbed_gene_count = X[i][column_index]
            avg_expression = condition_to_avg_control_vector[condition][column_index]
            efficiencies.append(perturbed_gene_count / avg_expression)
            if i % 5000 == 0:
                print(i, len(adata.obs["gene_target"]))
    adata.obs["knockdown_efficiencies"] = efficiencies
    return adata

def get_method_title(string):
    """
    Get title string of method and return
    """
    if  "SCGPT" in string.upper():
        return "scGPT"
    elif "SCIMILARITY" in string.upper():
        return "SCimilarity"
    elif "UMAP" in string.upper():
        return "UMAP"
    elif "STATE" in string.upper():
        return "STATE"
    elif "RANDOM" in string.upper():
        return "Random"
    elif "RAW" in string.upper():
        return "Raw"
    elif "CHATGPT5.0" in string.upper():
        return "ChatGPT 5.0"
    elif "COMPBIO_DIRECT" in string.upper():
        # return  "Traditional Comp Bio (Inflammatory)"
        return "DE (Inflammatory)"
    elif "COMPBIO_MANUAL" in string.upper():
        return "Comp Bio Manual"
    elif "COMPBIO_VOX" in string.upper():
        return "Comp Bio Vox"
    elif "COMPBIO_UNTREATED_DIRECT" in string.upper():
        # return  "Traditional Comp Bio (Basal)"
        return "DE (Basal)"
    else:
        raise Exception(f"method {string} not accounted for")

def get_dataset_title(string):
    """
    Get title string of dataset and return
    """
    if string.upper() == "FULL":
        return "Full Dataset"
    elif string.upper() in ["ALSAIGH", "PASS", "TELLIDES"]:
        return string.capitalize()
    else:
        raise Exception(f"no dataset found for string {string}")

def get_sorted_perts(dataset="full", method=None, knockdown_threshold=500.0, min_sample=10):
    """
    loads the sorted target map and count map, filters the target rankings by min_sample for cell counts
    """
    sorted_perts, count_map = pickle.load(open(f"pickles/{dataset}_{method}_sorted_perts_knockdown_threshold_{knockdown_threshold}_min_sample_0.pkl", "rb")) ##load ranking without any filtering and then apply cell count filter; dictionary with key: (anchor_state, treatment_type), value: ranked list of tuples(gene target, cosine distance) (ascending by cosine distance)
    if min_sample != 0: ##if cell count filtering needs to be applied
        for (anchor_state, treatment_type) in sorted_perts:
            entries = sorted_perts[(anchor_state, treatment_type)] ##ranked list of tuples(gene target, cosine distance)
            sorted_perts[(anchor_state, treatment_type)] = [(x[0], x[1]) for x in entries if count_map[treatment_type][x[0]] >= min_sample]
    return sorted_perts, count_map 

def plot_cosines(top_n=10, plot_type="distance", min_sample=0):
    """
    Plot distribution of cosine similarities / distances (depending on plot_type) for the different methods
    will also write CSVs with header ['Perturbation', f'Cosine Distance to {anchor_state}', 'Cell Count'] to outputs/cosines/
    """
    if not os.path.isdir("outputs/cosines/"):
        os.makedirs("outputs/cosines/")
    knockdown_threshold = 500.0
    methods = ["random", "raw", "umap", "scimilarity", "scgpt", "state"]
    for method in methods:
        sorted_perts, count_map = get_sorted_perts(method=method, min_sample=min_sample)
        for (anchor_state, treatment_type) in sorted_perts:
            perturbations = [x[0] for x in sorted_perts[(anchor_state, treatment_type)]]
            cosine_dist = [x[1] for x in sorted_perts[(anchor_state, treatment_type)]]
            cosine_sims = [1.0 - x[1] for x in sorted_perts[(anchor_state, treatment_type)]]
            if plot_type == "distance":
                cosine_plot = cosine_dist
                plt_label = "Distance"
            elif plot_type == "similarity":
                cosine_plot = cosine_sims
                plt_label = "Similarity"
            else:
                raise Exception()
            counts = [count_map[treatment_type][perturbation] for perturbation in perturbations]
            fig, ax = plt.subplots()
            x_labels = perturbations[0: top_n]
            x = np.array(range(0, len(x_labels)))
            y = cosine_plot[0: top_n]
            fontsize = 5 if top_n >= 25 else 9
            if top_n <= 50:
                ax.set_xticks(x)
                ax.set_xticklabels(x_labels, rotation=90, fontsize=fontsize)
            if plot_type == "similarity":
                ax.bar(x, y[::-1]) ##reverse for just the plot so ascending order
            else:
                ax.bar(x, y)
            ax.set_ylabel(f"Cosine {plt_label}")
            if top_n <= 100:
                plt.title(f"{get_method_title(method)}: Cosine {plt_label} for Top {min(top_n, len(perturbations))} Perturbations")
            else:
                plt.title(f"Distribution of Cosine {plt_label} for {get_method_title(method)}")
            if plot_type == "similarity":
                ax.set_ylim((min(cosine_plot) - 0.02, 1.0))
            plt.savefig(f"outputs/cosines/full_{method}_{treatment_type}_{anchor_state}_{knockdown_threshold}_{min_sample}_top_n.png", dpi=300, bbox_inches="tight")
            ##plot histogram of all cosines
            fig, ax = plt.subplots()
            ax.hist(cosine_plot, bins=50)
            ax.set_ylabel(f"Cosine {plt_label}")
            plt.title(f"{get_method_title(method)}: Histogram of Cosine Similarities for All Perturbations")
            plt.savefig(f"outputs/cosines/full_{method}_{treatment_type}_{anchor_state}_{knockdown_threshold}_{min_sample}_histogram.png", dpi=300, bbox_inches="tight")
            ##write csv of all cosine *DISTANCES*
            output_filename = f'outputs/cosines/full_{method}_{treatment_type}_perturbations_closest_to_{anchor_state}_{knockdown_threshold}_{min_sample}.csv'
            with open(output_filename, 'w', newline='') as csvfile:
                writer = csv.writer(csvfile)
                writer.writerow(['Perturbation', f'Cosine Distance to {anchor_state}', 'Cell Count'])
                for item1, item2, item3 in zip(perturbations, cosine_dist, counts):
                    writer.writerow([item1, item2, item3])
            print(f"CSV file '{output_filename}' created successfully.")

def get_library_abbrev(string):
    """
    given library string will return an abbreviated string
    """
    if string == "imru_internal":
        return "IMRU"
    elif string == "MSigDB_Hallmark_2020":
        return "MSig"
    elif "KEGG" in string.upper():
        return "KEGG"
    elif string == "GO_Biological_Process_2025":
        return "GO"
    elif "BIOCARTA" in string.upper():
        return "BioCarta"
    elif "REACTOME" in string.upper():
        return "Reactome"
    elif string == "curated":
        return "Curated"
    elif string == "curated filtered":
        return "Curated Filtered"
    elif string == "fgsea enrichment":
        return "FGSEA Enrichment"
    else:
        print(f"WARNING: no library abbreviation for {string}")
        return string

def enrichment_runner(min_sample=0):
    """
    runner call for enrichment analyses, saves to outputs/
    """
    library_sets = get_relevant_pathways(just_this_study=True, include_imru_internal=True, key_type="tuple")
    ##dictionary that will be used for plotting the main permutation of interest
    ##key: method, key: pathway, value: AUC for parameter-combo of interest: dataset=="full", kd=500.0 (no filtering), target="Untreated_SAFE_TARGET"
    plot_map = {} 
    bests = []
    knockdown_threshold = 500.0
    for library, pathway in library_sets:
        abbrev = f"{get_library_abbrev(library)}: {pathway.upper()} (n={len(library_sets[(library, pathway)])})" 
        print(library, pathway)
        meta_AUC, best_method = enrichment_plot(top_n=10, library=library, pathway=pathway, positive_controls=library_sets[(library, pathway)], save_dir=f"outputs/enrichment/{library}/{pathway}/", min_sample=min_sample, verbose=False)
        bests.append(best_method)
        for method in meta_AUC:
            if method not in plot_map:
                plot_map[method] = {abbrev: meta_AUC[method][knockdown_threshold]["Untreated_SAFE_TARGET"]}
            else:
                plot_map[method][abbrev] = meta_AUC[method][knockdown_threshold]["Untreated_SAFE_TARGET"]
    print(f"enrichment best methods on standard permutation: {Counter(bests)}")
    ##plot a summary bar graph of main permutation of interest
    fig, ax = plt.subplots()
    width = .09
    methods = ["random", "raw", "umap", "scimilarity", "scgpt", "state", "chatgpt5.0", "compbio_direct", "compbio_untreated_direct"]
    x_labels = list(plot_map["raw"].keys())
    print(x_labels)
    x_labels = sorted([x_ for x_ in x_labels if "INFLAMMATION REVERSAL GIVEN" not in x_]) ##don't include (+) control set in summary annotated pathway plot
    x = np.array(range(0, len(x_labels)))
    ax.set_xticks(x)
    for method in methods:
        y = [plot_map[method][pathway] for pathway in x_labels]
        ax.bar(x, y, width=width, label=get_method_title(method))
        x = x + width
    plt.yticks(fontsize=7)
    ax.set_xticklabels(x_labels, rotation=40, fontsize=6, ha="right", rotation_mode='anchor')
    ax.set_ylabel("AUC", fontsize=7)
    ax.set_ylim((0.0, 1.02))
    ax.set_title("Enrichment Comparison\n Across Relevant Pathways", fontsize=9)
    ax.set_yticklabels(['{:,.0%}'.format(y_) for y_ in ax.get_yticks()])
    box = ax.get_position()
    ax.set_position([box.x0, box.y0, box.width, box.height * 0.85])
    ax.legend(loc='upper right', prop={"size":7}, bbox_to_anchor=(1, 1.55))
    plt.gcf().subplots_adjust(top=.76)
    plt.savefig("outputs/enrichment/summary_enrichment_standard_permutation.png", dpi=300, bbox_inches="tight")

def enrichment_plot(top_n, library, pathway, positive_controls, save_dir, min_sample=0, verbose=False):
    """
    Given a set of positive_controls, a directory to save to, will save enrichment plots
    we can calculate enrichment scores of our ranked priority lists to see how early in the queue the (+) controls come up:
    will also print how many (+) controls come up in top_n for each method+dataset 
    returns map with key: method, key KD threshold, key: anchor_state, value: AUC for this pathway and set of positive_controls, and also which method was the best on the standard permutation
    """
    if not os.path.isdir(save_dir):
        os.makedirs(save_dir)
    methods = ["random", "raw", "umap", "scimilarity", "scgpt", "state", "chatgpt5.0", "compbio_direct", "compbio_untreated_direct"]
    knockdown_thresholds = [500.0]
    meta_AUC = {m: {kd: {} for kd in knockdown_thresholds} for m in methods}  ##key: method, key: anchor_state, key KD threshold, value: AUC for this pathway and set of positive_controls 
    global_max_AUC, global_max_params = 0, ""
    best_method, best_method_AUC = "", -1.0 ##the best method and best AUC 
    method_to_x_y = {} ##key: anchor_state, key: method, value: [x,y]
    for method in methods:
        for knockdown_threshold in knockdown_thresholds:
            sorted_perts, _ = get_sorted_perts(method=method, min_sample=min_sample)
            for (anchor_state, treatment_type) in sorted_perts:
                if anchor_state not in method_to_x_y: 
                    method_to_x_y[anchor_state] = {m: [] for m in methods}
                ranked_list = [x[0] for x in sorted_perts[(anchor_state, treatment_type)]]
                ##check if all positive controls in ranked_list, if missing remove from (+) control set local to this method's eval (don't touch positive_control global to this function because of multiple for loop iterations)
                target_set = set(ranked_list)
                missing = []
                for p in positive_controls:
                    if p not in target_set:
                        missing.append(p)
                if len(missing) > 0:
                    print(f"WARNING: {library} {pathway} {len(missing)} / {len(positive_controls)} (+) controls are missing: {missing}")
                positive_control_local = set([x for x in positive_controls if x not in missing])
                indices_of_positives = [ranked_list.index(p) for p in positive_control_local]
                if verbose:
                    hits_in_top_n = [(ranked_list[i], f"rank={i}") for i in indices_of_positives if i <= top_n]
                    print(f"    {method} enrichment {knockdown_threshold} {anchor_state}: {hits_in_top_n}\n.   {len(hits_in_top_n)} hit(s) found in top {top_n}")
                    print(f"    {method} indices of (+): {sorted(indices_of_positives)}")
                ##make the plot
                with warnings.catch_warnings():
                    warnings.simplefilter("ignore") 
                    fig, ax = plt.subplots()
                    x = np.linspace(0, 1, len(ranked_list))
                    y = []
                    running_sum = 0
                    for i, item in enumerate(ranked_list):
                        if i in indices_of_positives:
                            running_sum += 1 ##increment for a hit
                        y.append(running_sum / float(len(indices_of_positives)))
                    ax.plot(x, y)
                    method_to_x_y[anchor_state][method] = [x, y]
                    AUC = np.trapz(y, x) / ((max(x) - min(x)) * 1.0) ##global max will be 1.0
                    meta_AUC[method][knockdown_threshold][anchor_state] = AUC
                    if AUC > global_max_AUC:
                        global_max_AUC = AUC
                        global_max_params = f"{method} {anchor_state} {knockdown_threshold} {min_sample}"
                    if AUC > best_method_AUC and knockdown_threshold == 500.0 and anchor_state == "Untreated_SAFE_TARGET": ##keep track of best method for no-exclusion based off KD %, and anchor=Untreated_SAFE_TARGET
                        best_method_AUC = AUC 
                        best_method = method
                    plt.title(f"Enrichment Plot for {get_method_title(method)}\nAUC={AUC:.0%}")
                    plt.xlabel(f"Percentage in Priority Queue of {treatment_type} Perturbations")
                    plt.ylabel(f"Percentage of (+) Control Perturbations Found (n={len(positive_control_local)})")
                    ax.set_xticklabels(['{:,.0%}'.format(x_) for x_ in ax.get_xticks()])
                    ax.set_yticklabels(['{:,.0%}'.format(y_) for y_ in ax.get_yticks()])
                    plt.savefig(os.path.join(save_dir, f"full_{method}_anchor={anchor_state}_{knockdown_threshold}_{min_sample}.png"), dpi=300)
        print(f"    best method, no KD filtering, anchor=Untreated_SAFE_TARGET:\n    {best_method} with AUC={round(best_method_AUC, 3)}")
    ##plot all curves in one plot 
    max_AUC = 0.0
    for anchor_state in method_to_x_y:
        fig, ax = plt.subplots()
        for method in method_to_x_y[anchor_state]: 
            x, y = method_to_x_y[anchor_state][method]
            AUC = np.trapz(y, x) / ((max(x) - min(x)) * 1.0)
            max_AUC = max(AUC, max_AUC)
            ax.plot(x, y, label=get_method_title(method) + f", AUC={AUC:.0%}", linewidth=0.5, linestyle="dashed")
        plt.title(f"Enrichment Plots for {get_library_abbrev(library)}:\n{pathway}")
        plt.xlabel(f"Percentage in Priority Queue of {treatment_type} Perturbations")
        plt.ylabel(f"Percentage of (+) Control Perturbations Found (n={len(positive_control_local)})")
        ax.set_xticklabels(['{:,.0%}'.format(x_) for x_ in ax.get_xticks()])
        ax.set_yticklabels(['{:,.0%}'.format(y_) for y_ in ax.get_yticks()])
        box = ax.get_position()
        ax.set_position([box.x0, box.y0, box.width, box.height * 0.85])
        legend = ax.legend(loc='upper right', prop={"size":8}, bbox_to_anchor=(1, 1.62))
        legend_texts = legend.get_texts()
        legend_text_strings = [l.get_text() for l in legend.get_texts()]
        print(legend_text_strings)
        max_indices = [i for i in range(0, len(legend_texts)) if f"{max_AUC:.0%}" in legend_text_strings[i]]
        for max_index in max_indices:
            legend_texts[max_index].set_weight('bold')
        plt.gcf().subplots_adjust(top=.80)
        plt.savefig(f"outputs/enrichment/{library}/{pathway}/summary_{anchor_state}.png", dpi=300, bbox_inches="tight")
    return meta_AUC, best_method

def get_treatment_type(string):
    """
    given string, returns either Treated or Untreated if found in string
    """
    if "UNTREATED" in string.upper():
        return "Untreated"
    elif "TREATED" in string.upper(): ##if treated is a substring of untreated, but method will shor circuit if untreated
        return "Treated"
    else:
        raise Exception(f"treatment condition cannot be determined from {string}")

def replace_nth_char(s, old_char, new_char, n):
    count = 0
    start_index = 0
    while count < n:
        index = s.find(old_char, start_index)
        if index == -1:
            return s
        count += 1
        start_index = index + 1
    return s[:index] + new_char + s[index + 1:]

def plot_heatmap(matrix, label_matrix, label_type, x_labels, y_labels, vmin, vmax, rounding, title, title_fontsize=12, x_rotation=False, x_tick_fontsize=10, y_tick_fontsize=10, colorbar_label=None, fig_size=None, colorbar_shrink=1.0, reverse=False, add_plus_to_max=False, save_path=None):
    """
    standard heatmap plotting 
    """
    ##plot the result as heatmap
    if fig_size == None:
        fig, ax = plt.subplots()
    else:
        fig, ax = plt.subplots(figsize=fig_size) 
    color_spectrum = "viridis_r" if reverse else "viridis"
    cax = ax.imshow(matrix, aspect='auto', vmin=vmin, vmax=vmax, cmap=color_spectrum)
    ax.set_xticks(np.arange(len(x_labels)))
    ax.set_yticks(np.arange(len(y_labels)))
    if x_rotation:
        ax.set_xticklabels(x_labels, rotation=40, ha="right", rotation_mode='anchor', fontsize=x_tick_fontsize)
    else:
        ax.set_xticklabels(x_labels, fontsize=x_tick_fontsize)
    ax.set_yticklabels(y_labels, fontsize=y_tick_fontsize)
    cbar = fig.colorbar(cax, ax=ax, shrink=colorbar_shrink)

    ##add a "+" substring to max 
    if add_plus_to_max: 
        cbar_ticks = [int(t_.get_text()) for t_ in cbar.ax.get_yticklabels()]
        cbar.set_ticklabels([str(t_) if t_ < max(cbar_ticks) else str(t_) + "+" for t_ in cbar_ticks])

    if colorbar_label != None:
        cbar.set_label(colorbar_label, fontsize=title_fontsize, labelpad=12, rotation=360)
    if label_matrix != None:
        for i in range(len(label_matrix)): #add text annotations
            for j in range(len(label_matrix[i])):
                if label_type == "string":
                    ax.text(j, i, f'{label_matrix[i][j]}', ha='center', va='center', color='white', fontsize=5)
                else:
                    ax.text(j, i, f'{label_matrix[i][j]:.{rounding}%}', ha='center', va='center', color='white')
    plt.title(title, fontsize=title_fontsize)
    plt.savefig(save_path, dpi=300, bbox_inches="tight")

def generate_all_cosine_distances(adata_path, anchors=[]):
    """
    Given a path to adata with embeddings in adata.X, condition in .obs["condition"], and targets in .obs["gene_target"],
    create a CSV where each row is a cell, with columns: condition, gene_target, cosine_distance_to_{each anchor state}
    write to outputs/pairwise_cosines
    """
    adata = sc.read_h5ad(adata_path)
    ##calculate the centroid of each anchor
    anchor_to_centroid = {anchor_state: "" for anchor_state in anchors}
    embeddings = adata.X
    conditions = list(adata.obs["condition"])
    gene_targets = list(adata.obs["gene_target"])
    kd_ratio = list(adata.obs["knockdown_efficiencies"])
    kd_ratio = [x if x != -1 else 'NA' for x in kd_ratio]
    condition_targets = [conditions[i] + "_" + gene_targets[i] for i in range(0, len(conditions))]
    for anchor_state in anchors:
        anchor_indices = [i for i in range(0, len(condition_targets)) if condition_targets[i] == anchor_state]
        centroid = np.mean(embeddings[anchor_indices], axis=0)
        anchor_to_centroid[anchor_state] = centroid 
    ##iterate over each cell and calcualte distance to each anchor 
    anchor_to_distance = {anchor_state: [] for anchor_state in anchors} 
    for i in range(0, len(embeddings)):
        for anchor_state in anchor_to_centroid:
            distance = cosine_distance(embeddings[i], anchor_to_centroid[anchor_state])
            anchor_to_distance[anchor_state].append(distance)
    ##now that we have our lists: conditions, gene_targets, and cosine distances, write a CSV
    df = pd.DataFrame({"condition": conditions, "gene_target": gene_targets, "condition_gene_target": condition_targets, "kd_ratio (target expr / avg target expr in control)": kd_ratio}) 
    for anchor_state in anchor_to_distance:
        df[f"cosine_distance_to_{anchor_state}"] = anchor_to_distance[anchor_state]
    ##make sure anchor is closest to itself on average 
    for anchor in anchors:
        self_mean = 0
        other_means = []
        df_sub = df[df["condition_gene_target"] == anchor]
        for column in [f"cosine_distance_to_{anchor_state}" for anchor_state in anchors]:
            mean_ = np.mean(df_sub[column])
            if anchor in column:
                self_mean = mean_
            else:
                other_means.append(mean_)
        print(anchor, self_mean, sorted(other_means))
        if self_mean > min(other_means):
            raise Exception(f"Error: anchor {anchor} is not closest to itself!")
    df.to_csv(f"outputs/pairwise_cosines/{get_method_title(adata_path)}_cosine_distances_for_all_cells.csv", index=False)

def get_all_target_names():
    """
    will return a list of all target names from data/scgpt_imru_full_processed.h5ad
    if "pickles/target_names.pkl" does not exist, will save list to this file
    """
    if os.path.isfile("pickles/target_names.pkl"):
        l = pickle.load(open("pickles/target_names.pkl", "rb"))
        return sorted(l)
    else:
        adata = sc.read_h5ad("data/scgpt_imru_full_processed.h5ad")
        targets = set(adata.obs["gene_target"])
        l = sorted(targets)
        pickle.dump(l, open("pickles/target_names.pkl", "wb"))
        return l

class RankingSimilarity:
    """
    implementation of RBO from https://github.com/changyaochen/rbo
    This class will include some similarity measures between two different ranked lists.
    """
    from typing import List, Optional, Union
    def __init__(
        self,
        S: Union[List, np.ndarray],
        T: Union[List, np.ndarray],
        verbose: bool = False,
    ) -> None:
        """
        Initialize the object with the required lists.
        Examples of lists:
        S = ["a", "b", "c", "d", "e"]
        T = ["b", "a", 1, "d"]
        Both lists reflect the ranking of the items of interest, for example,
        list S tells us that item "a" is ranked first, "b" is ranked second,
        etc.
        Args:
            S, T (list or numpy array): lists with alphanumeric elements. They
                could be of different lengths. Both of the them should be
                ranked, i.e., each element"s position reflects its respective
                ranking in the list. Also we will require that there is no
                duplicate element in each list.
            verbose: If True, print out intermediate results. Default to False.
        """
        assert type(S) in [list, np.ndarray]
        assert type(T) in [list, np.ndarray]
        assert len(S) == len(set(S))
        assert len(T) == len(set(T))
        self.S, self.T = S, T
        self.N_S, self.N_T = len(S), len(T)
        self.verbose = verbose
        self.p = 0.5  # just a place holder

    def assert_p(self, p: float) -> None:
        """Make sure p is between (0, 1), if so, assign it to self.p.
        Args:
            p (float): The value p.
        """
        assert 0.0 < p < 1.0, "p must be between (0, 1)"
        self.p = p

    def _bound_range(self, value: float) -> float:
        """Bounds the value to [0.0, 1.0]."""
        try:
            assert (0 <= value <= 1 or np.isclose(1, value))
            return value
        except AssertionError:
            print("Value out of [0, 1] bound, will bound it.")
            larger_than_zero = max(0.0, value)
            less_than_one = min(1.0, larger_than_zero)
            return less_than_one

    def rbo(
        self,
        k: Optional[float] = None,
        p: float = 1.0,
        ext: bool = False,
    ) -> float:
        """
        This the weighted non-conjoint measures, namely, rank-biased overlap.
        Unlike Kendall tau which is correlation based, this is intersection
        based.
        The implementation if from Eq. (4) or Eq. (7) (for p != 1) from the
        RBO paper: http://www.williamwebber.com/research/papers/wmz10_tois.pdf
        If p = 1, it returns to the un-bounded set-intersection overlap,
        according to Fagin et al.
        https://researcher.watson.ibm.com/researcher/files/us-fagin/topk.pdf
        The fig. 5 in that RBO paper can be used as test case.
        Note there the choice of p is of great importance, since it
        essentially control the "top-weightness". Simply put, to an extreme,
        a small p value will only consider first few items, whereas a larger p
        value will consider more items. See Eq. (21) for quantitative measure.
        Args:
            k: The depth of evaluation.
            p: Weight of each agreement at depth d:
                p**(d-1). When set to 1.0, there is no weight, the rbo returns
                to average overlap.
            ext: If True, we will extrapolate the rbo, as in Eq. (23).
        Returns:
            The rbo at depth k (or extrapolated beyond).
        """
        if not self.N_S and not self.N_T:
            return 1  # both lists are empty
        if not self.N_S or not self.N_T:
            return 0  # one list empty, one non-empty
        if k is None:
            k = float("inf")
        k = min(self.N_S, self.N_T, k)
        # initialize the agreement and average overlap arrays
        A, AO = [0] * k, [0] * k
        if p == 1.0:
            weights = [1.0 for _ in range(k)]
        else:
            self.assert_p(p)
            weights = [1.0 * (1 - p) * p**d for d in range(k)]
        # using dict for O(1) look up
        S_running, T_running = {self.S[0]: True}, {self.T[0]: True}
        A[0] = 1 if self.S[0] == self.T[0] else 0
        AO[0] = weights[0] if self.S[0] == self.T[0] else 0
        for d in range(1, k):
            tmp = 0
            # if the new item from S is in T already
            if self.S[d] in T_running:
                tmp += 1
            # if the new item from T is in S already
            if self.T[d] in S_running:
                tmp += 1
            # if the new items are the same, which also means the previous
            # two cases did not happen
            if self.S[d] == self.T[d]:
                tmp += 1
            # update the agreement array
            A[d] = 1.0 * ((A[d - 1] * d) + tmp) / (d + 1)
            # update the average overlap array
            if p == 1.0:
                AO[d] = ((AO[d - 1] * d) + A[d]) / (d + 1)
            else:  # weighted average
                AO[d] = AO[d - 1] + weights[d] * A[d]
            # add the new item to the running set (dict)
            S_running[self.S[d]] = True
            T_running[self.T[d]] = True
        if ext and p < 1:
            return self._bound_range(AO[-1] + A[-1] * p**k)
        return self._bound_range(AO[-1])

    def rbo_ext(self, p=0.98):
        """
        This is the ultimate implementation of the rbo, namely, the
        extrapolated version. The corresponding formula is Eq. (32) in the rbo
        paper.
        """
        self.assert_p(p)
        if not self.N_S and not self.N_T:
            return 1  # both lists are empty
        if not self.N_S or not self.N_T:
            return 0  # one list empty, one non-empty
        # since we are dealing with un-even lists, we need to figure out the
        # long (L) and short (S) list first. The name S might be confusing
        # but in this function, S refers to short list, L refers to long list
        if len(self.S) > len(self.T):
            L, S = self.S, self.T
        else:
            S, L = self.S, self.T
        s, l = len(S), len(L)  # noqa
        # initialize the overlap and rbo arrays
        # the agreement can be simply calculated from the overlap
        X, A, rbo = [0] * l, [0] * l, [0] * l
        # first item
        S_running, L_running = {S[0]}, {L[0]}  # for O(1) look up
        X[0] = 1 if S[0] == L[0] else 0
        A[0] = X[0]
        rbo[0] = 1.0 * (1 - p) * A[0]
        # start the calculation
        disjoint = 0
        ext_term = A[0] * p
        for d in range(1, l):
            if d < s:  # still overlapping in length
                S_running.add(S[d])
                L_running.add(L[d])
                # again I will revoke the DP-like step
                overlap_incr = 0  # overlap increment at step d
                # if the new items are the same
                if S[d] == L[d]:
                    overlap_incr += 1
                else:
                    # if the new item from S is in L already
                    if S[d] in L_running:
                        overlap_incr += 1
                    # if the new item from L is in S already
                    if L[d] in S_running:
                        overlap_incr += 1
                X[d] = X[d - 1] + overlap_incr
                # Eq. (28) that handles the tie. len() is O(1)
                A[d] = 2.0 * X[d] / (len(S_running) + len(L_running))
                rbo[d] = rbo[d - 1] + 1.0 * (1 - p) * (p**d) * A[d]
                ext_term = 1.0 * A[d] * p**(d + 1)  # the extrapolate term
            else:  # the short list has fallen off the cliff
                L_running.add(L[d])  # we still have the long list
                # now there is one case
                overlap_incr = 1 if L[d] in S_running else 0
                X[d] = X[d - 1] + overlap_incr
                A[d] = 1.0 * X[d] / (d + 1)
                rbo[d] = rbo[d - 1] + 1.0 * (1 - p) * (p**d) * A[d]
                X_s = X[s - 1]  # this the last common overlap
                # second term in first parenthesis of Eq. (32)
                disjoint += 1.0 * (1 - p) * (p**d) * \
                    (X_s * (d + 1 - s) / (d + 1) / s)
                ext_term = 1.0 * ((X[d] - X_s) / (d + 1) + X[s - 1] / s) * \
                    p**(d + 1)  # last term in Eq. (32)
        return self._bound_range(rbo[-1] + disjoint + ext_term)

    def top_weightness(
            self,
            p: Optional[float] = None,
            d: Optional[int] = None):
        """
        This function will evaluate the degree of the top-weightness of the
        rbo. It is the implementation of Eq. (21) of the rbo paper.
        As a sanity check (per the rbo paper),
        top_weightness(p=0.9, d=10) should be 86%
        top_weightness(p=0.98, d=50) should be 86% too
        Args:
            p (float), default None: A value between zero and one.
            d (int), default None: Evaluation depth of the list.
        Returns:
            A float between [0, 1], that indicates the top-weightness.
        """
        # sanity check
        self.assert_p(p)
        if d is None:
            d = min(self.N_S, self.N_T)
        else:
            d = min(self.N_S, self.N_T, int(d))
        if d == 0:
            top_w = 1
        elif d == 1:
            top_w = 1 - 1 + 1.0 * (1 - p) / p * (np.log(1.0 / (1 - p)))
        else:
            sum_1 = 0
            for i in range(1, d):
                sum_1 += 1.0 * p**(i) / i
            top_w = 1 - p**(i) + 1.0 * (1 - p) / p * (i + 1) * \
                (np.log(1.0 / (1 - p)) - sum_1)  # here i == d-1
        if self.verbose:
            print("The first {} ranks have {:6.3%} of the weight of "
                  "the evaluation.".format(d, top_w))
        return self._bound_range(top_w)

def compare_method_rankings(ranking_method="RBO", min_sample=0):
    """
    compare the different triage ranks for each method for the case of Treated perturbations most similar to Untreated safe-target
    and plot heatmaps
    """
    methods = ["compbio_untreated_direct", "compbio_direct", "chatgpt5.0", "raw", "umap", "scimilarity", "scgpt", "state"] 
    method_to_rankings = {method: [] for method in methods}
    for method in methods: 
        ranking_map, _ = get_sorted_perts(method=method, min_sample=min_sample)
        ranking = [x[0] for x in ranking_map[('Untreated_SAFE_TARGET', 'Treated')]]
        method_to_rankings[method] = ranking
    combo_map = {} ##key: combo, value: ranking method score
    for m1, m2 in itertools.combinations(methods, 2):
        if ranking_method == "Spearman":
            score, _ = spearmanr(method_to_rankings[m1], method_to_rankings[m2])
        elif ranking_method == "RBO":
            score = RankingSimilarity(method_to_rankings[m1], method_to_rankings[m2]).rbo(p=0.9) 
        elif "Percent Overlap at Top" in ranking_method:
            n = int(ranking_method.replace("Percent Overlap at Top ", ""))
            score = len(set(method_to_rankings[m1][0:n]).intersection(set(method_to_rankings[m2][0:n]))) / float(n)
        else:
            raise Exception()
        combo_map[(m1, m2)] = score
        combo_map[(m2, m1)] = score
    for method in methods: ##similarity to self is by definition 1.0
        combo_map[(method, method)] = 1.0
    print(sorted(combo_map.items(), key=lambda x:x[1]))
    ##plot the result as heatmap
    k1_values = methods
    k2_values = methods[::-1]
    k1_labels = [get_method_title(k1) for k1 in k1_values]
    k2_labels = [get_method_title(k2) for k2 in k2_values]
    heatmap_matrix = [[combo_map[(m1,m2)] for m1 in k1_values] for m2 in k2_values]
    plot_heatmap(matrix=heatmap_matrix, label_matrix=heatmap_matrix, label_type="float", x_labels=k1_labels, y_labels=k2_labels, vmin=0.0, vmax=0.6, rounding=0, title=f'{ranking_method} Between Each Method', x_rotation=True, save_path=f"outputs/rank_similarity_{ranking_method}.png")

def method_ranking_heatmap(top_n, weighted=False, min_sample=0):
    """
    plot heatmap of x-axis method, y-axis target, value rank 
    if weighted:
        will display the top_n targets ranked by summation of rank for each method
    else:
        will display the intersection of all genes in top 50 of any ranking
    """
    methods = ["compbio_untreated_direct", "compbio_direct", "raw", "umap", "scimilarity", "scgpt", "state", "chatgpt5.0"] 
    intersection_set = set()
    target_to_sum = {} ##key: target, value: sum of rankings for all methods
    method_to_ranking = {} ##key: method, value: list of ranked targets 
    mapp = {m: {} for m in methods} ##key: method, key: target, value: ranking 
    for method in methods:
        sorted_perts, _ = get_sorted_perts(method=method, min_sample=min_sample)
        ranked_list = [x[0] for x in sorted_perts[("Untreated_SAFE_TARGET", "Treated")]]
        method_to_ranking[method] = ranked_list
        if weighted: 
            intersection_set = intersection_set.union(set(ranked_list))
        else:
            intersection_set = intersection_set.union(set(ranked_list[0:top_n]))
    intersection_set = list(intersection_set)
    ##for each target get the index of it in each method's rankings
    for target in intersection_set:
        for method in methods:
            index = method_to_ranking[method].index(target)
            mapp[method][target] = index
            if target not in target_to_sum:
                target_to_sum[target] = index 
            else:
                target_to_sum[target] += index 
    ##sort targets from lowest sum to highest sum 
    if weighted: ##just get top_n weighted genes as final list
        ordered_targets = [x[0] for x in sorted(target_to_sum.items(), key=lambda x:x[1])][0:top_n][::-1] 
        title = f"Target Rankings for Top {top_n} Weighted Targets"
    else: ##use the top_n gene intersection set as final list 
        ordered_targets = [x[0] for x in sorted(target_to_sum.items(), key=lambda x:x[1])][::-1]
        title = f"Target Rankings for Top {top_n} Targets of any Method"
    matrix = [[mapp[m][target] for m in methods] for target in ordered_targets]
    plot_heatmap(matrix=matrix, label_matrix=None, label_type="float", x_labels=[get_method_title(m) for m in methods], y_labels=ordered_targets, vmin=0.0, vmax=100, rounding=0, title=title, title_fontsize=8, x_rotation=True, x_tick_fontsize=4.5, y_tick_fontsize=4.5, colorbar_label="Rank", fig_size=(3,6), colorbar_shrink=0.4, reverse=True, add_plus_to_max=True, save_path=f"outputs/rank_similarity_heatmap.png")

def plot_replicate_count(adata_path, control_strings):
    """
    Plots stacked histogram of gene-target replicate counts (one histogram for each treatment condition) excluding controls
    """
    conditions = ["Untreated", "Treated"]
    adata = sc.read_h5ad(adata_path)
    adata = adata[~adata.obs["gene_target"].isin(control_strings)]
    condition_to_counts = {condition: "" for condition in conditions}
    for condition in conditions:
        sub_adata = adata[adata.obs["condition"] == condition]
        counts = Counter(sub_adata.obs["gene_target"])
        count_list = sorted(list(counts.values()))
        condition_to_counts[condition] = count_list
        print(f"gene target replicate stats for {condition}: min: {np.min(count_list)}, median: {np.median(count_list)}, max: {np.max(count_list)}")
    colors = {"Untreated": "lime", "Treated": "deepskyblue"}
    for condition in conditions:
        fig, ax = plt.subplots()
        x = np.array(range(0, len(count_list)))
        ax.hist(condition_to_counts[condition], bins=200, label=condition, color=colors[condition])
        ax.set_xlabel("Gene-target Replicate Count (# of Cells)")
        ax.set_ylabel("Frequency (# of Perturbations)")
        plt.title(f"Gene-target Replicate Counts for {condition}")
        plt.savefig(f"outputs/replicate_count_{condition}.png", dpi=300)

def create_random_adata():
    """
    Creates anndata with randomized embeddings of shape 128, will be l2 normed
    """
    adata = sc.read_h5ad(f"data/umap_imru_full_processed.h5ad")
    random_matrix = np.random.normal(0, 1, size=(adata.X.shape[0], 128)) ##need non-uniform (e.g. NOT random.rand) because: with uniform, sampling will always approach a value of 0.5 as we approach infinite samples --> targets with similar cell count will then have similar avg representations: SAFE_TARGET is highest population, so avg of this will be closest to 0.5 for all features, targets with high cell count will also approach 0.5 and targets with low cell count will be further from 0.5 due to chance. Won't be the case with normal distribution as samples will always be close to mean of normal regardless of sample size.
    np.random.shuffle(random_matrix)
    random_matrix = preprocess_matrix(random_matrix, standardize=False, normalize=True)
    ##need to keep condition + target pairs the same because some targets are present in only 1 condition, keep this consistent; shuffle the two lists synchronously
    random_indices = [i for i in range(0, len(adata.obs["gene_target"]))]
    random.Random(36).shuffle(random_indices)
    gene_targets = list(adata.obs["gene_target"])
    conditions = list(adata.obs["condition"])
    gene_targets = [gene_targets[i] for i in random_indices]
    conditions = [conditions[i] for i in random_indices]
    obs_df = pd.DataFrame({"gene_target": gene_targets, "condition": conditions, "knockdown_efficiencies": adata.obs["knockdown_efficiencies"]})
    ##create AnnData object and write to disk
    random_adata = AnnData(X=random_matrix, obs=obs_df)
    random_adata.obsm["umap"] = get_umap(random_matrix)
    print("wrote: data/random_imru_full_processed.h5ad")
    random_adata.write("data/random_imru_full_processed.h5ad")

def get_DE_genes_per_perturbation():
    """
    will read raw AnnData and 
    compute a dictionary with key: treatment condition, key: gene target, value: list of DE genes with p < 0.05
    saves file "pickles/raw_DE_gene_results.pkl"
    """
    conditions = ["Treated", "Untreated"]
    adata = sc.read_h5ad("data/imru_full.h5ad")
    sc.pp.log1p(adata) ##rank sum test expects logarithmic data 
    de_dict = {c: {} for c in conditions}
    for condition in conditions:
        adata_condition = adata[adata.obs["condition"] == condition]
        ##remove perturbations that only contain 1 sample 
        counts = Counter(adata_condition.obs["gene_target"])
        singles = set([key for key in counts if counts[key] == 1])
        adata_condition = adata_condition[~adata_condition.obs["gene_target"].isin(singles)]
        adata_condition = adata_condition.copy() ##change view to copy because rank sum test will write to this AnnData object
        ##run differential expression
        print("about to compute rank sum")
        sc.tl.rank_genes_groups(adata_condition, groupby="gene_target", reference="SAFE_TARGET", method="wilcoxon")
        print("finished compute rank sum")
        ##extract results
        names = adata_condition.uns["rank_genes_groups"]["names"]
        pvals_adj = adata_condition.uns["rank_genes_groups"]["pvals_adj"]
        for group in names.dtype.names:
            genes = []
            for gene, pval in zip(names[group], pvals_adj[group]):
                if pval < 0.05:
                    genes.append(gene)
            de_dict[condition][group] = genes
    pickle.dump(de_dict, open("pickles/raw_DE_gene_results.pkl", "wb"))

def plot_de_per_target():
    """
    plot scatterplot of each target and # DE genes in treated / untreated
    """
    if not os.path.isdir("outputs/DE/"):
        os.makedirs("outputs/DE/")
    ##plot all targets 
    fig, ax = plt.subplots()
    de_dict = pickle.load(open("pickles/raw_DE_gene_results.pkl", "rb"))
    targets = list(set(de_dict["Treated"].keys()).intersection(set(de_dict["Untreated"].keys()))) ##each point must have a DE count in treated and DE count in untreated
    x = [len(de_dict["Untreated"][t]) for t in targets]
    y = [len(de_dict["Treated"][t]) for t in targets]
    equals = [x[i] for i in range(0, len(x)) if y[i] == x[i]]
    x_greater = [x[i] for i in range(0, len(x)) if y[i] > x[i]]
    y_greater = [y[i] for i in range(0, len(x)) if y[i] > x[i]]
    global_ratio = round(len(x_greater) / len(x), 2)
    ax.scatter(x_greater, y_greater, s=1, color="red", label=f"DE Untreated < DE Treated, Ratio={global_ratio}")
    x_less = [x[i] for i in range(0, len(x)) if y[i] < x[i]]
    y_less = [y[i] for i in range(0, len(x)) if y[i] < x[i]]
    ax.scatter(x_less, y_less, s=1, color="green", label=f"DE Untreated > DE Treated, Ratio={round(len(x_less) / len(x), 2)}")
    x_equal = [x[i] for i in range(0, len(x)) if y[i] == x[i]]
    y_equal = [y[i] for i in range(0, len(x)) if y[i] == x[i]]
    ax.scatter(x_equal, y_equal, s=1, color="grey", label=f"DE Untreated == DE Treated, Ratio={round(len(x_equal) / len(x), 2)}")
    ax.plot([0, max(x + y)], [0, max(x + y)], linestyle='--', color='grey')
    print(f"{len(equals)} targets have equal DE expression in treated and untreated")
    print(f"{len(x_greater)} targets have Untreated > Treated")
    print(f"{len(y_greater)} targets have Treated > Unreated")
    print(f"for equal DE expression, avg = {np.mean(x_equal)}, median = {np.median(x_equal)}")
    # ##calculate ratios of # DE treated / # DE untreated for global and local top n 
    plt.title(f"DE Count Comparison for all Targets")
    ax.set_xlabel("# of DE Genes in Untreated")
    ax.set_ylabel("# of DE Genes in Treated")
    box = ax.get_position()
    ax.set_position([box.x0, box.y0, box.width, box.height * 0.80])
    ax.legend(loc='upper right', prop={"size":7}, bbox_to_anchor=(1, 1.23))
    plt.gcf().subplots_adjust(top=.80)
    plt.savefig(f"outputs/DE/scatter_all_targets.png", dpi=300)
    differences = [x[i] - y[i] for i in range(0, len(x))]
    differences = [x_ for x_ in differences if x_ != 0.0] ##exclude zero difference from histogram
    sorted_difference = sorted(differences)
    #histogram 
    fig, ax = plt.subplots()
    negative_differences = [x_ for x_ in differences if x_ < 0]
    positive_differences = [x_ for x_ in differences if x_ > 0]
    counts, bins, patches = ax.hist(negative_differences, bins=100, color="red", label=f"DE Treated > DE Untreated, n={len(negative_differences)}") 
    counts, bins, patches = ax.hist(positive_differences, bins=100, color="green", label=f"DE Untreated > DE Treated, n={len(positive_differences)}") 
    # zero_differences = [x_ for x_ in differences if x_ == 0]
    # ax.bar([0], [len(zero_differences)], color="grey", label=f"DE Untreated == DE Treated, n={len(zero_differences)}", width=55)
    ax.set_xlabel("Difference (# of DE Genes in Untreated - # of DE Genes in Treated)")
    ax.set_ylabel("Target Count")
    plt.title(f"Distribution of DE Untreated - DE Treated")
    box = ax.get_position()
    ax.set_position([box.x0, box.y0, box.width, box.height * 0.85])
    ax.legend(loc='upper right', prop={"size":7}, bbox_to_anchor=(1, 1.22))
    plt.gcf().subplots_adjust(top=.76)
    plt.savefig(f"outputs/DE/both.png", dpi=300)

def volcano_plot_target_genes(target_genes,logfc_threshold=1.0, pval_threshold=0.05):
    """
    Create a volcano plot of DE genes for each target_gene with given logfc and pval (FDR adjusted) thresholds 
    saves to outputs/volcanos
    """
    if not os.path.isdir("outputs/volcanos"):
        os.makedirs("outputs/volcanos")
    conditions = ["Treated", "Untreated"]
    adata = sc.read_h5ad("data/imru_full.h5ad")
    sc.pp.log1p(adata) ##rank sum test expects logarithmic data 
    for condition in conditions:
        ##subset anndata to just condition and target_genes
        adata_condition = adata[adata.obs["condition"] == condition]
        adata_condition = adata_condition[adata_condition.obs["gene_target"].isin(set(target_genes + ["SAFE_TARGET"]))] 
        adata_condition = adata_condition.copy() ##change view to copy because rank sum test will write to this AnnData object
        sc.tl.rank_genes_groups(adata_condition, groupby="gene_target", reference="SAFE_TARGET", method="wilcoxon")
        names = adata_condition.uns["rank_genes_groups"]["names"]
        pvals_adj = adata_condition.uns["rank_genes_groups"]["pvals_adj"]
        logfc = adata_condition.uns["rank_genes_groups"]["logfoldchanges"]
        for target_gene in names.dtype.names: ##iterate over perturbations
            #create DataFrame
            df = pd.DataFrame({
            'gene': names[target_gene],
            'log2FC': logfc[target_gene],
            'pval': pvals_adj[target_gene]})
            df['neg_log10_pval'] = -np.log10(df['pval'])
            ##exclude points with log2fold change magnitude greater than some threshold for the plot, but keep target gene 
            df = df[(df['gene'] == target_gene) | ((df['log2FC'] < 10) & (df['log2FC'] > -10))]
            #volcano plot of entire df 
            plt.figure(figsize=(8, 6))
            plt.scatter(df['log2FC'], df['neg_log10_pval'], color='gray', alpha=0.6)
            #highlight significant genes (and target gene)
            df_sig = df[(df['gene'] == target_gene) | ((abs(df['log2FC']) > logfc_threshold) & (df['pval'] < pval_threshold))] 
            plt.scatter(df_sig['log2FC'], df_sig['neg_log10_pval'], color='blue', alpha=0.8, s=20)
            #labeling the points - adjust text to avoid overlapping annotations
            texts = []
            for i, txt in enumerate(df_sig["gene"]):
                texts.append(plt.text(df_sig["log2FC"].iloc[i], df_sig["neg_log10_pval"].iloc[i], txt, ha='center', va='center', size=7))
            adjust_text(texts, arrowprops=dict(arrowstyle='->', color='gray', lw=0.5), min_arrow_len=5)
            #highlight target gene by plotting a red dot
            if target_gene in df['gene'].values:
                tg_row = df[df['gene'] == target_gene].iloc[0]
                plt.scatter(tg_row['log2FC'], tg_row['neg_log10_pval'], color='red', s=36)
            #add thresholds
            plt.axvline(x=logfc_threshold, color='black', linestyle='--')
            plt.axvline(x=-logfc_threshold, color='black', linestyle='--')
            plt.axhline(y=-np.log10(pval_threshold), color='black', linestyle='--')
            plt.title(f"Volcano Plot for Target {target_gene} in {condition} Condition")
            plt.xlabel(r'Log$_{2}$ Fold Change')
            plt.ylabel(r'-log$_{10}$ Adjusted p-value')
            plt.tight_layout()
            plt.savefig(f"outputs/volcanos/{target_gene}_{condition}.png", dpi=300)

def get_aggregated_representation_matrix(method, treatment_type, control_strings):
    """
    Get the X (matrix), y (list of targets) 
    for the target-aggregated embeddings for method and treatment_type
    """
    adata = sc.read_h5ad(f"data/{method}_imru_full_processed.h5ad")
    adata = adata[adata.obs["condition"] == treatment_type]
    adata = adata[~adata.obs["gene_target"].isin(control_strings)]
    all_targets = sorted(set(adata.obs["gene_target"]))
    ##build X = aggregated latent representation of target, y = target
    X, y = [], []
    for target in all_targets:
        sub_adata = adata[adata.obs["gene_target"] == target]
        aggregated_target = np.mean(sub_adata.X, axis=0)
        assert(len(set(sub_adata.obs["gene_target"])) == 1)
        X.append(aggregated_target)
        y.append(target)
    X = np.array(X)
    X = preprocess_matrix(X, standardize=False, normalize=True)
    return X, y

def GPT5_ranking(min_sample=0):
    """
    Write GPT5.0 target rankings into same data structure format as other competing methods
        (1) dictionary with key: (anchor_state, treatment_type), value: ranked list of tuples(gene target, cosine distance) (ascending by cosine distance)
        (2) count map with key: treatment_type, key: perturbation, value: # of cells
    saves to "pickles/full_gpt5.0_sorted_perts_knockdown_threshold_500.0_min_sample_{min_sample}.pkl"
    will have same target set as scgpt ranking 
    """
    uniques = set(['VWF', 'HNRNPA3', 'TRIM13', 'CORO1C', 'DDX27', 'PACSIN2']) ##remove treatment unique perturbations to match
    scgpt_ranking, count_map = get_sorted_perts(dataset="full", method="scgpt", min_sample=min_sample)
    scgpt_targets = set([x[0] for x in scgpt_ranking[('Untreated_SAFE_TARGET', 'Treated')]])
    gpt5 = pd.read_csv("data/GPT5.0_ranked_gene_targets.csv")
    gpt5_ranking = list([x for x in gpt5["gene target"] if x not in uniques])
    missing_targets = [x for x in scgpt_targets if x not in set(gpt5_ranking)]
    if len(missing_targets) > 0:
        print(f"WARNING: GPT5.0 ranking is missing {len(missing_targets)} targets that are present in scGPT ranking: {missing_targets}")
    gpt5_ranking = list([x for x in gpt5_ranking if x in scgpt_targets])
    return_dict = {('Untreated_SAFE_TARGET', 'Treated'): [(gpt5_ranking[i], -1) for i in range(0, len(gpt5_ranking))]}
    ##keep same counts map, doesn't change between methods 
    pickle.dump((return_dict, count_map), open(f"pickles/full_chatgpt5.0_sorted_perts_knockdown_threshold_500.0_min_sample_{min_sample}.pkl", "wb"))

def traditional_rankings(min_sample=0):
    """
    Write traditional comp bio approach (DE Inflammatory / DE Basal) target rankings into same data structure format as other competing methods
        (1) dictionary with key: (anchor_state, treatment_type), value: ranked list of tuples(gene target, cosine distance) (ascending by cosine distance)
        (2) count map with key: treatment_type, key: perturbation, value: # of cells
    saves to "pickles/full_compbio_sorted_perts_knockdown_threshold_500.0_min_sample_{min_sample}.pkl"
    """
    uniques = set(['SH3PXD2A', 'EMCN', 'RAB11FIP3', 'CKAP5', 'POLR2K'] + ['VWF', 'HNRNPA3', 'TRIM13', 'CORO1C', 'DDX27', 'PACSIN2']) ##set of perturbations that are only present in either treated OR untreated condition, but not both
    treated_final_set = set([]) ##set of treated targets, to make sure the final candidate lists are the same across all permutations 
    scgpt_ranking, count_map = get_sorted_perts(dataset="full", method="scgpt", min_sample=min_sample)
    scgpt_targets = [x[0] for x in scgpt_ranking[('Untreated_SAFE_TARGET', 'Treated')]]
    for permutation in ["direct", "manual", "vox", "untreated_direct"]:
        if permutation == "untreated_direct":
            compbio = pd.read_csv(f"data/20251208_perturbseq_ranked_list_direct_method_trad_all_cells_UNTREATED.txt", header=None, names=["gene target"])
        else:
            compbio = pd.read_csv(f"data/20251203_perturbseq_ranked_list_{permutation}_method_trad.txt", header=None, names=["gene target"])
        compbio_ranking = list([x for x in compbio["gene target"] if x not in uniques])
        compbio_ranking = [x for x in compbio_ranking if count_map["Treated"][x] >= min_sample]
        if permutation == "untreated_direct":
            compbio_ranking = [x for x in compbio_ranking if x in treated_final_set] ##for the untreated list, filter to a common set of targets we'll be using There are targets unique to the untreated condition that we can't use for comparison. 
        if permutation in ["direct", "manual", "vox"]:
            if len(treated_final_set) == 0: #unitialized
                treated_final_set = set(compbio_ranking)
            else:
                assert(set(compbio_ranking) == treated_final_set)
        missing_targets = [x for x in scgpt_targets if x not in set(compbio_ranking)]
        if len(missing_targets) > 0:
            print(f"WARNING: comp bio ranking is missing {len(missing_targets)} targets that are present in scGPT ranking with min_sample={min_sample}: {missing_targets}")
        return_dict = {('Untreated_SAFE_TARGET', 'Treated'): [(compbio_ranking[i], -1) for i in range(0, len(compbio_ranking))]}
        ##keep same counts map, doesn't change between methods 
        pickle.dump((return_dict, count_map), open(f"pickles/full_compbio_{permutation}_sorted_perts_knockdown_threshold_500.0_min_sample_{min_sample}.pkl", "wb"))

def get_unique_perturbations(adata_path):
    """
    given path to adata
    will return a dictionary with key: treatment, value: set of perturbations unique to that treatment 
    """
    adata = sc.read_h5ad(adata_path)
    adata_treated = adata[adata.obs["condition"] == "Treated"]
    treated_perts = set(adata_treated.obs["gene_target"])
    adata_untreated = adata[adata.obs["condition"] == "Untreated"]
    untreated_perts = set(adata_untreated.obs["gene_target"])
    treated_unique = set([x for x in treated_perts if x not in untreated_perts])
    untreated_unique = set([x for x in untreated_perts if x not in treated_perts])
    return {"Treated": treated_unique, "Untreated": untreated_unique}

def cell_count_representation_analysis(metric="cosine"):
    """
    calculate either cosine similarity or pearson (depending on metric) between (1) global target representation using all cells and (2) some subset of cells for that target
    saves a dictionary with key: treatment_type, key: method, key: n_cell, key: target, value: pearson correlation of target sample with global target population
    """
    methods = ["random", "raw", "umap", "scimilarity", "scgpt", "state"]
    treatment_types = ["Treated", "Untreated"]
    n_cells = list(range(1, 10)) + list(range(10, 22, 2)) + [50, 100]
    mapp = {t: {m: {n_cell: {} for n_cell in n_cells} for m in methods} for t in treatment_types} 
    ##key: treatment_type, key: method, key: n_cell, key: target, value: pearson correlation of target sample with global target population
    for treatment_type in treatment_types:
        for index, method in enumerate(methods):
            adata = sc.read_h5ad(f"data/{method}_imru_full_processed.h5ad")
            adata_condition = adata[adata.obs["condition"] == treatment_type]
            if metric == "cosine":
                X = preprocess_matrix(adata_condition.X, standardize=False, normalize=True)
            else:
                X = adata_condition.X
            target_list = list(adata_condition.obs["gene_target"])
            target_set = set(target_list)
            target_to_indices = {target: np.array([i for i in range(0, len(target_list)) if target_list[i] == target]) for target in target_set}
            target_to_centroid = {target: np.mean(X[target_to_indices[target]], axis=0) for target in target_set}
            for n_cell in n_cells:
                print(n_cell)
                for target in target_set:
                    if n_cell > len(target_to_indices[target]):
                        continue
                    else:
                        rng = np.random.default_rng(seed=42) ##need to re-init to keep consistency
                        sampled_indices = rng.choice(target_to_indices[target], size=n_cell, replace=False)
                        sample_mean =  np.mean(X[sampled_indices], axis=0)
                        if metric == "cosine":
                            mapp[treatment_type][method][n_cell][target] = 1.0 - cosine_distance(sample_mean, target_to_centroid[target])
                        elif metric == "pearson":
                            mapp[treatment_type][method][n_cell][target] = pearsonr(sample_mean, target_to_centroid[target])[0]
                        else:
                            raise Exception()
    pickle.dump(mapp, open(f"pickles/cell_count_rep_{metric}.pkl", "wb"))

def plot_cell_count_representation_analysis(metric):
    mapp = pickle.load(open(f"pickles/cell_count_rep_{metric}.pkl", "rb"))
    methods = ["raw", "umap", "scimilarity", "scgpt", "state"]
    for treatment_type in mapp:
        for max_count in [20, max(mapp[treatment_type]["scgpt"].keys())]:
            fig, ax = plt.subplots()
            x = sorted(mapp[treatment_type]["scgpt"].keys())
            x = [x_ for x_ in x if x_ <= max_count]
            for method in methods:
                y = np.array([np.mean(list(mapp[treatment_type][method][n_cell].values())) for n_cell in x])
                y_std = np.array([np.std(list(mapp[treatment_type][method][n_cell].values())) for n_cell in x])
                ax.plot(x, y, label=get_method_title(method))
                # ax.fill_between(x, y - y_std, y + y_std, color="grey", alpha=0.2)
                print(mapp[treatment_type][method][2])
            metric_title = "Pearson Correlation" if metric == "pearson" else "Cosine Similarity"
            ax.set_ylabel(f"{metric_title}")
            ax.set_xlabel("Subsample Cell Count")
            plt.title(f"{metric_title} Between Target Representations\nfrom Subsample and Entire Sample in {treatment_type} Condition")
            box = ax.get_position()
            ax.set_position([box.x0, box.y0, box.width, box.height * 0.85])
            ax.legend(loc='upper right', prop={"size":7}, bbox_to_anchor=(1, 1.38))
            plt.gcf().subplots_adjust(top=.76)
            plt.savefig(f"outputs/cell_count_{metric}_{treatment_type}_{max_count}.png", dpi=300)

def get_relevant_pathways(just_this_study=False, include_imru_internal=False, key_type="tuple"):
    """
    C2 canonical pathways (CP) collection here:
    https://www.gsea-msigdb.org/gsea/msigdb/human/collections.jsp#C2
    MSigDB v7.0 released on Aug 2019.
    if key_type == "string":
        dictionary key will be "library|pathway"
     if key_type == "tuple":
        dictionary key will be (library,pathway)
    Returns dictionary with key, value: list of genes based on content of "data/c2.cp.v7.0.symbols.gmt"
    if just_this_study:
        all the gene sets will be filtered to targets in just this study
    if include_imru_internal:
        adds custom gene sets, like the internal IMRU one
    """
    found = ["BIOCARTA_IL1R_PATHWAY", "BIOCARTA_TNFR1_PATHWAY", "KEGG_CYTOKINE_CYTOKINE_RECEPTOR_INTERACTION", "REACTOME_INTERLEUKIN_1_FAMILY_SIGNALING", "REACTOME_INTERLEUKIN_1_SIGNALING", "REACTOME_TNF_SIGNALING"]
    study_targets = get_all_target_names()
    filepath = "data/c2.cp.v7.0.symbols.gmt"
    pathway_to_genes = {}
    with open(filepath, 'r', encoding='utf-8') as f:
        for line in f:
            line = line.strip()
            if not line: 
                continue
            parts = line.split('\t')
            if len(parts) < 3:
                continue
            lib_pathway = parts[0].strip()
            if lib_pathway not in found:
                continue
            library, pathway = lib_pathway[0:lib_pathway.find("_")], lib_pathway[lib_pathway.find("_") + 1: ]
            genes = [g.strip() for g in parts[2:] if g.strip()]
            if just_this_study:
                genes = [x for x in genes if x in study_targets]
            if key_type == "string":
                key = f"{library}|{pathway}"
            if key_type == "tuple":
                key = (library, pathway)
            pathway_to_genes[key] = genes
    if include_imru_internal:
        if key_type == "string":
            key = "imru_internal|INFLAMMATION REVERSAL GIVEN EXPERIMENTAL CONTEXT"
        if key_type == "tuple":
            key = ("imru_internal", "INFLAMMATION REVERSAL GIVEN EXPERIMENTAL CONTEXT")
        pathway_to_genes[key] = ["TNFRSF1A", "TRADD", "JUNB","JUND", "NFKB1", "NFKB2"] ##internal IMRU expertise nominated 
    return pathway_to_genes

def enrichr_wrapper(gene_list, gene_sets=['KEGG_2021_Human', 'BioCarta_2016','Reactome_Pathways_2024', 'MSigDB_Hallmark_2020', 'GO_Biological_Process_2025'], organism="human", outdir=None, cutoff=0.05, background_type="universe", verbose=False):
    """
    Enrichr API sometimes crashes out, wrapper to call it again in case it fails
    returns enr object 
    if background_type == "universe" will use the enrichr's default background of all known genes
    if background_type == "targets" will use all the targets of this dataset 
    """
    warnings.simplefilter(action='ignore', category=FutureWarning)
    if background_type == "universe":
        background = None
    if background_type == "targets":
        background = pickle.load(open("pickles/target_names.pkl", "rb"))
    for i in range(0, 10): ##try up to 10 times, if it fails after 10 then raise an error
        try: 
            enr = gseapy.enrichr(gene_list=gene_list,
                gene_sets=gene_sets,
                organism=organism, 
                outdir=outdir,
                background=background,
                cutoff=cutoff) ##enr.results is a df 
            return enr ##short circuit 
        except: ##sometimes we need a pause, else enrichr will throw a JSON decoder error 
            if verbose: 
                print("Enrichr call failed, waited 1.0 seconds then tried again: ")
                print("succeeded on second try")
            time.sleep(1.0)
    traceback.print_exc()
    raise Exception("enrichr call failed 10x")

def call_gsea(statistic="FDR q-val", significance_thresh=0.05, rank_score="similarity", desired_anchor_state="SAFE_TARGET", run_networkx=False, methods=["scgpt"], min_sample=0, top_ns=[50]):
    """
    runner call for enrichr + gsea
    statistic = "FWER p-val": adjusted p-value (probably using Bonferroni or Holm's, downside is that it's overly strict)
    statistic = "FDR q-val": adjusted p-value for FDR (probably using Storey's method)
    significance_thresh will be p-value threshold
    """
    knockdown_threshold = 500.0
    for method in methods:
        sorted_perts, count_map = get_sorted_perts(method=method, min_sample=min_sample)
        for (anchor_state, treatment_type) in sorted_perts:
            if desired_anchor_state not in anchor_state: ##just look at one of the anchor states provided as argument to this function
                continue
            perturbations = [x[0] for x in sorted_perts[(anchor_state, treatment_type)]]
            distances = [x[1] for x in sorted_perts[(anchor_state, treatment_type)]]
            similarities = [1.0 - x for x in distances] ##cosine similarity is 1.0 - cosine distance
            save_name = f"full_{method}_{treatment_type}_perturbations_closest_to_{anchor_state}_{knockdown_threshold}"
            save_dir = "outputs/gsea/"
            #enrichr on top_n genes for multiple gene sets of interest - can combine multiple gene sets
            ##run enrichr
            for top_n in top_ns:
                pert_set_list = perturbations[0:top_n]
                if desired_anchor_state in pert_set_list:
                    print(f"WARNING: anchor state {desired_anchor_state} is in top {top_n}, removing from enrichr analysis and adding 1 more perturbation to get to {top_n}")
                    pert_set_list = [x_ for x_ in pert_set_list if x_ != desired_anchor_state]
                    pert_set_list.append(perturbations[top_n]) ##add 1 more to account for removing pert=anchor_state to get to total of top_n perturbations
                random.shuffle(pert_set_list)
                perform_gsea_enrichr(pert_set_list, significance_thresh=significance_thresh, title=f"Enrichr Plot for {get_method_title(method)}'s Top {len(pert_set_list)} Genes", save_dir=save_dir, save_name=save_name + f"_genes={top_n}", top_term=5) ##10 top_terms by default

def perform_gsea_enrichr(gene_list, significance_thresh=0.05, title=None, save_dir=None, save_name=None, top_term=5, gene_sets=['KEGG_2021_Human', 'BioCarta_2016','Reactome_Pathways_2024', 'MSigDB_Hallmark_2020', 'GO_Biological_Process_2025']):
    """
    Performs GSEA enrichr from https://gseapy.readthedocs.io/en/latest/gseapy_example.html#Enrichr-API 
    given gene_list as unranked/unsorted list find enriched pathways and gene_sets as a list of libraries to use
    top_term specifies maximal # of significant pathways to plot per gene_set (returned pathways = top_term - 1)
    """
    if not os.path.isdir(save_dir):
        os.makedirs(save_dir)
    try: 
        if len(gene_sets) <= 7:
            colors = ["salmon", "darkblue", "bisque", "forestgreen", "cornflowerblue", "salmon", "mediumslateblue"]
        else:
            prop_cycle = plt.rcParams['axes.prop_cycle']
            colors = prop_cycle.by_key()['color']
        ##gene set enrichment, given set of genes, get pathways, uses Enrichr web services
        time.sleep(5.0)
        enr = enrichr_wrapper(gene_list=gene_list, gene_sets=gene_sets, outdir=None, cutoff=significance_thresh, background_type="targets") ##if background is NOT default (e.g. is "universe"), then the enr.results df will not have an 'overlap' column; however we are working with a limited set of possible genes since this is a triage ranking so we should just use the dataset's target set as background
        enr.results.to_csv(os.path.join(save_dir, f"{save_name}_raw.csv"))
        sig_pathway_count = sum([1 for i in range(0, len(enr.results)) if enr.results["Adjusted P-value"].iloc[i] <= 0.05])
        ##barplot
        ax = gseapy.barplot(enr.results,
                    column="Adjusted P-value",
                    group='Gene_set', # set group, so you could do a multi-sample/library comparsion
                    size=10,
                    top_term=top_term,
                    figsize=(3,5),
                    title = title,
                    color={gene_sets[i]: colors[i] for i in range(0, len(gene_sets))}
                    )
        ax.tick_params(axis='x', labelsize='small')
        ax.tick_params(axis='y', labelsize='small')
        ytick_labels = [tick.get_text() for tick in ax.get_yticklabels()]
        ytick_labels = [y_ if y_.count(" ") < 8 else replace_nth_char(y_, " ", "\n", 6) for y_ in ytick_labels] ##if string is long break it up
        ax.set_yticklabels(ytick_labels)
        ax.set_title(title + f":\n{sig_pathway_count} Significantly Enriched Pathways", fontsize=11, fontweight="bold") 
        ax.set_xlabel(r'-log$_{10}$ (Adjusted P-value)', fontsize=11, fontweight="normal") ##default xlabel is this, but need to adjust the fontsize
        fig = ax.get_figure()
        fig.savefig(os.path.join(save_dir, f"{save_name}_bar.png"), dpi=300, bbox_inches="tight")    
    except Exception as e:
        print(f"caught exception for {save_name}: {e}")

def perform_gsea_prerank(gene_list, similarities, save_name=None, gene_sets=None, run_networkx=False, statistic="FWER p-val", min_size=15, max_size=500, significance_thresh=0.05, rank_score="similarity"):
    """
    Performs GSEA prerank from https://gseapy.readthedocs.io/en/latest/gseapy_example.html#Enrichr-API 
    gene_list: ranked/sorted list of genes
    similarities: corresponding cosine similarities for gene_list
    rank_score: if "similarity" will use cosine similarity as the score for the rankings, if "index" will simply be the (inverse) index in the queue e.g. top priority gene has score = # of total pertubations in the ranked gene_list, gene with lowest priority has score = 1 

    compute gene set enrichment plots

    will save an output CSV df: 
        some columns with ambiguous titles demystified:
        gene % = the position of the gene corresponding to enrichment score peak / number of genes in ranking list
        tag % = number of leading genes / numer of genes in certain pathway that overlapped with the input ranked list
        FWER p-val = adjusted p-value (probably using Bonferroni or Holm's? downside is that it's overly strict, controls strongly for false positives type I error but at cost of higher false negatives type II error)
        FDR q-val = adjusted p-value for FDR (probably using Storey's method? corrects for multiple testing but less strict than FWER p-val - will have lower type II error)
    """
    try: 
        if rank_score == "similarity":
            gene_df = pd.DataFrame({"genes": gene_list, "cosine sim": similarities}) ##top priority genes need a higher score for gseapy.prerank, let's use the cosine similarity 
        elif rank_score == "index":
            gene_df = pd.DataFrame({"genes": gene_list, "rank": range(len(gene_list), 0, -1)}) ##top priority genes need a higher score for gseapy.prerank, let's simply use the inverse index 
        else:
            raise Exception("rank_score must be one of similarity or index")
        pre_res = gseapy.prerank(rnk=gene_df,
                            gene_sets=gene_sets,
                            threads=4,
                            min_size=min_size, ##default 15/500, was 5/1000
                            max_size=max_size,
                            permutation_num=2000, # reduce number to speed up testing, increase number for greater p-value reliability
                            outdir=None, # don't write to disk
                            seed=6,
                            verbose=True, # see what's going on behind the scenes
                            )
        df = pre_res.res2d.sort_values(by=statistic) ##sort df by statistic (e.g. FWER p-val or FDR q-val) ascending
        df.to_csv(f"outputs/gsea/{save_name}_prerank_df.csv")
        all_terms = df.Term
        df_sig = df[df[statistic] < significance_thresh] ##only keep entries where adjusted p-val is less than 0.05 significance
        df_sig.to_csv(f"outputs/gsea/{save_name}_df_significant_{significance_thresh}.csv")
        ##plot all significant terms
        sig_terms = df_sig.Term
        if len(df_sig) == 0: ##if no pathways are significant, exit
            return 
        axs = pre_res.plot(terms=sig_terms, ##can subset the ranked terms if want less 
                        show_ranking=True, # whether to show the second yaxis
                        figsize=(3,4)
                        )
        fig = axs.get_figure()
        fig.savefig(f"outputs/gsea/{save_name}_GSEA_significant.png", dpi=300, bbox_inches="tight")    
        if run_networkx:
            import networkx as nx
            nodes, edges = gseapy.enrichment_map(df)
            G = nx.from_pandas_edgelist(edges,
                                source='src_idx',
                                target='targ_idx',
                                edge_attr=['jaccard_coef', 'overlap_coef', 'overlap_genes'])
            fig, ax = plt.subplots(figsize=(8, 8))
            pos=nx.layout.spiral_layout(G)
            nx.draw_networkx_nodes(G, # draw node
                                pos=pos,
                                cmap=plt.cm.RdYlBu,
                                node_color=list(nodes.NES),
                                node_size=list(nodes.Hits_ratio *1000))
            nx.draw_networkx_labels(G, # draw node label
                                    pos=pos,
                                    labels=nodes.Term.to_dict())
            edge_weight = nx.get_edge_attributes(G, 'jaccard_coef').values() # draw edge
            nx.draw_networkx_edges(G,
                                pos=pos,
                                width=list(map(lambda x: x*10, edge_weight)),
                                edge_color='#CDDBD4')
            plt.savefig(f"outputs/gsea/{save_name}_networkx.png", dpi=300, bbox_inches="tight")       
    except Exception as e:
        print(f"error with {save_name}: {e}")
        traceback.print_exc()

def task_get_sigs(sigs, gene_list):
    """
    sigs will be a mp.list to keep track of # of significantly enriched pathways, gene_list is the list of genes to perform enrichr on
    helper function for parallel_get_random_enrichments
    """
    enr = enrichr_wrapper(gene_list=gene_list, background_type="targets")
    sigs.append(sum([1 for i in range(0, len(enr.results)) if enr.results["Adjusted P-value"].iloc[i] <= 0.05]))

def parallel_get_random_enrichments(n_trials=50, max_n=None):
    """
    computationally expensive, will run enrichr on n_trials permutations of random genes
    used for a baseline so that we don't have to re-compute a null baseline every time 
    since each treatment condition has unique perturbations, need to create two dictionaries: one sampled from treated perturbation set, one sampled from just untreated perturbation set 
    if max_n not specified, will be the total # of unique targets per condition
    uses multiprocessing to run iterations in parallel to save time
    saves results to pickles/random_enrichment_trials={n_trials}.pkl: key: treatment type, key: n, value: number of significant pathways in random sample of length n
    """
    treatment_types = ["Untreated", "Treated"]
    adata = sc.read_h5ad(f"data/scgpt_imru_full_processed.h5ad") ##doesn't matter which, will be just be using gene target labels
    results_map = {t: {} for t in treatment_types} ##key: treatment type, key: n, value: number of significant pathways in random sample of length n
    for treatment_type in treatment_types: 
        adata_sub = adata[adata.obs["condition"] == treatment_type]
        gene_targets = list(set(adata_sub.obs["gene_target"]))
        if max_n == None:
            ns = range(1,  len(gene_targets) + 1)
        else:
            ns = range(1,  max_n + 1)
        ##if an existing null distribution with equal or fewer n_trials exists, use it and add more samples
        for n_ in range(0, n_trials + 1):
            if os.path.isfile(f"pickles/random_enrichment_trials={n_}.pkl"):
                loaded_results_map = pickle.load(open(f"pickles/random_enrichment_trials={n_}.pkl", "rb"))
                print(f"WARNING: loaded: pickles/random_enrichment_trials={n_}.pkl")
                for treatment_type in loaded_results_map:
                    for key in loaded_results_map[treatment_type]:
                        results_map[treatment_type][key] = loaded_results_map[treatment_type][key]
                break
        manager = mp.Manager() 
        for n in ns:
            if n not in results_map[treatment_type] or len(results_map[treatment_type][n]) <= n_trials:
                if n not in results_map[treatment_type]:
                    difference = n_trials
                else:
                    difference = n_trials - len(results_map[treatment_type][n]) ##if cached and n already has some trials, keep those and add more 
                print(n, difference)
                sigs = manager.list() ##will be the list of # of significantly enriched pathways for this n
                random_samples = [random.sample(gene_targets, n) for iter_ in range(0, difference)] ##list of lists of randomly sampled genes
                processes = [mp.Process(target=task_get_sigs, args=(sigs, random_samples[i])) for i in range(0, len(random_samples))]
                for process in processes: #start all processes
                    process.start()
                for process in processes: #wait for all processes to complete
                    process.join()
                for process in processes: #close the process and free up resources
                    process.close()
                results_map[treatment_type][n] = list(sigs)
            else:
                print(f"trial {n} has sufficient samples")
    pickle.dump(results_map, open(f"pickles/random_enrichment_trials={n_trials}.pkl", "wb"))

def plot_null_enrichr_distribution(n_trials=50):
    """
    plot median + std for each n of random null enrichr distribution
    """
    results_map = pickle.load(open(f"pickles/random_enrichment_trials={n_trials}.pkl", "rb"))
    for treatment_type in ["Treated", "Untreated"]:
        fig, ax = plt.subplots()
        x_values = sorted(results_map[treatment_type].keys())
        medians = [np.median(results_map[treatment_type][x]) for x in x_values]
        stds = [np.std(results_map[treatment_type][x]) for x in x_values]
        ax.scatter(x_values, medians, color='blue', label='Median', s=.1)
        lower_bound = np.array(medians) - np.array(stds)
        upper_bound = np.array(medians) + np.array(stds)
        ax.fill_between(x_values, lower_bound, upper_bound, color='blue', alpha=0.2)
        ax.set_xlabel("Genes in Set")
        ax.set_ylabel("Significantly Enriched Pathways\nMedian ± Std Dev")
        plt.title(f"Significantly Enriched Pathways for Random Distribution\nn trials = {n_trials}")
        plt.savefig(f"outputs/random_null_distribution_enrichr_{treatment_type}.png", dpi=300)

def rank_vs_enrichment(batch_size=64, min_sample=0, gene_set_type="de"):
    """
    if gene_set_type == top_n:
        saves mapp with key: "all_pathways" or "relevant_pathways", key: method, key: index, value: list of significantly enriched pathways for targets at rank:0 --> rank:index
    if gene_set_type == de:
        saves mapp with key: "all_pathways" or "relevant_pathways", key: method, key: index, value: list of significantly enriched pathways for de genes of target at index 

    to pickles/rank_vs_enrichment.pkl
    """
    def task_subtask(gene_dict, i, gene_list, gene_sets):
        """
        helper function to parallelize call to enrichr
        """
        ##find all pathways, not just relevant ones 
        enr = enrichr_wrapper(gene_list=gene_list, gene_sets=gene_sets, background_type="targets")
        significant_indices = [i for i in range(0, len(enr.results)) if enr.results["Adjusted P-value"].iloc[i] <= 0.05]
        if isinstance(gene_sets, list): ##if it's a list, then we are specifying the default behavior of enrichr where this is a list of libraries to use, enr.results["Gene_set"] will be library 
            significant_pathways = [enr.results["Gene_set"].iloc[i] + "_" + enr.results["Term"].iloc[i] for i in significant_indices]
        if isinstance(gene_sets, dict): ##if it's a dict, then we are specifying a custom mapping of library/pathway <--> gene sets,  enr.results["Gene_set"] will be a filler string "gs_ind_0", so let's not keep this
            significant_pathways = [enr.results["Term"].iloc[i] for i in significant_indices]
        gene_dict[i] = significant_pathways
    de_results = pickle.load(open("pickles/raw_DE_gene_results.pkl", "rb")) ##key: treatment_type, key: target, value: list of DE genes 
    methods = ["random", "raw", "umap", "scimilarity", "scgpt", "state", "chatgpt5.0", "compbio_direct", "compbio_untreated_direct"] 
    #we want to get enrichment results for (1) all pathways from a wide variety of libraries and (2) enirchment results for the relevant ones we selected
    pathway_granularities = ["relevant_pathways", "all_pathways"]
    mapp = {p_g: {m: {} for m in methods} for p_g in pathway_granularities} ##key: "all_pathways" or "relevant_pathways", key: method, key: index, value: list of significantly enriched pathways for targets 0 --> index
    for pathway_granularity in pathway_granularities:
        if pathway_granularity == "all_pathways":
            gene_sets = ['KEGG_2021_Human', 'BioCarta_2016','Reactome_Pathways_2024', 'MSigDB_Hallmark_2020', 'GO_Biological_Process_2025']
        elif pathway_granularity == "relevant_pathways":
            gene_sets = get_relevant_pathways(just_this_study=True, include_imru_internal=False, key_type="string")
        else:
            raise Exception()
        for method in methods:
            sorted_perts, _ = get_sorted_perts(method=method, min_sample=min_sample)
            ranked_list = [x[0] for x in sorted_perts[("Untreated_SAFE_TARGET", "Treated")]]
            ranked_list = [x for x in ranked_list if x not in ["SAFE_TARGET", "NO-TARGET"]] ##remove controls from ranked_list
            ##for feeding to multiprocessing function
            if gene_set_type == "top_n":
                target_batches = [(i, ranked_list[0:i]) for i in range(10, len(ranked_list), 1)] 
            if gene_set_type == "de":
                target_batches = [(i, de_results["Treated"][ranked_list[i]]) for i in range(0, len(ranked_list), 1)]
            for i in range(0, len(ranked_list)): ##init mapp with empty lists in case processes fail 
                mapp[pathway_granularity][method][i] = []
            ##parallel process target_batches
            manager = mp.Manager() 
            manager_dict = manager.dict()
            ##chunk it so we can reduce Process instances 
            for i in range(0, len(target_batches), batch_size):
                print(i, len(target_batches))
                mini_batch = target_batches[i: i + batch_size]
                processes = [mp.Process(target=task_subtask, args=(manager_dict, mini_batch[i][0], mini_batch[i][1], gene_sets)) for i in range(0, len(mini_batch))]
                for process in processes: #start all processes
                    process.start()
                for process in processes: #wait for all processes to complete
                    process.join()
                for process in processes: #close the process and free up resources
                    process.close()
                for key in manager_dict:
                    mapp[pathway_granularity][method][key] = manager_dict[key]
    pickle.dump(mapp, open(f"pickles/rank_vs_enrichment_{gene_set_type}.pkl", "wb"))

def plot_rank_vs_enrichment(gene_set_type="de"):
    """
    (1) plot rank vs # of sig pathways for each method
    (2) plot overlap of # of sig pathways of each method with traditional comp bio approach at each index 
    """
    mapp = pickle.load(open(f"pickles/rank_vs_enrichment_{gene_set_type}.pkl", "rb"))
    marker_list = [
    '.',  # point
    ',',  # pixel
    'o',  # circle
    'v',  # triangle_down
    '^',  # triangle_up
    '<',  # triangle_left
    '>',  # triangle_right
    '1',  # tri_down
    '2',  # tri_up
    '3',  # tri_left
    '4',  # tri_right
    's',  # square
    'p',  # pentagon
    '*',  # star
    ]
    if gene_set_type == "top_n":
        end_indices = [100, len(mapp["all_pathways"]["compbio_direct"].keys())]
    if gene_set_type == "de":
        end_indices = [10]
    ##plot fraction of relevant pathways found
    relevant_set = set(get_relevant_pathways(just_this_study=True, include_imru_internal=False, key_type="string").keys())
    for end_index in end_indices:
        fig, ax = plt.subplots()
        global_max = max([len(set(mapp["relevant_pathways"][method][i]).intersection(relevant_set)) / len(relevant_set) for method in mapp["relevant_pathways"] for i in sorted(mapp["relevant_pathways"][method].keys())[0:end_index]])
        max_AUC = 0.0
        for index, method in enumerate(mapp["relevant_pathways"]):
            if gene_set_type == "top_n":
                step_size = 5
                indices = np.array(sorted(mapp["relevant_pathways"][method].keys())[10: end_index + step_size - 1: step_size]) ##want to include the end_index
            if gene_set_type == "de":
                indices = sorted(mapp["relevant_pathways"][method].keys())[0:end_index]
            y = np.array([len(set(mapp["relevant_pathways"][method][i]).intersection(relevant_set)) / len(relevant_set) for i in indices])
            AUC = np.trapz(y, indices) / ((max(indices) - min(indices)) * global_max)
            max_AUC = max(AUC, max_AUC)
            indices_jitter = indices + (index * .1)
            ax.plot(indices_jitter, y, linewidth=0.5, linestyle="dashed") 
            ax.scatter(indices_jitter, y, label=f"{get_method_title(method)}, AUC={AUC:.0%}", marker=marker_list[index], s=7.0) ##add a little offset to avoid overlapping: (indices + index * .20, y)
        plt.title("Enrichment of Relevant Pathways at Different Rank Indices")
        ax.set_xlabel("Index Rank")
        ax.set_ylabel("Percentage of Relevant Pathways Found")
        ax.set_ylim((-.01, 1.01))
        ax.set_yticklabels(['{:,.0%}'.format(y_) for y_ in ax.get_yticks()])
        box = ax.get_position()
        ax.set_position([box.x0, box.y0, box.width, box.height * 0.85])
        legend = ax.legend(loc='upper right', prop={"size":7}, ncol=2, bbox_to_anchor=(1, 1.40))
        legend_texts = legend.get_texts()
        legend_text_strings = [l.get_text() for l in legend.get_texts()]
        max_indices = [i for i in range(0, len(legend_texts)) if f"{max_AUC:.0%}" in legend_text_strings[i]]
        for max_index in max_indices:
            legend_texts[max_index].set_weight('bold')
        plt.gcf().subplots_adjust(top=.76)
        plt.savefig(f"outputs/rank_vs_enrichment_{gene_set_type}_{end_index + 10}.png", dpi=300, bbox_inches="tight")