import json
import warnings
from datetime import datetime
from pathlib import Path

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import scanpy as sc
import seaborn as sns
import streamlit as st
from scipy.signal import find_peaks

from stile.nodes import (
    assign_cells_hdbscan,
    assign_cells_to_cores,
    compute_median_nn_distance,
    find_connected_components,
    merge_clusters_by_component,
    plot_cores,
    reassign_unassigned_to_nearest_center,
)

# ----------------------------
# Streamlit Page Setup
# ----------------------------
_ICON_PATH = Path(__file__).parent.parent / "docs" / "icon.svg"
_ICON_SVG = _ICON_PATH.read_text() if _ICON_PATH.exists() else None

st.set_page_config(
    page_title="STiLE Core Identifier",
    page_icon=_ICON_SVG or "🧬",
    layout="wide",
)
if _ICON_SVG:
    # Crop viewBox to the icon area (128x128 at offset 16,16) and scale to title height
    _inline_svg = _ICON_SVG.replace(
        'viewBox="0 0 400 160" width="400" height="160"',
        'viewBox="16 16 128 128"',
    ).replace("<svg ", '<svg style="height:1.8em;width:1.8em;vertical-align:middle;margin-right:0.3em" ')
    st.markdown(
        f'<h1 style="display:flex;align-items:center">{_inline_svg} STiLE: Identify Tissue Cores</h1>',
        unsafe_allow_html=True,
    )
else:
    st.title("🧬 STiLE: Identify Tissue Cores")
st.caption("Interactive UI equivalent of the `identify_tissue_cores` CLI command.")
warnings.filterwarnings("ignore", message=".*force_all_finite.*", category=FutureWarning)

# ----------------------------
# Session-state defaults
# ----------------------------
_DEFAULTS = {
    "min_cluster_size": 200,
    "min_samples": 7,
    "min_cells_per_center": 10,
    "num_rows": 7,
    "num_cols": 7,
    "buffer_distance": 0.0,
    "peak_prominence": 25,
    "peak_distance": 25,
    "density_bins": 200,  # legacy single value
    "density_bins_x": 200,
    "density_bins_y": 200,
    "rotation_deg": 0.0,
}


def _init_defaults():
    for key, val in _DEFAULTS.items():
        if key not in st.session_state:
            st.session_state[key] = val


_init_defaults()

if "regions_completed" not in st.session_state:
    st.session_state["regions_completed"] = []

if "col_options" not in st.session_state:
    st.session_state["col_options"] = []


def _numeric_obs_columns(adata: sc.AnnData) -> list[str]:
    return [c for c in adata.obs.columns if pd.api.types.is_numeric_dtype(adata.obs[c])]


def _set_column_options(adata: sc.AnnData) -> None:
    opts = _numeric_obs_columns(adata)
    st.session_state["col_options"] = opts


def _selectbox_cols(label: str, state_key: str):
    opts = st.session_state.get("col_options", [])
    if not opts:
        opts = ["(load data first)"]
    current = st.session_state.get(state_key, opts[0])
    if current not in opts:
        current = opts[0]
    val = st.sidebar.selectbox(label, options=opts, index=opts.index(current))
    st.session_state[state_key] = val
    return val


def _scatter_fig(
    adata, x_col, y_col, color_col=None, palette="tab20", title=None, unassigned_label=-1, dot_size=6, show_axes=False
):
    fig, ax = plt.subplots(figsize=(3, 3))
    df = adata.obs
    if color_col is None:
        sns.scatterplot(x=df[x_col], y=df[y_col], s=dot_size, alpha=0.45, color="#444444", edgecolor=None, ax=ax)
    else:
        if unassigned_label is not None:
            assigned_mask = df[color_col] != unassigned_label
            assigned_labels = sorted(df.loc[assigned_mask, color_col].unique())
            assigned_palette = sns.color_palette("tab20", n_colors=len(assigned_labels))
            palette_map = dict(zip(assigned_labels, assigned_palette))
            sns.scatterplot(
                x=df.loc[~assigned_mask, x_col],
                y=df.loc[~assigned_mask, y_col],
                s=dot_size,
                alpha=0.25,
                color="#b0b0b0",
                edgecolor=None,
                ax=ax,
            )
            sns.scatterplot(
                x=df.loc[assigned_mask, x_col],
                y=df.loc[assigned_mask, y_col],
                hue=df.loc[assigned_mask, color_col],
                palette=palette_map,
                hue_order=assigned_labels,
                s=dot_size,
                alpha=0.7,
                edgecolor=None,
                ax=ax,
                legend=False,
            )
        else:
            sns.scatterplot(
                x=df[x_col],
                y=df[y_col],
                hue=df[color_col],
                palette=palette,
                s=dot_size,
                alpha=0.7,
                edgecolor=None,
                ax=ax,
                legend=False,
            )
    ax.set_aspect("equal", adjustable="box")
    ax.set_xlabel(x_col)
    ax.set_ylabel(y_col)
    ax.invert_yaxis()
    if title:
        ax.set_title(title)
    if show_axes:
        ax.tick_params(labelsize=6)
    else:
        ax.axis("off")
    return fig


