"""Module describing the scattering matrix used in the PRISM algorithm."""
from __future__ import annotations
import inspect
import operator
import warnings
from abc import abstractmethod
from functools import partial, reduce

import dask.array as da
import numpy as np
from ase import Atoms
from dask.graph_manipulation import wait_on

from abtem.array import _validate_lazy, ArrayObject, ComputableList
from abtem.core.axes import (
    OrdinalAxis,
    AxisMetadata,
    ScanAxis,
    UnknownAxis,
    WaveVectorAxis,
)
from abtem.core.backend import get_array_module, cp, validate_device, copy_to_device
from abtem.core.chunks import chunk_ranges, validate_chunks, equal_sized_chunks, Chunks
from abtem.core.complex import complex_exponential
from abtem.core.energy import Accelerator
from abtem.core.ensemble import Ensemble, _wrap_with_array
from abtem.core.grid import Grid, GridUndefinedError
from abtem.core.utils import (
    safe_ceiling_int,
    expand_dims_to_broadcast,
    ensure_list,
    CopyMixin,
    EqualityMixin, tuple_range,
)
from abtem.detectors import (
    BaseDetector,
    _validate_detectors,
    WavesDetector,
    FlexibleAnnularDetector, AnnularDetector,
)
from abtem.measurements import BaseMeasurements
from abtem.multislice import (
    allocate_multislice_measurements,
    multislice_and_detect,
)
from abtem.potentials.iam import BasePotential, _validate_potential
from abtem.prism.utils import (
    plane_waves,
    wrapped_crop_2d,
    minimum_crop,
    batch_crop_2d,
)
from abtem.scan import BaseScan, _validate_scan, GridScan
from abtem.transfer import CTF
from abtem.waves import BaseWaves, _antialias_cutoff_gpts
from abtem.waves import Waves, Probe


def _extract_measurement(array, index):
    if array.size == 0:
        return array

    array = array.item()[index].array
    return array


def _wrap_measurements(measurements):
    return measurements[0] if len(measurements) == 1 else ComputableList(measurements)


def _finalize_lazy_measurements(
    arrays, waves, detectors, extra_ensemble_axes_metadata=None, chunks=None
):

    if extra_ensemble_axes_metadata is None:
        extra_ensemble_axes_metadata = []

    measurements = []
    for i, detector in enumerate(detectors):

        base_shape = detector._out_base_shape(waves)

        if isinstance(detector, AnnularDetector):
            # TODO
            base_shape = ()

        meta = detector._out_meta(waves)

        new_axis = tuple(range(len(arrays.shape), len(arrays.shape) + len(base_shape)))

        if chunks is None:
            chunks = arrays.chunks

        array = arrays.map_blocks(
            _extract_measurement,
            i,
            chunks=chunks + tuple((n,) for n in base_shape),
            new_axis=new_axis,
            meta=meta,
        )

        ensemble_axes_metadata = detector._out_ensemble_axes_metadata(waves)

        base_axes_metadata = detector._out_base_axes_metadata(waves)

        axes_metadata = ensemble_axes_metadata + base_axes_metadata

        metadata = detector._out_metadata(waves)

        cls = detector._out_type(waves)

        axes_metadata = extra_ensemble_axes_metadata + axes_metadata

        measurement = cls.from_array_and_metadata(
            array, axes_metadata=axes_metadata, metadata=metadata
        )

        if hasattr(measurement, "reduce_ensemble"):
            measurement = measurement.reduce_ensemble()

        measurements.append(measurement)

    return measurements


def _round_gpts_to_multiple_of_interpolation(
    gpts: tuple[int, int], interpolation: tuple[int, int]
) -> tuple[int, int]:
    return tuple(n + (-n) % f for f, n in zip(interpolation, gpts))  # noqa



[docs]
class BaseSMatrix(BaseWaves):
    """Base class for scattering matrices."""

    _device: str
    ensemble_axes_metadata: list[AxisMetadata]
    ensemble_shape: tuple[int, ...]
    _base_dims = 3

    @property
    def device(self):
        """The device where the S-Matrix is created and reduced."""
        return self._device

    @property
    @abstractmethod
    def interpolation(self):
        """Interpolation factor in the `x` and `y` directions"""
        pass

    @property
    @abstractmethod
    def wave_vectors(self) -> np.ndarray:
        """The wave vectors corresponding to each plane wave."""
        pass

    @property
    @abstractmethod
    def semiangle_cutoff(self) -> float:
        """The radial cutoff of the plane-wave expansion [mrad]."""
        pass

    @property
    @abstractmethod
    def window_extent(self):
        """The cropping window extent of the waves."""
        pass

    @property
    @abstractmethod
    def window_gpts(self):
        """The number of grid points describing the cropping window of the wave functions."""
        pass

    def __len__(self) -> int:
        return len(self.wave_vectors)

    @property
    def base_axes_metadata(self) -> list[AxisMetadata]:
        wave_axes_metadata = super().base_axes_metadata
        return [
            WaveVectorAxis(
                label="q",
                values=tuple(tuple(value) for value in self.wave_vectors),
            ),
            wave_axes_metadata[0],
            wave_axes_metadata[1],
        ]


[docs]
    def dummy_probes(
        self, scan: BaseScan = None, ctf: CTF = None, plane: str = "entrance", **kwargs
    ) -> Probe:
        # TODO
        """
        A probe or an ensemble of probes equivalent reducing the SMatrix at a single position.

        Parameters
        ----------
        scan : BaseScan
        ctf : CTF
        plane : str

        Returns
        -------
        dummy_probes : Probes
        """

        if ctf is None:
            ctf = CTF(energy=self.energy, semiangle_cutoff=self.semiangle_cutoff)
        elif isinstance(ctf, dict):
            ctf = CTF(energy=self.energy, semiangle_cutoff=self.semiangle_cutoff, **ctf)
        elif isinstance(ctf, CTF):
            ctf = ctf.copy()
        else:
            raise ValueError()

        if plane == "exit":
            defocus = 0.0
            if hasattr(self, "potential"):
                if self.potential is not None:
                    defocus = self.potential.thickness

            elif "accumulated_defocus" in self.metadata:
                defocus = self.metadata["accumulated_defocus"]

            ctf.defocus = ctf.defocus - defocus

        ctf.semiangle_cutoff = min(ctf.semiangle_cutoff, self.semiangle_cutoff)

        default_kwargs = {"device": self.device, "metadata": {**self.metadata}}
        kwargs = {**default_kwargs, **kwargs}

        probes = Probe._from_ctf(
            extent=self.window_extent,
            gpts=self.window_gpts,
            ctf=ctf,
            energy=self.energy,
            **kwargs,
        )

        if scan is not None:
            probes._positions = scan

        return probes




def _validate_interpolation(interpolation: int | tuple[int, int]):
    if isinstance(interpolation, int):
        interpolation = (interpolation,) * 2
    elif not len(interpolation) == 2:
        raise ValueError("Interpolation factor must be an integer.")
    return tuple(interpolation)


def _common_kwargs(a, b):
    a_kwargs = inspect.signature(a).parameters.keys()
    b_kwargs = inspect.signature(b).parameters.keys()
    return set(a_kwargs).intersection(b_kwargs)


def _pack_wave_vectors(wave_vectors):
    return tuple(
        (float(wave_vector[0]), float(wave_vector[1])) for wave_vector in wave_vectors
    )


def _chunked_axis(s_matrix_array):
    window_margin = s_matrix_array._window_margin

    argsort = np.argsort(
        (
            -s_matrix_array.gpts[0] // window_margin[0],
            -s_matrix_array.gpts[1] // window_margin[1],
        )
    )
    return int(argsort[0]), int(argsort[1])


def _chunks_for_multiple_rechunk_reduce(partitions):

    chunks_1 = ()
    chunk_indices_1 = ()
    for i in range(1, len(partitions) - 1, 3):
        chunks_1 += (sum(partitions[i - 1 : i + 2]),)
        chunk_indices_1 += (i - 1,)
    chunks_1 = chunks_1 + (sum(partitions[i + 2 :]),)

    assert sum(chunks_1) == sum(partitions)

    chunks_2 = (sum(partitions[:1]),)
    chunk_indices_2 = ()
    for i in range(2, len(partitions) - 1, 3):
        chunks_2 += (sum(partitions[i - 1 : i + 2]),)
        chunk_indices_2 += (i - 1,)
    chunks_2 = chunks_2 + (sum(partitions[i + 2 :]),)

    assert sum(chunks_2) == sum(partitions)

    chunks_3 = (sum(partitions[:2]),)
    chunk_indices_3 = ()
    for i in range(3, len(partitions) - 1, 3):
        chunks_3 += (sum(partitions[i - 1 : i + 2]),)
        chunk_indices_3 += (i - 1,)
    chunks_3 = chunks_3 + (sum(partitions[i + 2 :]),)

    assert sum(chunks_3) == sum(partitions)
    assert len(chunk_indices_1 + chunk_indices_2 + chunk_indices_3) == (
        len(partitions) - 2
    )
    return (chunks_1, chunks_2, chunks_3), (
        chunk_indices_1,
        chunk_indices_2,
        chunk_indices_3,
    )


