import logging

import SimpleITK as sitk
import numpy as np

import volume_utilities as vol


class VolumeProjector(object):

    def __init__(self, view_volume):
        logging.info('initializing volume projector')
        self.view_volume = view_volume
        

    def build_rotation_transform(self, from_axis, to_axis, angle):
        logging.info('constructing rotation')        

        transform = sitk.AffineTransform(3)
        transform.SetCenter((vol.sitk_get_center(self.view_volume)).tolist())
        transform.Rotate(to_axis, from_axis, angle, True)

        logging.info(transform.__str__())
        return transform

        
    def rotate(self, from_axis, to_axis, angle):
        logging.info('rotating from axis {0} to axis {1} '
                     'by {2:2.2f} radians'.format(from_axis, to_axis, angle))
        
        transform = self.build_rotation_transform(from_axis, to_axis, angle)
        rotated = sitk.Resample(self.view_volume, transform, sitk.sitkLinear, 
                                0.0, self.view_volume.GetPixelID())

        return rotated
    
    
    def extract(self, cb, volume=None):
        logging.info('extracting projection')
        
        if volume is None:
            volume=self.view_volume

        return cb(volume)

    
    def rotate_and_extract(self, from_axes, to_axes, angles, cb):
        
        for fax, tax, angle in zip(from_axes, to_axes, angles):
                
            rotated = self.rotate(fax, tax, angle)
            yield self.extract(cb, rotated)
    

    @classmethod
    def fixed_factory(cls, volume, size):

        view_volume = sitk.Image(size[0], size[1], size[2], volume.GetPixelID())
        view_volume = vol.sitk_paste_into_center(volume, view_volume)

        return cls(view_volume)


    @classmethod
    def safe_factory(cls, volume):

        max_extent = vol.sitk_get_diagonal_length(volume)
        max_extent = [np.ceil(max_extent).astype(int)] * 3

        vpar = vol.sitk_get_size_parity(volume)
        lpar = np.mod(max_extent, 2)

        for ax in xrange(volume.GetDimension()):
            if vpar[ax] != lpar[ax]:
                max_extent[ax] += 1
        
        return cls.fixed_factory(volume, max_extent)