def _apply_rotation(adata: sc.AnnData, x_col: str, y_col: str, rotation_deg: float) -> sc.AnnData:
    """Rotate xy coordinates (and obsm['spatial'] if present) by rotation_deg degrees."""
    if rotation_deg == 0:
        return adata
    theta = np.radians(rotation_deg)
    rot = np.array([[np.cos(theta), -np.sin(theta)], [np.sin(theta), np.cos(theta)]])

    if x_col in adata.obs.columns and y_col in adata.obs.columns:
        coords = adata.obs[[x_col, y_col]].to_numpy()
        adata.obs[[x_col, y_col]] = coords @ rot.T
    if "spatial" in adata.obsm:
        adata.obsm["spatial"] = adata.obsm["spatial"] @ rot.T
    return adata


def _marginal_peaks_fig(
    adata, x_col, y_col, density_bins_x, density_bins_y, x_peaks, y_peaks, color_col=None, dot_size=6
):
    from matplotlib import gridspec

    df = adata.obs
    fig = plt.figure(figsize=(3.8, 3.6))
    gs = gridspec.GridSpec(2, 2, width_ratios=[4, 1.2], height_ratios=[1.2, 4], wspace=0.05, hspace=0.05)
    ax_scatter = fig.add_subplot(gs[1, 0])
    ax_histx = fig.add_subplot(gs[0, 0], sharex=ax_scatter)
    ax_histy = fig.add_subplot(gs[1, 1], sharey=ax_scatter)

    # Compute shared data range with equal aspect padding
    x_vals = df[x_col].values
    y_vals = df[y_col].values
    x_min, x_max = float(x_vals.min()), float(x_vals.max())
    y_min, y_max = float(y_vals.min()), float(y_vals.max())
    x_range = x_max - x_min
    y_range = y_max - y_min
    margin = 0.02
    if x_range > y_range:
        pad = (x_range - y_range) / 2
        y_min -= pad
        y_max += pad
    else:
        pad = (y_range - x_range) / 2
        x_min -= pad
        x_max += pad
    x_margin = (x_max - x_min) * margin
    y_margin = (y_max - y_min) * margin

    # Histograms with explicit range matching the scatter limits
    ax_histx.hist(df[x_col], bins=density_bins_x, range=(x_min, x_max), color="#2563eb", alpha=0.6)
    for xp in x_peaks:
        ax_histx.axvline(xp, color="#e11d48", linestyle="--", linewidth=1)
    ax_histx.tick_params(labelbottom=False, labelleft=False, left=False, bottom=False)
    for spine in ax_histx.spines.values():
        spine.set_visible(False)

    ax_histy.hist(df[y_col], bins=density_bins_y, range=(y_min, y_max), color="#16a34a", alpha=0.6, orientation="horizontal")
    for yp in y_peaks:
        ax_histy.axhline(yp, color="#e11d48", linestyle="--", linewidth=1)
    ax_histy.tick_params(labelbottom=False, labelleft=False, left=False, bottom=False)
    for spine in ax_histy.spines.values():
        spine.set_visible(False)

    # Scatter
    if color_col and color_col in df.columns:
        assigned_labels = sorted(df[color_col].unique())
        assigned_palette = sns.color_palette("husl", n_colors=len(assigned_labels))
        palette_map = dict(zip(assigned_labels, assigned_palette))
        sns.scatterplot(
            x=df[x_col],
            y=df[y_col],
            hue=df[color_col],
            palette=palette_map,
            hue_order=assigned_labels,
            s=dot_size,
            alpha=0.7,
            edgecolor=None,
            ax=ax_scatter,
            legend=False,
        )
    else:
        sns.scatterplot(
            x=df[x_col], y=df[y_col], s=dot_size, alpha=0.45, color="#444444", edgecolor=None, ax=ax_scatter
        )
    for xp in x_peaks:
        ax_scatter.axvline(xp, color="#e11d48", linestyle="--", linewidth=1)
    for yp in y_peaks:
        ax_scatter.axhline(yp, color="#e11d48", linestyle="--", linewidth=1)

    # Set shared limits so histograms align with scatter
    ax_scatter.set_xlim(x_min - x_margin, x_max + x_margin)
    ax_scatter.set_ylim(y_max + y_margin, y_min - y_margin)  # inverted y
    ax_scatter.set_aspect("equal", adjustable="box")
    ax_scatter.axis("off")
    return fig


