"""
picasso.render
~~~~~~~~~~~~~~

Render single molecule localizations to a super-resolution image

:authors: Joerg Schnitzbauer 2015, Rafal Kowalewski 2023
:copyright: Copyright (c) 2015 Jungmann Lab, MPI of Biochemistry
"""

from typing import Literal

import numba
import numpy as np
import pandas as pd
from scipy import signal
from scipy.spatial.transform import Rotation


_DRAW_MAX_SIGMA = 3  # max. sigma from mean to render (mu +/- 3 sigma)


def render(
    locs: pd.DataFrame,
    info: dict | None = None,
    oversampling: float = 1,
    viewport: tuple[tuple[float, float], tuple[float, float]] | None = None,
    blur_method: (
        Literal["gaussian", "gaussian_iso", "smooth", "convolve"] | None
    ) = None,
    min_blur_width: float = 0,
    ang: tuple | None = None,
) -> tuple[int, np.ndarray]:
    """Render localizations given FOV and blur method.

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations to be rendered.
    info : dict, optional
        Contains localizations metadata. Needed only if no viewport
        specified.
    oversampling : float, optional
        Number of super-resolution pixels per camera pixel.
    viewport : tuple, optional
        Field of view to be rendered. The input is
        ``((y_min, x_min), (y_max, x_max))``. If None, all localizations
        are rendered.
    blur_method : {"gaussian", "gaussian_iso", "smooth", "convolve"} or None, \
            optional
        Defines localizations' blur. The string has to be one of
        'gaussian', 'gaussian_iso', 'smooth', 'convolve'. If None, no
        blurring is applied.
    min_blur_width : float, optional
        Minimum size of blur (camera pixels).
    ang : tuple, optional
        Rotation angles of locs around x, y and z axes in radians. If
        None, locs are not rotated.

    Raises
    ------
    Exception
        If blur_method not one of 'gaussian', 'gaussian_iso', 'smooth',
        'convolve' or None.

    Returns
    -------
    n : int
        Number of localizations rendered.
    image : np.ndarray
        Rendered image.
    """
    if viewport is None:
        try:
            # all locs
            viewport = [(0, 0), (info[0]["Height"], info[0]["Width"])]
        except TypeError:
            raise ValueError("Need info if no viewport is provided.")
    (y_min, x_min), (y_max, x_max) = viewport
    if blur_method is None:
        # no blur
        return render_hist(
            locs,
            oversampling,
            y_min,
            x_min,
            y_max,
            x_max,
            ang=ang,
        )
    elif blur_method == "gaussian":
        # individual localization precision
        return render_gaussian(
            locs,
            oversampling,
            y_min,
            x_min,
            y_max,
            x_max,
            min_blur_width,
            ang=ang,
        )
    elif blur_method == "gaussian_iso":
        # individual localization precision (same for x and y)
        return render_gaussian_iso(
            locs,
            oversampling,
            y_min,
            x_min,
            y_max,
            x_max,
            min_blur_width,
            ang=ang,
        )
    elif blur_method == "smooth":
        # one pixel blur
        return render_smooth(
            locs,
            oversampling,
            y_min,
            x_min,
            y_max,
            x_max,
            ang=ang,
        )
    elif blur_method == "convolve":
        # global localization precision
        return render_convolve(
            locs,
            oversampling,
            y_min,
            x_min,
            y_max,
            x_max,
            min_blur_width,
            ang=ang,
        )
    else:
        raise Exception("blur_method not understood.")


@numba.njit
def _render_setup(
    x: np.ndarray,
    y: np.ndarray,
    oversampling: float,
    y_min: float,
    x_min: float,
    y_max: float,
    x_max: float,
) -> tuple[np.ndarray, int, int, np.ndarray, np.ndarray, np.ndarray]:
    """Find coordinates to be rendered and sets up an empty image
    array.

    Parameters
    ----------
    x, y : np.ndarray
        x and y coordinates of the localizations to be rendered (1D
        arrays).
    oversampling : float
        Number of super-resolution pixels per camera pixel.
    y_min, x_min : float
        Minimum y and x coordinates to be rendered (camera pixels).
    y_max, x_max : float
        Maximum y and x coordinates to be rendered (camera pixels).

    Returns
    -------
    image : np.ndarray
        Empty image array.
    n_pixel_y : int
        Number of pixels in y.
    n_pixel_x : int
        Number of pixels in x.
    x : np.ndarray
        x coordinates to be rendered.
    y : np.ndarray
        y coordinates to be rendered.
    in_view : np.ndarray
        Indeces of the localizations to be rendered.
    """
    n_pixel_y = int(np.ceil(oversampling * (y_max - y_min)))
    n_pixel_x = int(np.ceil(oversampling * (x_max - x_min)))
    in_view = (x > x_min) & (y > y_min) & (x < x_max) & (y < y_max)
    x = x[in_view]
    y = y[in_view]
    x = oversampling * (x - x_min)
    y = oversampling * (y - y_min)
    image = np.zeros((n_pixel_y, n_pixel_x), dtype=np.float32)
    return image, n_pixel_y, n_pixel_x, x, y, in_view


