import functools
import logging

import numpy as np
import SimpleITK as sitk
from scipy.ndimage.interpolation import zoom

from image_sheet import ImageSheet
from projection_functions import max_projection
from volume_projector import VolumeProjector
import visualization_utilities as vis


def max_cb(max_sheet, depth_sheet, volume, axis, *a, **k):
    m, d = max_projection(volume, axis)
    max_sheet.append(m.T)
    depth_sheet.append(d.T)


def apply_colormap(image, colormap):
    color_image = colormap(image)
    color_image[:, :, -1] = image
    return color_image


def blend_with_background(image, background):

    for ii in xrange(3):
        image[:, :, ii] = np.multiply(np.squeeze(image[:, :, ii]), 
                                      np.squeeze(image[:, :, -1]))
        image[:, :, ii] += np.multiply(np.squeeze(background), 
                                       np.squeeze(1.0 - image[:, :, -1]))

    image[:, :, -1] = 1
    return image


def do_blur(image, blur):

    for ii in xrange(3):
        im = sitk.GetImageFromArray(image[:, :, ii])
        im = sitk.DiscreteGaussian(im, blur)
        image[:, :, ii] = sitk.GetArrayFromImage(im)
    
    return image


def handle_output_image(sheet, out_image, colormap, nsteps):

    sheet = sheet.copy()
    whole_sheet = sheet.get_output(-1)

    if out_image['blur'] > 0.0:
        logging.info('applying a gaussian blur with variance: {0:2.2f}'.format(out_image['blur']))
        whole_sheet = sitk.GetImageFromArray(whole_sheet)
        whole_sheet = sitk.DiscreteGaussian(whole_sheet, out_image['blur'])
        whole_sheet = sitk.GetArrayFromImage(whole_sheet)

    if out_image['scale'] != 1:
        whole_sheet = zoom(whole_sheet, zoom=out_image['scale'], order=1)

    whole_sheet = apply_colormap(whole_sheet, colormap)

    if out_image['background'] is not None:
        whole_sheet = blend_with_background(whole_sheet, out_image['background'])
    else:
        whole_sheet = blend_with_background(whole_sheet, np.zeros_like(whole_sheet)[:, :, -1])

    whole_sheet = np.around(whole_sheet * 255).astype(np.uint8)
    out_image['write'](whole_sheet[:, :, :-1])


def simple_rotation(from_axis, to_axis, start, end, nsteps):

    angles = np.linspace(start * np.pi, end * np.pi, nsteps, endpoint=False)
    from_axes = [from_axis] * nsteps
    to_axes = [to_axis] * nsteps

    return from_axes, to_axes, angles


def run(volume, imin, imax, rotations, colormap):

    volume = vis.sitk_safe_ln(volume)

    ln_imin = np.log(imin) if imin != 0 else -np.inf
    ln_imax = np.log(imax) if imax != 0 else np.inf

    volume = sitk.IntensityWindowing(volume, ln_imin, ln_imax, 0.0, 1.0)

    for rotation in rotations:
        
        max_sheet = ImageSheet()
        depth_sheet = ImageSheet()

        vp = VolumeProjector.fixed_factory(volume, rotation['window_size'])
        callback = functools.partial(max_cb, max_sheet, depth_sheet, **rotation['projection_parameters'])

        rot = rotation['rotation_parameters']
        from_axes, to_axes, angles = simple_rotation(**rot)

        for response in vp.rotate_and_extract(from_axes, to_axes, angles, callback):
            pass

        rotation['write_depth_sheet'](depth_sheet.get_output(-1))

        for out_image in rotation['output_images']:
            handle_output_image(max_sheet, out_image, colormap, rot['nsteps'])