# ----------------------------
# Data input
# ----------------------------
st.sidebar.header("Data Upload")
data_path_input = st.sidebar.text_input(
    "Full path to data file (.h5ad or .csv)",
    value="",
    placeholder="/path/to/adata.h5ad or /path/to/cells.csv",
)
st.sidebar.caption("Provide a full filesystem path; file is read locally.")
dot_size = st.sidebar.slider("Dot size (scatter)", min_value=0.1, max_value=5.0, value=1.0, step=0.1)

x_col = _selectbox_cols("x_col", "x_col")
y_col = _selectbox_cols("y_col", "y_col")
preliminary_cluster_col = "hdbscan_cluster"
core_label_col = "path_block_core"

st.sidebar.header("Region & Rotation")
rotation_deg = st.sidebar.number_input(
    "Rotation (degrees)",
    value=st.session_state.get("rotation_deg", 0.0),
    min_value=-180.0,
    max_value=180.0,
    step=1.0,
    key="rotation_deg",
)
# Apply pending bounds update: delete widget keys so fresh defaults take effect
_pending_bounds = st.session_state.pop("_pending_region_bounds", None)
if _pending_bounds:
    for _k in ["region_x_min", "region_x_max", "region_y_min", "region_y_max"]:
        st.session_state.pop(_k, None)
    _rb_xmin = _pending_bounds["x_min"]
    _rb_xmax = _pending_bounds["x_max"]
    _rb_ymin = _pending_bounds["y_min"]
    _rb_ymax = _pending_bounds["y_max"]
else:
    # Fall back to data coordinate range when available, otherwise 0.0
    if "adata_raw" in st.session_state:
        _ar = st.session_state["adata_raw"]
        _xc = st.session_state.get("x_col", "x_centroid")
        _yc = st.session_state.get("y_col", "y_centroid")
        _data_xmin = float(_ar.obs[_xc].min()) if _xc in _ar.obs.columns else 0.0
        _data_xmax = float(_ar.obs[_xc].max()) if _xc in _ar.obs.columns else 0.0
        _data_ymin = float(_ar.obs[_yc].min()) if _yc in _ar.obs.columns else 0.0
        _data_ymax = float(_ar.obs[_yc].max()) if _yc in _ar.obs.columns else 0.0
    else:
        _data_xmin = _data_xmax = _data_ymin = _data_ymax = 0.0
    _rb_xmin = st.session_state.get("region_x_min", _data_xmin)
    _rb_xmax = st.session_state.get("region_x_max", _data_xmax)
    _rb_ymin = st.session_state.get("region_y_min", _data_ymin)
    _rb_ymax = st.session_state.get("region_y_max", _data_ymax)
st.sidebar.number_input("x_min", value=_rb_xmin, step=1.0, key="region_x_min")
st.sidebar.number_input("x_max", value=_rb_xmax, step=1.0, key="region_x_max")
st.sidebar.number_input("y_min", value=_rb_ymin, step=1.0, key="region_y_min")
st.sidebar.number_input("y_max", value=_rb_ymax, step=1.0, key="region_y_max")


def _load_data(path: Path) -> sc.AnnData:
    if not path.exists():
        raise FileNotFoundError(f"File not found: {path}")
    if path.suffix.lower() in {".h5ad", ".h5"}:
        return sc.read_h5ad(path)
    if path.suffix.lower() in {".csv", ".tsv"}:
        import pandas as pd

        sep = "," if path.suffix.lower() == ".csv" else "\t"
        df = pd.read_csv(path, sep=sep)
        return sc.AnnData(df)
    raise ValueError("Unsupported file type. Provide .h5ad, .h5, .csv, or .tsv")


# ----------------------------
# Sidebar — Connectivity Analysis
# ----------------------------
st.sidebar.header("Connectivity Analysis")
buffer_distance = st.sidebar.number_input(
    "buffer_distance (0 = auto)", min_value=0.0, max_value=1000.0, format="%f", key="buffer_distance"
)

# ----------------------------
# Sidebar — Clustering
# ----------------------------
st.sidebar.header("Clustering")
min_cluster_size = st.sidebar.number_input("min_cluster_size", min_value=2, max_value=5000, key="min_cluster_size")
min_samples = st.sidebar.number_input("min_samples", min_value=1, max_value=5000, key="min_samples")

# ----------------------------
# Sidebar — Grid Detection & Refinement
# ----------------------------
st.sidebar.header("Grid Detection & Refinement")
num_rows = st.sidebar.number_input("num_rows (expected grid rows)", min_value=1, max_value=200, key="num_rows")
num_cols = st.sidebar.number_input("num_cols (expected grid cols)", min_value=1, max_value=200, key="num_cols")
density_bins_x = st.sidebar.number_input(
    "density_bins_x (x-axis)", min_value=10, max_value=20000, step=10, key="density_bins_x"
)
density_bins_y = st.sidebar.number_input(
    "density_bins_y (y-axis)", min_value=10, max_value=20000, step=10, key="density_bins_y"
)
peak_prominence = st.sidebar.number_input("peak_prominence", min_value=1, max_value=10000, key="peak_prominence")
peak_distance = st.sidebar.number_input("peak_distance", min_value=1, max_value=10000, key="peak_distance")
min_cells_per_center = st.sidebar.number_input(
    "min_cells_per_cluster", min_value=1, max_value=10000, key="min_cells_per_center"
)