@numba.njit
def _render_setup_anisotropic(
    x: np.ndarray,
    y: np.ndarray,
    oversampling_x: float,
    oversampling_y: float,
    y_min: float,
    x_min: float,
    y_max: float,
    x_max: float,
) -> tuple[np.ndarray, int, int, np.ndarray, np.ndarray, np.ndarray]:
    """Find coordinates to be rendered and sets up an empty image
    array. Allows for different pixel sizes in x and y (oversampling).

    Parameters
    ----------
    x, y : np.ndarray
        x and y coordinates of the localizations to be rendered (1D
        arrays).
    oversampling_x, oversampling_y : float
        Number of super-resolution pixels per camera pixel in x and y.
    y_min, x_min : float
        Minimum y and x coordinates to be rendered (camera pixels).
    y_max, x_max : float
        Maximum y and x coordinates to be rendered (camera pixels).

    Returns
    -------
    image : np.ndarray
        Empty image array.
    n_pixel_y : int
        Number of pixels in y.
    n_pixel_x : int
        Number of pixels in x.
    x : np.ndarray
        x coordinates to be rendered.
    y : np.ndarray
        y coordinates to be rendered.
    in_view : np.ndarray
        Indeces of the localizations to be rendered.
    """
    n_pixel_y = int(np.ceil(oversampling_y * (y_max - y_min)))
    n_pixel_x = int(np.ceil(oversampling_x * (x_max - x_min)))
    in_view = (x > x_min) & (y > y_min) & (x < x_max) & (y < y_max)
    x = x[in_view]
    y = y[in_view]
    x = oversampling_x * (x - x_min)
    y = oversampling_y * (y - y_min)
    image = np.zeros((n_pixel_y, n_pixel_x), dtype=np.float32)
    return image, n_pixel_y, n_pixel_x, x, y, in_view


@numba.njit
def _render_setup3d(
    x: np.ndarray,
    y: np.ndarray,
    z: np.ndarray,
    oversampling: float,
    y_min: float,
    x_min: float,
    y_max: float,
    x_max: float,
    z_min: float,
    z_max: float,
    pixelsize: float,
) -> tuple[
    np.ndarray,
    int,
    int,
    int,
    np.ndarray,
    np.ndarray,
    np.ndarray,
    np.ndarray,
]:
    """Find coordinates to be rendered in 3D and sets up an empty image
    array.

    Parameters
    ----------
    x, y, z : np.ndarray
        x, y and z coordinates of the localizations to be rendered (1D
        arrays).
    oversampling : float
        Number of super-resolution pixels per camera pixel.
    y_min, x_min : float
        Minimum y and x coordinate to be rendered (camera pixels).
    y_max, x_max : float
        Maximum y and x coordinate to be rendered (camera pixels).
    z_min : float
        Minimum z coordinate to be rendered (nm).
    z_max : float
        Maximum z coordinate to be rendered (nm).
    pixelsize : float
        Camera pixel size, used for converting z coordinates.

    Returns
    -------
    image : np.ndarray
        Empty image array.
    n_pixel_y, n_pixel_x, n_pixel_z : int
        Number of pixels in y, x, and z.
    x, y, z : np.ndarray
        x, y, z coordinates to be rendered.
    in_view : np.ndarray
        Indeces of the localizations to be rendered.
    """
    n_pixel_y = int(np.ceil(oversampling * (y_max - y_min)))
    n_pixel_x = int(np.ceil(oversampling * (x_max - x_min)))
    n_pixel_z = int(np.ceil(oversampling * (z_max - z_min)))
    z /= pixelsize
    in_view = (
        (x > x_min)
        & (y > y_min)
        & (z > z_min)
        & (x < x_max)
        & (y < y_max)
        & (z < z_max)
    )
    x = x[in_view]
    y = y[in_view]
    z = z[in_view]
    x = oversampling * (x - x_min)
    y = oversampling * (y - y_min)
    z = oversampling * (z - z_min)
    image = np.zeros((n_pixel_y, n_pixel_x, n_pixel_z), dtype=np.float32)
    return image, n_pixel_y, n_pixel_x, n_pixel_z, x, y, z, in_view