def _lazy_reduce(
    array,
    waves_partial,
    ensemble_axes_metadata,
    from_waves_kwargs,
    scan,
    ctf,
    detectors,
    max_batch_reduction,
):
    args = (array, ensemble_axes_metadata)
    waves = waves_partial(args).item()
    s_matrix = SMatrixArray._from_waves(waves, **from_waves_kwargs)

    measurements = s_matrix._batch_reduce_to_measurements(
        scan, ctf, detectors, max_batch_reduction
    )


    arr = np.zeros((1,) * (len(array.shape) - 1), dtype=object)
    arr.itemset(measurements)
    return arr


def _map_blocks(array, scans, block_indices, window_offset=(0, 0), **kwargs):

    ctf_chunks = tuple((n,) for n in kwargs["ctf"].ensemble_shape)

    blocks = ()
    for i, scan in zip(block_indices, scans):
        block = array.blocks[(slice(None),) * (len(array.shape) - 2) + i]

        new_chunks = array.chunks[:-3] + ctf_chunks + scan.shape

        kwargs["from_waves_kwargs"]["window_offset"] = (
            window_offset[0] + sum(array.chunks[-2][: i[0]]),
            window_offset[1] + sum(array.chunks[-1][: i[1]]),
        )

        if len(scan.shape) == 1:
            drop_axis = (len(array.shape) - 3, len(array.shape) - 1)
        elif len(scan.shape) == 2:
            drop_axis = (len(array.shape) - 3,)
        else:
            raise NotImplementedError

        block = da.map_blocks(
            _lazy_reduce,
            block,
            scan=scan,
            drop_axis=drop_axis,
            chunks=new_chunks,
            **kwargs,
            meta=np.array((), dtype=np.complex64),
        )

        if len(scan) == 0:
            block = da.zeros(
                (0,) * len(block.shape),
                dtype=np.complex64,
            )

        blocks += (block,)

    return blocks