# ----------------------------
# Step 1 — Upload Data
# ----------------------------
st.subheader("1️⃣ Upload Data")
load_btn = st.button("Load data")
if load_btn:
    if not data_path_input.strip():
        st.error("Please provide a data file path (.h5ad or .csv).")
        st.stop()
    adata_path = Path(data_path_input.strip())
    if not adata_path.exists():
        st.error(f"File not found: `{adata_path}`")
        st.stop()
    try:
        adata_raw = _load_data(adata_path)
        _set_column_options(adata_raw)
        if not st.session_state["col_options"]:
            st.error("No numeric columns found in obs; cannot plot coordinates.")
            st.stop()
        # Set preferred columns once options are known, then rerun to refresh dropdown defaults
        opts = st.session_state["col_options"]
        first = opts[0]
        x_default = "x_centroid" if "x_centroid" in opts else first
        if "y_centroid" in opts and "y_centroid" != x_default:
            y_default = "y_centroid"
        else:
            # pick a different column if available
            y_default = next((o for o in opts if o != x_default), first)
        prev_x = st.session_state.get("x_col")
        prev_y = st.session_state.get("y_col")
        st.session_state["x_col"] = prev_x if prev_x in opts else x_default
        st.session_state["y_col"] = prev_y if prev_y in opts else y_default

        # Apply optional rotation to coordinates (and spatial if present)
        adata_raw = _apply_rotation(
            adata_raw,
            x_col=st.session_state["x_col"],
            y_col=st.session_state["y_col"],
            rotation_deg=st.session_state.get("rotation_deg", 0.0),
        )

        st.session_state["data_path"] = str(adata_path)
        st.session_state["adata_raw"] = adata_raw
        st.session_state["adata_working"] = adata_raw.copy()
        st.session_state.pop("adata_conn", None)
        st.session_state.pop("adata_hdb", None)
        st.session_state.pop("adata_core", None)
        st.session_state["regions_completed"] = []
        # Initialize region bounds to data range
        _xc = st.session_state["x_col"]
        _yc = st.session_state["y_col"]
        st.session_state["_pending_region_bounds"] = {
            "x_min": float(adata_raw.obs[_xc].min()),
            "x_max": float(adata_raw.obs[_xc].max()),
            "y_min": float(adata_raw.obs[_yc].min()),
            "y_max": float(adata_raw.obs[_yc].max()),
        }
        st.success(f"Loaded `{adata_path}` with {adata_raw.n_obs} cells.")
        st.rerun()
    except Exception as e:
        st.error(f"Failed to load data: {e}")
        st.stop()

if "adata_raw" in st.session_state:
    adata_raw = st.session_state["adata_raw"]
    fig = _scatter_fig(
        adata_raw,
        st.session_state["x_col"],
        st.session_state["y_col"],
        title="All cells",
        dot_size=dot_size,
        show_axes=True,
    )
    st.pyplot(fig, clear_figure=True, width="content")
    total_cells = adata_raw.n_obs
    expected = total_cells / (num_rows * num_cols) if num_rows and num_cols else None
    st.markdown(
        f"**Summary**  \nTotal cells: `{total_cells}`"
        + (f" • Expected per cluster (~{num_rows}×{num_cols}): `{expected:.1f}`" if expected else "")
    )
else:
    st.info("Load data to enable downstream steps.")