@numba.njit
def _render_setup3d_anisotropic(
    x: np.ndarray,
    y: np.ndarray,
    z: np.ndarray,
    oversampling_x: float,
    oversampling_y: float,
    oversampling_z: float,
    y_min: float,
    x_min: float,
    y_max: float,
    x_max: float,
    z_min: float,
    z_max: float,
    pixelsize: float,
) -> tuple[
    np.ndarray,
    int,
    int,
    int,
    np.ndarray,
    np.ndarray,
    np.ndarray,
    np.ndarray,
]:
    """Find coordinates to be rendered in 3D and sets up an empty image
    array. Allows for different pixel sizes in x, y and z
    (oversampling).

    Parameters
    ----------
    x, y, z : np.ndarray
        x, y and z coordinates of the localizations to be rendered (1D
        arrays).
    oversampling : float
        Number of super-resolution pixels per camera pixel.
    y_min, x_min : float
        Minimum y and x coordinate to be rendered (camera pixels).
    y_max, x_max : float
        Maximum y and x coordinate to be rendered (camera pixels).
    z_min : float
        Minimum z coordinate to be rendered (nm).
    z_max : float
        Maximum z coordinate to be rendered (nm).
    pixelsize : float
        Camera pixel size, used for converting z coordinates.

    Returns
    -------
    image : np.ndarray
        Empty image array.
    n_pixel_y, n_pixel_x, n_pixel_z : int
        Number of pixels in y, x, and z.
    x, y, z : np.ndarray
        x, y, z coordinates to be rendered.
    in_view : np.ndarray
        Indeces of the localizations to be rendered.
    """
    n_pixel_y = int(np.ceil(oversampling_y * (y_max - y_min)))
    n_pixel_x = int(np.ceil(oversampling_x * (x_max - x_min)))
    n_pixel_z = int(np.ceil(oversampling_z * (z_max - z_min)))
    z /= pixelsize
    in_view = (
        (x > x_min)
        & (y > y_min)
        & (z > z_min)
        & (x < x_max)
        & (y < y_max)
        & (z < z_max)
    )
    x = x[in_view]
    y = y[in_view]
    z = z[in_view]
    x = oversampling_x * (x - x_min)
    y = oversampling_y * (y - y_min)
    z = oversampling_z * (z - z_min)
    image = np.zeros((n_pixel_y, n_pixel_x, n_pixel_z), dtype=np.float32)
    return image, n_pixel_y, n_pixel_x, n_pixel_z, x, y, z, in_view


@numba.njit
def _fill(image: np.ndarray, x: np.ndarray, y: np.ndarray) -> None:
    """Fill image with x and y coordinates. Image is not blurred.

    Parameters
    ----------
    image : np.ndarray
        Empty image array.
    x, y : np.ndarray
        x and y coordinates to be rendered.
    """
    x = x.astype(np.int32)
    y = y.astype(np.int32)
    for i, j in zip(x, y):
        image[j, i] += 1


@numba.njit
def _fill3d(
    image: np.ndarray, x: np.ndarray, y: np.ndarray, z: np.ndarray
) -> None:
    """Fill image with x, y and z coordinates. Image is not blurred.

    Parameters
    ----------
    image : np.ndarray
        Empty image array.
    x, y, z : np.ndarray
        x, y and z coordinates to be rendered.
    """
    x = x.astype(np.int32)
    y = y.astype(np.int32)
    z = z.astype(np.int32)
    z += np.min(z)  # because z takes also negative values
    for i, j, k in zip(x, y, z):
        image[j, i, k] += 1


@numba.njit
def _fill_gaussian(
    image: np.ndarray,
    x: np.ndarray,
    y: np.ndarray,
    sx: np.ndarray,
    sy: np.ndarray,
    n_pixel_x: int,
    n_pixel_y: int,
) -> None:
    """Fill image with blurred x and y coordinates. Each localization
    is rendered as a 2D Gaussian centered at (x, y) with standard
    deviations (sx, sy).

    Parameters
    ----------
    image : np.ndarray
        Empty image array.
    x, y : np.ndarray
        x and y coordinates to be rendered.
    sx, sy : np.ndarray
        Localization precision in x and y for each localization.
    n_pixel_x, n_pixel_y : int
        Number of pixels in x and y.
    """
    # render each localization separately
    for x_, y_, sx_, sy_ in zip(x, y, sx, sy):

        # get min and max indeces to draw the given localization
        max_y = _DRAW_MAX_SIGMA * sy_
        i_min = np.int32(y_ - max_y)
        if i_min < 0:
            i_min = 0
        i_max = np.int32(y_ + max_y + 1)
        if i_max > n_pixel_y:
            i_max = n_pixel_y
        max_x = _DRAW_MAX_SIGMA * sx_
        j_min = np.int32(x_ - max_x)
        if j_min < 0:
            j_min = 0
        j_max = np.int32(x_ + max_x) + 1
        if j_max > n_pixel_x:
            j_max = n_pixel_x

        # draw a localization as a 2D guassian PDF
        for i in range(i_min, i_max):
            for j in range(j_min, j_max):
                image[i, j] += np.exp(
                    -(
                        (j - x_ + 0.5) ** 2 / (2 * sx_**2)
                        + (i - y_ + 0.5) ** 2 / (2 * sy_**2)
                    )
                ) / (2 * np.pi * sx_ * sy_)


