"""
picasso.postprocess
~~~~~~~~~~~~~~~~~~~

Data analysis of localization lists.

:authors: Joerg Schnitzbauer, Maximilian Thomas Strauss, 2015-2018
:copyright: Copyright (c) 2015-2018 Jungmann Lab, MPI Biochemistry
"""

from __future__ import annotations

import itertools
import multiprocessing
from collections import OrderedDict
from collections.abc import Callable
from typing import Literal
from concurrent.futures import ThreadPoolExecutor as _ThreadPoolExecutor
from threading import Thread

import numba
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from scipy import interpolate
from scipy.optimize import curve_fit
from scipy.spatial import distance, KDTree
from tqdm import tqdm, trange

from . import lib, render, imageprocess, masking


def get_index_blocks(
    locs: pd.DataFrame,
    info: list[dict],
    size: float,
) -> tuple:
    """Split localizations into blocks of the given size. Used for fast
    localization indexing (e.g., for picking).

    Note: this function will be moved to ``picasso.lib`` in Picasso v0.11.0

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations.
    info : list of dicts
        Metadata of the localizations list.
    size : float
        Size of the blocks in camera pixels.

    Returns
    -------
    locs : pd.DataFrame
        Localizations in the specified blocks.
    size : float
        Size of the blocks in camera pixels.
    x_index : np.ndarray
        x indices of the localizations in the blocks.
    y_index : np.ndarray
        y indices of the localizations in the blocks.
    block_starts : np.ndarray
        Block start indices.
    block_ends : np.ndarray
        Block end indices.
    K : int
        Number of blocks in y direction.
    L : int
        Number of blocks in x direction.
    """
    locs = lib.ensure_sanity(locs, info)
    # Sort locs by indices
    x_index = np.uint32(locs["x"].to_numpy() / size)
    y_index = np.uint32(locs["y"].to_numpy() / size)
    sort_indices = np.lexsort([x_index, y_index])
    locs = locs.iloc[sort_indices]
    x_index = x_index[sort_indices]
    y_index = y_index[sort_indices]
    # Allocate block info arrays
    n_blocks_y, n_blocks_x = index_blocks_shape(info, size)
    block_starts = np.zeros((n_blocks_y, n_blocks_x), dtype=np.uint32)
    block_ends = np.zeros((n_blocks_y, n_blocks_x), dtype=np.uint32)
    K, L = block_starts.shape
    # Fill in block starts and ends
    thread = Thread(
        target=_fill_index_blocks,
        args=(block_starts, block_ends, x_index, y_index),
    )
    thread.start()
    thread.join()
    return locs, size, x_index, y_index, block_starts, block_ends, K, L


def index_blocks_shape(info: list[dict], size: float) -> tuple[int, int]:
    """Return the shape of the index grid, given the movie and grid
    sizes.

    Note: this function will be moved to ``picasso.lib`` in Picasso v0.11.0

    Parameters
    ----------
    info : list of dicts
        Metadata of the localizations list.
    size : float
        Size of the blocks.

    Returns
    -------
    n : tuple
        Number of blocks in y direction and x direction.
    """
    n_blocks_x = int(np.ceil(info[0]["Width"] / size))
    n_blocks_y = int(np.ceil(info[0]["Height"] / size))
    n = (n_blocks_y, n_blocks_x)
    return n


def get_block_locs_at(x: float, y: float, index_blocks: tuple) -> np.ndarray:
    """Return the localizations in the blocks around the given
    coordinates.

    Note: this function will be moved to ``picasso.lib`` in Picasso v0.11.0

    Parameters
    ----------
    x : float
        x coordinate.
    y : float
        y coordinate.
    index_blocks : tuple
        Index blocks information.

    Returns
    -------
    locs : np.ndarray
        Localizations in the blocks around the given coordinates.
    """
    locs, size, _, _, block_starts, block_ends, K, L = index_blocks
    x_index = np.uint32(x / size)
    y_index = np.uint32(y / size)
    indices = []
    for k in range(y_index - 1, y_index + 2):
        if 0 <= k < K:
            for ll in range(x_index - 1, x_index + 2):
                if 0 <= ll < L:
                    indices.append(
                        list(range(block_starts[k, ll], block_ends[k, ll]))
                    )
    indices = list(itertools.chain(*indices))
    return locs.iloc[indices]


@numba.jit(nopython=True, nogil=True)
def _fill_index_blocks(
    block_starts: np.ndarray,
    block_ends: np.ndarray,
    x_index: np.ndarray,
    y_index: np.ndarray,
) -> None:
    """Fill the block starts and ends arrays with the indices of
    localizations in the blocks.

    Note: this function will be moved to ``picasso.lib`` in Picasso v0.11.0
    """
    Y, X = block_starts.shape
    N = len(x_index)
    k = 0
    for i in range(Y):
        for j in range(X):
            k = _fill_index_block(
                block_starts, block_ends, N, x_index, y_index, i, j, k
            )


@numba.jit(nopython=True, nogil=True)
def _fill_index_block(
    block_starts: np.ndarray,
    block_ends: np.ndarray,
    N: int,
    x_index: np.ndarray,
    y_index: np.ndarray,
    i: int,
    j: int,
    k: int,
) -> int:
    """Fill the block starts and ends arrays for a single block.

    Note: this function will be moved to ``picasso.lib`` in Picasso v0.11.0
    """
    block_starts[i, j] = k
    while k < N and y_index[k] == i and x_index[k] == j:
        k += 1
    block_ends[i, j] = k
    return k


def picked_locs(
    locs: pd.DataFrame,
    info: list[dict],
    picks: list[tuple],
    pick_shape: Literal["Circle", "Rectangle", "Polygon", "Square"],
    pick_size: float = None,
    add_group: bool = True,
    index_blocks: tuple = None,
    callback: Callable[[int], None] | Literal["console"] | None = None,
) -> list[pd.DataFrame]:
    """Find picked localizations, i.e., localizations within the given
    regions of interest.

    Note: this function will be moved to ``picasso.lib`` in Picasso v0.11.0

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations.
    info : list of dicts
        Metadata of the localizations list.
    picks : list
        List of picks.
    pick_shape : {'Circle', 'Rectangle', 'Polygon', 'Square'}
        Shape of the pick.
    pick_size : float, optional
        Size of the pick. Radius for the circles, width for the
        rectangles, None for the polygons. Default is None.
    add_group : boolean, optional
        True if group id should be added to locs. Each pick will be
        assigned a different id. Default is True.
    index_blocks : tuple, optional
        Used only for circular picks. Precomputed index blocks for
        localizations, see  ``get_index_blocks``.If None, they will be
        calculated internally. Default is None.
    callback : function or "console" or None, optional
        Function to display progress. If "console", tqdm is used to
        display the progress. If None, no progress is displayed. Default
        is None.

    Returns
    -------
    picked_locs : list of pd.DataFrames
        List of pd.DataFrames, each containing locs from one pick.
    """
    _valid_shapes = ("Circle", "Rectangle", "Polygon", "Square")
    assert (
        pick_shape in _valid_shapes
    ), f"Invalid pick shape: {pick_shape}. Choose one of {_valid_shapes}."
    if len(picks):
        picked_locs = []
        if callback == "console":
            progress = tqdm(
                range(len(picks)),
                desc="Picking locs",
                unit="pick",
            )

        if pick_shape == "Circle":
            if index_blocks is None:
                index_blocks = get_index_blocks(locs, info, pick_size)
            locs_xy = index_blocks[0][["x", "y"]].to_numpy().T
            for i, pick in enumerate(picks):
                x, y = pick
                x_, y_ = int(x / pick_size), int(y / pick_size)
                block_locs_idx = _get_block_locs_at_numba(
                    x_,
                    y_,
                    index_blocks[4],
                    index_blocks[5],
                    index_blocks[6],
                    index_blocks[7],
                )
                block_locs = index_blocks[0].iloc[block_locs_idx]
                group_locs_idx = _locs_at_numba(
                    x, y, locs_xy[:, block_locs_idx], pick_size
                )
                group_locs = block_locs.iloc[group_locs_idx].copy()

                if add_group:
                    group_locs["group"] = i
                group_locs.sort_values(
                    by="frame",
                    kind="quicksort",
                    inplace=True,
                )
                picked_locs.append(group_locs)

                if callback == "console":
                    progress.update(1)
                elif callback is not None:
                    callback(i + 1)

        elif pick_shape == "Rectangle":
            for i, pick in enumerate(picks):
                (xs, ys), (xe, ye) = pick
                X, Y = lib.get_pick_rectangle_corners(
                    xs, ys, xe, ye, pick_size
                )
                x_min = min(X)
                x_max = max(X)
                y_min = min(Y)
                y_max = max(Y)
                group_locs = locs[locs["x"] > x_min]
                group_locs = group_locs[group_locs["x"] < x_max]
                group_locs = group_locs[group_locs["y"] > y_min]
                group_locs = group_locs[group_locs["y"] < y_max]
                group_locs = lib.locs_in_rectangle(group_locs, X, Y)
                # store rotated coordinates in x_rot and y_rot
                angle = 0.5 * np.pi - np.arctan2((ye - ys), (xe - xs))
                x_shifted = group_locs["x"] - xs
                y_shifted = group_locs["y"] - ys
                x_pick_rot = x_shifted * np.cos(angle) - y_shifted * np.sin(
                    angle
                )
                y_pick_rot = x_shifted * np.sin(angle) + y_shifted * np.cos(
                    angle
                )
                group_locs["x_pick_rot"] = x_pick_rot
                group_locs["y_pick_rot"] = y_pick_rot
                if add_group:
                    group_locs["group"] = i
                group_locs.sort_values(
                    by="frame", kind="quicksort", inplace=True
                )
                picked_locs.append(group_locs)

                if callback == "console":
                    progress.update(1)
                elif callback is not None:
                    callback(i + 1)

        elif pick_shape == "Polygon":
            for i, pick in enumerate(picks):
                X, Y = lib.get_pick_polygon_corners(pick)
                if X is None:
                    if callback == "console":
                        progress.update(1)
                    elif callback is not None:
                        callback(i + 1)
                    continue
                group_locs = locs[locs["x"] > min(X)]
                group_locs = group_locs[group_locs["x"] < max(X)]
                group_locs = group_locs[group_locs["y"] > min(Y)]
                group_locs = group_locs[group_locs["y"] < max(Y)]
                group_locs = lib.locs_in_polygon(group_locs, X, Y)
                if add_group:
                    group_locs["group"] = i
                group_locs.sort_values(
                    by="frame", kind="quicksort", inplace=True
                )
                picked_locs.append(group_locs)

                if callback == "console":
                    progress.update(1)
                elif callback is not None:
                    callback(i + 1)

        elif pick_shape == "Square":
            for i, pick in enumerate(picks):
                x, y = pick
                half_a = pick_size / 2
                x_min = x - half_a
                x_max = x + half_a
                y_min = y - half_a
                y_max = y + half_a
                group_locs = locs[locs["x"] > x_min]
                group_locs = group_locs[group_locs["x"] < x_max]
                group_locs = group_locs[group_locs["y"] > y_min]
                group_locs = group_locs[group_locs["y"] < y_max]
                if add_group:
                    group_locs["group"] = i
                group_locs.sort_values(
                    by="frame", kind="quicksort", inplace=True
                )
                picked_locs.append(group_locs)

                if callback == "console":
                    progress.update(1)
                elif callback is not None:
                    callback(i + 1)

        return picked_locs


