from __future__ import division
import logging
import os

from six import iteritems
import numpy as np
import SimpleITK as sitk
from PIL import Image, ImageDraw
from skimage.draw import polygon

from allensdk.config.manifest import Manifest


# use np_sitk_convert or sitk_np_convert to access
# TODO: check if this already exists. If not: add more dtypes
# it does not
NUMPY_SITK_TYPE_LOOKUP = {np.dtype(np.float32): sitk.sitkFloat32}
SITK_NUMPY_TYPE_LOOKUP = {v: k for k, v in iteritems(NUMPY_SITK_TYPE_LOOKUP)}


# ITK/Numpy


def set_image_spacing(image, spacing, origin=True):
    '''
    '''

    spacing = np.array(spacing)

    image.SetSpacing(spacing.tolist())

    if origin:
        image.SetOrigin( (0.5 * spacing).tolist() )


def new_image(dims, spacing, dtype, origin=True):
    '''
    '''
    
    if len(dims) == 2:
        image = sitk.Image(dims[0], dims[1], dtype)
    elif len(dims) == 3:
        image = sitk.Image(dims[0], dims[1], dims[2], dtype)
    set_image_spacing(image, spacing, origin)
    
    return image


def image_from_array(array, spacing, origin=True):
    '''
    '''

    image = sitk.GetImageFromArray(array)
    set_image_spacing(image, spacing, origin)
    
    return image


def np_sitk_convert(np_type):
    '''
    '''
    
    return NUMPY_SITK_TYPE_LOOKUP[np_type]
    
    
def sitk_np_convert(sitk_type):
    '''
    '''
    
    return SITK_NUMPY_TYPE_LOOKUP[sitk_type]
    
    
# Math
    
    
def compute_coarse_parameters(in_dims, in_spacing, out_spacing, reduce_level):
    '''
    '''
    
    reduce_factor = pow(2, reduce_level)
    fradius = np.divide(out_spacing, in_spacing) / 2.0 / reduce_factor
    
    coarse_grid_radius = np.round(fradius)
    coarse_grid_size = (coarse_grid_radius * 2 + 1) * reduce_factor
    
    coarse_grid_spacing = np.multiply(in_spacing, coarse_grid_size)
    coarse_grid_dims = np.ceil( np.divide(in_dims, coarse_grid_size) ).astype(int)
    
    return coarse_grid_dims, coarse_grid_spacing, coarse_grid_radius
    
    
def block_apply(in_image, out_shape, dtype, blocks, fn):
    '''
    '''
    
    out_image = np.zeros(out_shape, dtype=dtype)
    
    for ii, row_block in enumerate(blocks[0]):
        for jj, col_block in enumerate(blocks[1]):
        
            out_image[ii, jj] = fn(in_image[row_block[0]:row_block[1], 
                                            col_block[0]:col_block[1]])
                                            
    return out_image
    
    
def grid_image_blocks(in_shape, in_spacing, out_spacing):
    '''
    '''

    blocks = []
    out_shape = []
    for dim in range(len(in_shape)):
        in_px_centers = np.arange(in_spacing[dim]*0.5,
                                  in_shape[dim]*in_spacing[dim],
                                  in_spacing[dim])


        out_px_edges = np.arange(out_spacing[dim],
                                 (in_shape[dim]-0.5)*in_spacing[dim],
                                 out_spacing[dim])


        dig = np.digitize(in_px_centers, out_px_edges)        

        inds = np.where(np.diff(dig)>0)[0]+1  
        inds = [0] + inds.tolist() + [in_shape[dim]]

        dim_blocks = [ (int(inds[i]), int(inds[i+1])) for i in range(len(inds)-1) ]

        out_shape.append(len(dim_blocks))
        blocks.append(dim_blocks)
        
    return blocks, out_shape
    
    
# Polygons
    
  
def rasterize_polygons(shape, scale, polys):

    canvas = np.zeros(shape, dtype=np.uint8)
    for points in polys:
    
        rpts = np.array([int(np.around(item[1] * scale[1])) for item in points])
        cpts = np.array([int(np.around(item[0] * scale[0])) for item in points])

        poly = polygon(rpts, cpts)
        canvas[poly] = 1
        
    return canvas


# Transforms
    
    
def resample_into_volume(image, transform, z, vol, dtype=sitk.sitkFloat32):
    '''
    '''

    if transform is None:
        transform = sitk.Transform()

    timage = sitk.Resample(image, transform, sitk.sitkLinear, 0.0, dtype)
    tvol = sitk.JoinSeries(timage)
    return sitk.Paste(vol, tvol, tvol.GetSize(), destinationIndex=[0,0,z])
    
    
def build_affine_transform(aff_params):
    '''
    '''

    xfm = sitk.AffineTransform(3)
    xfm.SetParameters(aff_params)
    
    return xfm
    
    
def build_composite_transform(dfmfield=None, aff_params=None):
    '''
    '''

    if dfmfield is not None and dfmfield.GetPixelIDValue() != sitk.sitkVectorFloat64:
        dfmfield = sitk.Cast(dfmfield, sitk.sitkVectorFloat64)

    if dfmfield is None and aff_params is None:
        transform = sitk.Transform()
    elif dfmfield is not None and aff_params is None:
        transform = sitk.DisplacementFieldTransform(dfmfield)
    elif dfmfield is None and aff_params is not None:
        transform = build_affine_transform(aff_params)
    elif dfmfield is not None and aff_params is not None:
        dfmxfm = sitk.DisplacementFieldTransform(dfmfield)
        affxfm = build_affine_transform(aff_params)

        transform = sitk.Transform()
        transform.AddTransform(affxfm)
        transform.AddTransform(dfmxfm)

    return transform
    
    
def resample_volume(volume, dims, spacing, interpolator=None, transform=None):
    '''
    '''

    if transform is None:
        transform = sitk.Transform() 
    if interpolator is None:
        interpolator = sitk.sitkLinear
    
    ref = new_image(dims, spacing, sitk.sitkFloat32, False)
    return sitk.Resample(volume, ref, transform, interpolator)


def write_volume(volume, name, prefix=None, specify_resolution=None, extension='.nrrd', paths=None):

    if prefix is None:
        path = name
    else:
        path = os.path.join(prefix, name)

    if specify_resolution is not None:
        path = path + '_{0}'.format(specify_resolution)

    path = path + extension
    
    logging.info('writing {0} volume to {1}'.format(name, path))
    Manifest.safe_make_parent_dirs(path)
    volume.SetOrigin([0, 0, 0])
    sitk.WriteImage(volume, str(path), True)

    if paths is not None:
        paths.append(path)