@numba.njit
def _fill_gaussian_rot(
    image: np.ndarray,
    x: np.ndarray,
    y: np.ndarray,
    z: np.ndarray,
    sx: np.ndarray,
    sy: np.ndarray,
    sz: np.ndarray,
    n_pixel_x: int,
    n_pixel_y: int,
    ang: tuple[float, float, float],
) -> None:
    """Fill image with rotated gaussian-blurred localizations.

    Localization precisions (sx, sy and sz) are treated as standard
    deviations of the guassians to be rendered.

    See https://cs229.stanford.edu/section/gaussians.pdf

    Parameters
    ----------
    image : np.ndarray
        Empty image array.
    x, y, z : np.ndarray
        3D coordinates to be rendered.
    sx, sy, sz : np.ndarray
        Localization precision in x, y and z for each localization.
    n_pixel_x, n_pixel_y : int
        Number of pixels in x and y.
    ang : tuple
        Rotation angles of locs around x, y and z axes (radians).
    """
    (angx, angy, angz) = ang  # rotation angles

    rot_mat_x = np.array(
        [
            [1.0, 0.0, 0.0],
            [0.0, np.cos(angx), np.sin(angx)],
            [0.0, -np.sin(angx), np.cos(angx)],
        ],
        dtype=np.float32,
    )  # rotation matrix around x axis
    rot_mat_y = np.array(
        [
            [np.cos(angy), 0.0, np.sin(angy)],
            [0.0, 1.0, 0.0],
            [-np.sin(angy), 0.0, np.cos(angy)],
        ],
        dtype=np.float32,
    )  # rotation matrix around y axis
    rot_mat_z = np.array(
        [
            [np.cos(angz), -np.sin(angz), 0.0],
            [np.sin(angz), np.cos(angz), 0.0],
            [0.0, 0.0, 1.0],
        ],
        dtype=np.float32,
    )  # rotation matrix around z axis
    rot_matrix = rot_mat_x @ rot_mat_y @ rot_mat_z  # rotation matrix
    rot_matrixT = np.transpose(rot_matrix)  # ...and its transpose

    # draw each localization separately
    for x_, y_, z_, sx_, sy_, sz_ in zip(x, y, z, sx, sy, sz):

        # get min and max indeces to draw the given localization
        max_y = (_DRAW_MAX_SIGMA * 2.5) * sy_
        i_min = int(y_ - max_y)
        if i_min < 0:
            i_min = 0
        i_max = int(y_ + max_y + 1)
        if i_max > n_pixel_y:
            i_max = n_pixel_y
        max_x = (_DRAW_MAX_SIGMA * 2.5) * sx_
        j_min = int(x_ - max_x)
        if j_min < 0:
            j_min = 0
        j_max = int(x_ + max_x + 1)
        if j_max > n_pixel_x:
            j_max = n_pixel_x
        max_z = (_DRAW_MAX_SIGMA * 2.5) * sz_
        k_min = int(z_ - max_z)
        k_max = int(z_ + max_z + 1)

        # rotate localization precisions in 3D
        cov_matrix = np.array(
            [
                [sx_**2, 0, 0],
                [0, sy_**2, 0],
                [0, 0, sz_**2],
            ],
            dtype=np.float32,
        )  # covariance matrix (CM)
        cov_rot = rot_matrix @ cov_matrix @ rot_matrixT  # rotated CM
        cri = inverse_3x3(cov_rot)  # inverse of rotated CM
        dcr = determinant_3x3(cov_rot)  # determinant of rotated CM

        # draw a localization in 2D - sum z coordinates;
        # PDF of a rotated gaussian in 3D is calculated and
        # image is summed over z axis
        for i in range(i_min, i_max):
            b = np.float32(i + 0.5 - y_)
            for j in range(j_min, j_max):
                a = np.float32(j + 0.5 - x_)
                for k in range(k_min, k_max):
                    c = np.float32(k + 0.5 - z_)
                    exponent = (
                        a * a * cri[0, 0]
                        + a * b * cri[0, 1]
                        + a * c * cri[0, 2]
                        + a * b * cri[1, 0]
                        + b * b * cri[1, 1]
                        + b * c * cri[1, 2]
                        + a * c * cri[2, 0]
                        + b * c * cri[2, 1]
                        + c * c * cri[2, 2]
                    )  # Mahalanobis distance
                    image[i, j] += np.exp(-0.5 * exponent) / (
                        ((2 * np.pi) ** 3 * dcr) ** 0.5
                    )