# ----------------------------
# Step 2 — Region & Rotation
# ----------------------------
st.subheader("2️⃣ Region & Rotation")
if "adata_raw" in st.session_state:
    adata_raw = st.session_state["adata_raw"]
    xc = st.session_state["x_col"]
    yc = st.session_state["y_col"]
    data_xmin = float(adata_raw.obs[xc].min())
    data_xmax = float(adata_raw.obs[xc].max())
    data_ymin = float(adata_raw.obs[yc].min())
    data_ymax = float(adata_raw.obs[yc].max())
    st.markdown(
        f"**Data coordinate range:** x=[{data_xmin:.1f}, {data_xmax:.1f}], y=[{data_ymin:.1f}, {data_ymax:.1f}]"
    )

    col_crop, col_rotate, col_reset = st.columns(3)
    crop_btn = col_crop.button("✂️ Crop to bounds")
    rotate_btn = col_rotate.button("🔄 Apply rotation")
    reset_btn = col_reset.button("↩️ Reset to full data")

    if crop_btn:
        rx_min = st.session_state.get("region_x_min", data_xmin)
        rx_max = st.session_state.get("region_x_max", data_xmax)
        ry_min = st.session_state.get("region_y_min", data_ymin)
        ry_max = st.session_state.get("region_y_max", data_ymax)
        mask = (
            (adata_raw.obs[xc] >= rx_min)
            & (adata_raw.obs[xc] <= rx_max)
            & (adata_raw.obs[yc] >= ry_min)
            & (adata_raw.obs[yc] <= ry_max)
        )
        adata_cropped = adata_raw[mask].copy()
        st.session_state["adata_raw"] = adata_cropped
        st.session_state.pop("adata_conn", None)
        st.session_state.pop("adata_hdb", None)
        st.session_state.pop("adata_core", None)
        st.success(f"Cropped to {adata_cropped.n_obs} cells (from {adata_raw.n_obs}).")
        st.rerun()

    if rotate_btn:
        adata_rotated = _apply_rotation(
            adata_raw.copy(),
            x_col=st.session_state["x_col"],
            y_col=st.session_state["y_col"],
            rotation_deg=st.session_state.get("rotation_deg", 0.0),
        )
        st.session_state["adata_raw"] = adata_rotated
        st.session_state.pop("adata_conn", None)
        st.session_state.pop("adata_hdb", None)
        st.session_state.pop("adata_core", None)
        _xc = st.session_state["x_col"]
        _yc = st.session_state["y_col"]
        st.session_state["_pending_region_bounds"] = {
            "x_min": float(adata_rotated.obs[_xc].min()),
            "x_max": float(adata_rotated.obs[_xc].max()),
            "y_min": float(adata_rotated.obs[_yc].min()),
            "y_max": float(adata_rotated.obs[_yc].max()),
        }
        st.success(f"Applied {rotation_deg}° rotation to {adata_rotated.n_obs} cells.")
        st.rerun()

    if reset_btn:
        data_path = Path(st.session_state["data_path"])
        adata_reset = _load_data(data_path)
        adata_reset = _apply_rotation(
            adata_reset,
            x_col=st.session_state["x_col"],
            y_col=st.session_state["y_col"],
            rotation_deg=st.session_state.get("rotation_deg", 0.0),
        )
        st.session_state["adata_raw"] = adata_reset
        st.session_state.pop("adata_conn", None)
        st.session_state.pop("adata_hdb", None)
        st.session_state.pop("adata_core", None)
        _xc = st.session_state["x_col"]
        _yc = st.session_state["y_col"]
        st.session_state["_pending_region_bounds"] = {
            "x_min": float(adata_reset.obs[_xc].min()),
            "x_max": float(adata_reset.obs[_xc].max()),
            "y_min": float(adata_reset.obs[_yc].min()),
            "y_max": float(adata_reset.obs[_yc].max()),
        }
        st.success(f"Reset to full data: {adata_reset.n_obs} cells.")
        st.rerun()
else:
    st.info("Load data first.")

# ----------------------------
# Step 3 — Connectivity Analysis
# ----------------------------
st.subheader("3️⃣ Connectivity Analysis")
conn_btn = st.button("🔗 Compute connectivity")
if conn_btn:
    if "adata_raw" not in st.session_state:
        st.error("Load data first.")
        st.stop()
    adata_conn = st.session_state["adata_raw"].copy()
    try:
        median_nn = compute_median_nn_distance(
            adata_conn,
            x_col=st.session_state["x_col"],
            y_col=st.session_state["y_col"],
        )
        buf = buffer_distance if buffer_distance > 0 else None
        adata_conn = find_connected_components(
            adata_conn,
            buffer_distance=buf,
            x_col=st.session_state["x_col"],
            y_col=st.session_state["y_col"],
        )
        st.session_state["adata_conn"] = adata_conn
        st.session_state.pop("adata_hdb", None)
        st.session_state.pop("adata_core", None)
        effective_buf = buf if buf is not None else median_nn
        st.success("Connectivity analysis complete.")
        st.markdown(
            f"**Median NN distance:** `{median_nn:.4f}` • "
            f"**Effective buffer:** `{effective_buf:.4f}` • "
            f"**Connected components:** `{adata_conn.uns['n_connected_components']}`"
        )
    except Exception as e:
        st.error(f"Connectivity analysis failed: {e}")
        st.stop()

if "adata_conn" in st.session_state:
    adata_conn = st.session_state["adata_conn"]
    fig = _scatter_fig(
        adata_conn,
        st.session_state["x_col"],
        st.session_state["y_col"],
        color_col="connected_component",
        unassigned_label=None,
        title="Connected components",
        dot_size=dot_size,
    )
    st.pyplot(fig, clear_figure=True, width="content")
    st.markdown(
        f"**Components:** {adata_conn.uns['n_connected_components']} • "
        f"**Median NN dist:** `{adata_conn.uns.get('median_nn_distance', 'N/A')}`"
    )