def pick_similar(
    locs: pd.DataFrame,
    info: list[dict],
    picks: list[tuple],
    d: float,
    std_range: float = 2.0,
    index_blocks: tuple = None,
) -> list[tuple]:
    """Find similar picks based on the number of localizations and
    RMSD. Only implemented for circular picks.

    Mean number of localizations and RMSD in picks are calculated and
    the allowed range is defined by the given standard deviation range
    (``std_range``). Only picks with number of localizations and RMSD
    within these ranges are returned.

    Takes the grid of overlapping picks of the given size (defined by
    ``x``, ``y_shift`` and ``y_base``) and shifts each pick towards the
    center of mass of the localizations within the pick. If the picked
    localizations have the required number of localizations and the
    RMSD, it is added to the output list (``x_similar`` and
    ``y_similar``).

    This function calls ``_pick_similar`` which is implemented in numba
    for speed.

    Note: this function will be moved to ``picasso.lib`` in Picasso v0.11.0

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations.
    info : list of dicts
        Metadata of the localizations.
    picks : list
        List of picks (x, y) coordinates.
    d : float
        Pick diameter (in camera pixels).
    std_range : float, optional
        Standard deviation range for picking similar localizations.
        Default is 2.0.
    index_blocks : tuple, optional
        Precomputed index blocks for localizations, see
        ``get_index_blocks``. If None, they will be calculated
        internally. Default is None.

    Returns
    -------
    new_picks : list of tuples
        List of (x, y) coordinates of similar picks.
    """
    r = d / 2
    d2 = d**2
    # extract n_locs and rmsd from current picks
    if index_blocks is None:
        index_blocks = get_index_blocks(locs, info, r)
    locs_xy = index_blocks[0][["x", "y"]].to_numpy().T
    n_locs = []
    rmsd = []
    for i, pick in enumerate(picks):
        x, y = pick
        block_locs_xy = get_block_locs_at_numba(
            int(x / r),
            int(y / r),
            locs_xy,
            index_blocks[4],
            index_blocks[5],
            index_blocks[6],
            index_blocks[7],
        )
        pick_locs_xy = locs_at_numba(x, y, block_locs_xy, r)
        n_locs.append(pick_locs_xy.shape[1])
        rmsd.append(rmsd_at_com(pick_locs_xy))

    # calculate min and max n_locs and rmsd for picking similar
    mean_n_locs = np.mean(n_locs)
    mean_rmsd = np.mean(rmsd)
    std_n_locs = np.std(n_locs)
    std_rmsd = np.std(rmsd)
    min_n_locs = max(2, mean_n_locs - std_range * std_n_locs)
    max_n_locs = mean_n_locs + std_range * std_n_locs
    min_rmsd = mean_rmsd - std_range * std_rmsd
    max_rmsd = mean_rmsd + std_range * std_rmsd

    # x, y coordinates of found regions:
    x_similar = np.array([_[0] for _ in picks])
    y_similar = np.array([_[1] for _ in picks])

    # preparations for grid search
    x_range = np.arange(d / 2, info[0]["Width"], np.sqrt(3) * d / 2)
    y_range_base = np.arange(d / 2, info[0]["Height"] - d / 2, d)
    y_range_shift = y_range_base + d / 2
    locs_temp, size, _, _, block_starts, block_ends, K, L = index_blocks
    locs_xy = np.stack((locs_temp.x, locs_temp.y))
    x_r = np.uint64(x_range / size)
    y_r1 = np.uint64(y_range_shift / size)
    y_r2 = np.uint64(y_range_base / size)
    # pick similar
    x_similar, y_similar = _pick_similar(
        x_range,
        y_range_shift,
        y_range_base,
        min_n_locs,
        max_n_locs,
        min_rmsd,
        max_rmsd,
        x_r,
        y_r1,
        y_r2,
        locs_xy,
        block_starts,
        block_ends,
        K,
        L,
        x_similar,
        y_similar,
        r,
        d2,
    )
    new_picks = list(zip(x_similar, y_similar))
    return new_picks


@numba.jit(nopython=True, nogil=True, cache=True)
def _pick_similar(
    x: np.ndarray,
    y_shift: np.ndarray,
    y_base: np.ndarray,
    min_n_locs: int,
    max_n_locs: int,
    min_rmsd: float,
    max_rmsd: float,
    x_r: np.ndarray,
    y_r1: np.ndarray,
    y_r2: np.ndarray,
    locs_xy: np.ndarray,
    block_starts: np.ndarray,
    block_ends: np.ndarray,
    K: int,
    L: int,
    x_similar: np.ndarray,
    y_similar: np.ndarray,
    r: float,
    d2: float,
) -> tuple[np.ndarray, np.ndarray]:
    """Find similar picks based on the number of localizations and
    RMSD. Only implemented for circular picks.

    Takes the grid of overlapping picks of the given size (defined by
    ``x``, ``y_shift`` and ``y_base``) and shifts each pick towards the
    center of mass of the localizations within the pick. If the picked
    localizations have the required number of localizations and the
    RMSD, it is added to the output list (``x_similar`` and
    ``y_similar``).

    This function is implemented in numba for speed, called from
    ``pick_similar``. See that function for more user-friendly
    interface.

    Note: this function will be moved to ``picasso.lib`` in Picasso v0.11.0

    Parameters
    ----------
    x : np.ndarray
        x coordinates of the picks.
    y_shift : np.ndarray
        y coordinates of the picks, shifted for odd columns.
    y_base : np.ndarray
        y coordinates of the picks, not shifted.
    min_n_locs, max_n_locs : int
        Minimum and maximum number of localizations in the pick.
    min_rmsd, max_rmsd : float
        Minimum and maximum RMSD for the pick.
    x_r, y_r1, y_r2 : np.ndarray
        x and y ranges for the picks.
    locs_xy : np.ndarray
        Localizations in the blocks.
    block_starts : np.ndarray
        Block start indices.
    block_ends : np.ndarray
        Block end indices.
    K, L : int
        Number of blocks in y and x direction.
    x_similar, y_similar : np.ndarray
        Arrays to store the x and y coordinates of the similar picks.
    r : float
        Radius for the picks.
    d2 : float
        Squared distance threshold for the picks.

    Returns
    -------
    x_similar, y_similar : np.ndarray
        Arrays with the x and y coordinates of the similar picks.
    """
    for i, x_grid in enumerate(x):
        x_range = x_r[i]
        # y_grid is shifted for odd columns
        if i % 2:
            y = y_shift
            y_r = y_r1
        else:
            y = y_base
            y_r = y_r2
        for j, y_grid in enumerate(y):
            y_range = y_r[j]
            n_block_locs = n_block_locs_at(
                x_range,
                y_range,
                K,
                L,
                block_starts,
                block_ends,
            )
            if n_block_locs >= min_n_locs:
                block_locs_xy = get_block_locs_at_numba(
                    x_range,
                    y_range,
                    locs_xy,
                    block_starts,
                    block_ends,
                    K,
                    L,
                )
                picked_locs_xy = locs_at_numba(
                    x_grid, y_grid, block_locs_xy, r
                )
                if picked_locs_xy.shape[1] > 1:
                    # Move to COM peak
                    x_test_old = x_grid
                    y_test_old = y_grid
                    x_test = np.mean(picked_locs_xy[0])
                    y_test = np.mean(picked_locs_xy[1])
                    count = 0
                    while (
                        np.abs(x_test - x_test_old) > 1e-3
                        or np.abs(y_test - y_test_old) > 1e-3
                    ):
                        count += 1
                        # skip the locs if the loop is too long
                        if count > 500:
                            break
                        x_test_old = x_test
                        y_test_old = y_test
                        picked_locs_xy = locs_at_numba(
                            x_test, y_test, block_locs_xy, r
                        )
                        if picked_locs_xy.shape[1] > 1:
                            x_test = np.mean(picked_locs_xy[0])
                            y_test = np.mean(picked_locs_xy[1])
                        else:
                            break
                    if np.all(
                        (x_similar - x_test) ** 2 + (y_similar - y_test) ** 2
                        > d2
                    ):
                        if min_n_locs <= picked_locs_xy.shape[1] <= max_n_locs:
                            if (
                                min_rmsd
                                <= rmsd_at_com(picked_locs_xy)
                                <= max_rmsd
                            ):
                                x_similar = np.append(x_similar, x_test)
                                y_similar = np.append(y_similar, y_test)
    return x_similar, y_similar


@numba.jit(nopython=True, nogil=True)
def n_block_locs_at(
    x_range: int,
    y_range: int,
    K: int,
    L: int,
    block_starts: np.ndarray,
    block_ends: np.ndarray,
) -> int:
    """Return the number of localizations in the blocks around the
    given coordinates.

    Note: this function will be moved to ``picasso.lib`` in Picasso v0.11.0
    """
    step = 0
    for k in range(y_range - 1, y_range + 2):
        if 0 < k < K:
            for ll in range(x_range - 1, x_range + 2):
                if 0 < ll < L:
                    if step == 0:
                        n_block_locs = np.uint32(
                            block_ends[k][ll] - block_starts[k][ll]
                        )
                        step = 1
                    else:
                        n_block_locs += np.uint32(
                            block_ends[k][ll] - block_starts[k][ll]
                        )
    return n_block_locs


