neuroglancer_utils.py

"""
Neuroglancer integration utilities for viewing predictions vs ground truth.
"""

import os
import json
import re
import hashlib
import socket
import neuroglancer
import zarr
from upath import UPath
import pandas as pd

config = json.load(open(".config.json", "r"))
EM_S3_URL_TEMPLATE = config["EM_S3_URL_TEMPLATE"]

# Fixed color map for organelle layers. Each organelle type always gets the same
# color, and all segment IDs within a layer are rendered in that single color
# (via segment_default_color). Colors are grouped by biological family.
ORGANELLE_COLORS = {
    # Mitochondria - greens
    "mito": "#22cc44",
    "mito_lum": "#33dd55",
    "mito_mem": "#118833",
    "mito_ribo": "#44ee66",
    # ER - blues
    "er": "#3388ff",
    "er_lum": "#55aaff",
    "er_mem": "#1166cc",
    "er_mem_all": "#2277dd",
    # ER exit sites - dark blues
    "eres": "#0044aa",
    "eres_lum": "#1155bb",
    "eres_mem": "#003399",
    # Golgi - yellows/golds
    "golgi": "#ffcc00",
    "golgi_lum": "#ffdd33",
    "golgi_mem": "#ddaa00",
    # Nucleus - purples
    "nuc": "#9933ff",
    "nucpl": "#bb66ff",
    # Nuclear envelope - violets
    "ne": "#7722cc",
    "ne_lum": "#8833dd",
    "ne_mem": "#6611bb",
    "ne_mem_all": "#9944ee",
    # Nuclear pore - dark purples
    "np": "#550099",
    "np_in": "#6600aa",
    "np_out": "#440088",
    # Lipid droplets - oranges
    "ld": "#ff8800",
    "ld_lum": "#ff9933",
    "ld_mem": "#dd6600",
    # Lysosomes - reds
    "lyso": "#ff2200",
    "lyso_lum": "#ff4433",
    "lyso_mem": "#cc1100",
    # Endosomes - teals/cyans
    "endo": "#00ccbb",
    "endo_lum": "#33ddcc",
    "endo_mem": "#009988",
    # Vesicles - pinks
    "ves": "#ff66aa",
    "ves_lum": "#ff88bb",
    "ves_mem": "#dd4488",
    # Peroxisomes - browns/ambers
    "perox": "#cc6633",
    "perox_lum": "#dd7744",
    "perox_mem": "#aa4422",
    # Microtubules - lime/chartreuse
    "mt": "#88ee00",
    "mt_in": "#99ff11",
    "mt_out": "#66cc00",
    # Chromatin - magentas
    "chrom": "#ff00ff",
    "echrom": "#ff33cc",
    "hchrom": "#cc00cc",
    # Cell-level structures
    "cell": "#cccccc",
    "cyto": "#88bbdd",
    "ecs": "#aaddaa",
    "pm": "#666666",
}


def get_organelle_color(organelle_name):
    """Return a consistent hex color for the given organelle name.

    Uses the predefined ORGANELLE_COLORS map, falling back to a
    deterministic hash-based color for unknown organelle names.
    """
    if organelle_name in ORGANELLE_COLORS:
        return ORGANELLE_COLORS[organelle_name]
    # Deterministic fallback: hash the name to produce a stable color
    h = hashlib.md5(organelle_name.encode()).hexdigest()
    return f"#{h[:6]}"


# Global Neuroglancer server instance
_neuroglancer_server = None
_neuroglancer_port = None


def get_zarr_path(
    results_path,
    evaluation_id=None,  # None means GT
    crop=None,
    label=None,
    db_csv_filename: str = "evaluations.csv",
) -> str:
    """
    evaluations.csv should have columns: evaluation_id, username, submission_id, submission_name, data_path, results_file_path, status, created_at, last_updated, last_log_line
    """
    if evaluation_id is not None:
        db_csv_path = UPath(results_path).parent / db_csv_filename
        if not db_csv_path.exists():
            # Try working directory as fallback (in case results_path is relative and not rooted at the same level as the CSV)
            db_csv_path = UPath(os.getcwd()) / db_csv_filename
            if not db_csv_path.exists():
                raise FileNotFoundError(
                    f"Could not find {db_csv_filename} at {results_path} or current working directory"
                )
        db = pd.read_csv(db_csv_path.path)

        row = db[db["evaluation_id"] == evaluation_id]
        zarr_path = UPath(row["data_path"].values[0]).with_suffix(".zarr")
    else:
        zarr_path = UPath(results_path).with_suffix(".zarr")

    return os.path.join(
        *[x for x in (zarr_path.as_posix(), crop, label) if x is not None]
    )