@numba.njit
def inverse_3x3(a: np.ndarray) -> np.ndarray:
    """Calculate inverse of a 3x3 matrix. This function is faster than
    ``np.linalg.inv``.

    Parameters
    ----------
    a : np.ndarray
        3x3 matrix.

    Returns
    -------
    c : np.ndarray
        Inverse of ``a``.
    """
    c = np.zeros((3, 3), dtype=np.float32)
    det = determinant_3x3(a)

    c[0, 0] = (a[1, 1] * a[2, 2] - a[1, 2] * a[2, 1]) / det
    c[0, 1] = (a[0, 2] * a[2, 1] - a[0, 1] * a[2, 2]) / det
    c[0, 2] = (a[0, 1] * a[1, 2] - a[0, 2] * a[1, 1]) / det

    c[1, 0] = (a[1, 2] * a[2, 0] - a[1, 0] * a[2, 2]) / det
    c[1, 1] = (a[0, 0] * a[2, 2] - a[0, 2] * a[2, 0]) / det
    c[1, 2] = (a[0, 2] * a[1, 0] - a[0, 0] * a[1, 2]) / det

    c[2, 0] = (a[1, 0] * a[2, 1] - a[1, 1] * a[2, 0]) / det
    c[2, 1] = (a[0, 1] * a[2, 0] - a[0, 0] * a[2, 1]) / det
    c[2, 2] = (a[0, 0] * a[1, 1] - a[0, 1] * a[1, 0]) / det

    return c


@numba.njit
def determinant_3x3(a: np.ndarray) -> np.float32:
    """Calculate determinant of a 3x3 matrix. This function is faster
    than ``np.linalg.det``.

    Parameters
    ----------
    a : np.ndarray
        3x3 matrix.

    Returns
    -------
    det : float
        Determinant of ``a``.
    """
    det = np.float32(
        a[0, 0] * (a[1, 1] * a[2, 2] - a[1, 2] * a[2, 1])
        - a[0, 1] * (a[1, 0] * a[2, 2] - a[2, 0] * a[1, 2])
        + a[0, 2] * (a[1, 0] * a[2, 1] - a[2, 0] * a[1, 1])
    )
    return det


def render_hist(
    locs: pd.DataFrame,
    oversampling: float,
    y_min: float,
    x_min: float,
    y_max: float,
    x_max: float,
    ang: tuple[float, float, float] | None = None,
) -> tuple[int, np.ndarray]:
    """Render localizations with no blur by assigning them to pixels.

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations to be rendered.
    oversampling : float
        Number of super-resolution pixels per camera pixel.
    y_min, x_min : float
        Minimum y and x coordinates to be rendered (camera pixels)
    y_max, x_max : float
        Maximum y and x coordinates to be rendered (camera pixels)
    ang : tuple (default=None)
        Rotation angles of locs around x, y and z axes in radians. If
        None, locs are not rotated.

    Returns
    -------
    n : int
        Number of localizations rendered.
    image : np.ndarray
        Rendered image.
    """
    image, n_pixel_y, n_pixel_x, x, y, in_view = _render_setup(
        locs["x"].to_numpy(),
        locs["y"].to_numpy(),
        oversampling,
        y_min,
        x_min,
        y_max,
        x_max,
    )
    if ang:
        x, y, _, _ = locs_rotation(
            locs,
            oversampling,
            x_min,
            x_max,
            y_min,
            y_max,
            ang,
        )
    _fill(image, x, y)
    n = len(x)
    return n, image


@numba.jit(nopython=True, nogil=True)
def render_hist3d(
    x: np.ndarray,
    y: np.ndarray,
    z: np.ndarray,
    oversampling: float,
    y_min: float,
    x_min: float,
    y_max: float,
    x_max: float,
    z_min: float,
    z_max: float,
    pixelsize: float,
) -> tuple[int, np.ndarray]:
    """Render localizations in 3D with no blur by assigning them to
    pixels.

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations to be rendered.
    oversampling : float (default=1)
        Number of super-resolution pixels per camera pixel.
    y_min, x_min : float
        Minimum y and x coordinates to be rendered (camera pixels).
    y_max, x_max : float
        Maximum y and x coordinates to be rendered (camera pixels).
    z_min : float
        Minimum z coordinate to be rendered (nm).
    z_max : float
        Maximum z coordinate to be rendered (nm).
    pixelsize : float
        Camera pixel size in nm, used for converting z coordinates.

    Returns
    -------
    n : int
        Number of localizations rendered.
    image : np.ndarray
        Rendered 3D image.
    """
    z_min = z_min / pixelsize
    z_max = z_max / pixelsize

    image, n_pixel_y, n_pixel_x, n_pixel_z, x, y, z, in_view = _render_setup3d(
        x,
        y,
        z,
        oversampling,
        y_min,
        x_min,
        y_max,
        x_max,
        z_min,
        z_max,
        pixelsize,
    )
    _fill3d(image, x, y, z)
    n = len(x)
    z *= pixelsize  # convert back to nm
    return n, image


