""" Contains CTImagesAugmentedBatch: masked ct-batch with some augmentation actions """

import numpy as np

from .ct_masked_batch import CTImagesMaskedBatch
from ..batchflow import action, Sampler  # pylint: disable=no-name-in-module
from .mask import insert_cropped

class CTImagesAugmentedBatch(CTImagesMaskedBatch):
    """ Masked ct-batch with augmenting actions.

    Adds cutout, additive/multiplicative noise - augmentations.
    """
    @action
    def init_with_ones(self, shape=(32, 64, 64)):
        """ Loader for tests, fills batch with ones.
        """
        self.images = np.ones(shape=(len(self) * shape[0], *shape[1:]))
        self._bounds = np.cumsum((0, ) + (shape[0], ) * len(self))
        return self

    @action
    def cutout(self, positions, sizes, components='images', fill_with=0):
        """ Fill a box from each scan with some density-value.

        Parameters:
        -----------
        positions : ndarray
            array of starting positions of boxes, has shape (len(batch), 3).
        sizes : ndarray
            array of box-sizes, has shape (len(batch), 3).
        components : str or list
            names of components to apply cutout
        fill_with : ndarray, float or string
            value or filling scheme. Value can be float or an array of the shape,
            that can be broadcasted to box-shape. When string, can be either scan-wise
            mean ('mean') or scan-wise minimum/maximum ('min', 'max').
        """
        if isinstance(components, str):
            components = [components]
        for i in range(len(self)):
            size, position = sizes[i].astype(np.int64), positions[i].astype(np.int64)
            for component in components:
                item = self.get(i, component)

                # parse filling scheme
                fill_with = getattr(np, fill_with)(item) if isinstance(fill_with, str) else fill_with
                filled = np.ones(shape=size) * fill_with

                # perform insertion
                insert_cropped(item, filled, position)

        return self

    @action
    def apply_noise(self, noise, op='+', component='images'):
        """ For each item apply the noise to the item using op.

        Parameters:
        -----------
        noise : Sampler/ndarray
            1d-sampler/ndarray of shape=(len(batch), item.shape).
        op : str
            operation to perform on item. Can be either '+', '-', '*'.
        component : str
            component to add noise to.
        """
        # prepare noise-array
        all_items = getattr(self, component)
        noise = noise.sample(size=all_items.size).reshape(all_items.shape) if isinstance(noise, Sampler) else noise

        # parse and apply op in-place
        op_dict = {'+': '__add__', '*': '__mul__', '-': '__sub__'}
        op = op_dict[op]
        all_items[:] = getattr(all_items, op)(noise)

        return self