def _find_free_port(start_port, bind_address="0.0.0.0"):
    port = start_port
    while True:
        with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
            try:
                s.bind((bind_address, port))
                return port
            except OSError:
                port += 1


def initialize_neuroglancer_server(bind_address="0.0.0.0", port=8080):
    """
    Initialize the shared Neuroglancer server.

    Args:
        bind_address: IP address to bind to
        port: Starting port number; increments if already in use

    Returns:
        Neuroglancer server instance
    """
    global _neuroglancer_server, _neuroglancer_port

    if _neuroglancer_server is None:
        port = _find_free_port(port, bind_address)
        neuroglancer.set_server_bind_address(bind_address, bind_port=port)
        # Create a viewer to force the HTTP server to start immediately.
        # Without this, the server only starts on the first Viewer() call
        # inside create_viewer_url(), causing a race condition where the URL
        # is opened before the server is ready to accept connections.
        _viewer = neuroglancer.Viewer()
        _neuroglancer_server = True
        _neuroglancer_port = port
        print(f"Neuroglancer server initialized on {bind_address}:{port}")

    return _neuroglancer_server


def get_neuroglancer_port():
    """Return the port the Neuroglancer server is listening on."""
    return _neuroglancer_port


def load_zarr_array(zarr_path):
    """
    Load a zarr array from the given path.

    Args:
        zarr_path: Path to zarr array

    Returns:
        tuple: (array, voxel_size, offset) or (None, None, None) if not found
    """
    try:
        # Open the zarr array
        arr = zarr.open(zarr_path, mode="r")

        # Get voxel size (try different attribute names)
        voxel_size = None
        for attr_name in ["voxel_size", "resolution", "scale"]:
            if attr_name in arr.attrs:
                voxel_size = arr.attrs[attr_name]
                break

        if voxel_size is not None:
            # Ensure it's a list
            voxel_size = list(voxel_size)

        # Get offset/translation
        offset = None
        for attr_name in ["offset", "translation"]:
            if attr_name in arr.attrs:
                offset = arr.attrs[attr_name]
                break

        if offset is not None:
            offset = list(offset)
            if voxel_size is None:
                voxel_size = [
                    1.0,
                    1.0,
                    1.0,
                ]  # Assume isotropic 1nm voxels if not specified

            # Convert offset to voxel units
            offset = [int(o / v) for o, v in zip(offset, voxel_size)]

        return arr, voxel_size, offset

    except Exception as e:
        print(f"Error loading zarr array from {zarr_path}: {e}")
        return None, None, None


def get_dataset_for_crop(crop, crop_to_dataset):
    """
    Look up the dataset name for a given crop.

    Args:
        crop: Crop name (e.g., 'crop982')
        crop_to_dataset: Dict mapping crop number strings to dataset names

    Returns:
        str: Dataset name (e.g., 'jrc_mus-nacc-1') or None if not found
    """
    # Strip 'crop' prefix to get the number
    crop_num = re.sub(r"^crop", "", crop)
    return crop_to_dataset.get(crop_num)


def get_available_organelles(eval_id, crop, results_df=None):
    """
    Get list of available organelles for a given evaluation and crop.

    Args:
        eval_id: Evaluation ID
        crop: Crop name
        results_df: DataFrame with evaluation results (optional, will try to discover from files if not provided)

    Returns:
        list: List of organelle names
    """
    if results_df is not None:
        # Filter by eval_id and crop
        filtered = results_df[
            (results_df["evaluation_id"] == eval_id) & (results_df["crop"] == crop)
        ]
        organelles = filtered["organelle"].unique().tolist()
        return sorted(organelles)

    # Fallback: try to discover from filesystem
    # This would require scanning the zarr files, which we'll skip for now
    return []


