"""Provides an adapted version of the gunpowder ShiftAugment node
"""
import logging
import random

import numpy as np

from gunpowder.roi import Roi
from gunpowder.coordinate import Coordinate
from gunpowder import BatchFilter

logger = logging.getLogger(__name__)


class ShiftAugment(BatchFilter):
    """Modification of Gunpowder ShiftAugment, only shift/slip center

    This class differs from the gunpowder ShiftAugment by only
    slipping/shifting at the center frame.
    """
    def __init__(
            self,
            prob_slip=0,
            prob_shift=0,
            sigma=0,
            shift_axis=0):
        '''
        Arguments:

            prob_slip (``float``):

                Probability that a single section will move independently
                of neighboring sections.

            prob_shift (``float``):

                Probability that a section and all following sections will
                shift.

            sigma (``float`` or ``float[]``):

                The standard deviation of shifting in each direction.
                If a single value, will be extended to all dimensions
                except the shift axis. If an array, must have length
                equal to the number of dimensions, and the value at
                shift_axis must be zero.

            shift_axis (``int``):

                The axis along which to shift. The algorithm will slice
                along this axis and shift sections relative to each other.
        '''

        self.prob_slip = prob_slip
        self.prob_shift = prob_shift
        self.sigma = sigma
        self.shift_axis = shift_axis

        self.ndim = None
        self.shift_sigmas = None
        self.shift_array = None
        self.lcm_voxel_size = None

    def prepare(self, request):
        random.seed(request.random_seed)

        self.ndim = request.get_total_roi().dims
        assert self.shift_axis in range(self.ndim)

        try:
            self.shift_sigmas = tuple(self.sigma)
        except TypeError:
            self.shift_sigmas = [float(self.sigma)] * self.ndim
            self.shift_sigmas[self.shift_axis] = 0.0
            self.shift_sigmas = tuple(self.shift_sigmas)

        assert len(self.shift_sigmas) == self.ndim
        assert self.shift_sigmas[self.shift_axis] == 0.0

        has_nonzero = False
        for sigma in self.shift_sigmas:
            if sigma != 0.0:
                has_nonzero = True
                break
        assert has_nonzero

        if not request.array_specs:
            raise ValueError(
                    "Request passed to ShiftAugment node must contain at least"
                    " one array key. Check to make sure that ShiftAugment"
                    " node is not upstream of a RandomLocation node "
                    "with an ensure_nonempty argument.")

        self.lcm_voxel_size = self.spec.get_lcm_voxel_size(
                array_keys=request.array_specs.keys())
        assert self.lcm_voxel_size

        roi_shape = request.get_total_roi().get_shape()
        assert (roi_shape // self.lcm_voxel_size * self.lcm_voxel_size ==
                roi_shape), ("total roi shape {} must be divisible by least"
                             " common voxel size {}".format(
                                 roi_shape, self.lcm_voxel_size))
        roi_shape_adjusted = roi_shape // self.lcm_voxel_size
        shift_axis_len = roi_shape_adjusted[self.shift_axis]

        self.shift_array = self.construct_global_shift_array(
                shift_axis_len,
                self.shift_sigmas,
                self.prob_slip,
                self.prob_shift,
                self.lcm_voxel_size)

        for key, spec in request.items():
            sub_shift_array = self.get_sub_shift_array(
                    request.get_total_roi(),
                    spec.roi,
                    self.shift_array,
                    self.shift_axis,
                    self.lcm_voxel_size)
            updated_roi = self.compute_upstream_roi(spec.roi, sub_shift_array)
            spec.roi.set_offset(updated_roi.get_offset())
            spec.roi.set_shape(updated_roi.get_shape())
            request[key] = spec

        deps = request
        return deps

    def process(self, batch, request):
        for array_key, array in batch.arrays.items():
            sub_shift_array = self.get_sub_shift_array(
                    request.get_total_roi(),
                    array.spec.roi,
                    self.shift_array,
                    self.shift_axis,
                    self.lcm_voxel_size)

            input_data = array.data
            ndims = array.spec.roi.dims
            logger.debug("Data has shape: %s" % str(input_data.shape))
            if input_data.ndim > ndims:
                logger.debug("Multichannel data (%d channels)"
                             % input_data.shape[0])
                channels = []
                assert input_data.ndim - ndims == 1,\
                    "Expected a single extra dimension for channels"
                for channel_arr in input_data:
                    channels.append(self.shift_and_crop(
                        channel_arr,
                        request[array_key].roi.get_shape(),
                        sub_shift_array,
                        array.spec.voxel_size))
                array.data = np.stack(channels, axis=0)
            else:
                array.data = self.shift_and_crop(
                        array.data,
                        request[array_key].roi.get_shape(),
                        sub_shift_array,
                        array.spec.voxel_size)
            array.spec.roi = request[array_key].roi
            assert (request[array_key].roi.get_shape() ==
                    Coordinate(array.data.shape[-ndims:])
                    * self.lcm_voxel_size), \
                ("request roi shape {} is not the same as "
                 "generated array shape {}").format(
                     request[array_key].roi.get_shape(), array.data.shape)
            batch[array_key] = array

        for points_key, points in batch.graphs.items():
            sub_shift_array = self.get_sub_shift_array(
                    request.get_total_roi(),
                    points.spec.roi,
                    self.shift_array,
                    self.shift_axis,
                    self.lcm_voxel_size)
            points = self.shift_points(
                    points,
                    request[points_key].roi,
                    sub_shift_array,
                    self.shift_axis,
                    self.lcm_voxel_size)
            batch[points_key] = points

    def shift_and_crop(self, arr, roi_shape, sub_shift_array, voxel_size):
        """ Shift an array received from upstream and crop it to the target
            downstream region

        :param arr: an array of upstream data to be shifted and cropped
        :param roi_shape: the shape of the downstream ROI
        :param sub_shift_array: the cropped section of the global shift array
            that applies to this specific request
        :param voxel_size: the voxel sizes of the data in the array
        :return: an array of shape roi_shape that contains the array
            to be passed downstream
        """
        array_shift_axis_len = arr.shape[self.shift_axis]
        sub_shift_array_len = len(sub_shift_array)
        assert array_shift_axis_len % sub_shift_array_len == 0, \
            ("array shift axis length {} is not divisible by "
             "the sub_shift_array length {}").format(
                arr.shape[self.shift_axis], sub_shift_array.shape[0])

        voxel_ratio = array_shift_axis_len // sub_shift_array_len

        # assumption: each sub shift array element divides evenly by voxel size
        rescaled_sub_shift_array = sub_shift_array // np.array(
                voxel_size, dtype=int)

        max_shift = rescaled_sub_shift_array.max(axis=0)
        batch = arr.copy()
        batch_view = np.moveaxis(batch, self.shift_axis, 0)
        for index, plane in enumerate(batch_view):
            adjusted_index = index // voxel_ratio
            shift = rescaled_sub_shift_array[adjusted_index, :] - max_shift
            shift = np.delete(shift, self.shift_axis, axis=0)
            assert(len(shift) == plane.ndim)
            plane = np.roll(plane, shift, axis=tuple(range(len(shift))))
            batch_view[index] = plane

        adjusted_roi_shape = Coordinate(roi_shape) // Coordinate(voxel_size)

        sl = tuple(slice(0, adjusted_roi_shape[index])
                   for index in range(self.ndim))
        return batch[sl]

    @staticmethod
    def shift_points(
            points,
            request_roi,
            sub_shift_array,
            shift_axis,
            lcm_voxel_size):
        """ Shift a set of points received from upstream and crop out those
            not the the target downstream region

        :param points: the points from upstream
        :param request_roi: the downstream ROI
        :param sub_shift_array: the cropped section of the global shift array
            that applies to this specific request
        :param shift_axis: the axis to perform the shift along
        :param lcm_voxel_size: the least common voxel size for the arrays
            in the request
        :return a Points object with the updated point locations and ROI
        """

        nodes = list(points.nodes)
        spec = points.spec
        shift_axis_start_pos = spec.roi.get_offset()[shift_axis]

        for node in nodes:
            loc = node.location
            shift_axis_position = loc[shift_axis]
            shift_array_index = int(
                (shift_axis_position - shift_axis_start_pos)
                // lcm_voxel_size[shift_axis])
            assert(shift_array_index >= 0)
            shift = Coordinate(sub_shift_array[shift_array_index])
            loc += shift
            if not request_roi.contains(loc):
                points.remove_node(node)

        points.spec.roi = request_roi
        return points

    @staticmethod
    def get_sub_shift_array(
            total_roi,
            item_roi,
            shift_array,
            shift_axis,
            lcm_voxel_size):
        """ Slices the global shift array to return the sub-shift array
            to shift an item in the request

        :param total_roi: the total roi of the request
        :param item_roi: the roi of the item (array or points) being shifted
        :param shift_array: the shift array for the total_roi
        :param shift_axis: the axis along which we are shifting
        :param lcm_voxel_size: the least common voxel size for the arrays
            in the request
        :return: the portion of the global shift array that should be used
            to shift the item
        """
        item_offset_from_total = item_roi.get_offset() - total_roi.get_offset()
        offset_in_shift_axis = (item_offset_from_total[shift_axis]
                                // lcm_voxel_size[shift_axis])
        len_in_shift_axis = (item_roi.get_shape()[shift_axis]
                             // lcm_voxel_size[shift_axis])
        return shift_array[offset_in_shift_axis:
                           offset_in_shift_axis + len_in_shift_axis]

    @staticmethod
    def construct_global_shift_array(
            shift_axis_len,
            shift_sigmas,
            prob_slip,
            prob_shift,
            lcm_voxel_size):
        """ Sets the attribute variable self.shift_array

        :param shift_axis_len: the length of the shift axis
        :param shift_sigmas: the sigma to generate the normal distribution
            of shift amounts in each direction
        :param prob_slip: the probability of the slice shifting independently
            of all other slices
        :param prob_shift: the probability of the slice and all following
            slices shifting
        :param lcm_voxel_size: the least common voxel size of all the arrays
            in the request
        :return: the shift_array for the total_roi
        """

        assert shift_axis_len == 7,\
            "Request contains %d frames, expected 7" % shift_axis_len

        # each row is one slice along shift axis
        shift_array = np.zeros(
                shape=(shift_axis_len, len(shift_sigmas)),
                dtype=int)
        base_shift = np.zeros(shape=len(shift_sigmas), dtype=int)
        assert(prob_slip + prob_shift <= 1)

        for shift_axis_position in range(shift_axis_len):
            # only shift at the center frame
            if shift_axis_position == 3:
                r = random.random()
                slip = np.array([np.random.normal(scale=sigma
                                 / lcm_voxel_size[dim])
                                 for dim, sigma in enumerate(shift_sigmas)])
                slip = np.rint(slip).astype(int)
                slip = slip * np.array(lcm_voxel_size, dtype=int)

                if r <= prob_slip:
                    shift_array[shift_axis_position] = base_shift + slip
                elif r <= prob_slip + prob_shift:
                    base_shift += slip
                    shift_array[shift_axis_position] = base_shift
                else:
                    shift_array[shift_axis_position] = base_shift
            else:
                shift_array[shift_axis_position] = base_shift

        center_shift = shift_array[3]
        shift_array -= center_shift
        return shift_array

    @staticmethod
    def compute_upstream_roi(request_roi, sub_shift_array):
        """ Compute the ROI to pass upstream for a specific item
            (array or points) in a request

        :param request_roi: the downstream ROI passed to the Jitter node
        :param sub_shift_array: the portion of the global shift array that
            should be used to shift the item
        :return: the expanded ROI to pass upstream
        """

        max_shift = Coordinate(sub_shift_array.max(axis=0))
        min_shift = Coordinate(sub_shift_array.min(axis=0))

        downstream_offset = request_roi.get_offset()
        upstream_offset = downstream_offset - max_shift
        upstream_shape = request_roi.get_shape() + max_shift - min_shift
        return Roi(offset=upstream_offset, shape=upstream_shape)
