#!/usr/bin/python
# -*- coding: utf-8 -*-
# \file rasterizer.py
# \brief TODO
# \author Florent Guiotte <florent.guiotte@gmail.com>
# \version 0.1
# \date 03 juil. 2018
#
# TODO details

import numpy as np
from scipy.interpolate import griddata

def rasterize(spatial, values, resolution=1., method='linear', reverse_alt=False, dtype=None):
    '''
        Rasterize spatialised data.

        :param spatial: Coordinates of the points [x, y, z] shape (n, 3)
        :type spatial: ndarray
        :param values: Values of the points
        :type values: ndarray
        :param resolution: Resolution of the raster in meters
        :type resolution: float
        :param method: Method of interpolation for NoData values {'linear', 'nearest', 'cubic'}
        :type method: string
        :return: 2D raster of the values along z axis
        :rtype: ndarray

        .. warning:: WIP
    '''
    
    clip = method == 'cubic-clip'
    method = 'cubic' if method == 'cubic-clip' else method
    
    raster = vhr(spatial, values, resolution, reverse_alt, dtype)
    
    print('Interpolate NoData...')
    if method != 'nearest':
        raster = interpolate(raster, method)
        
    if clip:
        print('Clipping interpolation...')
        np.clip(raster, values.min(), values.max(), out=raster)
        
    print('Final touches...')
    raster = interpolate(raster, 'nearest')
    
    return raster

def vhr(spatial, values, resolution=1., reverse_alt=False, dtype=None):
    dtype = values.dtype if dtype is None else dtype
    
    bins = np.rint((spatial.max(axis=0) - spatial.min(axis=0)) / resolution).astype(int)
     
    print('Compute histogram...')
    voxels_histo, edges = np.histogramdd(spatial, bins=bins)
    voxels_histo = voxels_histo.astype(dtype)
    
    print('Compute weighted histogram...')
    voxels_histo_v, edges = np.histogramdd(spatial, bins=bins, weights=values)
    voxels_histo_v = voxels_histo_v.astype(dtype)

    layers = voxels_histo.shape[2]
        
    print('Compute voxels value...')
    voxels = (voxels_histo_v / voxels_histo).T
    
    raster = np.zeros_like(voxels_histo.T[0], dtype=dtype)
    raster[:] = np.nan

    del voxels_histo_v, voxels_histo

    print('Stack voxels...')
    for i in reversed(range(layers)) if not reverse_alt else range(layers):
        nan_filter = np.isnan(raster)
        raster[nan_filter] = voxels[i][nan_filter]
    
    del voxels
    
    return np.flip(raster, 0)

def interpolate(raster, method='linear'):
    nan_filter = np.isnan(raster)
    val_filter = np.logical_not(nan_filter)
    coords = np.argwhere(val_filter)
    values = raster[val_filter]
    grid   = np.argwhere(nan_filter)
    
    raster[nan_filter] = griddata(coords, values, grid, method=method)
    return raster