# ----------------------------
# Step 4 — Run clustering (HDBSCAN + Merge)
# ----------------------------
st.subheader("4️⃣ Run clustering (HDBSCAN + Merge)")
run_hdb_btn = st.button("▶️ Run HDBSCAN + Merge")
if run_hdb_btn:
    if "adata_conn" not in st.session_state:
        st.error("Run connectivity analysis first.")
        st.stop()
    adata_hdb = st.session_state["adata_conn"].copy()
    try:
        adata_hdb = assign_cells_hdbscan(
            adata_hdb,
            x_col=st.session_state["x_col"],
            y_col=st.session_state["y_col"],
            cluster_label_col=preliminary_cluster_col,
            min_cluster_size=min_cluster_size,
            min_samples=min_samples,
        )
        adata_hdb = merge_clusters_by_component(
            adata_hdb,
            cluster_label_col=preliminary_cluster_col,
        )
        st.session_state["adata_hdb"] = adata_hdb
        st.session_state.pop("adata_core", None)
        st.success("Clustering and merge complete.")
    except Exception as e:
        st.error(f"HDBSCAN + Merge failed: {e}")
        st.stop()

if "adata_hdb" in st.session_state:
    adata_hdb = st.session_state["adata_hdb"]
    fig = _scatter_fig(
        adata_hdb,
        st.session_state["x_col"],
        st.session_state["y_col"],
        color_col="merged_cluster",
        unassigned_label=-1,
        title="Merged clusters (grey = unassigned)",
        dot_size=dot_size,
    )
    st.pyplot(fig, clear_figure=True, width="content")
    labels = adata_hdb.obs["merged_cluster"].values
    n_clusters = len(np.unique(labels[labels != -1]))
    assigned = np.sum(labels != -1)
    total = len(labels)
    pct = assigned / total * 100
    st.markdown(
        f"**Merged clusters:** {n_clusters} • **Assigned:** {assigned} ({pct:.1f}%) "
        f"• **Unassigned:** {total - assigned}"
    )

    # Option to skip grid detection and use HDBSCAN clusters directly
    use_hdb_btn = st.button("✅ Use HDBSCAN clusters as cores (skip grid)")
    if use_hdb_btn:
        adata_core = adata_hdb.copy()
        merged = adata_core.obs["merged_cluster"]
        adata_core.obs[core_label_col] = merged.apply(lambda v: f"core_{v}" if v != -1 else "unassigned")
        st.session_state["adata_core"] = adata_core
        st.rerun()

# ----------------------------
# Step 5 — Detect Grid & Reassignment (optional)
# ----------------------------
st.subheader("5️⃣ Detect grid & assign cores (optional)")
grid_btn = st.button("📐 Detect grid & assign cores")
if grid_btn:
    if "adata_hdb" not in st.session_state:
        st.error("Run clustering first.")
        st.stop()
    adata_core = st.session_state["adata_hdb"].copy()
    try:
        # Peak detection for display
        x = adata_core.obs[st.session_state["x_col"]].values
        y = adata_core.obs[st.session_state["y_col"]].values
        x_hist, x_edges = np.histogram(x, bins=density_bins_x)
        y_hist, y_edges = np.histogram(y, bins=density_bins_y)
        x_idx, _ = find_peaks(x_hist, prominence=peak_prominence, distance=peak_distance)
        y_idx, _ = find_peaks(y_hist, prominence=peak_prominence, distance=peak_distance)
        x_peaks = (x_edges[x_idx] + x_edges[x_idx + 1]) / 2 if len(x_idx) else []
        y_peaks = (y_edges[y_idx] + y_edges[y_idx + 1]) / 2 if len(y_idx) else []

        adata_core = assign_cells_to_cores(
            adata_core,
            max_x_peaks=num_cols,
            max_y_peaks=num_rows,
            x_col=st.session_state["x_col"],
            y_col=st.session_state["y_col"],
            peak_prominence=peak_prominence,
            peak_distance=peak_distance,
            core_label_col=core_label_col,
            density_bins_x=density_bins_x,
            density_bins_y=density_bins_y,
            min_cells_per_center=min_cells_per_center,
            preliminary_cluster_col="merged_cluster",
        )
        adata_core, _ = reassign_unassigned_to_nearest_center(
            adata_core,
            core_label_col=core_label_col,
            x_col=st.session_state["x_col"],
            y_col=st.session_state["y_col"],
        )
        st.session_state["adata_core"] = adata_core

        # Plot marginals with peaks
        fig_peaks = _marginal_peaks_fig(
            adata_core,
            st.session_state["x_col"],
            st.session_state["y_col"],
            density_bins_x,
            density_bins_y,
            x_peaks,
            y_peaks,
            color_col="merged_cluster",
            dot_size=dot_size,
        )
        st.pyplot(fig_peaks, clear_figure=True, width="content")
        st.success("Grid detection and refinement complete.")
    except Exception as e:
        st.error(f"Grid detection failed: {e}")
        st.stop()

