from __future__ import division
import itertools as it
import logging

from skimage.measure import block_reduce
from skimage.util import view_as_windows
from scipy.ndimage.filters import convolve
import numpy as np


def downsample_average(volume, current_spacing, target_spacing):
    
    factor = target_spacing / current_spacing

    if factor == 1:
        return volume

    if factor - np.floor(factor) == 0:
        volume = block_average(volume, factor)
    elif factor - np.floor(factor) == 0.5:
        volume = window_average(volume, factor)
    else:
        raise ValueError('voxels cannot be unevenly split!')

    return volume


def block_average(volume, factor):
    logging.info('downsampling by block averaging with a factor of {0}'.format(factor))
    factor = np.around(factor).astype(int)
    return block_reduce(volume, tuple([factor, factor, factor]), np.mean, 0)


def apply_divisions(image, window_size):

    for axis in xrange(image.ndim):

        slc = [slice(window_size-1, None, window_size) 
               if ii == axis else slice(0, None) 
               for ii in xrange(image.ndim)]

        image[slc] = image[slc] / 2


def window_average(volume, factor):
    logging.info('downsampling by window averaging with a factor of {0}'.format(factor))
    volume = volume.copy()

    window_size = np.ceil(factor).astype(int)
    window_step = 2 * window_size - 1
    output_size = np.ceil([sh / factor for sh in volume.shape]).astype(int)

    apply_divisions(volume, window_size)
    volume = conv(volume, factor, window_size)

    return extract(volume, factor, window_size, window_step, output_size)


def conv(image, factor, window_size):
    kernel = np.ones([window_size for ii in image.shape])  
    return convolve(image, kernel, mode='constant', cval=0.0) / factor ** image.ndim


def extract(image, factor, window_size, window_step, output_shape):
  
    output = np.zeros( output_shape )

    for case in it.product(*([[0, 1]] * image.ndim)):     

        inp = [slice(window_size - 2, None, window_step) 
               if not ii else slice(window_size, None, window_step) for ii in case]
        out = [slice(0, None, 2) if not ii else slice(1, None, 2) for ii in case]

        output[out] = image[inp]

    return output

