from __future__ import division, print_function

import logging
import random

import numpy as np

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

from .batch_filter import BatchFilter

logger = logging.getLogger(__name__)


class ShiftAugment(BatchFilter):
    def __init__(self, prob_slip=0, prob_shift=0, sigma=0, shift_axis=0, p=1.0):
        self.prob_slip = prob_slip
        self.prob_shift = prob_shift
        self.sigma = sigma
        self.shift_axis = shift_axis
        self.p = p

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

    def skip_node(self, request):
        return random.random() > self.p

    def prepare(self, request):
        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 Jitter node must contain at least one array key. "
                + "Check to make sure that Jitter 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().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.offset = updated_roi.offset
            spec.roi.shape = updated_roi.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,
            )
            array.data = self.shift_and_crop(
                array.data,
                request[array_key].roi.shape,
                sub_shift_array,
                array.spec.voxel_size,
            )
            array.spec.roi = request[array_key].roi
            assert (
                request[array_key].roi.shape
                == Coordinate(array.data.shape) * self.lcm_voxel_size
            ), "request roi shape {} is not the same as generated array shape {}".format(
                request[array_key].roi.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 the 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 Graph object with the updated point locations and ROI
        """

        nodes = list(points.nodes)
        spec = points.spec
        shift_axis_start_pos = spec.roi.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.offset - total_roi.offset
        offset_in_shift_axis = (
            item_offset_from_total[shift_axis] // lcm_voxel_size[shift_axis]
        )
        len_in_shift_axis = item_roi.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
        """
        # 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):
            r = random.random()
            slip = np.array(
                [
                    np.random.normal(scale=sigma / lcm_voxel_size[dimension])
                    for dimension, 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

        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.offset
        upstream_offset = downstream_offset - max_shift
        upstream_shape = request_roi.shape + max_shift - min_shift
        return Roi(offset=upstream_offset, shape=upstream_shape)