@numba.jit(nopython=True, nogil=True)
def render_hist3d_anisotropic(
    x: np.ndarray,
    y: np.ndarray,
    z: np.ndarray,
    oversampling_x: float,
    oversampling_y: float,
    oversampling_z: float,
    y_min: float,
    x_min: float,
    y_max: float,
    x_max: float,
    z_min: float,
    z_max: float,
    pixelsize: float,
) -> tuple[int, np.ndarray]:
    """Render localizations in 3D with no blur by assigning them to
    pixels. Allows for different pixel sizes in x, y and z
    (oversampling).

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations to be rendered.
    oversampling_x, oversampling_y, oversampling_z : float (default=1)
        Number of super-resolution pixels per camera pixel in x, y, and
        z directions.
    y_min, x_min : float
        Minimum y and x coordinates to be rendered (camera pixels).
    y_max, x_max : float
        Maximum y and x coordinates to be rendered (camera pixels).
    z_min : float
        Minimum z coordinate to be rendered (nm).
    z_max : float
        Maximum z coordinate to be rendered (nm).
    pixelsize : float
        Camera pixel size in nm, used for converting z coordinates.

    Returns
    -------
    n : int
        Number of localizations rendered.
    image : np.ndarray
        Rendered 3D image.
    """
    z_min = z_min / pixelsize
    z_max = z_max / pixelsize

    image, n_pixel_y, n_pixel_x, n_pixel_z, x, y, z, in_view = (
        _render_setup3d_anisotropic(
            x,
            y,
            z,
            oversampling_x,
            oversampling_y,
            oversampling_z,
            y_min,
            x_min,
            y_max,
            x_max,
            z_min,
            z_max,
            pixelsize,
        )
    )
    _fill3d(image, x, y, z)
    n = len(x)
    z *= pixelsize  # convert back to nm
    return n, image


def render_gaussian(
    locs: pd.DataFrame,
    oversampling: float,
    y_min: float,
    x_min: float,
    y_max: float,
    x_max: float,
    min_blur_width: float,
    ang: tuple[float, float, float] | None = None,
) -> tuple[int, np.ndarray]:
    """Render localizations with with individual localization precision
    which differs in x and y.

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations to be rendered.
    oversampling : float
        Number of super-resolution pixels per camera pixel.
    y_min, y_max : float
        Minimum and maximum y coordinates to be rendered (camera pixels).
    x_min, x_max : float
        Minimum and maximum x coordinates to be rendered (camera pixels).
    min_blur_width : float
        Minimum localization precision (camera pixels).
    ang : tuple, optional
        Rotation angles of localizations around x, y and z axes in
        radians. If None, localizations are not rotated.

    Returns
    -------
    n : int
        Number of localizations rendered.
    image : np.ndarray
        Rendered image.
    """
    image, n_pixel_y, n_pixel_x, x, y, in_view = _render_setup(
        locs["x"].to_numpy(),
        locs["y"].to_numpy(),
        oversampling,
        y_min,
        x_min,
        y_max,
        x_max,
    )

    if not ang:  # not rotated
        blur_width = oversampling * np.maximum(
            locs["lpx"].to_numpy(), min_blur_width
        )
        blur_height = oversampling * np.maximum(
            locs["lpy"].to_numpy(), min_blur_width
        )
        sy = blur_height[in_view]
        sx = blur_width[in_view]

        _fill_gaussian(image, x, y, sx, sy, n_pixel_x, n_pixel_y)

    else:  # rotated
        x, y, in_view, z = locs_rotation(
            locs,
            oversampling,
            x_min,
            x_max,
            y_min,
            y_max,
            ang,
        )
        blur_width = oversampling * np.maximum(
            locs["lpx"].to_numpy(), min_blur_width
        )
        blur_height = oversampling * np.maximum(
            locs["lpy"].to_numpy(), min_blur_width
        )
        # if lpz not found, make it twice the mean of lpx and lpy
        if "lpz" in locs:
            lpz = locs["lpz"].to_numpy()
        else:
            lpz = 2 * locs[["lpx", "lpy"]].to_numpy().mean(axis=1)
        blur_depth = oversampling * np.maximum(lpz, min_blur_width)

        sy = blur_height[in_view]
        sx = blur_width[in_view]
        sz = blur_depth[in_view]

        _fill_gaussian_rot(
            image, x, y, z, sx, sy, sz, n_pixel_x, n_pixel_y, ang
        )

    n = len(x)
    return n, image