def get_seg_layer(
    array=None,
    voxel_offset=None,
    dimensions=None,
    local_path=None,
    voxel_size=None,
    color=None,
) -> neuroglancer.SegmentationLayer:
    if array is not None:
        layer = neuroglancer.SegmentationLayer(
            source=neuroglancer.LocalVolume(
                data=array,
                voxel_offset=voxel_offset,
                dimensions=dimensions,
            ),
        )
    elif local_path is not None and voxel_size is not None and voxel_offset is not None:
        path = config["SEMI-PUBLIC_URL"] + local_path.replace(
            "/nrs/cellmap/", "/nrs/"
        ).replace("/prfs/cellmap/cellmap/", "/prfs/")
        # Affine: scale by voxel_size, translate by translation
        matrix = [
            [1, 0, 0, voxel_offset[0]],
            [0, 1, 0, voxel_offset[1]],
            [0, 0, 1, voxel_offset[2]],
        ]
        input_dims = neuroglancer.CoordinateSpace(
            names=(
                "d0",
                "d1",
                "d2",
            ),  # or ("x","y","z") depending on the Zarr axis order
            units=("nm", "nm", "nm"),
            scales=voxel_size,
        )
        transform = neuroglancer.CoordinateSpaceTransform(
            matrix=matrix,
            output_dimensions=dimensions,
            input_dimensions=input_dims,
        )
        layer = neuroglancer.SegmentationLayer(
            source=neuroglancer.LayerDataSource(
                url=path,
                transform=transform,
            )
        )
    else:
        raise ValueError("Either array or local_path and voxel_size must be provided")

    if color is not None:
        layer.segment_default_color = color

    return layer


