import multiprocessing as mp
import logging

from six import iteritems
import SimpleITK as sitk
import numpy as np

from base_subimage import run_subimage
from classic_subimage import ClassicSubImage
import image_utilities as iu
from downsampling_utilities import block_average, window_average


#==============================================================================


class ImageSeriesGridder(object):

    @property
    def transform(self):

        if not hasattr(self, '_transform'):
            dfmfld = sitk.ReadImage(str(self.dfmfld_path))
            self._transform = iu.build_composite_transform(dfmfld, self.affine_params)
            del dfmfld
  
        return self._transform


    def __init__(self, in_dims, in_spacing, 
                       out_dims, out_spacing,
                       reduce_level,  
                       subimages, 
                       subimage_kwargs, 
                       nprocesses, 
                       affine_params, 
                       dfmfld_path):
        
        self.in_dims = np.array(in_dims)
        self.in_spacing = np.array(in_spacing)
        
        self.out_dims = np.array(out_dims)
        self.out_spacing = np.array(out_spacing)
        
        self.reduce_level = reduce_level
    
        self.nprocesses = nprocesses
        
        self.affine_params = affine_params
        self.dfmfld_path = dfmfld_path
        
        self.volumes = {}
        
        self.subimages = subimages
        self.subimage_kwargs = subimage_kwargs
        
        
    def set_coarse_grid_parameters(self):
        
        self.coarse_dims, self.coarse_spacing, self.coarse_grid_radius = \
            iu.compute_coarse_parameters(self.in_dims, self.in_spacing, 
                                         self.out_spacing[::-1], 
                                         self.reduce_level)
                                         
        self.coarse_dims[-1] = self.in_dims[-1]
        self.coarse_spacing[-1] = self.in_spacing[-1]
        self.coarse_grid_radius = self.coarse_grid_radius[0]
        
        
    def setup_subimages(self):
        
        if not hasattr(self, 'coarse_grid_radius'):
            self.set_coarse_grid_parameters()
        
        dc = {'in_dims': self.in_dims[:2], 
              'in_spacing': self.in_spacing[:2], 
              'coarse_dims': self.coarse_dims[:2], 
              'coarse_spacing': self.coarse_spacing[:2], 
              'reduce_level': self.reduce_level}
        dc.update(self.subimage_kwargs)
        
        for si in self.subimages:
            si.update(dc)
        
    
    def initialize_coarse_volume(self, key, dtype):
        logging.info('initializing {0} coarse grid volume'.format(key))
        self.volumes[key] = iu.new_image(self.coarse_dims, self.coarse_spacing, dtype, True)

        origin = list(self.volumes[key].GetOrigin())
        origin[2] = 0
        self.volumes[key].SetOrigin(origin)
        
        
    def paste_slice(self, key, index, slice_array):
        '''
        '''
        
        if not key in self.volumes:
            sitk_type = iu.np_sitk_convert(slice_array.dtype)
            self.initialize_coarse_volume(key, sitk_type)
        
        logging.info('resampling data from index {0} into {1} coarse grid volume'.format(index, key))
        slice_image = iu.image_from_array(slice_array.T, self.coarse_spacing[:2], True)
        self.volumes[key] = iu.resample_into_volume(slice_image, None, index, self.volumes[key])
        
    
    def paste_subimage(self, index, output):
        '''Inserts planar accumulators into coarse grid volumes
        '''
        
        for key, array in iteritems(output):
            self.paste_slice(key, index, array)
            output[key] = None
            
        del output
    
    
    def build_coarse_grids(self):
    
        pool = mp.Pool(processes=self.nprocesses)
        mapper = pool.imap_unordered(run_subimage, self.subimages)
        
        logging.info('building coarse grids')
        for index, output in mapper:
            
            logging.info('received coarse planar data from subimage at index {0}'.format(index))
            self.paste_subimage(index, output)
        

    def resample_volume(self, key):
        logging.info('resampling {0} volume'.format(key))
        self.volumes[key] = iu.resample_volume(self.volumes[key], self.out_dims, 
                                               self.out_spacing, None, 
                                               self.transform)


    def consume_volume(self, key, cb):
        logging.info('consuming {0} volume'.format(key))
        self.resample_volume(key)
        cb(self.volumes[key])
        del self.volumes[key]


    def accumulator_to_numpy(self, key, cb):
        self.resample_volume(key)
        cb(self.volumes[key])
        logging.info('converting {0} volume to ndarray'.format(key))
        self.volumes[key] = sitk.GetArrayFromImage(self.volumes[key])


    def make_ratio_volume(self, num_key, den_key, ratio_key):
        '''assume parents numpified
        '''

        self.volumes[ratio_key] = np.divide(self.volumes[num_key], self.volumes[den_key])
        self.volumes[ratio_key][np.isnan(self.volumes[ratio_key])] = 0