@numba.jit(nopython=True, nogil=True, cache=True)
def _get_block_locs_at_numba(
    x_index: int,
    y_index: int,
    block_starts: np.ndarray,
    block_ends: np.ndarray,
    K: int,
    L: int,
) -> np.ndarray:
    """Numba implementation of ``get_block_locs_at``. Return the indices
    of localizations in the blocks around the given coordinates.

    Note: this function will be moved to ``picasso.lib`` in Picasso v0.11.0
    """
    step = 0
    for k in range(y_index - 1, y_index + 2):
        if 0 <= k < K:
            for ll in range(x_index - 1, x_index + 2):
                if 0 <= ll < L:
                    if block_ends[k, ll] - block_starts[k, ll] > 0:
                        # numba does not work if you attach concatenate
                        # to an empty list so the first step is
                        # different
                        if step == 0:
                            indices = np.arange(
                                float(block_starts[k, ll]),
                                float(block_ends[k, ll]),
                                dtype=np.uint32,
                            )
                            step = 1
                        else:
                            indices = np.concatenate(
                                (
                                    indices,
                                    np.arange(
                                        float(block_starts[k, ll]),
                                        float(block_ends[k, ll]),
                                        dtype=np.uint32,
                                    ),
                                )
                            )
    return indices


@numba.jit(nopython=True, nogil=True, cache=True)
def get_block_locs_at_numba(
    x_index: int,
    y_index: int,
    locs_xy: np.ndarray,
    block_starts: np.ndarray,
    block_ends: np.ndarray,
    K: int,
    L: int,
) -> np.ndarray:
    """Numba implementation of ``get_block_locs_at. Return the
    localizations in the blocks around the given coordinates.

    Note: this function will be moved to ``picasso.lib`` in Picasso v0.11.0
    """
    indices = _get_block_locs_at_numba(
        x_index,
        y_index,
        block_starts,
        block_ends,
        K,
        L,
    )
    return locs_xy[:, indices]


@numba.jit(nopython=True, nogil=True, cache=True)
def _locs_at_numba(
    x: float,
    y: float,
    locs_xy: np.ndarray,
    r: float,
) -> np.ndarray:
    """Numba implementation of ``lib.locs_at``. Return the indices of
    localizations at the given coordinates within radius ``r``.

    Note: this function will be moved to ``picasso.lib`` in Picasso v0.11.0"""
    dx = locs_xy[0] - x
    dy = locs_xy[1] - y
    r2 = r**2
    # print(x, y, r, dx, dy, )
    is_picked = dx**2 + dy**2 < r2
    return is_picked


@numba.jit(nopython=True, nogil=True, cache=True)
def locs_at_numba(
    x: float,
    y: float,
    locs_xy: np.ndarray,
    r: float,
) -> np.ndarray:
    """Numba implementation of ``lib.locs_at``. Return the localizations
    at the given coordinates within radius ``r``.

    Note: this function will be moved to ``picasso.lib`` in Picasso v0.11.0
    """
    is_picked = _locs_at_numba(x, y, locs_xy, r)
    return locs_xy[:, is_picked]


@numba.jit(nopython=True, nogil=True)
def rmsd_at_com(locs_xy: np.ndarray) -> float:
    """Calculate the RMSD of the localizations at the center of mass
    (COM) of the localizations.

    Note: this function will be moved to ``picasso.lib`` in Picasso v0.11.0"""
    com_x = np.mean(locs_xy[0])
    com_y = np.mean(locs_xy[1])
    return np.sqrt(
        np.mean((locs_xy[0] - com_x) ** 2 + (locs_xy[1] - com_y) ** 2)
    )


@numba.jit(nopython=True, nogil=True)
def _distance_histogram(
    x: np.ndarray,
    y: np.ndarray,
    bin_size: float,
    r_max: float,
    x_index: np.ndarray,
    y_index: np.ndarray,
    block_starts: np.ndarray,
    block_ends: np.ndarray,
    start: int,
    chunk: int,
) -> np.ndarray:
    """Calculate the distance histogram for a chunk of localizations."""
    dh_len = np.uint32(r_max / bin_size)
    dh = np.zeros(dh_len, dtype=np.uint32)
    r_max_2 = r_max**2
    K, L = block_starts.shape
    end = min(start + chunk, len(x))
    for i in range(start, end):
        xi = x[i]
        yi = y[i]
        ki = y_index[i]
        li = x_index[i]
        for k in range(ki, ki + 2):
            if k < K:
                for ll in range(li, li + 2):
                    if ll < L:
                        for j in range(block_starts[k, ll], block_ends[k, ll]):
                            if j > i:
                                dx2 = (xi - x[j]) ** 2
                                if dx2 < r_max_2:
                                    dy2 = (yi - y[j]) ** 2
                                    if dy2 < r_max_2:
                                        d = np.sqrt(dx2 + dy2)
                                        if d < r_max:
                                            bin = np.uint32(d / bin_size)
                                            if bin < dh_len:
                                                dh[bin] += 1
    return dh


def distance_histogram(
    locs: pd.DataFrame,
    info: list[dict],
    bin_size: float,
    r_max: float,
) -> np.ndarray:
    """Calculate the distance histogram for the given localizations,
    i.e., the pairwise distances between localizations.

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations.
    info : list[dict]
        Metadata of the localizations.
    bin_size : float
        Size of the bins for the histogram.
    r_max : float
        Maximum distance probed in the histogram.

    Returns
    -------
    dh : np.ndarray
        Distance histogram.
    """
    locs, size, x_index, y_index, b_starts, b_ends, K, L = get_index_blocks(
        locs, info, r_max
    )
    N = len(locs)
    n_threads = min(
        60, max(1, int(0.75 * multiprocessing.cpu_count()))
    )  # Python crashes when using >64 cores
    chunk = int(N / n_threads)
    starts = range(0, N, chunk)
    args = [
        (
            locs["x"].to_numpy(),
            locs["y"].to_numpy(),
            bin_size,
            r_max,
            x_index,
            y_index,
            b_starts,
            b_ends,
            start,
            chunk,
        )
        for start in starts
    ]
    with _ThreadPoolExecutor() as executor:
        futures = [executor.submit(_distance_histogram, *_) for _ in args]
    results = [future.result() for future in futures]
    dh = np.sum(results, axis=0)
    return dh


def nena(
    locs: pd.DataFrame,
    info: any = None,
    callback: Callable[[int], None] | None = None,
) -> tuple[dict, float]:
    """Calculate NeNA - experimental estimate of localization
    precision. Please refer to the original paper for details:
    Endesfelder, et al. Histochemistry and Cell Biology, 2014.
    DOI: 10.1007/s00418-014-1192-3.

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations.
    info : any
        Metadata of the localizations. Not used.
    callback : function or None
        Function to display progress (takes in an integer). If None, no
        progress is displayed.

    Returns
    -------
    result : dict
        Data on the results, including the distances probed, best fit
        and fitted parameters.
    s : float
        Estimated localization precision.
    """
    bin_centers, dnfl = next_frame_neighbor_distance_histogram(locs, callback)

    def func(d, delta_a, s, ac, dc, sc):
        a = ac + delta_a  # make sure a >= ac
        p_single = a * (d / (2 * s**2)) * np.exp(-(d**2) / (4 * s**2))
        p_short = (
            ac
            / (sc * np.sqrt(2 * np.pi))
            * np.exp(-0.5 * ((d - dc) / sc) ** 2)
        )
        return p_single + p_short

    area = np.trapezoid(dnfl, bin_centers)
    median_lp = np.mean([np.median(locs["lpx"]), np.median(locs["lpy"])])
    p0 = [0.8 * area, median_lp, 0.1 * area, 2 * median_lp, median_lp]
    bounds = ([0, 0, 0, 0, 0], [np.inf, np.inf, np.inf, np.inf, np.inf])
    popt, _ = curve_fit(func, bin_centers, dnfl, p0=p0, bounds=bounds)
    s = popt[1]  # NeNA
    result = {
        "d": bin_centers,  # distances probed
        "data": dnfl,
        "best_fit": func(bin_centers, *popt),
        "best_values": {
            "delta_a": popt[0],
            "s": popt[1],
            "ac": popt[2],
            "dc": popt[3],
            "sc": popt[4],
        },
    }
    return result, s