# ----------------------------
# Step 6 — Final assignment + Save Region
# ----------------------------
st.subheader("6️⃣ Final core assignment")
if "adata_core" in st.session_state:
    adata_final = st.session_state["adata_core"]
    fig_final = plot_cores(
        adata_final,
        x_col=st.session_state["x_col"],
        y_col=st.session_state["y_col"],
        core_label_col=core_label_col,
    )
    fig_final.set_size_inches(5, 5, forward=True)
    st.pyplot(fig_final, clear_figure=True, width="content")

    labels = adata_final.obs[core_label_col].values
    assigned = np.sum(labels != "unassigned")
    total = len(labels)
    st.markdown(f"**Assignment stats:** Total `{total}` • Assigned `{assigned}` • Unassigned `{total - assigned}`")

    # --- Save Region ---
    st.markdown("---")
    st.markdown("**Save region results**")
    n_regions = len(st.session_state.get("regions_completed", []))
    region_name = st.text_input("Region name", value=f"R{n_regions + 1}", key="region_name_input")
    save_region_btn = st.button("💾 Save region results")
    if save_region_btn:
        adata_c = st.session_state["adata_core"]
        obs_result = adata_c.obs.copy()
        # Prefix core labels with region name
        obs_result[core_label_col] = obs_result[core_label_col].apply(
            lambda v: f"{region_name}_{v}" if v != "unassigned" else v
        )
        # Save region to disk
        base_path = Path(st.session_state.get("data_path", ".")) if st.session_state.get("data_path") else Path(".")
        base_dir = base_path.parent if base_path.is_file() else base_path
        base_name = base_path.stem if base_path.is_file() else "adata"
        region_dir = base_dir / f"regions_{base_name}"
        region_dir.mkdir(exist_ok=True)
        region_csv = region_dir / f"{region_name}.csv"
        export_cols = ["cell_id", st.session_state["x_col"], st.session_state["y_col"], core_label_col]
        df_export = obs_result.copy()
        if "cell_id" not in df_export.columns:
            df_export.insert(0, "cell_id", df_export.index)
        export_cols = [c for c in export_cols if c in df_export.columns]
        df_export[export_cols].to_csv(region_csv, index=False)

        region_entry = {
            "name": region_name,
            "x_min": st.session_state.get("region_x_min", 0.0),
            "x_max": st.session_state.get("region_x_max", 0.0),
            "y_min": st.session_state.get("region_y_min", 0.0),
            "y_max": st.session_state.get("region_y_max", 0.0),
            "obs_result": obs_result,
            "csv_path": str(region_csv),
        }
        st.session_state["regions_completed"].append(region_entry)
        n_saved_cores = len(obs_result[obs_result[core_label_col] != "unassigned"][core_label_col].unique())
        st.success(
            f"Saved region **{region_name}** ({assigned} assigned cells, {n_saved_cores} cores).  \nCSV: `{region_csv}`"
        )

    # Display saved regions table
    if st.session_state.get("regions_completed"):
        st.markdown("**Saved regions:**")
        for i, reg in enumerate(st.session_state["regions_completed"]):
            obs_r = reg["obs_result"]
            n_assigned = int((obs_r[core_label_col] != "unassigned").sum())
            n_cores = len(obs_r[obs_r[core_label_col] != "unassigned"][core_label_col].unique())
            col1, col2 = st.columns([4, 1])
            csv_info = f"  \n`{reg['csv_path']}`" if reg.get("csv_path") else ""
            col1.markdown(
                f"**{reg['name']}** — bounds: x=[{reg['x_min']:.1f}, {reg['x_max']:.1f}], "
                f"y=[{reg['y_min']:.1f}, {reg['y_max']:.1f}] — "
                f"{n_assigned} cells, {n_cores} cores{csv_info}"
            )
            if col2.button("🗑️ Delete", key=f"del_region_{i}"):
                st.session_state["regions_completed"].pop(i)
                st.rerun()
else:
    st.info("Run grid detection to see final assignments.")