def render_gaussian_iso(
    locs: pd.DataFrame,
    oversampling: float,
    y_min: float,
    x_min: float,
    y_max: float,
    x_max: float,
    min_blur_width: float,
    ang: tuple[float, float, float] | None = None,
) -> tuple[int, np.ndarray]:
    """Same as ``render_gaussian``, but uses the same localization
    precision in x and y."""
    image, n_pixel_y, n_pixel_x, x, y, in_view = _render_setup(
        locs["x"].to_numpy(),
        locs["y"].to_numpy(),
        oversampling,
        y_min,
        x_min,
        y_max,
        x_max,
    )

    if not ang:  # not rotated
        blur_width = oversampling * np.maximum(
            locs["lpx"].to_numpy(), min_blur_width
        )
        blur_height = oversampling * np.maximum(
            locs["lpy"].to_numpy(), min_blur_width
        )
        sy = (blur_height[in_view] + blur_width[in_view]) / 2
        sx = sy

        _fill_gaussian(image, x, y, sx, sy, n_pixel_x, n_pixel_y)

    else:  # rotated
        x, y, in_view, z = locs_rotation(
            locs,
            oversampling,
            x_min,
            x_max,
            y_min,
            y_max,
            ang,
        )
        blur_width = oversampling * np.maximum(
            locs["lpx"].to_numpy(), min_blur_width
        )
        blur_height = oversampling * np.maximum(
            locs["lpy"].to_numpy(), min_blur_width
        )
        # for now, let lpz be twice the mean of lpx and lpy
        if "lpz" in locs:
            lpz = locs["lpz"].to_numpy()
        else:
            lpz = 2 * locs[["lpx", "lpy"]].to_numpy().mean(axis=1)
        blur_depth = oversampling * np.maximum(lpz, min_blur_width)

        sy = (blur_height[in_view] + blur_width[in_view]) / 2
        sx = sy
        sz = blur_depth[in_view]

        _fill_gaussian_rot(
            image, x, y, z, sx, sy, sz, n_pixel_x, n_pixel_y, ang
        )

    return len(x), image


def render_convolve(
    locs: pd.DataFrame,
    oversampling: float,
    y_min: float,
    x_min: float,
    y_max: float,
    x_max: float,
    min_blur_width: float,
    ang: tuple[float, float, float] | None = None,
) -> tuple[int, np.ndarray]:
    """Render localizations with with global localization precision,
    i.e. each localization is blurred by the median localization
    precision in x and y.

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations to be rendered.
    oversampling : float
        Number of super-resolution pixels per camera pixel.
    y_min, x_min : float
        Minimum y and x coordinates to be rendered (camera pixels).
    y_max, x_max : float
        Maximum y and x coordinates to be rendered (camera pixels).
    min_blur_width : float
        Minimum localization precision (camera pixels).
    ang : tuple, optional
        Rotation angles of localizations around x, y and z axes in
        radians. If None, localizations are not rotated.

    Returns
    -------
    n : int
        Number of localizations rendered.
    image : np.ndarray
        Rendered image.
    """
    image, n_pixel_y, n_pixel_x, x, y, in_view = _render_setup(
        locs["x"].to_numpy(),
        locs["y"].to_numpy(),
        oversampling,
        y_min,
        x_min,
        y_max,
        x_max,
    )
    if ang:  # rotate
        x, y, in_view, _ = locs_rotation(
            locs,
            oversampling,
            x_min,
            x_max,
            y_min,
            y_max,
            ang,
        )

    n = len(x)
    if n == 0:
        return 0, image
    else:
        _fill(image, x, y)
        blur_width = oversampling * max(
            np.median(locs["lpx"].to_numpy()[in_view]), min_blur_width
        )
        blur_height = oversampling * max(
            np.median(locs["lpy"].to_numpy()[in_view]), min_blur_width
        )
        return n, _fftconvolve(image, blur_width, blur_height)


def render_smooth(
    locs: pd.DataFrame,
    oversampling: float,
    y_min: float,
    x_min: float,
    y_max: float,
    x_max: float,
    ang: tuple[float, float, float] | None = None,
) -> tuple[int, np.ndarray]:
    """Render localizations with with blur of one display pixel (set by
    oversampling).

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations to be rendered.
    oversampling : float
        Number of super-resolution pixels per camera pixel.
    y_min, x_min : float
        Minimum y and x coordinates to be rendered (camera pixels).
    y_max, x_max : float
        Maximum y and x coordinates to be rendered (camera pixels).
    ang : tuple, optional
        Rotation angles of localizations around x, y and z axes in
        radians. If None, localizations are not rotated.

    Returns
    -------
    n : int
        Number of localizations rendered.
    image : np.array
        Rendered image.
    """
    image, n_pixel_y, n_pixel_x, x, y, in_view = _render_setup(
        locs["x"].to_numpy(),
        locs["y"].to_numpy(),
        oversampling,
        y_min,
        x_min,
        y_max,
        x_max,
    )

    if ang:
        x, y, _, _ = locs_rotation(
            locs,
            oversampling,
            x_min,
            x_max,
            y_min,
            y_max,
            ang,
        )

    n = len(x)
    if n == 0:
        return 0, image
    else:
        _fill(image, x, y)
        return n, _fftconvolve(image, 1, 1)