def create_viewer_url(
    eval_id,
    crop,
    organelle,
    gt_data_path,
    evaluations_df=None,
    results_df=None,
    crop_to_dataset=None,
    return_viewer=False,
):
    """
    Create a Neuroglancer viewer URL with raw EM, GT, and prediction layers.

    Args:
        eval_id: Evaluation ID
        crop: Crop name (e.g., 'crop982')
        organelle: Organelle name (e.g., 'mito', 'ld', 'nuc'), or None/empty to load all organelles
        gt_data_path: Base path to ground truth data
        evaluations_df: DataFrame with evaluation metadata (to get data_path)
        results_df: DataFrame with evaluation results (to get available organelles)
        crop_to_dataset: Dict mapping crop number strings to dataset names (for EM data)
        return_viewer: If True, return the Neuroglancer viewer instance instead of the URL

    Returns:
        str: Neuroglancer viewer URL or None if error
        OR
        neuroglancer.Viewer: Viewer instance if return_viewer is True
    """
    try:
        # If organelle is None or empty, load all available organelles
        # If comma-separated string, split into list
        if not organelle or organelle == "":
            if results_df is not None:
                organelles = get_available_organelles(eval_id, crop, results_df)
                if not organelles:
                    print(f"No organelles found for eval_id={eval_id}, crop={crop}")
                    return None
                print(f"Loading all available organelles: {organelles}")
            else:
                print("Cannot load all organelles without results_df")
                return None
        elif "," in organelle:
            organelles = [o.strip() for o in organelle.split(",") if o.strip()]
            print(f"Loading selected organelles: {organelles}")
        else:
            organelles = [organelle]

        # Get prediction data path
        if evaluations_df is not None:
            eval_row = evaluations_df[evaluations_df["evaluation_id"] == eval_id]
            if not eval_row.empty:
                data_path = eval_row.iloc[0]["data_path"]
                results_path = data_path
            else:
                print(f"Evaluation {eval_id} not found in metadata")
                return None
        else:
            results_path = None

        # Create viewer (this creates a new viewer with fresh state, resetting field of view)
        viewer = neuroglancer.Viewer()

        # Track dimensions and position info from first organelle
        first_dimensions = None
        first_offset = None
        first_voxel_size = None
        first_shape = None

        with viewer.txn() as s:
            # --- Load raw EM data from S3 ---
            em_loaded = False
            if crop_to_dataset:
                dataset_name = get_dataset_for_crop(crop, crop_to_dataset)
                if dataset_name:
                    em_s3_path = EM_S3_URL_TEMPLATE.format(dataset_name=dataset_name)
                    print(f"Loading raw EM from: {em_s3_path}")

                    try:
                        s.layers["em"] = neuroglancer.ImageLayer(
                            source=em_s3_path,
                        )
                        em_loaded = True
                        print(f"Added EM layer from S3 path: {em_s3_path}")
                    except Exception as e:
                        print(f"Could not load EM data from {em_s3_path}")

                else:
                    print(
                        f"No dataset mapping found for crop '{crop}', "
                        f"skipping EM layer"
                    )

            # --- Load GT and prediction segmentation layers ---
            # Track which organelles have GT vs predictions only
            gt_layers = []
            pred_layers = []
            has_any_gt = False

            for org in organelles:
                # Get zarr paths
                gt_path = None
                gt_array, gt_voxel_size, gt_offset = None, None, None
                if gt_data_path:
                    gt_path = get_zarr_path(
                        results_path=gt_data_path,
                        evaluation_id=None,
                        crop=crop,
                        label=org,
                    )
                    print(f"GT path ({org}): {gt_path}")
                    gt_array, gt_voxel_size, gt_offset = load_zarr_array(gt_path)
                else:
                    print(f"GT path not configured, skipping GT layer for {org}")

                pred_path = get_zarr_path(
                    results_path=results_path,
                    evaluation_id=eval_id,
                    crop=crop,
                    label=org,
                )

                print(f"Prediction path ({org}): {pred_path}")

                # Load prediction array
                pred_array, pred_voxel_size, pred_offset = load_zarr_array(pred_path)

                if pred_array is None:
                    print(
                        f"Could not load prediction array for {org} from {pred_path}, skipping"
                    )
                    continue

                # Set dimensions from first successful prediction load
                if first_dimensions is None:
                    voxel_size = gt_voxel_size if gt_voxel_size else pred_voxel_size
                    offset = gt_offset if gt_offset else pred_offset
                    shape = (
                        tuple(gt_array.shape)
                        if gt_array is not None
                        else tuple(pred_array.shape)
                    )
                    first_dimensions = neuroglancer.CoordinateSpace(
                        names=["z", "y", "x"],
                        units=["nm", "nm", "nm"],
                        scales=voxel_size,
                    )
                    first_offset = offset
                    first_voxel_size = voxel_size
                    first_shape = shape
                    s.dimensions = first_dimensions

                # Get the fixed color for this organelle
                org_color = get_organelle_color(org)

                # Add GT layer if available
                if gt_array is not None:
                    gt_layer = get_seg_layer(
                        # array=gt_array,
                        local_path=gt_path,
                        voxel_offset=gt_offset,
                        dimensions=first_dimensions,
                        voxel_size=gt_voxel_size,
                        color=org_color,
                    )
                    s.layers[f"gt_{org}"] = gt_layer
                    gt_layers.append(f"gt_{org}")
                    has_any_gt = True
                else:
                    print(
                        f"Could not load GT array for {org} from {gt_path}, showing predictions only"
                    )

                # Add prediction layer
                pred_layer = get_seg_layer(
                    # array=pred_array,
                    local_path=pred_path,
                    voxel_offset=pred_offset if pred_offset else gt_offset,
                    dimensions=first_dimensions,
                    voxel_size=pred_voxel_size if pred_voxel_size else gt_voxel_size,
                    color=org_color,
                )
                s.layers[f"pred_{org}"] = pred_layer
                pred_layers.append(f"pred_{org}")

            if (
                first_dimensions is None
                or first_offset is None
                or first_shape is None
                or first_voxel_size is None
            ):
                print("No valid organelle data could be loaded")
                return None

            # Set layout based on GT availability
            if has_any_gt:
                # Side-by-side view: GT + EM on left, Predictions + EM on right
                left_layers = []
                right_layers = []
                if em_loaded:
                    left_layers.append("em")
                    right_layers.append("em")
                left_layers.extend(gt_layers)
                right_layers.extend(pred_layers)
                s.layout = neuroglancer.row_layout(
                    [
                        neuroglancer.LayerGroupViewer(layers=left_layers),
                        neuroglancer.LayerGroupViewer(layers=right_layers),
                    ]
                )
                print("Using side-by-side layout: GT (left) vs Predictions (right)")
            else:
                # Single pane: just EM + Predictions
                single_pane_layers = []
                if em_loaded:
                    single_pane_layers.append("em")
                single_pane_layers.extend(pred_layers)
                s.layout = neuroglancer.row_layout(
                    [neuroglancer.LayerGroupViewer(layers=single_pane_layers)]
                )
                print("Using single pane layout (no GT available)")

            # Set initial position to center of volume (in world units / nm)
            # first_offset is in voxel units, so convert to world units
            center_position = [
                ((o * v) + (s * v) / 2) / v
                for o, v, s in zip(first_offset, first_voxel_size, first_shape)
            ]
            s.position = center_position

            # Set a reasonable cross-section scale based on the volume size
            # This ensures the field of view shows a reasonable portion of the crop
            max_dimension = max(first_shape[i] * first_voxel_size[i] for i in range(3))
            s.crossSectionScale = max_dimension / 1000  # Adjust this factor as needed

        url = viewer.get_viewer_url()
        print(f"Created viewer URL: {url}")
        print(neuroglancer.to_url(viewer.state))

        if return_viewer:
            return viewer
        else:
            return url

    except Exception as e:
        print(f"Error creating Neuroglancer viewer: {e}")
        import traceback

        traceback.print_exc()
        return None