# ============================================================
# HiMaLAYAS — Cluster-specific zoom analysis (full workflow)
# ============================================================

# ------------------------------------------------------------
# Hierarchical clustering options (quick reference)
# ------------------------------------------------------------
#
# Linkage_method:
#   - "Ward"      : compact, variance-minimizing clusters (requires euclidean)
#   - "Average"   : average pairwise distance between clusters
#   - "Complete"  : maximum pairwise distance (tight clusters)
#   - "Single"    : minimum pairwise distance (chaining; usually discouraged)
#
# Linkage_metric:
#   - "Euclidean"   : required for ward; default and safest
#   - "Correlation" : distance = 1 − Pearson correlation
#   - "Cosine"      : angle-based similarity
#   - "Cityblock"   : Manhattan distance
#
# Recommended default:
#   Linkage_method="ward"
#   Linkage_metric="euclidean"
# ------------------------------------------------------------


# ============================================================
# Step 1: Extract cluster-specific submatrix
# ============================================================

CLUSTER_ID = 9

zoom_view = results.subset(cluster=CLUSTER_ID)
zoom_matrix = zoom_view.matrix

print(f"Zoomed matrix shape: {zoom_matrix.values.shape}")


# ============================================================
# Step 2: Build annotations filtered to the subset matrix
# ============================================================

zoom_annotations = Annotations(
    go_bp,
    zoom_matrix,
)


# ============================================================
# Step 3: Re-run clustering + enrichment (hypergeom)
# ============================================================

ZOOM_LINKAGE_THRESHOLD = 6

zoom_analysis = (
    Analysis(zoom_matrix, zoom_annotations)
    .cluster(
        linkage_method="ward",
        linkage_metric="euclidean",
        linkage_threshold=ZOOM_LINKAGE_THRESHOLD,
        min_cluster_size=6,
    )
    .enrich(
        min_overlap=2,
        background=results.matrix,
    )
    .finalize(
        col_cluster=True,
        add_qvalues=True,
    )
)
zoom_results = zoom_analysis.results


# ============================================================
# Step 4: Map GO term IDs to human-readable names
# ============================================================

zoom_results.df["term_name"] = (
    zoom_results.df["term"].map(go_id_to_name).fillna(zoom_results.df["term"])
)


# ============================================================
# Step 5: Filter to significant enrichments
# ============================================================

zoom_results_sig = zoom_results.filter("qval <= 0.05")

print(f"All enriched rows: {len(zoom_results.df):,}")
print(f"Significant rows (q<=0.05): {len(zoom_results_sig.df):,}")
print(zoom_results_sig.df.head(10))


# ============================================================
# Step 6: Summarize clusters into plot-ready labels
# ============================================================

zoom_cluster_labels = summarize_clusters(
    zoom_results_sig.df,
    label_mode="top_term",
    label_col="term_name",
)

print("\nCluster-level labels (zoomed):")
print(zoom_cluster_labels)


# ============================================================
# Step 7: Plot zoomed clustered matrix
# ============================================================

# Universal colors
LABEL_COLOR = "white"
BACKGROUND_COLOR = "black"

# Robust heatmap scaling
vals = zoom_matrix.values
mask = np.isfinite(vals) & (vals != 0)
if not np.any(mask):
    raise ValueError("Matrix has no finite nonzero entries for heatmap scaling")
vlim = float(np.percentile(np.abs(vals[mask]), 99))

# Create the plotter with your enrichment results
plotter = Plotter(zoom_results)

(
    plotter.plot_dendrogram(
        axes=[0.06, 0.05, 0.09, 0.90],
        data_pad=0.75,
        color="#888888",
        linewidth=0.75,
    )
    .plot_matrix(
        cmap=matrix_cmap,
        center=0,
        vmin=-vlim,
        vmax=vlim,
        figsize=(13, 7),
        subplots_adjust={
            "left": 0.15,
            "right": 0.70,
            "bottom": 0.05,
            "top": 0.95,
        },
    )
    .plot_matrix_axis_labels(
        xlabel="Array gene",
        ylabel="Query gene",
        fontsize=16,
        color=LABEL_COLOR,
        xlabel_pad=4.5,
        ylabel_pad=0.005,
    )
    .plot_row_ticks(
        max_labels=120,
        fontsize=8,
        color=LABEL_COLOR,
    )
    .plot_col_ticks(
        max_labels=300,
        fontsize=8,
        color=LABEL_COLOR,
        rotation=90,
    )
    .plot_cluster_labels(
        zoom_cluster_labels,
        max_words=24,
        wrap_text=True,  # enable wrapping
        wrap_width=36,  # characters per line (tune if needed)
        overflow="wrap",  # never ellipsize
        font="Helvetica",
        fontsize=18,
        color=LABEL_COLOR,
        skip_unlabeled=False,
        label_fields=("label", "n", "p"),
        omit_words={},
        boundary_color=LABEL_COLOR,
        boundary_lw=1,
        boundary_alpha=0.8,
        dendro_boundary_alpha=0.0,
        label_text_pad=0.012,
        label_sep_xmin=None,
        label_sep_xmax=0.5,
        label_sep_color="#bbbbbb",
        label_sep_lw=1,
        label_sep_alpha=0.0,
        label_gutter_color=BACKGROUND_COLOR,
        axes=[0.70, 0.05, 0.29, 0.90],
    )
    .plot_cluster_bar(
        norm=Normalize(0, 30),
        name="sigbar",
        title="Enrichment",
        values=zoom_cluster_labels,
        pval_col="pval",
        width=0.05,
        left_pad=0.00,
        right_pad=0.01,
    )
    .plot_gene_bar(
        values=gene_essential_map,
        mode="categorical",
        colors=gene_essential_colors,
        placement="label_panel",
        gene_bar_left_pad=0.06,
        gene_bar_width=0.05,
        gene_bar_right_pad=0.02,
        name="essentiality",
        title="Essential",
    )
    .plot_gene_bar(
        values=gene_characterization_map,
        mode="categorical",
        colors=gene_characterization_colors,
        placement="label_panel",
        gene_bar_left_pad=0.0,
        gene_bar_width=0.05,
        gene_bar_right_pad=0.02,
        name="unchar",
        title="Uncharacterized",
    )
    .plot_bar_labels(
        font="Helvetica",
        fontsize=14,
        color=LABEL_COLOR,
        pad=4,
        rotation=90,
    )
    .set_label_track_order(("essentiality", "unchar", "sigbar"))
    .set_background(color=BACKGROUND_COLOR)
    .plot_title(
        f"Yeast SGA Matrix (Cluster {CLUSTER_ID})",
        color=LABEL_COLOR,
        fontsize=20,
    )
    .show()
)
