from __future__ import division

import SimpleITK as sitk
import numpy as np


def convert_axis(axis):
    return 2 - axis


def max_projection(volume, axis, *a, **k):
    volume = sitk.GetArrayFromImage(volume)
    axis = convert_axis(axis)

    return np.amax(volume, axis), np.argmax(volume, axis)


def template_projection(volume, axis, gain=2, maxv=1, *a, **k):
    volume = sitk.GetArrayFromImage(volume)
    axis = convert_axis(axis)

    output_shape = list(volume.shape)
    del output_shape[axis]
    output = np.zeros(output_shape, dtype=float)

    for ii in xrange(volume.shape[axis]):
        current = volume.take(ii, axis)

        output = np.multiply(output, (maxv - current) / maxv)
        output += gain * np.multiply(current, current) / maxv

    return output



