""" Utility functions for template matching.

    Copyright (c) 2023 European Molecular Biology Laboratory

    Author: Valentin Maurer <valentin.maurer@embl-hamburg.de>
"""
import os
import traceback
import pickle
from shutil import move
from tempfile import mkstemp
from itertools import product
from typing import Tuple, Dict, Callable
from concurrent.futures import ThreadPoolExecutor

import numpy as np
from numpy.typing import NDArray
from scipy.spatial import ConvexHull
from scipy.spatial.transform import Rotation

from .helpers import quaternion_to_rotation_matrix, load_quaternions_by_angle
from .extensions import max_euclidean_distance
from .matching_memory import estimate_ram_usage


def handle_traceback(last_type, last_value, last_traceback):
    """
    Handle sys.exc_info().

    Parameters
    ----------
    last_type : type
        The type of the last exception.
    last_value :
        The value of the last exception.
    last_traceback : traceback
        The traceback object encapsulating the call stack at the point
        where the exception originally occurred.

    Raises
    ------
    Exception
        Re-raises the last exception.
    """
    if last_type is None:
        return None
    traceback.print_tb(last_traceback)
    raise Exception(last_value)
    # raise last_type(last_value)


def generate_tempfile_name(suffix=None):
    """
    Returns the path to a potential temporary file location. If the environment
    variable TME_TMPDIR is defined, the temporary file will be created there.
    Otherwise the default tmp directory will be used.

    Parameters
    ----------
    suffix : str, optional
        File suffix. By default the file has no suffix.

    Returns
    -------
    str
        The generated filename
    """
    tmp_dir = os.environ.get("TMPDIR", None)
    _, filename = mkstemp(suffix=suffix, dir=tmp_dir)
    return filename


def array_to_memmap(arr: NDArray, filename: str = None) -> str:
    """
    Converts a numpy array to a np.memmap.

    Parameters
    ----------
    arr : np.ndarray
        The numpy array to be converted.
    filename : str, optional
        Desired filename for the memmap. If not provided, a temporary
        file will be created.

    Notes
    -----
        If the environment variable TME_TMPDIR is defined, the temporary
        file will be created there. Otherwise the default tmp directory
        will be used.

    Returns
    -------
    str
        The filename where the memmap was written to.
    """
    if filename is None:
        filename = generate_tempfile_name()

    shape, dtype = arr.shape, arr.dtype
    arr_memmap = np.memmap(filename, mode="w+", dtype=dtype, shape=shape)

    arr_memmap[:] = arr[:]
    arr_memmap.flush()

    return filename


def memmap_to_array(arr: NDArray) -> NDArray:
    """
    Converts a np.memmap into an numpy array.

    Parameters
    ----------
    arr : np.memmap
        The numpy array to be converted.

    Returns
    -------
    np.ndarray
        The converted array.
    """
    if type(arr) == np.memmap:
        memmap_filepath = arr.filename
        arr = np.array(arr)
        os.remove(memmap_filepath)
    return arr


def close_memmap(arr: np.ndarray) -> None:
    """
    Remove the file associated with a numpy memmap array.

    Parameters
    ----------
    arr : np.ndarray
        The numpy array which might be a memmap.
    """
    try:
        os.remove(arr.filename)
        # arr._mmap.close()
    except Exception:
        pass


def write_pickle(data: object, filename: str) -> None:
    """
    Serialize and write data to a file invalidating the input data in
    the process. This function  uses type-specific serialization for
    certain objects, such as np.memmap, for optimized storage. Other
    objects are serialized using standard pickle.

    Parameters
    ----------
    data : iterable or object
        The data to be serialized.
    filename : str
        The name of the file where the serialized data will be written.

    See Also
    --------
    :py:meth:`load_pickle`
    """
    if type(data) not in (list, tuple):
        data = (data,)

    dirname = os.path.dirname(filename)
    with open(filename, "wb") as ofile, ThreadPoolExecutor() as executor:
        for i in range(len(data)):
            futures = []
            item = data[i]
            if isinstance(item, np.memmap):
                _, new_filename = mkstemp(suffix=".mm", dir=dirname)
                new_item = ("np.memmap", item.shape, item.dtype, new_filename)
                futures.append(executor.submit(move, item.filename, new_filename))
                item = new_item
            pickle.dump(item, ofile)
        for future in futures:
            future.result()