# ----------------------------
# Step 7 — Merge Regions
# ----------------------------
st.subheader("7️⃣ Merge Regions")
if st.session_state.get("regions_completed"):
    st.markdown(f"**{len(st.session_state['regions_completed'])} region(s)** ready to merge.")
    merge_btn = st.button("🔀 Merge all regions")
    if merge_btn:
        # Re-read full original data and apply rotation
        data_path = Path(st.session_state["data_path"])
        adata_full = _load_data(data_path)
        adata_full = _apply_rotation(
            adata_full,
            x_col=st.session_state["x_col"],
            y_col=st.session_state["y_col"],
            rotation_deg=st.session_state.get("rotation_deg", 0.0),
        )
        # Initialize all cells as unassigned
        adata_full.obs[core_label_col] = "unassigned"
        # Apply each region's results by matching cell indices
        for reg in st.session_state["regions_completed"]:
            obs_r = reg["obs_result"]
            assigned_mask = obs_r[core_label_col] != "unassigned"
            assigned_obs = obs_r.loc[assigned_mask, core_label_col]
            common_idx = adata_full.obs.index.intersection(assigned_obs.index)
            adata_full.obs.loc[common_idx, core_label_col] = assigned_obs.loc[common_idx]
        st.session_state["adata_core"] = adata_full
        # Show merged plot
        fig_merged = plot_cores(
            adata_full,
            x_col=st.session_state["x_col"],
            y_col=st.session_state["y_col"],
            core_label_col=core_label_col,
        )
        fig_merged.set_size_inches(5, 5, forward=True)
        st.pyplot(fig_merged, clear_figure=True, width="content")
        merged_labels = adata_full.obs[core_label_col].values
        m_assigned = np.sum(merged_labels != "unassigned")
        m_total = len(merged_labels)
        m_cores = len(np.unique(merged_labels[merged_labels != "unassigned"]))
        st.success(
            f"Merged {len(st.session_state['regions_completed'])} regions into full dataset: "
            f"{m_total} total cells, {m_assigned} assigned, {m_cores} cores."
        )
else:
    st.info("Save one or more region results above, then merge them here.")

# ----------------------------
# Step 8 — Export assignments (no layout needed)
# ----------------------------
st.subheader("8️⃣ Export assignments (no layout)")
if "adata_core" in st.session_state:
    export_assign_format = st.radio(
        "Export format (assignments only)",
        options=["CSV", "AnnData (.h5ad)"],
        index=0,
        horizontal=True,
        key="export_assign_format",
    )
    export_assign_btn = st.button("💾 Export assignments")
    if export_assign_btn:
        # Build output paths with timestamp and base filename to avoid overwrite
        base_path = Path(st.session_state.get("data_path", ".")) if st.session_state.get("data_path") else Path(".")
        if base_path.is_file():
            base_dir = base_path.parent
            base_name = base_path.stem
        else:
            base_dir = base_path
            base_name = "adata"
        ts = datetime.now().strftime("%Y%m%d-%H%M%S")
        adata_export = st.session_state["adata_core"]
        df = adata_export.obs.copy()
        if "cell_id" not in df.columns:
            df.insert(0, "cell_id", df.index)
        cols = ["cell_id", st.session_state["x_col"], st.session_state["y_col"], core_label_col]
        if preliminary_cluster_col in df.columns:
            cols.append(preliminary_cluster_col)
        cols = [c for c in cols if c in df.columns]
        # Save parameters as JSON
        params = {
            "data_path": st.session_state.get("data_path", ""),
            "x_col": st.session_state.get("x_col", ""),
            "y_col": st.session_state.get("y_col", ""),
            "rotation_deg": st.session_state.get("rotation_deg", 0.0),
            "region_x_min": st.session_state.get("region_x_min", 0.0),
            "region_x_max": st.session_state.get("region_x_max", 0.0),
            "region_y_min": st.session_state.get("region_y_min", 0.0),
            "region_y_max": st.session_state.get("region_y_max", 0.0),
            "buffer_distance": st.session_state.get("buffer_distance", 0.0),
            "min_cluster_size": st.session_state.get("min_cluster_size", 200),
            "min_samples": st.session_state.get("min_samples", 7),
            "num_rows": st.session_state.get("num_rows", 7),
            "num_cols": st.session_state.get("num_cols", 7),
            "density_bins_x": st.session_state.get("density_bins_x", 200),
            "density_bins_y": st.session_state.get("density_bins_y", 200),
            "peak_prominence": st.session_state.get("peak_prominence", 25),
            "peak_distance": st.session_state.get("peak_distance", 25),
            "min_cells_per_center": st.session_state.get("min_cells_per_center", 10),
        }
        params_path = base_dir / f"{base_name}_params_{ts}.json"
        params_path.write_text(json.dumps(params, indent=2))

        if export_assign_format.startswith("CSV"):
            out_path = base_dir / f"{base_name}_core_assignments_{ts}.csv"
            df[cols].to_csv(out_path, index=False)
            st.success(f"Exported core assignments CSV to `{out_path}`")
        else:
            out_path = base_dir / f"{base_name}_core_assignments_{ts}.h5ad"
            adata_export.write_h5ad(out_path)
            st.success(f"Exported AnnData with assignments to `{out_path}`")
        st.info(f"Parameters saved to `{params_path}`")
else:
    st.caption("Run grid detection to export assignments.")


# ----------------------------
# Footer
# ----------------------------
st.markdown("---")
st.caption("This Streamlit UI runs the same workflow as your CLI command `identify_tissue_cores`.")