def _tuple_from_indices(*args):
    temp_list = [None] * (len(args) // 2)

    for arg1, arg2 in zip(args[::2], args[1::2]):
        temp_list[arg1] = arg2

    return tuple(temp_list)


def _multiple_rechunk_reduce(s_matrix_array, scan, detectors, ctf, max_batch_reduction):
    assert np.all(s_matrix_array.periodic)

    window_margin = s_matrix_array._window_margin
    chunked_axis, nochunks_axis = _chunked_axis(s_matrix_array)

    pad_amounts = _tuple_from_indices(
        chunked_axis, (window_margin[chunked_axis],) * 2, nochunks_axis, (0, 0)
    )
    s_matrix_array = s_matrix_array._pad(pad_amounts)

    chunk_size = window_margin[chunked_axis]

    size = s_matrix_array.shape[-2:][chunked_axis] - window_margin[chunked_axis] * 2

    num_chunks = -(size // -chunk_size)

    partitions = _tuple_from_indices(
        chunked_axis,
        (chunk_size,) * num_chunks,
        nochunks_axis,
        (s_matrix_array.shape[len(s_matrix_array.shape) - 2 + chunked_axis],),
    )

    chunk_extents = tuple(
        tuple(((cc[0]) * d, (cc[1]) * d) for cc in c)
        for c, d in zip(chunk_ranges(partitions), s_matrix_array.sampling)
    )

    scan, scan_chunks = scan._sort_into_extents(chunk_extents)

    scans = [(indices, scan.item()) for indices, _, scan in scan.generate_blocks(scan_chunks)]

    partitions = (pad_amounts[chunked_axis][0],) + partitions[chunked_axis]
    partitions = partitions + (
        s_matrix_array.shape[len(s_matrix_array.shape) - 2 + chunked_axis]
        - sum(partitions),
    )

    (chunks_1, chunks_2, chunks_3), (
        scan_indices_1,
        scan_indices_2,
        scan_indices_3,
    ) = _chunks_for_multiple_rechunk_reduce(partitions)

    chunks_1 = (
        s_matrix_array.array.chunks[:-3]
        + (-1,)
        + _tuple_from_indices(chunked_axis, chunks_1, nochunks_axis, -1)
    )
    chunks_2 = (
        s_matrix_array.array.chunks[:-3]
        + (-1,)
        + _tuple_from_indices(chunked_axis, chunks_2, nochunks_axis, -1)
    )
    chunks_3 = (
        s_matrix_array.array.chunks[:-3]
        + (-1,)
        + _tuple_from_indices(chunked_axis, chunks_3, nochunks_axis, -1)
    )

    shape = tuple(len(c) for c in scan_chunks)
    blocks = np.zeros(shape, dtype=object)

    kwargs = {
        "waves_partial": s_matrix_array.waves._from_partitioned_args(),
        "ensemble_axes_metadata": s_matrix_array.waves.ensemble_axes_metadata,
        "from_waves_kwargs": s_matrix_array._copy_kwargs(exclude=("array", "extent")),
        "ctf": ctf,
        "detectors": detectors,
        "max_batch_reduction": max_batch_reduction,
    }

    array = s_matrix_array.array.rechunk(chunks_1)

    window_offset = s_matrix_array.window_offset

    block_indices = [
        _tuple_from_indices(chunked_axis, i, nochunks_axis, 0)
        for i in range(len(scan_indices_1))
    ]

    new_blocks = _map_blocks(
        array,
        [scans[i][1] for i in scan_indices_1],
        block_indices,
        window_offset=window_offset,
        **kwargs,
    )

    for i, block in zip(scan_indices_1, new_blocks):
        blocks.itemset(scans[i][0], block)

    if s_matrix_array.ensemble_shape:
        fp_arrays = []
        for i in np.ndindex(s_matrix_array.ensemble_shape):
            try:
                fp_new_blocks = tuple(block[i] for block in new_blocks)
                fp_array = wait_on(array[i], *fp_new_blocks)[0]
                fp_arrays.append(fp_array)
            except IndexError:
                fp_arrays.append(array[i])

        array = da.stack(fp_arrays, axis=0)

    array = array.rechunk(chunks_2)

    block_indices = [
        _tuple_from_indices(chunked_axis, i, nochunks_axis, 0)
        for i in range(1, len(scan_indices_2) + 1)
    ]

    new_blocks = _map_blocks(
        array,
        [scans[i][1] for i in scan_indices_2],
        block_indices,
        window_offset=window_offset,
        **kwargs,
    )

    for i, block in zip(scan_indices_2, new_blocks):
        blocks.itemset(scans[i][0], block)

    if s_matrix_array.ensemble_shape:
        fp_arrays = []
        for i in np.ndindex(s_matrix_array.ensemble_shape):
            try:
                fp_new_blocks = tuple(block[i] for block in new_blocks)
                fp_array = wait_on(array[i], *fp_new_blocks)[0]
                fp_arrays.append(fp_array)
            except IndexError:
                fp_arrays.append(array[i])

        array = da.stack(fp_arrays, axis=0)

    array = array.rechunk(chunks_3)

    block_indices = [
        _tuple_from_indices(chunked_axis, i, nochunks_axis, 0)
        for i in range(1, len(scan_indices_3) + 1)
    ]

    new_blocks = _map_blocks(
        array,
        [scans[i][1] for i in scan_indices_3],
        block_indices,
        window_offset=window_offset,
        **kwargs,
    )

    for i, block in zip(scan_indices_3, new_blocks):
        blocks.itemset(scans[i][0], block)

    array = da.block(blocks.tolist())

    dummy_probes = s_matrix_array.dummy_probes(scan=scan, ctf=ctf)

    measurements = _finalize_lazy_measurements(
        array,
        waves=dummy_probes,
        detectors=detectors,
        extra_ensemble_axes_metadata=s_matrix_array.ensemble_axes_metadata,
    )

    return measurements


def _single_rechunk_reduce(
    s_matrix_array: "SMatrixArray",
    scan: BaseScan,
    detectors: list[BaseDetector],
    ctf: CTF,
    max_batch_reduction: int,
):

    chunked_axis, nochunks_axis = _chunked_axis(s_matrix_array)

    num_chunks = (
        s_matrix_array.gpts[chunked_axis] // s_matrix_array._window_margin[chunked_axis]
    )

    chunks = equal_sized_chunks(
        s_matrix_array.shape[-2:][chunked_axis], num_chunks=num_chunks
    )

    assert np.all(np.array(chunks) > s_matrix_array._window_margin[chunked_axis])

    chunks = (
        s_matrix_array.array.chunks[:-3]
        + (-1,)
        + _tuple_from_indices(chunked_axis, chunks, nochunks_axis, -1)
    )

    array = s_matrix_array._array.rechunk(chunks)

    assert all(s_matrix_array.periodic)

    # chunk_extents = tuple(
    #     tuple(((cc[0]) * d, (cc[1]) * d) for cc in c)
    #     for c, d in zip(chunk_ranges(array.chunks[-2:]), s_matrix_array.sampling)
    # )
    chunk_extents_x = tuple(
        ((cc[0]) * s_matrix_array.sampling[0], (cc[1]) * s_matrix_array.sampling[0])
        for cc in array.chunks[-2]
    )
    chunk_extents_y = tuple(
        ((cc[0]) * s_matrix_array.sampling[1], (cc[1]) * s_matrix_array.sampling[1])
        for cc in array.chunks[-1]
    )

    chunk_extents = (chunk_extents_x, chunk_extents_y)
    scan, scan_chunks = scan._sort_into_extents(chunk_extents)

    ctf_chunks = tuple((n,) for n in ctf.ensemble_shape)
    chunks = array.chunks[:-3] + ctf_chunks

    shape = tuple(len(c) for c, p in zip(scan_chunks, s_matrix_array.periodic))
    blocks = np.zeros((1,) * len(array.shape[:-3]) + shape, dtype=object)

    kwargs = {
        "waves_partial": s_matrix_array.waves._from_partitioned_args(),
        "ensemble_axes_metadata": s_matrix_array.waves.ensemble_axes_metadata,
        "from_waves_kwargs": s_matrix_array._copy_kwargs(exclude=("array", "extent")),
        "ctf": ctf,
        "detectors": detectors,
        "max_batch_reduction": max_batch_reduction,
    }

    for indices, _, sub_scan in scan.generate_blocks(scan_chunks):

        sub_scan = sub_scan.item()

        if len(sub_scan) == 0:
            blocks.itemset(
                (0,) * len(array.shape[:-3]) + indices,
                da.zeros(
                    (0,) * len(blocks.shape),
                    dtype=np.complex64,
                ),
            )
            continue

        slics = (slice(None),) * (len(array.shape) - 2)
        window_offset = ()
        for i, k in enumerate(indices):

            if len(array.chunks[-2:][i]) > 1:
                slics += ([k - 1, k, (k + 1) % len(array.chunks[-2:][i])],)
                window_offset += (
                    sum(array.chunks[-2:][i][:k]) - array.chunks[-2:][i][k - 1],
                )

            else:
                slics += (slice(None),)
                window_offset += (0,)

        new_block = array.blocks[slics]
        new_block = new_block.rechunk(array.chunks[:-2] + (-1, -1))
        new_chunks = chunks + sub_scan.shape

        kwargs["from_waves_kwargs"]["window_offset"] = tuple(window_offset)

        if len(scan.shape) == 1:
            drop_axis = (len(array.shape) - 3, len(array.shape) - 1)
        elif len(scan.shape) == 2:
            drop_axis = (len(array.shape) - 3,)
        else:
            raise NotImplementedError

        new_block = da.map_blocks(
            _lazy_reduce,
            new_block,
            scan=sub_scan,
            drop_axis=drop_axis,
            chunks=new_chunks,
            **kwargs,
            meta=np.array((), dtype=np.complex64),
        )

        blocks.itemset((0,) * len(array.shape[:-3]) + indices, new_block)

    array = da.block(blocks.tolist())

    dummy_probes = s_matrix_array.dummy_probes(scan=scan, ctf=ctf)

    measurements = _finalize_lazy_measurements(
        array,
        waves=dummy_probes,
        detectors=detectors,
        extra_ensemble_axes_metadata=s_matrix_array.ensemble_axes_metadata,
    )

    return measurements


def _no_chunks_reduce(
    s_matrix_array: "SMatrixArray",
    scan: BaseScan,
    detectors: list[BaseDetector],
    ctf: CTF,
    max_batch_reduction: int,
):
    array = s_matrix_array._array.rechunk(
        s_matrix_array.array.chunks[:-3] + (-1, -1, -1)
    )

    kwargs = {
        "waves_partial": s_matrix_array.waves._from_partitioned_args(),
        "ensemble_axes_metadata": s_matrix_array.waves.ensemble_axes_metadata,
        "from_waves_kwargs": s_matrix_array._copy_kwargs(exclude=("array", "extent")),
        "ctf": ctf,
        "detectors": detectors,
        "max_batch_reduction": max_batch_reduction,
    }

    ctf_chunks = tuple((n,) for n in ctf.ensemble_shape)

    chunks = array.chunks[:-3] + ctf_chunks + scan.shape

    if len(scan.shape) == 1:
        drop_axis = (len(array.shape) - 3, len(array.shape) - 1)
    elif len(scan.shape) == 2:
        drop_axis = (len(array.shape) - 3,)
    else:
        raise NotImplementedError

    array = da.map_blocks(
        _lazy_reduce,
        array,
        scan=scan,
        drop_axis=drop_axis,
        chunks=chunks,
        **kwargs,
        meta=np.array((), dtype=np.complex64),
    )

    dummy_probes = s_matrix_array.dummy_probes(scan=scan, ctf=ctf)

    measurements = _finalize_lazy_measurements(
        array,
        waves=dummy_probes,
        detectors=detectors,
        extra_ensemble_axes_metadata=s_matrix_array.ensemble_axes_metadata,
    )
    return measurements



[docs]
class SMatrixArray(BaseSMatrix, ArrayObject):
    """
    A scattering matrix defined by a given array of dimension 3, where the first indexes the probe plane waves and the
    latter two are the `y` and `x` scan directions.

    Parameters
    ----------
    array : np.ndarray
        Array defining the scattering matrix. Must be 3D or higher, dimensions before the last three dimensions should
        represent ensemble dimensions, the next dimension indexes the plane waves and the last two dimensions
        represent the spatial extent of the plane waves.
    wave_vectors : np.ndarray
        Array defining the wave vectors corresponding to each plane wave. Must have shape Nx2, where N is equal to the
        number of plane waves.
    semiangle_cutoff : float
        The radial cutoff of the plane-wave expansion [mrad].
    energy : float
        Electron energy [eV].
    sampling : one or two float, optional
        Lateral sampling of wave functions [1 / Å]. Provide only if potential is not given. Will be ignored if 'gpts'
        is also provided.
    extent : one or two float, optional
        Lateral extent of wave functions [Å]. Provide only if potential is not given.
    interpolation : one or two int, optional
        Interpolation factor in the `x` and `y` directions (default is 1, ie. no interpolation). If a single value is
        provided, assumed to be the same for both directions.
    window_gpts : tuple of int
        The number of grid points describing the cropping window of the wave functions.
    window_offset : tuple of int
        The number of grid points from the origin the cropping windows of the wave functions is displaced.
    periodic: tuple of bool
        Specifies whether the SMatrix should be assumed to be periodic along the x and y-axis.
    device : str, optional
        The calculations will be carried out on this device ('cpu' or 'gpu'). Default is 'cpu'. The default is
        determined by the user configuration.
    ensemble_axes_metadata : list of AxesMetadata
        Axis metadata for each ensemble axis. The axis metadata must be compatible with the shape of the array.
    metadata : dict
        A dictionary defining wave function metadata. All items will be added to the metadata of measurements derived
        from the waves.
    """


[docs]
    def __init__(
        self,
        array: np.ndarray,
        wave_vectors: np.ndarray,
        semiangle_cutoff: float,
        energy: float = None,
        interpolation: int | tuple[int, int] = (1, 1),
        sampling: float | tuple[float, float] = None,
        extent: float | tuple[float, float] = None,
        window_gpts: tuple[int, int] = (0, 0),
        window_offset: tuple[int, int] = (0, 0),
        periodic: tuple[bool, bool] = (True, True),
        device: str = None,
        ensemble_axes_metadata: list[AxisMetadata] = None,
        metadata: dict = None,
    ):

        self._grid = Grid(
            extent=extent, gpts=array.shape[-2:], sampling=sampling, lock_gpts=True
        )
        self._accelerator = Accelerator(energy=energy)
        self._wave_vectors = wave_vectors

        super().__init__(
            array=array,
            ensemble_axes_metadata=ensemble_axes_metadata,
            metadata=metadata,
        )

        self._semiangle_cutoff = semiangle_cutoff
        self._window_gpts = tuple(window_gpts)
        self._window_offset = tuple(window_offset)
        self._interpolation = _validate_interpolation(interpolation)
        self._device = device
        self._periodic = periodic


    @classmethod
    def _pack_kwargs(cls, kwargs):
        kwargs["wave_vectors"] = _pack_wave_vectors(kwargs["wave_vectors"])
        return super()._pack_kwargs(kwargs)

    @classmethod
    def _unpack_kwargs(cls, attrs):
        kwargs = super()._unpack_kwargs(attrs)
        kwargs["wave_vectors"] = np.array(kwargs["wave_vectors"], dtype=np.float32)
        return kwargs

        # kwargs["wave_vectors"] = _pack_wave_vectors(kwargs["wave_vectors"])


[docs]
    def copy_to_device(self, device: str) -> "SMatrixArray":
        """Copy SMatrixArray to specified device."""
        s_matrix = super().copy_to_device(device)
        s_matrix._wave_vectors = copy_to_device(self._wave_vectors, device)
        return s_matrix


    @staticmethod
    def _packed_wave_vectors(wave_vectors):
        return _pack_wave_vectors(wave_vectors)

    @property
    def device(self):
        """The device on which the SMatrixArray is reduced."""
        return self._device

    @property
    def storage_device(self):
        """The device on which the SMatrixArray is stored."""
        return super().device

    @classmethod
    def _from_waves(cls, waves: Waves, **kwargs):
        common_kwargs = _common_kwargs(cls, Waves)
        kwargs.update({key: getattr(waves, key) for key in common_kwargs})
        kwargs["ensemble_axes_metadata"] = kwargs["ensemble_axes_metadata"][:-1]

        return cls(**kwargs)

    @property
    def waves(self) -> Waves:
        """The wave vectors describing each plane wave."""
        kwargs = {
            key: getattr(self, key) for key in _common_kwargs(self.__class__, Waves)
        }
        kwargs["ensemble_axes_metadata"] = (
            kwargs["ensemble_axes_metadata"] + self.base_axes_metadata[:-2]
        )
        return Waves(**kwargs)

    def _copy_with_new_waves(self, waves):
        keys = set(
            inspect.signature(self.__class__).parameters.keys()
        ) - _common_kwargs(self.__class__, Waves)
        kwargs = {key: getattr(self, key) for key in keys}
        return self._from_waves(waves, **kwargs)

    @property
    def periodic(self) -> tuple[bool, bool]:
        """If True the SMatrix is assumed to be periodic along corresponding axis."""
        return self._periodic

    @property
    def metadata(self) -> dict:
        self._metadata["energy"] = self.energy
        return self._metadata

    @property
    def ensemble_axes_metadata(self) -> list[AxisMetadata]:
        """Axis metadata for each ensemble axis."""
        return self._ensemble_axes_metadata

    @property
    def ensemble_shape(self) -> tuple[int, int]:
        return self.array.shape[:-3]

    @property
    def interpolation(self) -> tuple[int, int]:

        return self._interpolation


[docs]
    def rechunk(self, chunks: Chunks, in_place: bool = True):

        array = self.array.rechunk(chunks)

        if in_place:
            self._array = array
            return self
        else:
            kwargs = self._copy_kwargs(exclude=("array",))
            return self.__class__(array, **kwargs)


    @property
    def semiangle_cutoff(self) -> float:
        """The cutoff semiangle of the plane wave expansion."""
        return self._semiangle_cutoff

    @property
    def wave_vectors(self) -> np.ndarray:
        return self._wave_vectors

    @property
    def window_gpts(self) -> tuple[int, int]:
        return self._window_gpts

    @property
    def window_extent(self) -> tuple[float, float]:
        return (
            self.window_gpts[0] * self.sampling[0],
            self.window_gpts[1] * self.sampling[1],
        )

    @property
    def window_offset(self) -> tuple[float, float]:
        """The number of grid points from the origin the cropping windows of the wave functions is displaced."""
        return self._window_offset


[docs]
    def multislice(self, potential: BasePotential = None) -> "SMatrixArray":
        """


        Parameters
        ----------
        potential :

        Returns
        -------

        """
        waves = self.waves.multislice(potential)
        return self._copy_with_new_waves(waves)


    def _reduce_to_waves(
        self,
        array,
        positions,
        position_coefficients,
    ):
        xp = get_array_module(self._device)

        if self._device == "gpu" and isinstance(array, np.ndarray):
            array = xp.asarray(array)

        position_coefficients = xp.array(position_coefficients, dtype=xp.complex64)

        # return xp.zeros(position_coefficients.shape[:2] + self.window_gpts, dtype=xp.complex64)

        if self.window_gpts != self.gpts:
            pixel_positions = positions / xp.array(self.waves.sampling) - xp.asarray(
                self.window_offset
            )

            crop_corner, size, corners = minimum_crop(pixel_positions, self.window_gpts)

            array = wrapped_crop_2d(array, crop_corner, size)

            array = xp.tensordot(position_coefficients, array, axes=[-1, -3])

            if len(self.waves.shape) > 3:
                array = xp.moveaxis(array, -3, 0)

            array = batch_crop_2d(array, corners, self.window_gpts)

        else:
            array = xp.tensordot(position_coefficients, array, axes=[-1, -3])

            if len(self.waves.shape) > 3:
                array = xp.moveaxis(array, -3, 0)

        return array

    def _calculate_positions_coefficients(self, scan):
        xp = get_array_module(self.wave_vectors)

        if isinstance(scan, GridScan):
            x = xp.asarray(scan._x_coordinates())
            y = xp.asarray(scan._y_coordinates())
            coefficients = complex_exponential(
                -2.0 * xp.pi * x[:, None, None] * self.wave_vectors[None, None, :, 0]
            ) * complex_exponential(
                -2.0 * xp.pi * y[None, :, None] * self.wave_vectors[None, None, :, 1]
            )
        else:
            positions = xp.asarray(scan.get_positions())
            coefficients = complex_exponential(
                -2.0 * xp.pi * positions[..., 0, None] * self.wave_vectors[:, 0][None]
                - 2.0 * xp.pi * positions[..., 1, None] * self.wave_vectors[:, 1][None]
            )

        return coefficients

    def _calculate_ctf_coefficients(self, ctf):
        wave_vectors = self.wave_vectors
        xp = get_array_module(wave_vectors)

        alpha = (
            xp.sqrt(wave_vectors[:, 0] ** 2 + wave_vectors[:, 1] ** 2) * ctf.wavelength
        )
        phi = xp.arctan2(wave_vectors[:, 1], wave_vectors[:, 0])
        array = ctf._evaluate_from_angular_grid(alpha, phi)
        array = array / xp.sqrt((array**2).sum(axis=-1, keepdims=True))
        return array

    def _batch_reduce_to_measurements(
        self,
        scan: BaseScan,
        ctf: CTF,
        detectors: list[BaseDetector],
        max_batch_reduction: int,
    ) -> tuple[BaseMeasurements | Waves, ...]:

        dummy_probes = self.dummy_probes(scan=scan, ctf=ctf)
        #print(self.waves.ensemble_axes_metadata[:-1])

        measurements = allocate_multislice_measurements(
            dummy_probes,
            detectors,
            extra_ensemble_axes_shape=self.waves.ensemble_shape[:-1],
            extra_ensemble_axes_metadata=self.waves.ensemble_axes_metadata[:-1],
        )

        xp = get_array_module(self._device)

        if self._device == "gpu" and isinstance(self.waves.array, np.ndarray):
            array = cp.asarray(self.waves.array)
        else:
            array = self.waves.array

        for _, ctf_slics, sub_ctf in ctf.generate_blocks(1):
            sub_ctf = sub_ctf.item()
            ctf_coefficients = self._calculate_ctf_coefficients(sub_ctf)

            for _, slics, sub_scan in scan.generate_blocks(max_batch_reduction):
                sub_scan = sub_scan.item()
                positions = xp.asarray(sub_scan.get_positions())

                positions_coefficients = self._calculate_positions_coefficients(
                    sub_scan
                )

                if ctf_coefficients is not None:
                    (
                        expanded_ctf_coefficients,
                        positions_coefficients,
                    ) = expand_dims_to_broadcast(
                        ctf_coefficients,
                        positions_coefficients,
                        match_dims=[(-1,), (-1,)],
                    )
                    coefficients = positions_coefficients * expanded_ctf_coefficients
                else:
                    coefficients = positions_coefficients

                ensemble_shape = len(array.shape[:-3]) + len(sub_ctf.ensemble_shape)

                ensemble_axes_metadata = []
                ensemble_axes_metadata.extend(
                    [UnknownAxis() for _ in range(ensemble_shape)]
                )
                ensemble_axes_metadata.extend(
                    [ScanAxis() for _ in range(len(scan.shape))]
                )

                waves_array = self._reduce_to_waves(array, positions, coefficients)

                waves = Waves(
                    waves_array,
                    sampling=self.sampling,
                    energy=self.energy,
                    ensemble_axes_metadata=ensemble_axes_metadata,
                    metadata=self.metadata,
                )

                indices = (
                    (slice(None),) * (len(self.waves.shape) - 3) + ctf_slics + slics
                )

                for detector, measurement in zip(detectors, measurements):
                    measurement.array[indices] = detector.detect(waves).array

        return tuple(measurements)

    @property
    def _window_margin(self):
        return -(self.window_gpts[0] // -2), -(self.window_gpts[1] // -2)

    def _pad(self, pad_width):
        array = self.array

        pad_width = ((0,) * 2,) * len(array.shape[:-2]) + tuple(pad_width)

        pad_amounts = sum(pad_width[-2]), sum(pad_width[-1])

        pad_chunks = array.chunks[:-2] + (
            array.shape[-2] + pad_amounts[-2],
            array.shape[-1] + pad_amounts[-1],
        )

        array = array.map_blocks(
            np.pad,
            pad_width=pad_width,
            meta=array._meta,
            chunks=pad_chunks,
            mode="wrap",
        )

        kwargs = self._copy_kwargs(exclude=("array", "extent"))

        kwargs["periodic"] = tuple(
            False if pad_amount else periodic
            for periodic, pad_amount in zip(kwargs["periodic"], pad_amounts)
        )

        kwargs["window_offset"] = tuple(
            window_offset - pad_amount[0]
            for window_offset, pad_amount in zip(
                kwargs["window_offset"], pad_width[-2:]
            )
        )

        return self.__class__(array, **kwargs)

    def _chunks_for_reduction(self):
        chunks = -(self.gpts[0] // -(self.interpolation[0] * 2)), -(
            self.gpts[1] // -(self.interpolation[1] * 2)
        )

        num_chunks = self.gpts[0] // chunks[0], self.gpts[1] // chunks[1]

        if num_chunks[1] > num_chunks[0]:
            num_chunks = (1, num_chunks[1])
        else:
            num_chunks = (num_chunks[0], 1)

        chunks = tuple(
            equal_sized_chunks(n, num_chunks=nsc)
            for n, nsc in zip(self.shape[-2:], num_chunks)
        )

        if chunks is None:
            chunks = self.array.chunks[-2:]
        else:
            chunks = validate_chunks(self.shape[-2:], chunks)

        return chunks

    def _validate_max_batch_reduction(
        self, scan, max_batch_reduction: int | str = "auto"
    ):

        shape = (len(scan),) + self.window_gpts
        chunks = (max_batch_reduction, -1, -1)

        return validate_chunks(shape, chunks, dtype=np.dtype("complex64"))[0][0]

    def _validate_reduction_scheme(self, reduction_scheme):

        if self.interpolation == (1, 1) and reduction_scheme == "no-chunks":
            raise NotImplementedError

        if reduction_scheme == "auto" and max(self.interpolation) <= 2:
            return "no-chunks"
        elif reduction_scheme == "auto":
            return "multiple-rechunk"

        return reduction_scheme


[docs]
    def reduce(
        self,
        scan: BaseScan = None,
        ctf: CTF = None,
        detectors: BaseDetector | list[BaseDetector] = None,
        max_batch_reduction: int | str = "auto",
        reduction_scheme: str = "auto",
    ) -> BaseMeasurements | Waves | list[BaseMeasurements | Waves]:

        """
        Scan the probe across the potential and record a measurement for each detector.

        Parameters
        ----------
        detectors : list of Detector objects
            The detectors recording the measurements.
        scan : Scan object
            Scan defining the positions of the probe wave functions.
        ctf: CTF object, optional
            The probe contrast transfer function. Default is None (aperture is set by the planewave cutoff).
        max_batch_reduction : int or str, optional
            Number of positions per reduction operation. A large number of positions better utilize thread
            parallelization, but requires more memory and floating point operations. If 'auto' (default), the batch size
            is automatically chosen based on the abtem user configuration settings "dask.chunk-size" and
            "dask.chunk-size-gpu".
        rechunk : two int or str, optional
            Partitioning of the scan. The scattering matrix will be reduced in similarly partitioned chunks.
            Should be equal to or greater than the interpolation.
        """

        self.accelerator.check_is_defined()

        if ctf is None:
            ctf = CTF(semiangle_cutoff=self.semiangle_cutoff)

        ctf.grid.match(self.dummy_probes())
        ctf.accelerator.match(self)

        if ctf.semiangle_cutoff == np.inf:
            ctf.semiangle_cutoff = self.semiangle_cutoff

        if not isinstance(scan, BaseScan):
            squeeze = (-3,)
        else:
            squeeze = ()

        if scan is None:
            squeeze_scan = True
            scan = self.extent[0] / 2, self.extent[1] / 2

        scan = _validate_scan(
            scan, Probe._from_ctf(extent=self.extent, ctf=ctf, energy=self.energy)
        )

        detectors = _validate_detectors(detectors)

        max_batch_reduction = self._validate_max_batch_reduction(
            scan, max_batch_reduction
        )

        reduction_scheme = self._validate_reduction_scheme(reduction_scheme)

        if self.is_lazy:

            if reduction_scheme == "multiple-rechunk":
                measurements = _multiple_rechunk_reduce(
                    self, scan, detectors, ctf, max_batch_reduction
                )
            elif reduction_scheme == "single-rechunk":
                measurements = _single_rechunk_reduce(
                    self, scan, detectors, ctf, max_batch_reduction
                )
            elif reduction_scheme == "no-chunks":
                measurements = _no_chunks_reduce(
                    self, scan, detectors, ctf, max_batch_reduction
                )
            else:
                raise ValueError()
        else:
            measurements = self._batch_reduce_to_measurements(
                scan, ctf, detectors, max_batch_reduction
            )

        measurements = [measurement.squeeze(squeeze) for measurement in measurements]

        return _wrap_measurements(measurements)



[docs]
    def scan(
        self,
        scan: BaseScan = None,
        detectors: BaseDetector | list[BaseDetector] = None,
        ctf: CTF = None,
        max_batch_reduction: int | str = "auto",
        rechunk: tuple[int, int] | str = "auto",
    ):
        """
        Reduce the SMatrix using coefficients calculated by a BaseScan and a CTF, to obtain the exit wave functions
        at given initial probe positions and aberrations.

        Parameters
        ----------
        scan : BaseScan
            Positions of the probe wave functions. If not given, scans across the entire potential at Nyquist sampling.
        detectors : BaseDetector, list of BaseDetector, optional
            A detector or a list of detectors defining how the wave functions should be converted to measurements after
            running the multislice algorithm. See abtem.measurements.detect for a list of implemented detectors.
        ctf : CTF
            Contrast transfer function from used for calculating the expansion coefficients in the reduction of the
            SMatrix.
        max_batch_reduction : int or str, optional
            Number of positions per reduction operation. A large number of positions better utilize thread
            parallelization, but requires more memory and floating point operations. If 'auto' (default), the batch size
            is automatically chosen based on the abtem user configuration settings "dask.chunk-size" and
            "dask.chunk-size-gpu".
        rechunk : str or tuple of int, optional
            Parallel reduction of the SMatrix requires rechunking the Dask array from chunking along the expansion axis
            to chunking over the spatial axes. If given as a tuple of int of length the SMatrix is rechunked to have
            those chunks. If 'auto' (default) the chunks are taken to be identical to the interpolation factor.

        Returns
        -------
        detected_waves : BaseMeasurements or list of BaseMeasurement
            The detected measurement (if detector(s) given).
        exit_waves : Waves
            Wave functions at the exit plane(s) of the potential (if no detector(s) given).
        """
        if scan is None:
            scan = GridScan()

        if detectors is None:
            detectors = [FlexibleAnnularDetector()]

        return self.reduce(
            scan=scan,
            ctf=ctf,
            detectors=detectors,
            max_batch_reduction=max_batch_reduction,
            rechunk=rechunk,
        )





[docs]
class SMatrix(BaseSMatrix, Ensemble, CopyMixin, EqualityMixin):
    """
    The scattering matrix is used for simulating STEM experiments using the PRISM algorithm.

    Parameters
    ----------
    semiangle_cutoff : float
        The radial cutoff of the plane-wave expansion [mrad].
    energy : float
        Electron energy [eV].
    potential : Atoms or AbstractPotential, optional
        Atoms or a potential that the scattering matrix represents. If given as atoms, a default potential will be
        created. If nothing is provided the scattering matrix will represent a vacuum potential, in which case the
        sampling and extent must be provided.
    gpts : one or two int, optional
        Number of grid points describing the scattering matrix. Provide only if potential is not given.
    sampling : one or two float, optional
        Lateral sampling of scattering matrix [1 / Å]. Provide only if potential is not given. Will be ignored if 'gpts'
        is also provided.
    extent : one or two float, optional
        Lateral extent of scattering matrix [Å]. Provide only if potential is not given.
    interpolation : one or two int, optional
        Interpolation factor in the `x` and `y` directions (default is 1, ie. no interpolation). If a single value is
        provided, assumed to be the same for both directions.
    downsample : {'cutoff', 'valid'} or float or bool
        Controls whether to downsample the scattering matrix after running the multislice algorithm.

            ``cutoff`` :
                Downsample to the antialias cutoff scattering angle (default).

            ``valid`` :
                Downsample to the largest rectangle that fits inside the circle with a radius defined by the antialias
                cutoff scattering angle.

            float :
                Downsample to a specified maximum scattering angle [mrad].
    device : str, optional
        The calculations will be carried out on this device ('cpu' or 'gpu'). Default is 'cpu'. The default is
        determined by the user configuration.
    store_on_host : bool, optional
        If True, store the scattering matrix in host (cpu) memory so that the necessary memory is transferred as chunks
        to the device to run calculations (default is False).
    """


[docs]
    def __init__(
        self,
        semiangle_cutoff: float,
        energy: float,
        potential: Atoms | BasePotential = None,
        gpts: int | tuple[int, int] = None,
        sampling: float | tuple[float, float] = None,
        extent: float | tuple[float, float] = None,
        interpolation: int | tuple[int, int] = 1,
        downsample: bool | str = "cutoff",
        # tilt: Tuple[float, float] = (0.0, 0.0),
        device: str = None,
        store_on_host: bool = False,
    ):

        if downsample is True:
            downsample = "cutoff"

        self._device = validate_device(device)
        self._grid = Grid(extent=extent, gpts=gpts, sampling=sampling)

        if potential is None:
            try:
                self.grid.check_is_defined()
            except GridUndefinedError:
                raise ValueError("Provide a potential or provide 'extent' and 'gpts'.")
        else:
            potential = _validate_potential(potential)
            self.grid.match(potential)
            self._grid = potential.grid

        self._potential = potential
        self._interpolation = _validate_interpolation(interpolation)
        self._semiangle_cutoff = semiangle_cutoff
        self._downsample = downsample

        self._accelerator = Accelerator(energy=energy)
        # self._beam_tilt = BeamTilt(tilt=tilt)

        self._store_on_host = store_on_host

        assert semiangle_cutoff > 0.0

        if not all(n % f == 0 for f, n in zip(self.interpolation, self.gpts)):
            warnings.warn(
                "The interpolation factor does not exactly divide 'gpts', normalization may not be exactly preserved."
            )


    @property
    def base_shape(self) -> tuple[int, int, int]:
        """Shape of the base axes of the SMatrix."""
        return len(self), self.gpts[0], self.gpts[1]

    @property
    def tilt(self):
        """The small-angle tilt of applied to the Fresnel propagator [mrad]."""
        return 0.0, 0.0


[docs]
    def round_gpts_to_interpolation(self) -> SMatrix:
        """
        Round the gpts of the SMatrix to the closest multiple of the interpolation factor.

        Returns
        -------
        s_matrix_with_rounded_gpts : SMatrix
        """

        rounded = _round_gpts_to_multiple_of_interpolation(
            self.gpts, self.interpolation
        )
        if rounded == self.gpts:
            return self

        self.gpts = rounded
        return self


    @property
    def downsample(self) -> str | bool:
        """How to downsample the scattering matrix after running the multislice algorithm."""
        return self._downsample

    @property
    def store_on_host(self) -> bool:
        """Store the SMatrix in host memory. The reduction may still be calculated on the device."""
        return self._store_on_host

    @property
    def metadata(self):
        return {"energy": self.energy}

    @property
    def shape(self) -> tuple[int, ...]:
        """Shape of the SMatrix."""
        return self.ensemble_shape + (len(self),) + self.gpts

    @property
    def ensemble_shape(self) -> tuple[int, ...]:
        """Shape of the SMatrix ensemble axes."""
        if self.potential is None:
            return ()
        else:
            return self.potential.ensemble_shape

    @property
    def ensemble_axes_metadata(self):
        """Axis metadata for each ensemble axis."""
        if self.potential is None:
            return []
        else:
            return self.potential.ensemble_axes_metadata

    @property
    def wave_vectors(self) -> np.ndarray:
        self.grid.check_is_defined()
        self.accelerator.check_is_defined()
        dummy_probes = self.dummy_probes(device="cpu")

        aperture = dummy_probes.aperture._evaluate_kernel(dummy_probes)

        indices = np.where(aperture > 0.0)

        n = np.fft.fftfreq(aperture.shape[0], d=1 / aperture.shape[0])[indices[0]]
        m = np.fft.fftfreq(aperture.shape[1], d=1 / aperture.shape[1])[indices[1]]

        w, h = self.extent

        kx = n / w * np.float32(self.interpolation[0])
        ky = m / h * np.float32(self.interpolation[1])

        xp = get_array_module(self.device)
        return xp.asarray([kx, ky]).T

    @property
    def potential(self) -> BasePotential:
        """The potential described by the SMatrix."""
        return self._potential

    @potential.setter
    def potential(self, potential: BasePotential):
        self._potential = potential
        self._grid = potential.grid

    @property
    def semiangle_cutoff(self) -> float:
        """Plane-wave expansion cutoff."""
        return self._semiangle_cutoff

    @semiangle_cutoff.setter
    def semiangle_cutoff(self, value: float):
        self._semiangle_cutoff = value

    @property
    def interpolation(self) -> tuple[int, int]:
        return self._interpolation

    def _wave_vector_chunks(self, max_batch):
        if isinstance(max_batch, int):
            max_batch = max_batch * reduce(operator.mul, self.gpts)

        chunks = validate_chunks(
            shape=(len(self),) + self.gpts,
            chunks=("auto", -1, -1),
            limit=max_batch,
            dtype=np.dtype("complex64"),
            device=self.device,
        )
        return chunks

    @property
    def downsampled_gpts(self) -> tuple[int, int]:
        """The gpts of the SMatrix after downsampling."""
        if self.downsample:
            downsampled_gpts = self._gpts_within_angle(self.downsample)
            rounded = _round_gpts_to_multiple_of_interpolation(
                downsampled_gpts, self.interpolation
            )
            return rounded
        else:
            return self.gpts

    @property
    def window_gpts(self):
        return (
            safe_ceiling_int(self.downsampled_gpts[0] / self.interpolation[0]),
            safe_ceiling_int(self.downsampled_gpts[1] / self.interpolation[1]),
        )

    @property
    def window_extent(self):
        sampling = (
            self.extent[0] / self.downsampled_gpts[0],
            self.extent[1] / self.downsampled_gpts[1],
        )

        return (
            self.window_gpts[0] * sampling[0],
            self.window_gpts[1] * sampling[1],
        )

    # @staticmethod
    # def _wrapped_build_s_matrix(*args, s_matrix_partial):
    #     s_matrix = s_matrix_partial(*tuple(arg.item() for arg in args[:-1]))
    #
    #     wave_vector_range = slice(*np.squeeze(args[-1]))
    #     array = s_matrix._build_s_matrix(wave_vector_range).array
    #     return array
    #
    # def _s_matrix_partial(self):
    #     def s_matrix(*args, potential_partial, **kwargs):
    #         if potential_partial is not None:
    #             potential = potential_partial(*args + (np.array([None], dtype=object),))
    #         else:
    #             potential = None
    #         return SMatrix(potential=potential, **kwargs)
    #
    #     potential_partial = (
    #         self.potential._from_partitioned_args()
    #         if self.potential is not None
    #         else None
    #     )
    #     return partial(
    #         s_matrix,
    #         potential_partial=potential_partial,
    #         **self._copy_kwargs(exclude=("potential",)),
    #     )


[docs]
    def multislice(
        self,
        potential=None,
        lazy: bool = None,
        max_batch: int | str = "auto",
    ):
        """


        Parameters
        ----------
        potential
        lazy : bool, optional
            If True, create the wave functions lazily, otherwise, calculate instantly. If not given, defaults to the
            setting in the user configuration file.
        max_batch : int or str, optional
            The number of expansion plane waves in each run of the multislice algorithm.

        Returns
        -------

        """
        s_matrix = self.__class__(
            potential=potential, **self._copy_kwargs(exclude=("potential",))
        )
        return s_matrix.build(lazy=lazy, max_batch=max_batch)


    @property
    def _default_ensemble_chunks(self):
        return self.potential._default_ensemble_chunks

    def _partition_args(self, chunks=(1,), lazy: bool = True):
        if self.potential is not None:
            return self.potential._partition_args(chunks, lazy=lazy)
        else:
            array = np.empty((1,), dtype=object)
            if lazy:
                array = da.from_array(array, chunks=1)
            return (array,)

    @staticmethod
    def _s_matrix(*args, potential_partial, **kwargs):
        potential = potential_partial(*args).item()
        s_matrix = SMatrix(potential=potential, **kwargs)
        return _wrap_with_array(s_matrix)

    def _from_partitioned_args(self, *args, **kwargs):
        if self.potential is not None:
            potential_partial = self.potential._from_partitioned_args()
            kwargs = self._copy_kwargs(exclude=("potential", "sampling", "extent"))
        else:
            potential_partial = lambda *args, **kwargs: _wrap_with_array(None, 1)
            kwargs = self._copy_kwargs(exclude=("potential",))

        return partial(self._s_matrix, potential_partial=potential_partial, **kwargs)

    @staticmethod
    def _wave_vector_blocks(wave_vector_chunks, lazy: bool = True):

        wave_vector_blocks = chunk_ranges(wave_vector_chunks)[0]

        array = np.zeros(len(wave_vector_blocks), dtype=object)
        for i, wave_vector_block in enumerate(wave_vector_blocks):
            array.itemset(i, wave_vector_block)

        if lazy:
            array = da.from_array(array, chunks=1)
        return array

    @staticmethod
    def _build_s_matrix(s_matrix, wave_vector_range=slice(None)):

        if isinstance(s_matrix, np.ndarray):
            s_matrix = s_matrix.item()

        if isinstance(wave_vector_range, np.ndarray):
            wave_vector_range = slice(*wave_vector_range.item())

        xp = get_array_module(s_matrix.device)

        wave_vectors = xp.asarray(s_matrix.wave_vectors, dtype=xp.float32)

        array = plane_waves(
            wave_vectors[wave_vector_range], s_matrix.extent, s_matrix.gpts
        )

        array *= np.prod(s_matrix.interpolation) / np.prod(array.shape[-2:])

        waves = Waves(
            array,
            energy=s_matrix.energy,
            extent=s_matrix.extent,
            ensemble_axes_metadata=[
                OrdinalAxis(values=wave_vectors[wave_vector_range])
            ],
        )

        if s_matrix.potential is not None:
            waves = multislice_and_detect(waves, s_matrix.potential, [WavesDetector()])[
                0
            ]

        if s_matrix.downsampled_gpts != s_matrix.gpts:
            waves.metadata[
                "adjusted_antialias_cutoff_gpts"
            ] = waves.antialias_cutoff_gpts

            waves = waves.downsample(
                gpts=s_matrix.downsampled_gpts,
                normalization="intensity",
            )

        if s_matrix.store_on_host and s_matrix.device == "gpu":
            waves = waves.to_cpu()

        return waves.array


[docs]
    def build(
        self, lazy: bool = None, max_batch: int | str = "auto", bound: bool = None
    ) -> SMatrixArray:
        """
        Build the plane waves of the scattering matrix and propagate them through the potential using the
        multislice algorithm.

        Parameters
        ----------
        lazy : bool, optional
            If True, create the wave functions lazily, otherwise, calculate instantly. If not given, defaults to the
            setting in the user configuration file.
        max_batch : int or str, optional
            The number of expansion plane waves in each run of the multislice algorithm.

        Returns
        -------
        s_matrix_array : SMatrixArray
            The built scattering matrix.
        """
        lazy = _validate_lazy(lazy)

        downsampled_gpts = self.downsampled_gpts

        s_matrix_blocks = self.ensemble_blocks(1)
        xp = get_array_module(self.device)

        wave_vector_chunks = self._wave_vector_chunks(max_batch)

        if lazy:
            wave_vector_blocks = self._wave_vector_blocks(
                wave_vector_chunks, lazy=False
            )

            wave_vector_blocks = np.tile(
                wave_vector_blocks[None], (len(s_matrix_blocks), 1)
            )

            wave_vector_blocks = da.from_array(wave_vector_blocks, chunks=1)

            from dask.graph_manipulation import bind

            if bound is not None:
                wave_vector_blocks = bind(wave_vector_blocks, bound)

            adjust_chunks = {
                1: wave_vector_chunks[0],
                2: (downsampled_gpts[0],),
                3: (downsampled_gpts[1],),
            }

            symbols = (0, 1, 2, 3)
            if self.potential is None or not self.potential.ensemble_shape:
                symbols = symbols[1:]

            array = da.blockwise(
                self._build_s_matrix,
                symbols,
                s_matrix_blocks,
                (0,),
                wave_vector_blocks[..., None, None],
                (0, 1, 2, 3),
                concatenate=True,
                adjust_chunks=adjust_chunks,
                meta=xp.array((), dtype=np.complex64),
            )

        else:
            wave_vector_blocks = self._wave_vector_blocks(
                wave_vector_chunks, lazy=False
            )

            if self.store_on_host:
                array = np.zeros(
                    self.ensemble_shape + (len(self),) + self.downsampled_gpts,
                    dtype=np.complex64,
                )
            else:
                array = xp.zeros(
                    self.ensemble_shape + (len(self),) + self.downsampled_gpts,
                    dtype=np.complex64,
                )

            for i, _, s_matrix in self.generate_blocks(1):
                s_matrix = s_matrix.item()
                for start, stop in wave_vector_blocks:
                    items = (slice(start, stop),)
                    if self.ensemble_shape:
                        items = i + items

                    new_array = self._build_s_matrix(s_matrix, slice(start, stop))

                    if self.store_on_host:
                        new_array = xp.asnumpy(new_array)

                    array[items] = new_array

        waves = Waves(
            array,
            energy=self.energy,
            extent=self.extent,
            ensemble_axes_metadata=self.ensemble_axes_metadata
            + self.base_axes_metadata[:1],
        )

        if self.downsampled_gpts != self.gpts:
            waves.metadata["adjusted_antialias_cutoff_gpts"] = _antialias_cutoff_gpts(
                self.window_gpts, self.sampling
            )

        s_matrix_array = SMatrixArray._from_waves(
            waves,
            wave_vectors=self.wave_vectors,
            interpolation=self.interpolation,
            semiangle_cutoff=self.semiangle_cutoff,
            window_gpts=self.window_gpts,
            device=self.device,
        )

        return s_matrix_array



[docs]
    def scan(
        self,
        scan: np.ndarray | BaseScan = None,
        detectors: BaseDetector | list[BaseDetector] = None,
        ctf: CTF | dict = None,
        max_batch_multislice: str | int = "auto",
        max_batch_reduction: str | int = "auto",
        reduction_scheme: str = "auto",
        disable_s_matrix_chunks: bool = "auto",
        lazy: bool = None,
    ) -> BaseMeasurements | Waves | list[BaseMeasurements | Waves]:

        """
        Run the multislice algorithm, then reduce the SMatrix using coefficients calculated by a BaseScan and a CTF,
        to obtain the exit wave functions at given initial probe positions and aberrations.

        Parameters
        ----------
        scan : BaseScan
            Positions of the probe wave functions. If not given, scans across the entire potential at Nyquist sampling.
        detectors : BaseDetector, list of BaseDetector, optional
            A detector or a list of detectors defining how the wave functions should be converted to measurements after
            running the multislice algorithm. See abtem.measurements.detect for a list of implemented detectors.
        ctf : CTF
            Contrast transfer function from used for calculating the expansion coefficients in the reduction of the
            SMatrix.
        max_batch_multislice : int, optional
            The number of wave functions in each chunk of the Dask array. If 'auto' (default), the batch size is
            automatically chosen based on the abtem user configuration settings "dask.chunk-size" and
            "dask.chunk-size-gpu".
        max_batch_reduction : int or str, optional
            Number of positions per reduction operation. A large number of positions better utilize thread
            parallelization, but requires more memory and floating point operations. If 'auto' (default), the batch size
            is automatically chosen based on the abtem user configuration settings "dask.chunk-size" and
            "dask.chunk-size-gpu".
        reduction_scheme : str or tuple of int, optional
            Parallel reduction of the SMatrix requires rechunking the Dask array from chunking along the expansion axis
            to chunking over the spatial axes. If given as a tuple of int of length the SMatrix is rechunked to have
            those chunks. If 'auto' (default) the chunks are taken to be identical to the interpolation factor.
        disable_s_matrix_chunks : bool, optional
            If True, each S-Matrix is kept as a single chunk, thus lowering the communication overhead, but providing
            fewer opportunities for parallelization.
        lazy : bool, optional
            If True, create the measurements lazily, otherwise, calculate instantly. If None, this defaults to the value
            set in the configuration file.

        Returns
        -------
        detected_waves : BaseMeasurements or list of BaseMeasurement
            The detected measurement (if detector(s) given).
        exit_waves : Waves
            Wave functions at the exit plane(s) of the potential (if no detector(s) given).
        """

        if scan is None:
            scan = GridScan(
                start=(0, 0),
                end=self.extent,
                sampling=self.dummy_probes().aperture.nyquist_sampling,
            )

        if detectors is None:
            detectors = FlexibleAnnularDetector()

        return self.reduce(
            scan=scan,
            detectors=detectors,
            max_batch_reduction=max_batch_reduction,
            max_batch_multislice=max_batch_multislice,
            ctf=ctf,
            reduction_scheme=reduction_scheme,
            disable_s_matrix_chunks=disable_s_matrix_chunks,
            lazy=lazy,
        )


    def _eager_build_s_matrix_detect(self, scan, ctf, detectors, squeeze):
        extra_ensemble_axes_shape = ()
        extra_ensemble_axes_metadata = []
        for shape, axis_metadata in zip(
            self.ensemble_shape, self.ensemble_axes_metadata
        ):

            extra_ensemble_axes_metadata += [axis_metadata]
            extra_ensemble_axes_shape += (shape,)

            if axis_metadata._ensemble_mean:
                extra_ensemble_axes_shape = (1,) + extra_ensemble_axes_shape[1:]

        if self.potential is not None and len(self.potential.exit_planes) > 1:
            extra_ensemble_axes_shape = extra_ensemble_axes_shape + (
                len(self.potential.exit_planes),
            )
            extra_ensemble_axes_metadata = extra_ensemble_axes_metadata + [
                self.potential.base_axes_metadata[0]
            ]

        detectors = _validate_detectors(detectors)

        if self.ensemble_shape:
            measurements = allocate_multislice_measurements(
                self.build(lazy=True).dummy_probes(scan, ctf),
                detectors,
                extra_ensemble_axes_shape,
                extra_ensemble_axes_metadata,
            )
        else:
            measurements = None

        for i, _, s_matrix in self.generate_blocks(1):
            s_matrix = s_matrix.item()
            s_matrix_array = s_matrix.build(lazy=False)

            new_measurements = s_matrix_array.reduce(
                scan=scan, detectors=detectors, ctf=ctf
            )

            new_measurements = ensure_list(new_measurements)

            if measurements is None:
                measurements = new_measurements
            else:
                for measurement, new_measurement in zip(measurements, new_measurements):
                    if measurement.axes_metadata[0]._ensemble_mean:
                        measurement.array[:] += new_measurement.array
                    else:
                        measurement.array[i] = new_measurement.array

        # measurements = list(measurements.values())

        for i, measurement in enumerate(measurements):
            if (
                hasattr(measurement.axes_metadata[0], "_ensemble_mean")
                and measurement.axes_metadata[0]._ensemble_mean
            ) and squeeze:
                measurements[i] = measurement.squeeze((0,))

        return measurements

    @staticmethod
    def _lazy_build_s_matrix_detect(s_matrix, scan, ctf, detectors):
        s_matrix = s_matrix.item()
        measurements = s_matrix._eager_build_s_matrix_detect(
            scan=scan, ctf=ctf, detectors=detectors, squeeze=False
        )
        # measurements = ensure_list(measurements)

        array = np.zeros((1,) + (1,) * len(scan.shape), dtype=object)
        array.itemset(0, measurements)

        return array


[docs]
    def reduce(
        self,
        scan: np.ndarray | BaseScan = None,
        detectors: BaseDetector | list[BaseDetector] = None,
        ctf: CTF | dict = None,
        reduction_scheme: str = "auto",
        max_batch_multislice: str | int = "auto",
        max_batch_reduction: str | int = "auto",
        disable_s_matrix_chunks: bool = "auto",
        lazy: bool = None,
    ) -> BaseMeasurements | Waves | list[BaseMeasurements | Waves]:

        """
        Run the multislice algorithm, then reduce the SMatrix using coefficients calculated by a BaseScan and a CTF,
        to obtain the exit wave functions at given initial probe positions and aberrations.

        Parameters
        ----------
        scan : BaseScan
            Positions of the probe wave functions. If not given, scans across the entire potential at Nyquist sampling.
        detectors : BaseDetector, list of BaseDetector, optional
            A detector or a list of detectors defining how the wave functions should be converted to measurements after
            running the multislice algorithm. See abtem.measurements.detect for a list of implemented detectors.
        ctf : CTF
            Contrast transfer function from used for calculating the expansion coefficients in the reduction of the
            SMatrix.
        max_batch_multislice : int, optional
            The number of wave functions in each chunk of the Dask array. If 'auto' (default), the batch size is
            automatically chosen based on the abtem user configuration settings "dask.chunk-size" and
            "dask.chunk-size-gpu".
        max_batch_reduction : int or str, optional
            Number of positions per reduction operation. A large number of positions better utilize thread
            parallelization, but requires more memory and floating point operations. If 'auto' (default), the batch size
            is automatically chosen based on the abtem user configuration settings "dask.chunk-size" and
            "dask.chunk-size-gpu".
        reduction_scheme : str, optional
            Parallel reduction of the SMatrix requires rechunking the Dask array from chunking along the expansion axis
            to chunking over the spatial axes. If given as a tuple of int of length the SMatrix is rechunked to have
            those chunks. If 'auto' (default) the chunks are taken to be identical to the interpolation factor.
        disable_s_matrix_chunks : bool, optional
            If True, each S-Matrix is kept as a single chunk, thus lowering the communication overhead, but providing
            fewer opportunities for parallelization.
        lazy : bool, optional
            If True, create the measurements lazily, otherwise, calculate instantly. If None, this defaults to the value
            set in the configuration file.

        Returns
        -------
        measurements : BaseMeasurements or Waves or list of BaseMeasurements or list of Waves
            The detected measurement (if detector(s) given).
        """

        detectors = _validate_detectors(detectors)

        if scan is None:
            scan = (self.extent[0] / 2, self.extent[1] / 2)

        lazy = _validate_lazy(lazy)

        if self.device == "gpu" and disable_s_matrix_chunks == "auto":
            disable_s_matrix_chunks = True
        elif disable_s_matrix_chunks == "auto":
            disable_s_matrix_chunks = False

        if not lazy:
            scan = _validate_scan(scan, self)

            measurements = self._eager_build_s_matrix_detect(
                scan, ctf, detectors, squeeze=True
            )
            return _wrap_measurements(measurements)

        if disable_s_matrix_chunks:
            scan = _validate_scan(scan, self)

            blocks = self.ensemble_blocks(1)
            chunks = blocks.chunks + scan.shape

            new_axis = tuple_range(offset=len(blocks.shape), length=len(scan.shape))
            drop_axis = ()
            if len(self.ensemble_shape) == 0:
                drop_axis = (0,)
                chunks = chunks[1:]
                new_axis = tuple_range(offset=len(blocks.shape) - 1, length=len(scan.shape))

            arrays = blocks.map_blocks(
                self._lazy_build_s_matrix_detect,
                drop_axis=drop_axis,
                new_axis=new_axis,
                chunks=chunks,
                scan=scan,
                ctf=ctf,
                detectors=detectors,
                meta=np.array((), dtype=object),
            )

            waves = self.build(lazy=True).dummy_probes(scan=scan)

            extra_axes_metadata = []
            if self.potential is not None:
                extra_axes_metadata = self.potential.ensemble_axes_metadata

            measurements = _finalize_lazy_measurements(
                arrays, waves, detectors, extra_axes_metadata
            )

            return _wrap_measurements(measurements)

        s_matrix_array = self.build(max_batch=max_batch_multislice, lazy=lazy)

        return s_matrix_array.reduce(
            scan=scan,
            detectors=detectors,
            reduction_scheme=reduction_scheme,
            max_batch_reduction=max_batch_reduction,
            ctf=ctf,
        )