def load_pickle(filename: str) -> object:
    """
    Load and deserialize data written by :py:meth:`write_pickle`.

    Parameters
    ----------
    filename : str
        The name of the file to read and deserialize data from.

    Returns
    -------
    object or iterable
        The deserialized data.

    See Also
    --------
    :py:meth:`write_pickle`
    """

    def _load_pickle(file_handle):
        try:
            while True:
                yield pickle.load(file_handle)
        except EOFError:
            pass

    def _is_pickle_memmap(data):
        ret = False
        if type(data[0]) == str:
            if data[0] == "np.memmap":
                ret = True
        return ret

    items = []
    with open(filename, "rb") as ifile:
        for data in _load_pickle(ifile):
            if isinstance(data, tuple):
                if _is_pickle_memmap(data):
                    _, shape, dtype, filename = data
                    data = np.memmap(filename, shape=shape, dtype=dtype)
            items.append(data)
    return items[0] if len(items) == 1 else items


def compute_parallelization_schedule(
    shape1: NDArray,
    shape2: NDArray,
    max_cores: int,
    max_ram: int,
    matching_method: str,
    backend: str = None,
    split_only_outer: bool = False,
    shape1_padding: NDArray = None,
    analyzer_method: str = None,
    max_splits: int = 256,
    float_nbytes: int = 4,
    complex_nbytes: int = 8,
    integer_nbytes: int = 4,
) -> Tuple[Dict, int, int]:
    """
    Computes a parallelization schedule for a given computation.

    This function estimates the amount of memory that would be used by a computation
    and breaks down the computation into smaller parts that can be executed in parallel
    without exceeding the specified limits on the number of cores and memory.

    Parameters
    ----------
    shape1 : NDArray
        The shape of the first input tensor.
    shape1_padding : NDArray, optional
        Padding for shape1 used for each split. None by defauly
    shape2 : NDArray
        The shape of the second input tensor.
    max_cores : int
        The maximum number of cores that can be used.
    max_ram : int
        The maximum amount of memory that can be used.
    matching_method : str
        The metric used for scoring the computations.
    backend : str, optional
        Backend used for computations.
    split_only_outer : bool, optional
        Whether only outer splits sould be considered.
    analyzer_method : str
        The method used for score analysis.
    max_splits : int, optional
        The maximum number of parts that the computation can be split into,
        by default 256.
    float_nbytes : int
        Number of bytes of the used float, e.g. 4 for float32.
    complex_nbytes : int
        Number of bytes of the used complex, e.g. 8 for complex64.
    integer_nbytes : int
        Number of bytes of the used integer, e.g. 4 for int32.

    Notes
    -----
        This function assumes that no residual memory remains after each split,
        which not always holds true, e.g. when using
        :py:class:`tme.analyzer.MaxScoreOverRotations`.

    Returns
    -------
    dict
        The optimal splits for each axis of the first input tensor.
    int
        The number of outer jobs.
    int
        The number of inner jobs per outer job.
    """
    shape1, shape2 = np.array(shape1), np.array(shape2)
    if shape1_padding is None:
        shape1_padding = np.zeros_like(shape1)
    core_assignments = []
    for i in range(1, int(max_cores**0.5) + 1):
        if max_cores % i == 0:
            core_assignments.append((i, max_cores // i))
            core_assignments.append((max_cores // i, i))

    if split_only_outer:
        core_assignments = [(1, max_cores)]

    possible_params, split_axis = [], np.argmax(shape1)
    split_factor, n_splits = [1 for _ in range(len(shape1))], 0
    while n_splits <= max_splits:
        splits = {k: split_factor[k] for k in range(len(split_factor))}
        array_slices = split_numpy_array_slices(shape=shape1, splits=splits)
        array_widths = [
            tuple(x.stop - x.start for x in split) for split in array_slices
        ]
        n_splits = np.prod(list(splits.values()))

        for inner_cores, outer_cores in core_assignments:
            if outer_cores > n_splits:
                continue
            ram_usage = [
                estimate_ram_usage(
                    shape1=np.add(shp, shape1_padding),
                    shape2=shape2,
                    matching_method=matching_method,
                    analyzer_method=analyzer_method,
                    backend=backend,
                    ncores=inner_cores,
                    float_nbytes=float_nbytes,
                    complex_nbytes=complex_nbytes,
                    integer_nbytes=integer_nbytes,
                )
                for shp in array_widths
            ]
            max_usage = 0
            for i in range(0, len(ram_usage), outer_cores):
                usage = np.sum(ram_usage[i : (i + outer_cores)])
                if usage > max_usage:
                    max_usage = usage

            inits = n_splits // outer_cores
            if max_usage < max_ram:
                possible_params.append(
                    (*split_factor, outer_cores, inner_cores, n_splits, inits)
                )
        split_factor[split_axis] += 1
        split_axis += 1
        if split_axis == shape1.size:
            split_axis = 0

    possible_params = np.array(possible_params)
    if not len(possible_params):
        print(
            "No suitable assignment found. Consider increasing "
            "max_ram or decrease max_cores."
        )
        return None, None

    init = possible_params.shape[1] - 1
    possible_params = possible_params[
        np.lexsort((possible_params[:, init], possible_params[:, (init - 1)]))
    ]
    splits = {k: possible_params[0, k] for k in range(shape1.size)}
    core_assignment = (
        possible_params[0, shape1.size],
        possible_params[0, (shape1.size + 1)],
    )

    return splits, core_assignment


def centered(arr: NDArray, newshape: Tuple[int]) -> NDArray:
    """
    Extract the centered portion of an array based on a new shape.

    Parameters
    ----------
    arr : NDArray
        Input array.
    newshape : tuple
        Desired shape for the central portion.

    Returns
    -------
    NDArray
        Central portion of the array with shape `newshape`.

    References
    ----------
    .. [1] https://github.com/scipy/scipy/blob/v1.11.2/scipy/signal/_signaltools.py#L388
    """
    new_shape = np.asarray(newshape)
    current_shape = np.array(arr.shape)
    starts = (current_shape - new_shape) // 2
    stops = starts + newshape
    box = tuple(slice(start, stop) for start, stop in zip(starts, stops))
    return arr[box]


def centered_mask(arr: NDArray, newshape: Tuple[int]) -> NDArray:
    """
    Mask the centered portion of an array based on a new shape.

    Parameters
    ----------
    arr : NDArray
        Input array.
    newshape : tuple
        Desired shape for the mask.

    Returns
    -------
    NDArray
        Array with central portion unmasked and the rest set to 0.
    """
    new_shape = np.asarray(newshape)
    current_shape = np.array(arr.shape)
    starts = (current_shape - new_shape) // 2
    stops = starts + newshape
    box = tuple(slice(start, stop) for start, stop in zip(starts, stops))
    mask = np.zeros_like(arr)
    mask[box] = 1
    arr *= mask
    return arr


def apply_convolution_mode(
    arr: NDArray,
    convolution_mode: str,
    s1: Tuple[int],
    s2: Tuple[int],
    mask_output: bool = False,
) -> NDArray:
    """
    Applies convolution_mode to arr.

    Parameters
    ----------
    arr : NDArray
        Numpy array containing convolution result of arrays with shape s1 and s2.
    convolution_mode : str
        Analogous to mode in ``scipy.signal.convolve``:

        +---------+----------------------------------------------------------+
        | 'full'  | returns full template matching result of the inputs.     |
        +---------+----------------------------------------------------------+
        | 'valid' | returns elements that do not rely on zero-padding..      |
        +---------+----------------------------------------------------------+
        | 'same'  | output is the same size as s1.                           |
        +---------+----------------------------------------------------------+
    s1 : tuple
        Tuple of integers corresponding to shape of convolution array 1.
    s2 : tuple
        Tuple of integers corresponding to shape of convolution array 2.
    mask_output : bool, optional
        Whether to mask values outside of convolution_mode rather than
        removing them. Defaults to False.

    Returns
    -------
    NDArray
        The numpy array after applying the convolution mode.

    References
    ----------
    .. [1] https://github.com/scipy/scipy/blob/v1.11.2/scipy/signal/_signaltools.py#L519
    """
    # This removes padding to next fast fourier length
    arr = arr[tuple(slice(s1[i] + s2[i] - 1) for i in range(len(s1)))]

    if convolution_mode not in ("full", "same", "valid"):
        raise ValueError("Supported convolution_mode are 'full', 'same' and 'valid'.")

    func = centered_mask if mask_output else centered
    if convolution_mode == "full":
        return arr
    elif convolution_mode == "same":
        return func(arr, s1)
    elif convolution_mode == "valid":
        valid_shape = [s1[i] - s2[i] + s2[i] % 2 for i in range(arr.ndim)]
        return func(arr, valid_shape)


def compute_full_convolution_index(
    outer_shape: Tuple[int],
    inner_shape: Tuple[int],
    outer_split: Tuple[slice],
    inner_split: Tuple[slice],
) -> Tuple[slice]:
    """
    Computes the position of the convolution of pieces in the full convolution.

    Parameters
    ----------
    outer_shape : tuple
        Tuple of integers corresponding to the shape of the outer array.
    inner_shape : tuple
        Tuple of integers corresponding to the shape of the inner array.
    outer_split : tuple
        Tuple of slices used to split outer array
        (see :py:meth:`split_numpy_array_slices`).
    inner_split : tuple
        Tuple of slices used to split inner array
        (see :py:meth:`split_numpy_array_slices`).

    Returns
    -------
    tuple
        Tuple of slices corresponding to the position of the given convolution
        in the full convolution.
    """
    outer_shape = np.asarray(outer_shape)
    inner_shape = np.asarray(inner_shape)

    outer_width = np.array([outer.stop - outer.start for outer in outer_split])
    inner_width = np.array([inner.stop - inner.start for inner in inner_split])
    convolution_shape = outer_width + inner_width - 1

    end_inner = np.array([inner.stop for inner in inner_split]).astype(int)
    start_outer = np.array([outer.start for outer in outer_split]).astype(int)

    offsets = start_outer + inner_shape - end_inner

    score_slice = tuple(
        (slice(offset, offset + shape))
        for offset, shape in zip(offsets, convolution_shape)
    )

    return score_slice


def split_numpy_array_slices(
    shape: NDArray, splits: Dict, margin: NDArray = None
) -> Tuple[slice]:
    """
    Returns a tuple of slices to subset a numpy array into pieces along multiple axes.

    Parameters
    ----------
    shape : NDArray
        Shape of the array to split.
    splits : dict
        A dictionary where the keys are the axis numbers and the values
        are the number of splits along that axis.
    margin : NDArray, optional
        Padding on the left hand side of the array.

    Returns
    -------
    tuple
        A tuple of slices, where each slice corresponds to a split along an axis.
    """
    ndim = len(shape)
    if margin is None:
        margin = np.zeros(ndim, dtype=int)
    splits = {k: max(splits.get(k, 0), 1) for k in range(ndim)}
    new_shape = np.divide(shape, [splits.get(i, 1) for i in range(ndim)]).astype(int)

    slice_list = [
        tuple(
            (slice(max((n_splits * length) - margin[axis], 0), (n_splits + 1) * length))
            if n_splits < splits.get(axis, 1) - 1
            else (slice(max((n_splits * length) - margin[axis], 0), shape[axis]))
            for n_splits in range(splits.get(axis, 1))
        )
        for length, axis in zip(new_shape, splits.keys())
    ]

    splits = tuple(product(*slice_list))

    return splits


def get_rotation_matrices(
    angular_sampling: float, dim: int = 3, use_optimized_set: bool = True
) -> NDArray:
    """
    Returns rotation matrices in format k x dim x dim, where k is determined
    by ``angular_sampling``.

    Parameters
    ----------
    angular_sampling : float
        The angle in degrees used for the generation of rotation matrices.
    dim : int, optional
        Dimension of the rotation matrices.
    use_optimized_set : bool, optional
        Whether to use pre-computed rotational sets with more optimal sampling.
        Currently only available when dim=3.

    Notes
    -----
        For the case of dim = 3 optimized rotational sets are used, otherwise
        QR-decomposition.

    Returns
    -------
    NDArray
        Array of shape (k, dim, dim) containing k rotation matrices.
    """
    if dim == 3 and use_optimized_set:
        quaternions, *_ = load_quaternions_by_angle(angular_sampling)
        ret = quaternion_to_rotation_matrix(quaternions)
    else:
        num_rotations = dim * (dim - 1) // 2
        k = int((360 / angular_sampling) ** num_rotations)
        As = np.random.randn(k, dim, dim)
        ret, _ = np.linalg.qr(As)
        dets = np.linalg.det(ret)
        neg_dets = dets < 0
        ret[neg_dets, :, -1] *= -1
    return ret


def minimum_enclosing_box(
    coordinates: NDArray,
    margin: NDArray = None,
    use_geometric_center: bool = False,
) -> Tuple[int]:
    """
    Computes the minimal enclosing box around coordinates with margin.

    Parameters
    ----------
    coordinates : NDArray
        Coordinates of which the enclosing box should be computed. The shape
        of this array should be [d, n] with d dimensions and n coordinates.
    margin : NDArray, optional
        Box margin. Defaults to None.
    use_geometric_center : bool, optional
        Whether the box should accommodate the geometric or the coordinate
        center. Defaults to False.

    Returns
    -------
    tuple
        Integers corresponding to the minimum enclosing box shape.
    """
    point_cloud = np.asarray(coordinates)
    dim = point_cloud.shape[0]
    point_cloud = point_cloud - point_cloud.min(axis=1)[:, None]

    margin = np.zeros(dim) if margin is None else margin
    margin = np.asarray(margin).astype(int)

    norm_cloud = point_cloud - point_cloud.mean(axis=1)[:, None]
    # Adding one avoids clipping during scipy.ndimage.affine_transform
    shape = np.repeat(
        np.ceil(2 * np.linalg.norm(norm_cloud, axis=0).max()) + 1, dim
    ).astype(int)
    if use_geometric_center:
        hull = ConvexHull(point_cloud.T)
        distance, _ = max_euclidean_distance(point_cloud[:, hull.vertices].T)
        distance += np.linalg.norm(np.ones(dim))
        shape = np.repeat(np.rint(distance).astype(int), dim)

    return shape


def crop_input(
    target: "Density",
    template: "Density",
    target_mask: "Density" = None,
    template_mask: "Density" = None,
    map_cutoff: float = 0,
    template_cutoff: float = 0,
) -> Tuple[int]:
    """
    Crop target and template maps for efficient fitting. Input densities
    are cropped in place.

    Parameters
    ----------
    target : Density
        Target to be fitted on.
    template : Density
        Template to fit onto the target.
    target_mask : Density, optional
        Path to mask of target. Will be croppped like target.
    template_mask : Density, optional
        Path to mask of template. Will be cropped like template.
    map_cutoff : float, optional
        Cutoff value for trimming the target Density. Default is 0.
    map_cutoff : float, optional
        Cutoff value for trimming the template Density. Default is 0.

    Returns
    -------
    Tuple[int]
        Tuple containing reference fit index
    """
    convolution_shape_init = np.add(target.shape, template.shape) - 1
    # If target and template are aligned, fitting should return this index
    reference_fit = np.subtract(template.shape, 1)

    target_box = tuple(slice(0, x) for x in target.shape)
    if map_cutoff is not None:
        target_box = target.trim_box(cutoff=map_cutoff)

    target_mask_box = target_box
    if target_mask is not None and map_cutoff is not None:
        target_mask_box = target_mask.trim_box(cutoff=map_cutoff)
    target_box = tuple(
        slice(min(arr.start, mask.start), max(arr.stop, mask.stop))
        for arr, mask in zip(target_box, target_mask_box)
    )

    template_box = tuple(slice(0, x) for x in template.shape)
    if template_cutoff is not None:
        template_box = template.trim_box(cutoff=template_cutoff)

    template_mask_box = template_box
    if template_mask is not None and template_cutoff is not None:
        template_mask_box = template_mask.trim_box(cutoff=template_cutoff)
    template_box = tuple(
        slice(min(arr.start, mask.start), max(arr.stop, mask.stop))
        for arr, mask in zip(template_box, template_mask_box)
    )

    cut_right = np.array(
        [shape - x.stop for shape, x in zip(template.shape, template_box)]
    )
    cut_left = np.array([x.start for x in target_box])

    origin_difference = np.divide(target.origin - template.origin, target.sampling_rate)
    origin_difference = origin_difference.astype(int)

    target.adjust_box(target_box)
    template.adjust_box(template_box)

    if target_mask is not None:
        target_mask.adjust_box(target_box)
    if template_mask is not None:
        template_mask.adjust_box(template_box)

    reference_fit -= cut_right + cut_left + origin_difference

    convolution_shape = np.array(target.shape)
    convolution_shape += np.array(template.shape) - 1

    print(f"Cropped volume of target is: {target.shape}")
    print(f"Cropped volume of template is: {template.shape}")
    saving = 1 - (np.prod(convolution_shape)) / np.prod(convolution_shape_init)
    saving *= 100

    print(
        "Cropping changed array size from "
        f"{round(4*np.prod(convolution_shape_init)/1e6, 3)} MB "
        f"to {round(4*np.prod(convolution_shape)/1e6, 3)} MB "
        f"({'-' if saving > 0 else ''}{abs(round(saving, 2))}%)"
    )
    return reference_fit


def euler_to_rotationmatrix(angles: Tuple[float]) -> NDArray:
    """
    Convert Euler angles to a rotation matrix.

    Parameters
    ----------
    angles : tuple
        A tuple representing the Euler angles in degrees.

    Returns
    -------
    NDArray
        The generated rotation matrix.
    """
    if len(angles) == 1:
        angles = (angles, 0, 0)
    rotation_matrix = (
        Rotation.from_euler("zyx", angles, degrees=True).as_matrix().astype(np.float32)
    )
    return rotation_matrix


def euler_from_rotationmatrix(rotation_matrix: NDArray) -> Tuple:
    """
    Convert a rotation matrix to euler angles.

    Parameters
    ----------
    rotation_matrix : NDArray
        A 2 x 2 or 3 x 3 rotation matrix in z y x form.

    Returns
    -------
    Tuple
        The generate euler angles in degrees
    """
    if rotation_matrix.shape[0] == 2:
        temp_matrix = np.eye(3)
        temp_matrix[:2, :2] = rotation_matrix
        rotation_matrix = temp_matrix
    euler_angles = (
        Rotation.from_matrix(rotation_matrix)
        .as_euler("zyx", degrees=True)
        .astype(np.float32)
    )
    return euler_angles


def rigid_transform(
    coordinates: NDArray,
    rotation_matrix: NDArray,
    out: NDArray,
    translation: NDArray,
    use_geometric_center: bool = False,
    coordinates_mask: NDArray = None,
    out_mask: NDArray = None,
    center: NDArray = None,
) -> None:
    """
    Apply a rigid transformation (rotation and translation) to given coordinates.

    Parameters
    ----------
    coordinates : NDArray
        An array representing the coordinates to be transformed [d x N].
    rotation_matrix : NDArray
        The rotation matrix to be applied [d x d].
    translation : NDArray
        The translation vector to be applied [d].
    out : NDArray
        The output array to store the transformed coordinates.
    coordinates_mask : NDArray, optional
        An array representing the mask for the coordinates [d x T].
    out_mask : NDArray, optional
        The output array to store the transformed coordinates mask.
    use_geometric_center : bool, optional
        Whether to use geometric or coordinate center.

    Returns
    -------
    None
    """
    coordinate_dtype = coordinates.dtype
    center = coordinates.mean(axis=1) if center is None else center
    if not use_geometric_center:
        coordinates = coordinates - center[:, None]

    np.matmul(rotation_matrix, coordinates, out=out)
    if use_geometric_center:
        axis_max, axis_min = out.max(axis=1), out.min(axis=1)
        axis_difference = axis_max - axis_min
        translation = np.add(translation, center - axis_max + (axis_difference // 2))
    else:
        translation = np.add(translation, np.subtract(center, out.mean(axis=1)))

    out += translation[:, None]
    if coordinates_mask is not None and out_mask is not None:
        if not use_geometric_center:
            coordinates_mask = coordinates_mask - center[:, None]
        np.matmul(rotation_matrix, coordinates_mask, out=out_mask)
        out_mask += translation[:, None]

    if not use_geometric_center and coordinate_dtype != out.dtype:
        np.subtract(out.mean(axis=1), out.astype(int).mean(axis=1), out=translation)
        out += translation[:, None]


def _format_string(string: str) -> str:
    """
    Formats a string by adding quotation marks if it contains white spaces.

    Parameters
    ----------
    string : str
        Input string to be formatted.

    Returns
    -------
    str
        Formatted string with added quotation marks if needed.
    """
    if " " in string:
        return f"'{string}'"
    # Occurs e.g. for C1' atoms. The trailing whitespace is necessary.
    if string.count("'") == 1:
        return f'"{string}"'
    return string


def _format_mmcif_colunns(subdict: Dict) -> Dict:
    """
    Formats the columns of a mmcif dictionary.

    Parameters
    ----------
    subdict : dict
        Input dictionary where each key corresponds to a column and the
        values are iterables containing the column values.

    Returns
    -------
    dict
        Formatted dictionary with the columns of the mmcif file.
    """
    subdict = {k: [_format_string(s) for s in v] for k, v in subdict.items()}
    key_length = {
        key: len(max(value, key=lambda x: len(x), default=""))
        for key, value in subdict.items()
    }
    padded_subdict = {
        key: [s.ljust(key_length[key] + 1) for s in values]
        for key, values in subdict.items()
    }
    return padded_subdict


def create_mask(mask_type: str, **kwargs) -> NDArray:
    """
    Creates a mask of the specified type.

    Parameters
    ----------
    mask_type : str
        Type of the mask to be created. Can be "ellipse", "box", or "tube".
    kwargs : dict
        Additional parameters required by the mask creating functions.

    Returns
    -------
    NDArray
        The created mask.

    Raises
    ------
    ValueError
        If the mask_type is invalid.

    See Also
    --------
    :py:meth:`elliptical_mask`
    :py:meth:`box_mask`
    :py:meth:`tube_mask`
    """
    mapping = {"ellipse": elliptical_mask, "box": box_mask, "tube": tube_mask}
    if mask_type not in mapping:
        raise ValueError(f"mask_type has to be one of {','.join(mapping.keys())}")

    return mapping[mask_type](**kwargs)


def elliptical_mask(
    shape: Tuple[int], radius: Tuple[float], center: Tuple[int]
) -> NDArray:
    """
    Creates an ellipsoidal mask.

    Parameters
    ----------
    shape : tuple
        Shape of the mask to be created.
    radius : tuple
        Radius of the ellipse.
    center : tuple
        Center of the ellipse.

    Returns
    -------
    NDArray
        The created ellipsoidal mask.

    Raises
    ------
    ValueError
        If the length of center and radius is not one or the same as shape.

    Examples
    --------
    >>> mask = elliptical_mask(shape = (20,20), radius = (5,5), center = (10,10))
    """
    center, shape, radius = np.asarray(center), np.asarray(shape), np.asarray(radius)

    radius = np.repeat(radius, shape.size // radius.size)
    center = np.repeat(center, shape.size // center.size)

    if radius.size != shape.size:
        raise ValueError("Length of radius has to be either one or match shape.")
    if center.size != shape.size:
        raise ValueError("Length of center has to be either one or match shape.")

    n = shape.size
    center = center.reshape((-1,) + (1,) * n)
    radius = radius.reshape((-1,) + (1,) * n)

    mask = np.linalg.norm((np.indices(shape) - center) / radius, axis=0)
    mask = (mask <= 1).astype(int)

    return mask


def box_mask(shape: Tuple[int], center: Tuple[int], height: Tuple[int]) -> np.ndarray:
    """
    Creates a box mask centered around the provided center point.

    Parameters
    ----------
    shape : Tuple[int]
        Shape of the output array.
    center : Tuple[int]
        Center point coordinates of the box.
    height : Tuple[int]
        Height (side length) of the box along each axis.

    Returns
    -------
    NDArray
        The created box mask.
    """
    if len(shape) != len(center) or len(center) != len(height):
        raise ValueError("The length of shape, center, and height must be consistent.")

    # Calculate min and max coordinates for the box using the center and half-heights
    center, height = np.array(center, dtype=int), np.array(height, dtype=int)

    half_heights = height // 2
    starts = np.maximum(center - half_heights, 0)
    stops = np.minimum(center + half_heights + np.mod(height, 2) + 1, shape)
    slice_indices = tuple(slice(*coord) for coord in zip(starts, stops))

    out = np.zeros(shape)
    out[slice_indices] = 1
    return out


def tube_mask(
    shape: Tuple[int],
    symmetry_axis: int,
    base_center: Tuple[int],
    inner_radius: float,
    outer_radius: float,
    height: int,
) -> NDArray:
    """
    Creates a tube mask.

    Parameters
    ----------
    shape : tuple
        Shape of the mask to be created.
    symmetry_axis : int
        The axis of symmetry for the tube.
    base_center : tuple
        Center of the base circle of the tube.
    inner_radius : float
        Inner radius of the tube.
    outer_radius : float
        Outer radius of the tube.
    height : int
        Height of the tube.

    Returns
    -------
    NDArray
        The created tube mask.

    Raises
    ------
    ValueError
        If the inner radius is larger than the outer radius. Or height is larger
        than the symmetry axis shape.
    """
    if inner_radius > outer_radius:
        raise ValueError("inner_radius should be smaller than outer_radius.")

    if height > shape[symmetry_axis]:
        raise ValueError(f"Height can be no larger than {shape[symmetry_axis]}.")

    if symmetry_axis > len(shape):
        raise ValueError(f"symmetry_axis can be not larger than {len(shape)}.")

    circle_shape = tuple(b for ix, b in enumerate(shape) if ix != symmetry_axis)
    base_center = tuple(b for ix, b in enumerate(base_center) if ix != symmetry_axis)

    inner_circle = create_mask(
        mask_type="ellipse",
        shape=circle_shape,
        radius=inner_radius,
        center=base_center,
    )
    outer_circle = create_mask(
        mask_type="ellipse",
        shape=circle_shape,
        radius=outer_radius,
        center=base_center,
    )
    circle = outer_circle - inner_circle
    circle = np.expand_dims(circle, axis=symmetry_axis)

    center = shape[symmetry_axis] // 2
    start_idx = center - height // 2
    stop_idx = center + height // 2 + height % 2

    slice_indices = tuple(
        slice(None) if i != symmetry_axis else slice(start_idx, stop_idx)
        for i in range(len(shape))
    )
    tube = np.zeros(shape)
    tube[slice_indices] = np.repeat(circle, height, axis=symmetry_axis)

    return tube


def scramble_phases(
    arr: NDArray, noise_proportion: float = 0.5, seed: int = 42
) -> NDArray:
    """
    Applies random phase scrambling to a given array.

    This function takes an input array, applies a Fourier transform, then scrambles the
    phase with a given proportion of noise, and finally applies an
    inverse Fourier transform to the scrambled data. The phase scrambling
    is controlled by a random seed.

    Parameters
    ----------
    arr : NDArray
        The input array to be scrambled.
    noise_proportion : float, optional
        The proportion of noise in the phase scrambling, by default 0.5.
    seed : int, optional
        The seed for the random phase scrambling, by default 42.

    Returns
    -------
    NDArray
        The array with scrambled phases.

    Raises
    ------
    ValueError
        If noise_proportion is not within [0, 1].
    """
    if noise_proportion < 0 or noise_proportion > 1:
        raise ValueError("noise_proportion has to be within [0, 1].")

    arr_fft = np.fft.fftn(arr)

    amp = np.abs(arr_fft)
    ph = np.angle(arr_fft)

    np.random.seed(seed)
    ph_noise = np.random.permutation(ph)
    ph_new = ph * (1 - noise_proportion) + ph_noise * noise_proportion
    ret = np.real(np.fft.ifftn(amp * np.exp(1j * ph_new)))
    return ret


def conditional_execute(func: Callable, execute_operation: bool = True) -> Callable:
    """
    Return the given function or a no-operation function based on execute_operation.

    Parameters
    ----------
    func : callable
        The function to be executed if execute_operation is True.
    execute_operation : bool, optional
        A flag that determines whether to return `func` or a no-operation function.
        Default is True.

    Returns
    -------
    callable
        Either the given function `func` or a no-operation function.

    Examples
    --------
    >>> def greet(name):
    ...     return f"Hello, {name}!"
    ...
    >>> operation = conditional_execute(greet, False)
    >>> operation("Alice")
    >>> operation = conditional_execute(greet, True)
    >>> operation("Alice")
    'Hello, Alice!'
    """

    def noop(*args, **kwargs):
        """No operation function."""
        pass

    return func if execute_operation else noop