def _fftconvolve(
    image: np.ndarray,
    blur_width: float,
    blur_height: float,
) -> np.ndarray:
    """Blur (convolves) 2D image using fast fourier transform.

    Parameters
    ----------
    image : np.ndarray
        Image with rendered but not blurred localizations.
    blur_width, blur_height : float
        Blur width and height in pixels.

    Returns
    -------
    image : np.ndarray
        Blurred image.
    """
    kernel_width = 10 * int(np.round(blur_width)) + 1
    kernel_height = 10 * int(np.round(blur_height)) + 1
    kernel_y = signal.windows.gaussian(kernel_height, blur_height)
    kernel_x = signal.windows.gaussian(kernel_width, blur_width)
    kernel = np.outer(kernel_y, kernel_x)
    kernel /= kernel.sum()
    image = signal.fftconvolve(image, kernel, mode="same")
    return image.astype(np.float32)


def rotation_matrix(angx: float, angy: float, angz: float) -> Rotation:
    """Find rotation matrix given rotation angles around axes.

    Parameters
    ----------
    angx, angy, angz : float
        Rotation angles around x, y and z axes in radians.

    Returns
    -------
    scipy.spatial.transform.Rotation
        Scipy class that can be applied to rotate an Nx3 np.ndarray
    """
    rot_mat_x = np.array(
        [
            [1.0, 0.0, 0.0],
            [0.0, np.cos(angx), np.sin(angx)],
            [0.0, -np.sin(angx), np.cos(angx)],
        ]
    )  # rotation matrix around x axis
    rot_mat_y = np.array(
        [
            [np.cos(angy), 0.0, np.sin(angy)],
            [0.0, 1.0, 0.0],
            [-np.sin(angy), 0.0, np.cos(angy)],
        ]
    )  # rotation matrix around y axis
    rot_mat_z = np.array(
        [
            [np.cos(angz), -np.sin(angz), 0.0],
            [np.sin(angz), np.cos(angz), 0.0],
            [0.0, 0.0, 1.0],
        ]
    )  # rotation matrix around z axis
    rot_mat = rot_mat_x @ rot_mat_y @ rot_mat_z
    return Rotation.from_matrix(rot_mat)


def locs_rotation(
    locs: pd.DataFrame,
    oversampling: float,
    x_min: float,
    x_max: float,
    y_min: float,
    y_max: float,
    ang: tuple[float, float, float],
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    """Rotate localizations within a FOV.

    Parameters
    ----------
    locs : pd.DataFrame
        Localizations to be rotated.
    oversampling : float
        Number of super-resolution pixels per camera pixel.
    y_min, x_min : float
        Minimum y and x coordinate to be rendered (camera pixels).
    y_max, x_max : float
        Maximum y and x coordinate to be rendered (camera pixels).
    ang : tuple
        Rotation angles of localizations around x, y and z axes in
        radians.

    Returns
    -------
    x : np.ndarray
        New (rotated) x coordinates
    y : np.ndarray
        New y coordinates
    in_view : np.ndarray
        Indeces of locs that are rendered
    z : np.ndarray
        New z coordinates
    """
    # z is translated to pixels
    locs_coord = locs[["x", "y", "z"]].to_numpy()

    # x and y are in range (x_min/y_min, x_max/y_max) so they need to be
    # shifted (scipy rotation is around origin)
    locs_coord[:, 0] -= x_min + (x_max - x_min) / 2
    locs_coord[:, 1] -= y_min + (y_max - y_min) / 2

    # rotate locs
    R = rotation_matrix(ang[0], ang[1], ang[2])
    locs_coord = R.apply(locs_coord)

    # unshift locs
    locs_coord[:, 0] += x_min + (x_max - x_min) / 2
    locs_coord[:, 1] += y_min + (y_max - y_min) / 2

    # output
    x = locs_coord[:, 0]
    y = locs_coord[:, 1]
    z = locs_coord[:, 2]
    in_view = (x > x_min) & (y > y_min) & (x < x_max) & (y < y_max)
    x = x[in_view]
    y = y[in_view]
    z = z[in_view]
    x = oversampling * (x - x_min)
    y = oversampling * (y - y_min)
    z *= oversampling
    return x, y, in_view, z