def next_frame_neighbor_distance_histogram(
    locs: pd.DataFrame,
    callback: Callable[[int], None] | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    """Calculate the next frame neighbor distance histogram (NFNDH).

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations.
    callback : function or None
        Function to display progress. If None, no progress is displayed.

    Returns
    -------
    bin_centers : np.ndarray
        Centers of the bins for the histogram.
    dnfl : np.ndarray
        Distance histogram of next frame neighbors.
    """
    locs.sort_values(kind="quicksort", by="frame", inplace=True)
    frame = locs["frame"].to_numpy()
    x = locs["x"].to_numpy()
    y = locs["y"].to_numpy()
    if "group" in locs.columns:
        group = locs["group"].to_numpy()
    else:
        group = np.zeros(len(locs), dtype=np.int32)
    bin_size = 0.001
    d_max = 1.0
    return _nfndh(frame, x, y, group, d_max, bin_size, callback)


def _nfndh(
    frame: np.ndarray,
    x: np.ndarray,
    y: np.ndarray,
    group: np.ndarray,
    d_max: float,
    bin_size: float,
    callback: Callable[[int], None] | None = None,
) -> tuple[np.ndarray, np.ndarray]:
    """Calculate the next frame neighbor distance histogram (NFNDH)."""
    N = len(frame)
    bins = np.arange(0, d_max, bin_size)
    dnfl = np.zeros(len(bins))
    one_percent = int(N / 100)
    starts = one_percent * np.arange(100)
    for k, start in enumerate(starts):
        for i in range(start, start + one_percent):
            _fill_dnfl(N, frame, x, y, group, i, d_max, dnfl, bin_size)
        if callback is not None:
            callback(k + 1)
    bin_centers = bins + bin_size / 2
    return bin_centers, dnfl


@numba.jit(nopython=True)
def _fill_dnfl(
    N: int,
    frame: np.ndarray,
    x: np.ndarray,
    y: np.ndarray,
    group: np.ndarray,
    i: int,
    d_max: float,
    dnfl: np.ndarray,
    bin_size: float,
) -> None:
    """Fill the next frame neighbor distance histogram (NFNDH) for a
    single localization."""
    frame_i = frame[i]
    x_i = x[i]
    y_i = y[i]
    group_i = group[i]
    min_frame = frame_i + 1
    for min_index in range(i + 1, N):
        if frame[min_index] >= min_frame:
            break
    max_frame = frame_i + 1
    for max_index in range(min_index, N):
        if frame[max_index] > max_frame:
            break
    d_max_2 = d_max**2
    for j in range(min_index, max_index):
        if group[j] == group_i:
            dx2 = (x_i - x[j]) ** 2
            if dx2 <= d_max_2:
                dy2 = (y_i - y[j]) ** 2
                if dy2 <= d_max_2:
                    d = np.sqrt(dx2 + dy2)
                    if d <= d_max:
                        bin = int(d / bin_size)
                        dnfl[bin] += 1


def frc(
    locs: pd.DataFrame,
    info: list[dict],
    viewport: tuple[tuple[float, float], tuple[float, float]],
    *,
    random_seed: int = 42,
) -> dict:
    """Calculate the Fourier Ring Correlation (FRC) resolution.

    See Nieuwenhuizen et al., Nat. Methods 10, 557–562 (2013).

    Parameters
    ----------
    locs : pd.DataFrame
        Localization list.
    info : list of dicts
        Metadata of the localizations list.
    viewport : tuple of floats
        Viewport ((y_min, x_min), (y_max, x_max)) for rendering the
        images. Note that the origin of the image is in the top-left
        corner.
    random_seed : int, optional
        Random seed for splitting the data into halves. Default is 42.

    Returns
    -------
    frc_result : dict
        Dictionary with keys "frc_curve", "frc_curve_smooth",
        "frequencies" (for spatial frequencies probed (nm^-1)),
        "resolution" (estimated resolution in nm) and "images"
        (2 grayscale images rendered and masked).
    """
    pixelsize = lib.get_from_metadata(info, "Pixelsize", raise_error=True)
    lp = nena(locs, info)[1]
    # correct for the viewport to be square
    viewport_width = viewport[1][1] - viewport[0][1]
    viewport_height = viewport[1][0] - viewport[0][0]
    if viewport_width < viewport_height:
        y_center = 0.5 * (viewport[0][0] + viewport[1][0])
        y_min = y_center - viewport_width / 2
        y_max = y_center + viewport_width / 2
        viewport = ((y_min, viewport[0][1]), (y_max, viewport[1][1]))
    elif viewport_height < viewport_width:
        x_center = 0.5 * (viewport[0][1] + viewport[1][1])
        x_min = x_center - viewport_height / 2
        x_max = x_center + viewport_height / 2
        viewport = ((viewport[0][0], x_min), (viewport[1][0], x_max))

    # select the locs within the viewport
    (y_min, x_min), (y_max, x_max) = viewport
    in_view = (
        (locs["x"] > x_min)
        & (locs["y"] > y_min)
        & (locs["x"] < x_max)
        & (locs["y"] < y_max)
    )
    locs = locs.loc[in_view]

    np.random.seed(random_seed)
    # split locs randomly into two halves
    r_idx = np.random.permutation(len(locs))
    locs1 = locs.iloc[r_idx[: len(r_idx) // 2]]
    locs2 = locs.iloc[r_idx[len(r_idx) // 2 :]]
    # run FRC
    frc_curve, frc_curve_smooth, frequencies, resolution, images = _frc(
        locs1, locs2, pixelsize, lp, viewport
    )
    # smooth the FRC curve and summarize findings
    frc_result = {
        "frc_curve": frc_curve,
        "frc_curve_smooth": frc_curve_smooth,
        "frequencies": frequencies,
        "resolution": resolution,
        "images": images,
    }
    return frc_result


def _frc(
    locs1: pd.DataFrame,
    locs2: pd.DataFrame,
    pixelsize: float,
    lp: float,
    viewport: tuple[tuple[float, float], tuple[float, float]],
) -> tuple[
    np.ndarray,
    np.ndarray,
    np.ndarray,
    float | None,
    tuple[np.ndarray, np.ndarray],
]:
    """Calculate the Fourier Ring Correlation (FRC) resolution once.

    Generates images from the two sets of localizations and calculates
    the FRC curve and resolution (at 1/7 threshold).

    See Nieuwenhuizen et al., Nat. Methods 10, 557–562 (2013).

    Note: do not use this function directly, use ``frc`` instead.

    Parameters
    ----------
    locs1, locs2 : pd.DataFrame
        Localization lists already split to render images.
    pixelsize: float
        Camera pixel size in nm.
    lp : float
        Average localization precision (NeNA), used for bin size of the
        rendered images (binsize = lp / 2).
    viewport : tuple of floats
        Viewport ((y_min, x_min), (y_max, x_max)) for rendering the
        images. Note that the origin of the image is in the top-left
        corner.

    Returns
    -------
    frc_curve : np.ndarray
        FRC curve.
    frc_curve_smooth : np.ndarray
        Smoothed FRC curve (LOESS).
    frequencies : np.ndarray
        Spatial frequencies corresponding to the FRC curve (nm^-1).
    resolution : float or None
        Estimated resolution in nm, given the 1/7 threshold. None if
        resolution could not be determined.
    images : tuple of 2 np.ndarrays
        2 grayscale images used for calculating FRC. Already masked.
    """
    # render images
    binsize = lp / 2
    oversampling = 1 / binsize
    im1 = render.render(locs1, None, oversampling, viewport, None)[1]
    im2 = render.render(locs2, None, oversampling, viewport, None)[1]

    # ensure the images are odd-sized and mask them (tukey)
    if im1.shape[0] % 2 == 0:
        im1 = im1[:-1, :-1]
        im2 = im2[:-1, :-1]
    mask = masking.threshold_tukey(im1)
    im1 *= mask
    im2 *= mask

    # fourier
    f1 = np.fft.fftshift(np.fft.fft2(im1))
    f2 = np.fft.fftshift(np.fft.fft2(im2))

    # frc curve
    frc_num = np.real(imageprocess.radial_sum(f1 * np.conj(f2)))
    frc_denom = np.sqrt(
        np.abs(
            imageprocess.radial_sum(np.abs(f1) ** 2)
            * imageprocess.radial_sum(np.abs(f2) ** 2)
        )
    )
    with np.errstate(divide="ignore", invalid="ignore"):
        frc_curve = frc_num / frc_denom
    frc_curve[np.isnan(frc_curve)] = 0

    # smooth the frc curve
    sspan = max(int(np.ceil(int(im1.shape[0] / 2) / 20)), 5)
    frc_curve_smooth = masking.loess_smooth(frc_curve, sspan)

    # find the frequencies (q) and resolution
    frequencies = (
        np.arange(len(frc_curve)) / im1.shape[0] / (pixelsize * binsize)
    )
    threshold = 1 / 7  # resolution at 1/7 threshold
    resolution = None
    for i in range(1, len(frc_curve_smooth)):
        if (
            frc_curve_smooth[i - 1] >= threshold
            and frc_curve_smooth[i] < threshold
        ):
            # linear interpolation
            f1 = frequencies[i - 1]
            f2 = frequencies[i]
            r1 = frc_curve_smooth[i - 1]
            r2 = frc_curve_smooth[i]
            f_res = f1 + (threshold - r1) * (f2 - f1) / (r2 - r1)
            resolution = 1 / f_res  # in nm
            break
    return frc_curve, frc_curve_smooth, frequencies, resolution, (im1, im2)


def pair_correlation(
    locs: pd.DataFrame,
    info: list[dict],
    bin_size: float,
    r_max: float,
) -> tuple[np.ndarray, np.ndarray]:
    """Calculate the pair correlation function for the given
    localizations.

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations.
    info : list of dicts
        Metadata of the localizations.
    bin_size : float
        Size of the bins for the histogram.
    r_max : float
        Maximum distance for the histogram.

    Returns
    -------
    bins_lower : np.ndarray
        Lower bounds of the bins for the histogram.
    pc : np.ndarray
        Pair correlation function.
    """
    dh = distance_histogram(locs, info, bin_size, r_max)
    # Start with r-> otherwise area will be 0
    bins_lower = np.arange(bin_size, r_max + bin_size, bin_size)

    if bins_lower.shape[0] > dh.shape[0]:
        bins_lower = bins_lower[:-1]
    area = np.pi * bin_size * (2 * bins_lower + bin_size)
    pc = dh / area
    return bins_lower, pc


@numba.jit(nopython=True, nogil=True)
def _local_density(
    x: np.ndarray,
    y: np.ndarray,
    radius: float,
    x_index: np.ndarray,
    y_index: np.ndarray,
    block_starts: np.ndarray,
    block_ends: np.ndarray,
    start: int,
    chunk: int,
) -> np.ndarray:
    """Calculate densities in blocks around each localization."""
    N = len(x)
    r2 = radius**2
    end = min(start + chunk, N)
    density = np.zeros(N, dtype=np.uint32)
    for i in range(start, end):
        yi = y[i]
        xi = x[i]
        ki = y_index[i]
        li = x_index[i]
        di = 0
        for k in range(ki - 1, ki + 2):
            for ll in range(li - 1, li + 2):
                j_min = block_starts[k, ll]
                j_max = block_ends[k, ll]
                for j in range(j_min, j_max):
                    dx2 = (xi - x[j]) ** 2
                    if dx2 < r2:
                        dy2 = (yi - y[j]) ** 2
                        if dy2 < r2:
                            d2 = dx2 + dy2
                            if d2 < r2:
                                di += 1
        density[i] = di
    return density


def compute_local_density(
    locs: pd.DataFrame,
    info: list[dict],
    radius: float,
) -> pd.DataFrame:
    """Compute the local density of localizations in blocks around
    each localization.

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations.
    info : list of dicts
        Metadata of the localizations.
    radius : float
        Radius for local density computation.

    Returns
    -------
    locs : pd.DataFrame
        Localizations with added 'density' field/column.
    """
    locs, size, x_index, y_index, block_starts, block_ends, K, L = (
        get_index_blocks(locs, info, radius)
    )
    N = len(locs)
    n_threads = min(
        60, max(1, int(0.75 * multiprocessing.cpu_count()))
    )  # Python crashes when using >64 cores
    chunk = int(N / n_threads)
    starts = range(0, N, chunk)
    args = [
        (
            locs["x"].to_numpy(),
            locs["y"].to_numpy(),
            radius,
            x_index,
            y_index,
            block_starts,
            block_ends,
            start,
            chunk,
        )
        for start in starts
    ]
    with _ThreadPoolExecutor() as executor:
        futures = [executor.submit(_local_density, *_) for _ in args]
    density = np.sum([future.result() for future in futures], axis=0)
    locs["density"] = density
    return locs


def compute_dark_times(
    locs: pd.DataFrame,
    group: np.ndarray | None = None,
) -> pd.DataFrame:
    """Compute dark time for each binding event.

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations that were linked, i.e., binding events.
    group : np.ndarray, optional
        Grouping array for binding events. If None, all binding events
        are considered to be in the same group.

    Returns
    -------
    locs : pd.DataFrame
        Binding events with added 'dark' field/column, which contains
        the dark time for each binding event. If a binding event is not
        followed by another binding event in the same group, the dark
        time is set to -1.
    """
    if "len" not in locs.columns:
        raise AttributeError(
            "Length not found. Please link localizations first."
        )
    dark = dark_times(locs, group)
    locs["dark"] = np.int32(dark)
    locs = locs[locs.dark != -1]
    return locs


def dark_times(
    locs: pd.DataFrame,
    group: np.ndarray | None = None,
) -> np.ndarray:
    """Calculate dark times for each binding event.

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations that were linked, i.e., binding events.
    group : np.ndarray, optional
        Grouping array for binding events. If None, all binding events
        are considered to be in the same group.

    Returns
    -------
    dark : np.ndarray
        Array of dark times for each binding event. If a binding event
        is not followed by another binding event in the same group, the
        dark time is set to -1.
    """
    frame = locs["frame"].to_numpy()
    lens = locs["len"].to_numpy()
    last_frame = frame + lens - 1
    if group is None:
        if "group" in locs.columns:
            group = locs["group"].to_numpy()
        else:
            group = np.zeros(len(locs))
    dark = _dark_times(frame, group, last_frame)
    return dark


@numba.jit(nopython=True)
def _dark_times(
    frame: np.ndarray,
    group: np.ndarray,
    last_frame: np.ndarray,
) -> np.ndarray:
    """Calculate dark times for each binding event."""
    N = len(frame)
    max_frame = frame.max()
    dark = max_frame * np.ones(len(frame), dtype=np.int32)
    for i in range(N):
        for j in range(N):
            if (group[i] == group[j]) and (i != j):
                dark_ij = frame[i] - last_frame[j]
                if (dark_ij > 0) and (dark_ij < dark[i]):
                    dark[i] = dark_ij
    for i in range(N):
        if dark[i] == max_frame:
            dark[i] = -1
    return dark


def link(
    locs: pd.DataFrame,
    info: list[dict],
    r_max: float = 0.05,
    max_dark_time: int = 3,
    combine_mode: Literal["average", "refit"] = "average",
    remove_ambiguous_lengths: bool = True,
) -> pd.DataFrame:
    """Link localizations, i.e., group them into binding events based
    on their spatiotemporal proximity.

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations.
    info : list of dicts
        Metadata of the localizations.
    r_max : float, optional
        Maximum distance for linking localizations. Default is 0.05.
    max_dark_time : int, optional
        Maximum dark time for linking localizations. Default is 1.
    combine_mode : {'average', 'refit'}, optional
        Mode for combining linked localizations. 'average' calculates
        the average position and properties of the linked localizations,
        while 'refit' would refit the linked localizations to a model.
        'refit' is not implemented yet. Default is 'average'.
    remove_ambiguous_lengths : bool, optional
        If True, removes linked localizations with ambiguous lengths,
        i.e., localizations that are linked to multiple binding events
        with different lengths. Default is True.

    Returns
    -------
    linked_locs : pd.DataFrame
        Linked localizations, i.e., binding events with their
        properties.
    """
    if len(locs) == 0:  # special case of an empty localization list
        linked_locs = locs.copy()
        if "frame" in locs.columns:
            linked_locs["len"] = np.array([], dtype=np.int32)
            linked_locs["n"] = np.array([], dtype=np.int32)
        if "photons" in locs.columns:
            linked_locs["photon_rate"] = np.array([], dtype=np.float32)
    else:
        locs = locs.sort_values(kind="quicksort", by="frame")
        if "group" in locs.columns:
            group = locs["group"].to_numpy()
        else:
            group = np.zeros(len(locs), dtype=np.int32)
        frame = locs["frame"].to_numpy()
        x = locs["x"].to_numpy()
        y = locs["y"].to_numpy()
        link_group = get_link_groups(frame, x, y, r_max, max_dark_time, group)
        if combine_mode == "average":
            linked_locs = link_loc_groups(
                locs,
                info,
                link_group,
                remove_ambiguous_lengths=remove_ambiguous_lengths,
            )
        elif combine_mode == "refit":
            raise NotImplementedError(
                "Refit mode is not implemented yet. Please use 'average' mode."
            )
    return linked_locs


# def weighted_variance(locs):
#     n = len(locs)
#     w = locs.photons
#     x = locs.x
#     y = locs.y
#     xWbarx = np.average(locs.x, weights=w)
#     xWbary = np.average(locs.y, weights=w)
#     wbarx = np.mean(locs.lpx)
#     wbary = np.mean(locs.lpy)
#     variance_x = (
#         n
#         / ((n - 1) * sum(w) ** 2)
#         * (
#             sum((w * x - wbarx * xWbarx) ** 2)
#             - 2 * xWbarx * sum((w - wbarx) * (w * x - wbarx * xWbarx))
#             + xWbarx**2 * sum((w - wbarx) ** 2)
#         )
#     )
#     variance_y = (
#         n
#         / ((n - 1) * sum(w) ** 2)
#         * (
#             sum((w * y - wbary * xWbary) ** 2)
#             - 2 * xWbary * sum((w - wbary) * (w * y - wbary * xWbary))
#             + xWbary**2 * sum((w - wbary) ** 2)
#         )
#     )
#     return variance_x, variance_y


# Combine localizations: calculate the properties of the group
def cluster_combine(locs: pd.DataFrame) -> pd.DataFrame:
    """Combine localizations into clusters and calculate their
    properties such as center of mass, standard deviation, and number
    of localizations in each cluster.

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations with 'group' and 'cluster' columns.

    Returns
    -------
    combined_locs : pd.DataFrame
        Combined localizations with calculated properties for each
        cluster.
    """
    combined_locs = []
    if "z" in locs.columns:
        for group in tqdm(np.unique(locs["group"])):
            temp = locs[locs["group"] == group]
            cluster = np.unique(temp["cluster"].to_numpy())
            n_cluster = len(cluster)
            mean_frame = np.zeros(n_cluster)
            std_frame = np.zeros(n_cluster)
            com_x = np.zeros(n_cluster)
            com_y = np.zeros(n_cluster)
            com_z = np.zeros(n_cluster)
            std_x = np.zeros(n_cluster)
            std_y = np.zeros(n_cluster)
            std_z = np.zeros(n_cluster)
            group_id = np.zeros(n_cluster)
            n = np.zeros(n_cluster, dtype=np.int32)
            for i, clusterval in enumerate(cluster):
                cluster_locs = temp[temp["cluster"] == clusterval]
                mean_frame[i] = cluster_locs["frame"].mean()
                com_x[i] = np.average(
                    cluster_locs["x"].to_numpy(),
                    weights=cluster_locs["photons"].to_numpy(),
                )
                com_y[i] = np.average(
                    cluster_locs["y"].to_numpy(),
                    weights=cluster_locs["photons"].to_numpy(),
                )
                com_z[i] = np.average(
                    cluster_locs["z"].to_numpy(),
                    weights=cluster_locs["photons"].to_numpy(),
                )
                std_frame[i] = cluster_locs["frame"].std()
                std_x[i] = cluster_locs["x"].std() / np.sqrt(len(cluster_locs))
                std_y[i] = cluster_locs["y"].std() / np.sqrt(len(cluster_locs))
                std_z[i] = cluster_locs["z"].std() / np.sqrt(len(cluster_locs))
                n[i] = len(cluster_locs)
                group_id[i] = group
            clusters = pd.DataFrame(
                {
                    "group": group_id,
                    "cluster": cluster,
                    "mean_frame": mean_frame.astype(np.float32),
                    "x": com_x.astype(np.float32),
                    "y": com_y.astype(np.float32),
                    "z": com_z.astype(np.float32),
                    "std_frame": std_frame.astype(np.float32),
                    "lpx": std_x.astype(np.float32),
                    "lpy": std_y.astype(np.float32),
                    "lpz": std_z.astype(np.float32),
                    "n": n.astype(np.int32),
                }
            )
            combined_locs.append(clusters)
    else:
        for group in tqdm(np.unique(locs["group"])):
            temp = locs[locs["group"] == group]
            cluster = np.unique(temp["cluster"].to_numpy())
            n_cluster = len(cluster)
            mean_frame = np.zeros(n_cluster)
            std_frame = np.zeros(n_cluster)
            com_x = np.zeros(n_cluster)
            com_y = np.zeros(n_cluster)
            std_x = np.zeros(n_cluster)
            std_y = np.zeros(n_cluster)
            group_id = np.zeros(n_cluster)
            n = np.zeros(n_cluster, dtype=np.int32)
            for i, clusterval in enumerate(cluster):
                cluster_locs = temp[temp["cluster"] == clusterval]
                mean_frame[i] = cluster_locs["frame"].mean()
                com_x[i] = np.average(
                    cluster_locs["x"].to_numpy(),
                    weights=cluster_locs["photons"].to_numpy(),
                )
                com_y[i] = np.average(
                    cluster_locs["y"].to_numpy(),
                    weights=cluster_locs["photons"].to_numpy(),
                )
                std_frame[i] = cluster_locs["frame"].std()
                std_x[i] = cluster_locs["x"].std() / np.sqrt(len(cluster_locs))
                std_y[i] = cluster_locs["y"].std() / np.sqrt(len(cluster_locs))
                n[i] = len(cluster_locs)
                group_id[i] = group
            clusters = pd.DataFrame(
                {
                    "group": group_id,
                    "cluster": cluster,
                    "mean_frame": mean_frame.astype(np.float32),
                    "x": com_x.astype(np.float32),
                    "y": com_y.astype(np.float32),
                    "std_frame": std_frame.astype(np.float32),
                    "lpx": std_x.astype(np.float32),
                    "lpy": std_y.astype(np.float32),
                    "n": n.astype(np.int32),
                }
            )
            combined_locs.append(clusters)

    combined_locs = pd.concat(combined_locs, ignore_index=True)
    return combined_locs


def cluster_combine_dist(
    locs: pd.DataFrame, pixelsize: float | None = None
) -> pd.DataFrame:
    """Similar to ``cluster_combine``, but also calculates the distance
    to the nearest neighbor in the same group and the distance to the
    nearest neighbor in the same cluster in the same group.

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations with 'group' and 'cluster' fields.
    pixelsize : float or None, optional
        Pixel size in nm for z-scaling. If None, defaults to 130 nm.

    Returns
    -------
    combined_locs : pd.DataFrame
        Combined localizations with calculated properties for each
        cluster, including distances to nearest neighbors.
    """
    if "z" in locs.columns:
        pixelsize = 130 if pixelsize is None else pixelsize
        combined_locs = []
        for group in tqdm(np.unique(locs["group"])):
            temp = locs[locs["group"] == group]
            cluster = np.unique(temp["cluster"])
            n_cluster = len(cluster)
            mean_frame = temp["mean_frame"].to_numpy()
            std_frame = temp["std_frame"].to_numpy()
            com_x = temp["x"].to_numpy()
            com_y = temp["y"].to_numpy()
            com_z = temp["z"].to_numpy()
            std_x = temp["lpx"].to_numpy()
            std_y = temp["lpy"].to_numpy()
            std_z = temp["lpz"].to_numpy()
            group_id = temp["group"].to_numpy()
            n = temp["n"].to_numpy()
            min_dist = np.zeros(n_cluster)
            min_dist_xy = np.zeros(n_cluster)
            for i, clusterval in enumerate(cluster):
                # find nearest neighbor in xyz
                group_locs = temp[temp["cluster"] != clusterval]
                cluster_locs = temp[temp["cluster"] == clusterval]
                ref_point = np.stack(
                    (
                        cluster_locs["x"].to_numpy(),
                        cluster_locs["y"].to_numpy(),
                        cluster_locs["z"].to_numpy() / pixelsize,
                    )
                )
                all_points = np.stack(
                    (
                        group_locs["x"].to_numpy(),
                        group_locs["y"].to_numpy(),
                        group_locs["z"].to_numpy() / pixelsize,
                    )
                )
                distances = distance.cdist(
                    ref_point.transpose(), all_points.transpose()
                )
                min_dist[i] = np.amin(distances)
                # find nearest neighbor in xy
                ref_point_xy = np.array(cluster_locs[["x", "y"]])
                all_points_xy = np.array(group_locs[["x", "y"]])
                distances_xy = distance.cdist(
                    ref_point_xy.transpose(), all_points_xy.transpose()
                )
                min_dist_xy[i] = np.amin(distances_xy)

            clusters = pd.DataFrame(
                {
                    "group": group_id,
                    "cluster": cluster,
                    "mean_frame": mean_frame.astype(np.float32),
                    "x": com_x.astype(np.float32),
                    "y": com_y.astype(np.float32),
                    "z": com_z.astype(np.float32),
                    "std_frame": std_frame.astype(np.float32),
                    "lpx": std_x.astype(np.float32),
                    "lpy": std_y.astype(np.float32),
                    "lpz": std_z.astype(np.float32),
                    "n": n.astype(np.int32),
                    "min_dist": min_dist.astype(np.float32),
                    "mind_dist_xy": min_dist_xy.astype(np.float32),
                }
            )
            combined_locs.append(clusters)

    else:  # 2D case
        combined_locs = []
        for group in tqdm(np.unique(locs["group"])):
            temp = locs[locs["group"] == group]
            cluster = np.unique(temp["cluster"])
            n_cluster = len(cluster)
            mean_frame = temp["mean_frame"].to_numpy()
            std_frame = temp["std_frame"].to_numpy()
            com_x = temp["x"].to_numpy()
            com_y = temp["y"].to_numpy()
            std_x = temp["lpx"].to_numpy()
            std_y = temp["lpy"].to_numpy()
            group_id = temp["group"].to_numpy()
            n = temp["n"].to_numpy()
            min_dist = np.zeros(n_cluster)

            for i, clusterval in enumerate(cluster):
                # find nearest neighbor in xyz
                group_locs = temp[temp["cluster"] != clusterval]
                cluster_locs = temp[temp["cluster"] == clusterval]
                ref_point_xy = np.array(cluster_locs[["x", "y"]])
                all_points_xy = np.array(group_locs[["x", "y"]])
                distances_xy = distance.cdist(
                    ref_point_xy.transpose(), all_points_xy.transpose()
                )
                min_dist[i] = np.amin(distances_xy)

            clusters = pd.DataFrame(
                {
                    "group": group_id,
                    "cluster": cluster,
                    "mean_frame": mean_frame.astype(np.float32),
                    "x": com_x.astype(np.float32),
                    "y": com_y.astype(np.float32),
                    "std_frame": std_frame.astype(np.float32),
                    "lpx": std_x.astype(np.float32),
                    "lpy": std_y.astype(np.float32),
                    "n": n.astype(np.int32),
                    "min_dist": min_dist.astype(np.float32),
                }
            )
            combined_locs.append(clusters)

    combined_locs = pd.concat(combined_locs, ignore_index=True)
    return combined_locs


@numba.jit(nopython=True)
def get_link_groups(
    frame: np.ndarray,
    x: np.ndarray,
    y: np.ndarray,
    d_max: float,
    max_dark_time: int,
    group: np.ndarray,
) -> np.ndarray:
    """Find the groups for linking localizations into binding events.
    Assumes that ``locs`` are sorted by frame.

    Parameters
    ----------
    frame : np.ndarray
        Frame numbers of localizations.
    x, y : np.ndarray
        Coordinates of localizations.
    d_max : float
        Maximum distance for linking localizations.
    max_dark_time : int
        Maximum number of frames between localizations to be considered
        as originating from the same binding event.
    group : np.ndarray
        Grouping array for binding events. If None, all binding events
        are considered to be in the same group.

    Returns
    -------
    link_group : np.ndarray
        Array of link groups for each localization. Each group is
        represented by a unique integer. Localizations that are not
        linked to any other localization are assigned -1.
    """
    N = len(x)
    link_group = -np.ones(N, dtype=np.int32)
    current_link_group = -1
    for i in range(N):
        if link_group[i] == -1:  # loc has no group yet
            current_link_group += 1
            link_group[i] = current_link_group
            current_index = i
            next_loc_index_in_group = _get_next_loc_index_in_link_group(
                current_index,
                link_group,
                N,
                frame,
                x,
                y,
                d_max,
                max_dark_time,
                group,
            )
            while next_loc_index_in_group != -1:
                link_group[next_loc_index_in_group] = current_link_group
                current_index = next_loc_index_in_group
                next_loc_index_in_group = _get_next_loc_index_in_link_group(
                    current_index,
                    link_group,
                    N,
                    frame,
                    x,
                    y,
                    d_max,
                    max_dark_time,
                    group,
                )
    return link_group


@numba.jit(nopython=True)
def _get_next_loc_index_in_link_group(
    current_index: int,
    link_group: np.ndarray,
    N: int,
    frame: np.ndarray,
    x: np.ndarray,
    y: np.ndarray,
    d_max: float,
    max_dark_time: float,
    group: np.ndarray,
) -> int:
    """Find the next localization index in the link group for a given
    current localization index. The next localization is the one that
    is in the same group, has a frame greater than the current frame
    plus one, and is within the maximum distance defined by d_max.
    If no such localization is found, returns -1."""
    current_frame = frame[current_index]
    current_x = x[current_index]
    current_y = y[current_index]
    current_group = group[current_index]
    min_frame = current_frame + 1
    for min_index in range(current_index + 1, N):
        if frame[min_index] >= min_frame:
            break
    max_frame = current_frame + max_dark_time + 1
    for max_index in range(min_index, N):
        if frame[max_index] > max_frame:
            break
    else:
        max_index = N
    d_max_2 = d_max**2
    for j in range(min_index, max_index):
        if group[j] == current_group:
            if link_group[j] == -1:
                dx2 = (current_x - x[j]) ** 2
                if dx2 <= d_max_2:
                    dy2 = (current_y - y[j]) ** 2
                    if dy2 <= d_max_2:
                        if dx2 + dy2 <= d_max_2:
                            return j
    return -1


@numba.jit(nopython=True)
def _link_group_count(
    link_group: np.ndarray,
    n_locs: int,
    n_groups: int,
) -> np.ndarray:
    """Count the number of localizations in each link group."""
    result = np.zeros(n_groups, dtype=np.uint32)
    for i in range(n_locs):
        i_ = link_group[i]
        result[i_] += 1
    return result


@numba.jit(nopython=True)
def _link_group_sum(
    column: np.ndarray,
    link_group: np.ndarray,
    n_locs: int,
    n_groups: int,
) -> np.ndarray:
    """Sum the values of a column for each link group."""
    result = np.zeros(n_groups, dtype=column.dtype)
    for i in range(n_locs):
        i_ = link_group[i]
        result[i_] += column[i]
    return result


@numba.jit(nopython=True)
def _link_group_mean(
    column: np.ndarray,
    link_group: np.ndarray,
    n_locs: int,
    n_groups: int,
    n_locs_per_group: np.ndarray,
) -> np.ndarray:
    """Calculate the mean of a column for each link group."""
    group_sum = _link_group_sum(column, link_group, n_locs, n_groups)
    result = np.empty(
        n_groups, dtype=np.float32
    )  # this ensures float32 after the division
    result[:] = group_sum / n_locs_per_group
    return result


@numba.jit(nopython=True)
def _link_group_weighted_mean(
    column: np.ndarray,
    weights: np.ndarray,
    link_group: np.ndarray,
    n_locs: int,
    n_groups: int,
    n_locs_per_group: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
    """Calculate the mean of a column for each link group and the sum
    of the weights."""
    sum_weights = _link_group_sum(weights, link_group, n_locs, n_groups)
    return (
        _link_group_mean(
            column * weights,
            link_group,
            n_locs,
            n_groups,
            sum_weights,
        ),
        sum_weights,
    )


@numba.jit(nopython=True)
def _link_group_min_max(
    column: np.ndarray,
    link_group: np.ndarray,
    n_locs: int,
    n_groups: int,
) -> tuple[np.ndarray, np.ndarray]:
    """Calculate the minimum and maximum of a column for each link
    group."""
    min_ = np.empty(n_groups, dtype=column.dtype)
    max_ = np.empty(n_groups, dtype=column.dtype)
    min_[:] = column.max()
    max_[:] = column.min()
    for i in range(n_locs):
        i_ = link_group[i]
        value = column[i]
        if value < min_[i_]:
            min_[i_] = value
        if value > max_[i_]:
            max_[i_] = value
    return min_, max_


@numba.jit(nopython=True)
def _link_group_last(
    column: np.ndarray,
    link_group: np.ndarray,
    n_locs: int,
    n_groups: int,
) -> np.ndarray:
    """Return the last value of a column for each link group."""
    result = np.zeros(n_groups, dtype=column.dtype)
    for i in range(n_locs):
        i_ = link_group[i]
        result[i_] = column[i]
    return result


def link_loc_groups(
    locs: pd.DataFrame,
    info: list[dict],
    link_group: np.ndarray,
    remove_ambiguous_lengths: bool = True,
) -> pd.DataFrame:
    """Combine localizations into binding events based on the
    spatiotemporal proximity defined by the ``link_group``. Takes the
    average position to calculate the coordinates of the binding events.

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations.
    info : list of dicts
        Metadata of the localization list.
    link_group : np.ndarray
        Array that defines the link groups for the localizations.
    remove_ambiguous_lengths : bool, optional
        If True, removes linked localizations with ambiguous lengths,
        i.e., localizations that are linked to multiple binding events
        with different lengths. Default is True.

    Returns
    -------
    linked_locs : pd.DataFrame
        Linked localizations, i.e., binding events with their
        properties.
    """
    n_locs = len(link_group)
    n_groups = link_group.max() + 1
    n_ = _link_group_count(link_group, n_locs, n_groups)
    columns = OrderedDict()
    if "frame" in locs.columns:
        first_frame_, last_frame_ = _link_group_min_max(
            locs["frame"].to_numpy(), link_group, n_locs, n_groups
        )
        columns["frame"] = first_frame_
    if "x" in locs.columns:
        weights_x = 1 / locs["lpx"].to_numpy() ** 2
        columns["x"], sum_weights_x_ = _link_group_weighted_mean(
            locs["x"].to_numpy(), weights_x, link_group, n_locs, n_groups, n_
        )
    if "y" in locs.columns:
        weights_y = 1 / locs["lpy"].to_numpy() ** 2
        columns["y"], sum_weights_y_ = _link_group_weighted_mean(
            locs["y"].to_numpy(), weights_y, link_group, n_locs, n_groups, n_
        )
    if "photons" in locs.columns:
        columns["photons"] = _link_group_sum(
            locs["photons"].to_numpy(),
            link_group,
            n_locs,
            n_groups,
        )
    if "sx" in locs.columns:
        columns["sx"] = _link_group_mean(
            locs["sx"].to_numpy(),
            link_group,
            n_locs,
            n_groups,
            n_,
        )
    if "sy" in locs.columns:
        columns["sy"] = _link_group_mean(
            locs["sy"].to_numpy(),
            link_group,
            n_locs,
            n_groups,
            n_,
        )
    if "bg" in locs.columns:
        columns["bg"] = _link_group_sum(
            locs["bg"].to_numpy(),
            link_group,
            n_locs,
            n_groups,
        )
    if "x" in locs.columns:
        columns["lpx"] = np.sqrt(1 / sum_weights_x_)
    if "y" in locs.columns:
        columns["lpy"] = np.sqrt(1 / sum_weights_y_)
    if "ellipticity" in locs.columns:
        columns["ellipticity"] = _link_group_mean(
            locs["ellipticity"].to_numpy(), link_group, n_locs, n_groups, n_
        )
    if "net_gradient" in locs.columns:
        columns["net_gradient"] = _link_group_mean(
            locs["net_gradient"].to_numpy(), link_group, n_locs, n_groups, n_
        )
    if "likelihood" in locs.columns:
        columns["likelihood"] = _link_group_mean(
            locs["likelihood"].to_numpy(), link_group, n_locs, n_groups, n_
        )
    if "iterations" in locs.columns:
        columns["iterations"] = _link_group_mean(
            locs["iterations"].to_numpy(), link_group, n_locs, n_groups, n_
        )
    if "z" in locs.columns:
        columns["z"] = _link_group_mean(
            locs["z"].to_numpy(),
            link_group,
            n_locs,
            n_groups,
            n_,
        )
    if "d_zcalib" in locs.columns:
        columns["d_zcalib"] = _link_group_mean(
            locs["d_zcalib"].to_numpy(), link_group, n_locs, n_groups, n_
        )
    if "group" in locs.columns:
        columns["group"] = _link_group_last(
            locs["group"].to_numpy(),
            link_group,
            n_locs,
            n_groups,
        )
    if "frame" in locs.columns:
        columns["len"] = last_frame_ - first_frame_ + 1
    columns["n"] = n_
    if "photons" in locs.columns:
        columns["photon_rate"] = np.float32(columns["photons"] / n_)
    linked_locs = pd.DataFrame(columns)
    if remove_ambiguous_lengths:
        valid = np.logical_and(
            first_frame_ > 0,
            last_frame_ < info[0]["Frames"],
        )
        linked_locs = linked_locs[valid]
    return linked_locs


def n_segments(info: list[dict], segmentation: int) -> int:
    """Calculate the number of segments for the given segmentation
    for undrifting.

    Parameters
    ----------
    info : list of dicts
        Metadata of the localizations list.
    segmentation : int
        Number of segments to divide the data into.

    Returns
    -------
    n_segments : int
        Number of segments based on the total number of frames and the
        segmentation value.
    """
    n_frames = info[0]["Frames"]
    n_segments = int(np.round(n_frames / segmentation))
    return n_segments


def segment(
    locs: pd.DataFrame,
    info: list[dict],
    segmentation: int,
    kwargs: dict = {},
    callback: Callable[[int], None] = None,
) -> tuple[np.ndarray, np.ndarray]:
    """Split localizations into temporal segments (number of segments
    is defined by the segmentation parameter) and render each segment
    into a 2D image.

    Parameters
    ----------
    locs : np.ndarray
        Localization list.
    info : list of dicts
        Metadata of the localization list.
    segmentation : int
        Number of segments to divide the data into.
    kwargs : dict, optional
        Additional keyword arguments for the rendering function.
        Default is an empty dictionary.
    callback : Callable[[int], None], optional
        Callback function to report progress. It should accept an
        integer argument representing the current segment index.
        Default is None, which means no callback is used.

    Returns
    -------
    bounds : np.ndarray
        Array of bounds for each segment, where each bound is the
        starting frame of the segment.
    segments : np.ndarray
        3D array of segments, where each segment is a 2D image of the
        localizations in that segment.
    """
    Y = info[0]["Height"]
    X = info[0]["Width"]
    n_frames = info[0]["Frames"]
    n_seg = n_segments(info, segmentation)
    bounds = np.linspace(0, n_frames - 1, n_seg + 1, dtype=np.uint32)
    segments = np.zeros((n_seg, Y, X))
    if callback is None:
        it = trange(n_seg, desc="Generating segments", unit="segments")
    else:
        callback(0)
        it = range(n_seg)
    for i in it:
        segment_locs = locs[
            (locs["frame"] >= bounds[i]) & (locs["frame"] < bounds[i + 1])
        ]
        _, segments[i] = render.render(segment_locs, info, **kwargs)
        if callback is not None:
            callback(i + 1)
    return bounds, segments


def undrift(
    locs: pd.DataFrame,
    info: list[dict],
    segmentation: int,
    display: bool = True,
    segmentation_callback: Callable[[int], None] = None,
    rcc_callback: Callable[[int], None] = None,
) -> tuple[pd.DataFrame, pd.DataFrame]:
    """Undrift by RCC. See Wang, Schnitzbauer, et al. Optics Express,
    2014.

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations to undrift.
    info : list of dicts
        Metadata of the localization list.
    segmentation : int
        Number of segments to divide the data into for undrifting.
    display : bool, optional
        If True, displays the estimated drift. Default is True.
    segmentation_callback : Callable[[int], None], optional
        Callback function to report progress during segmentation. It
        should accept an integer argument representing the current
        segment index. Default is None, which means no callback is used.
    rcc_callback : Callable[[int], None], optional
        Callback function to report progress during RCC calculation.
        It should accept an integer argument representing the current
        segment index. Default is None, which means no callback is used.

    Returns
    -------
    drift : pd.DataFrame
        Estimated drift as a DataFrame with columns 'x' and 'y'.
    locs : pd.DataFrame
        Undrifted localization list with the drift applied to the 'x'
        and 'y' coordinates.
    """
    bounds, segments = segment(
        locs,
        info,
        segmentation,
        {"blur_method": "gaussian", "min_blur_width": 1},
        segmentation_callback,
    )
    shift_y, shift_x = imageprocess.rcc(segments, 32, rcc_callback)
    t = (bounds[1:] + bounds[:-1]) / 2
    drift_x_pol = interpolate.InterpolatedUnivariateSpline(t, shift_x, k=3)
    drift_y_pol = interpolate.InterpolatedUnivariateSpline(t, shift_y, k=3)
    t_inter = np.arange(info[0]["Frames"])
    drift_ = (drift_x_pol(t_inter), drift_y_pol(t_inter))
    drift = pd.DataFrame({"x": drift_[0], "y": drift_[1]})
    if display:
        fig1 = plt.figure(figsize=(10, 6), constrained_layout=True)
        plt.suptitle("Estimated drift")
        plt.subplot(1, 2, 1)
        plt.plot(drift["x"], label="x interpolated")
        plt.plot(drift["y"], label="y interpolated")
        t = (bounds[1:] + bounds[:-1]) / 2
        plt.plot(
            t,
            shift_x,
            "o",
            color=list(plt.rcParams["axes.prop_cycle"])[0]["color"],
            label="x",
        )
        plt.plot(
            t,
            shift_y,
            "o",
            color=list(plt.rcParams["axes.prop_cycle"])[1]["color"],
            label="y",
        )
        plt.legend(loc="best")
        plt.xlabel("Frame")
        plt.ylabel("Drift (pixel)")
        plt.subplot(1, 2, 2)
        plt.plot(
            drift["x"],
            drift["y"],
            color=list(plt.rcParams["axes.prop_cycle"])[2]["color"],
        )
        plt.plot(
            shift_x,
            shift_y,
            "o",
            color=list(plt.rcParams["axes.prop_cycle"])[2]["color"],
        )
        plt.axis("equal")
        plt.xlabel("x")
        plt.ylabel("y")
        plt.show()
    locs["x"] -= drift["x"][locs["frame"]].to_numpy()
    locs["y"] -= drift["y"][locs["frame"]].to_numpy()
    return drift, locs


def undrift_from_picked(
    picked_locs: list[pd.DataFrame], info: list[dict]
) -> pd.DataFrame:
    """Find drift from picked localizations. Note that unlike other
    undrifting functions, this function does not return undrifted
    localizations but only drift.

    Parameters
    ----------
    picked_locs : list of pd.DataFrames
        List of picked localizations, where each element is a data frame
        of localizations for a single pick.
    info : list of dicts
        Metadata of the localization list, where each element
        corresponds to the metadata of the localizations in
        ``picked_locs``.

    Returns
    -------
    drift : pd.DataFrame
        Estimated drift as a DataFrame with columns 'x', 'y', and
        optionally 'z' if the z coordinate exists in the picked
        localizations.
    """
    drift_x = _undrift_from_picked_coordinate(picked_locs, info, "x")
    drift_y = _undrift_from_picked_coordinate(picked_locs, info, "y")

    # A data frame to store the applied drift
    drift = pd.DataFrame({"x": drift_x, "y": drift_y})
    # If z coordinate exists, also apply drift there
    if all(["z" in _.columns for _ in picked_locs]):
        drift_z = _undrift_from_picked_coordinate(picked_locs, info, "z")
        drift["z"] = drift_z
    return drift


def _undrift_from_picked_coordinate(
    picked_locs: list[pd.DataFrame],
    info: list[dict],
    coordinate: Literal["x", "y", "z"],
) -> np.ndarray:
    """Calculate drift in a given coordinate from picked localizations.
    Uses the center of mass of each pick to find the drift in the
    specified coordinate across all frames. The drift is calculated as
    the average of the localizations' coordinates minus the mean of the
    coordinates for each pick.

    Parameters
    ----------
    picked_locs : list of pd.DataFrames
        List of pd.DataFrames with locs for each pick.
    info : list of dicts
        Localizations' metadata.
    coordinate : {"x", "y", "z"}
        Spatial coordinate where drift is to be found.

    Returns
    -------
    drift_mean : np.ndarray
        Average drift across picks for all frames
    """
    n_picks = len(picked_locs)
    n_frames = info[0]["Frames"]

    # Drift per pick per frame
    drift = np.empty((n_picks, n_frames))
    drift.fill(np.nan)

    # Remove center of mass offset
    for i, locs in enumerate(picked_locs):
        coordinates = locs[coordinate].to_numpy()
        drift[i, locs["frame"].to_numpy()] = coordinates - np.mean(coordinates)

    # Mean drift over picks
    drift_mean = np.nanmean(drift, 0)
    # Square deviation of each pick's drift to mean drift along frames
    sd = (drift - drift_mean) ** 2
    # Mean of square deviation for each pick
    msd = np.nanmean(sd, 1)
    # New mean drift over picks
    # where each pick is weighted according to its msd
    nan_mask = np.isnan(drift)
    drift = np.ma.MaskedArray(drift, mask=nan_mask)
    drift_mean = np.ma.average(drift, axis=0, weights=1 / msd)
    drift_mean = drift_mean.filled(np.nan)

    # Linear interpolation for frames without localizations
    def nan_helper(y):
        return np.isnan(y), lambda z: z.nonzero()[0]

    nans, nonzero = nan_helper(drift_mean)
    drift_mean[nans] = np.interp(
        nonzero(nans), nonzero(~nans), drift_mean[~nans]
    )
    return drift_mean


def align(
    locs: list[pd.DataFrame],
    infos: list[dict],
    display: bool = False,
) -> pd.DataFrame:
    """Align localizations from multiple channels (one per each element
    in `locs`) by calculating the shifts between the rendered images
    using RCC.

    Parameters
    ----------
    locs : list of pd.DataFrames
        List of localization datasets, where each element is a
        DataFrame of localizations for a single image.
    infos : list of dicts
        List of metadata dictionaries corresponding to each
        localization array in `locs`.
    display : bool, optional
        Not used.

    Returns
    -------
    locs : list of pd.DataFrames
        Aligned localizations with the shifts applied to the 'x' and
        'y' coordinates.
    """
    images = []
    for i, (locs_, info_) in enumerate(zip(locs, infos)):
        _, image = render.render(locs_, info_, blur_method="smooth")
        images.append(image)
    shift_y, shift_x = imageprocess.rcc(images)
    for i, (locs_, dx, dy) in enumerate(zip(locs, shift_x, shift_y)):
        locs_["y"] -= dy
        locs_["x"] -= dx
    return locs


def groupprops(
    locs: pd.DataFrame,
    callback: Callable[[int], None] | None = None,
) -> pd.DataFrame:
    """Calculate group statistics for localizations, such as mean and
    standard deviation.

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations with a 'group' field that defines the groups.
    callback : Callable[[int], None], optional
        Callback function to report progress. It should accept an
        integer argument representing the current group index.
        Default is None, which means no callback is used.

    Returns
    -------
    groups : pd.DataFrame
        Group statistics for each group in the localization list.
    """
    try:
        locs = locs[locs["dark"] != -1]
    except AttributeError:
        pass
    group_ids = np.unique(locs["group"])
    n = len(group_ids)
    names = ["group", "n_events"] + list(
        itertools.chain(*[(_ + "_mean", _ + "_std") for _ in locs.columns])
    )
    groups = pd.DataFrame(np.empty((n, len(names))), columns=names)
    if callback is not None:
        callback(0)
        it = enumerate(group_ids)
    else:
        it = enumerate(
            tqdm(group_ids, desc="Calculating group statistics", unit="Groups")
        )
    for i, group_id in it:
        group_locs = locs[locs["group"] == group_id]
        groups.loc[i, "group"] = group_id
        groups.loc[i, "n_events"] = len(group_locs)
        for name in locs.columns:
            groups.loc[i, name + "_mean"] = group_locs[name].mean()
            groups.loc[i, name + "_std"] = group_locs[name].std()
        if callback is not None:
            callback(i + 1)
    # set dtypes
    groups = groups.astype(
        {
            "group": np.int32,
            "n_events": np.int32,
            **{name + "_mean": np.float32 for name in locs.columns},
            **{name + "_std": np.float32 for name in locs.columns},
        }
    )
    # add qpaint idx
    if "dark_mean" in groups.columns:
        groups["qpaint_idx"] = 1 / groups["dark_mean"]
    return groups


def calculate_fret(
    acc_locs: pd.DataFrame,
    don_locs: pd.DataFrame,
) -> tuple[dict, pd.DataFrame]:
    """Calculate the FRET efficiency in picked regions, this is for one
    trace."""
    fret_dict = {}
    if len(acc_locs) == 0:
        max_frames = don_locs["frame"].max()
    elif len(don_locs) == 0:
        max_frames = acc_locs["frame"].max()
    else:
        max_frames = max([acc_locs["frame"].max(), don_locs["frame"].max()])

    # Initialize a vector filled with zeros for the duration of the movie
    xvec = np.arange(max_frames + 1)
    yvec = xvec[:] * 0
    acc_trace = yvec.copy()
    don_trace = yvec.copy()
    # Fill vector with the photon numbers of events that happend
    acc_trace[acc_locs["frame"]] = acc_locs["photons"] - acc_locs["bg"]
    don_trace[don_locs["frame"]] = don_locs["photons"] - don_locs["bg"]

    # Calculate the FRET efficiency
    fret_trace = acc_trace / (acc_trace + don_trace)
    # Only select FRET values between 0 and 1
    selector = np.logical_and(fret_trace > 0, fret_trace < 1)

    # Select the final fret events based on the 0 to 1 range
    fret_events = fret_trace[selector]
    fret_timepoints = np.arange(len(fret_trace))[selector]

    f_locs = []
    if len(fret_timepoints) > 0:
        # Calculate FRET locs: Select the locs when FRET happens
        sel_locs = []
        for element in fret_timepoints:
            sel_locs.append(don_locs[don_locs["frame"] == element])

        f_locs = pd.concat(sel_locs, ignore_index=True)
        f_locs["fret"] = np.array(fret_events)

    fret_dict["fret_events"] = np.array(fret_events)
    fret_dict["fret_timepoints"] = fret_timepoints
    fret_dict["acc_trace"] = acc_trace
    fret_dict["don_trace"] = don_trace
    fret_dict["frames"] = xvec
    fret_dict["maxframes"] = max_frames

    return fret_dict, f_locs


def nn_analysis(
    X1: np.ndarray,
    X2: np.ndarray,
    nn_count: int,
) -> np.ndarray:
    """Find the nearest neighbors between two sets of localizations.

    Parameters
    ----------
    X1, X2 : np.ndarray
        Arrays of shape (N, D) and (M, D) representing the coordinates
        of the two sets of localizations, where N and M are the number
        of localizations in each set, and D is the number of spatial
        dimensions (2 or 3).
    nn_count : int
        Number of nearest neighbors to find for each localization in
        the second set.

    Returns
    -------
    nnd : np.ndarray
        Array of nearest neighbors distances, where each row corresponds
        to a localization in the first set and contains the distances to
        its nearest neighbors in the second set.
    """
    if X1.shape[1] != X2.shape[1]:
        raise ValueError("X1 and X2 must have the same number of dimensions.")
    tree = KDTree(X2)
    if np.array_equal(X1, X2):
        distances, indices = tree.query(X1, k=nn_count + 1)
        nn = distances[:, 1:]
    else:
        distances, indices = tree.query(X1, k=nn_count)
        nn = distances
    nn.reshape(-1, nn_count)  # ensure the shape is (N, nn_count)
    return nn
