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

import sys
from pathlib import Path
import numpy as np
import matplotlib.pyplot as plt
import laspy
sys.path.append('../triskele/python/')
import triskele
import rasterizer

def rasterize_cache(field, data, resolution=.5, interpolation='nearest', reverse_direction=False, cache_dir='/tmp'):
    """Cache layer for rasterize"""
    
    cache_dir = Path(cache_dir)
    name = '{}_{}_f{}_c{}_r{}_{}{}'.format(field, data.name, data.filter_treshold, data.clip_treshold, resolution,
                                       interpolation, '_reversed' if reverse_direction else '')
    png_file = cache_dir.joinpath(Path(name + '.png'))
    tif_file = cache_dir.joinpath(Path(name + '.tif'))
    
    if tif_file.exists() :
        print ('WARNING: Loading cached result {}'.format(tif_file))
        raster = triskele.read(tif_file)
    else:
        raster = rasterizer.rasterize(data.spatial, getattr(data, field), resolution, interpolation, reverse_direction, np.float32)
        triskele.write(tif_file, raster)
        plt.imsave(png_file, raster)
        
    return raster

def find_las(path):
    path = Path(path)
    las = list()
    
    if path.is_dir():
        for child in path.iterdir():
            las.extend(find_las(child))
    
    if path.is_file() and path.suffix == '.las':
        las.append(path)
    
    return las

def extremum_filter(data, treshold=.5):
    """Return boolean filter of extremums in data, treshold in percentil [0;100]"""
    tresholds = np.percentile(data, [treshold, 100 - treshold])
    return np.logical_or(data < tresholds[0], data > tresholds[1])

def extremum_clip(data, treshold=.5):
    """Clip extremums in data, treshold in percentil [0;100]"""
    tresholds = np.percentile(data, [treshold, 100 - treshold])
    np.clip(data, *tresholds, out=data)
    
def bulk_load(path, name=None, filter_treshold=.1, clip_treshold=.5, dtype=None):
    data = {'file': path, 
            'name': name,
            'filter_treshold': filter_treshold,
            'clip_treshold': clip_treshold}

    attributes = ['x', 'y', 'z', 'intensity', 'num_returns']#, 'scan_angle_rank', 'pt_src_id', 'gps_time']
    for a in attributes:
        data[a] = list()
        
    paths = [path] if isinstance(path, str) else path
    
    print('Load data...')
    for path in paths: 
        for f in find_las(path):
            print('{}: '.format(f), end='')
            infile = laspy.file.File(f)
            for i, a in enumerate(attributes):
                print('\r {}: [{:3d}%]'.format(f, int(i/len(attributes) * 100)), end='')
                data[a].extend(getattr(infile, a))
            infile.close()
            print('\r {}: [Done]'.format(f))

    
    print('Create matrices...', end='')
    for i, a in enumerate(attributes):
        print('\rCreate matrices: [{:3d}%]'.format(int(i/len(attributes) * 100)), end='')
        data[a] = np.array(data[a], dtype=dtype)
    print('\rCreate matrices: [Done]')
        
    print('Filter data...', end='')
    z_or_i = np.logical_or(extremum_filter(data['z'], filter_treshold),
                           extremum_filter(data['intensity'], filter_treshold))
    for i, a in enumerate(attributes):
        print('\rFilter data: [{:3d}%]'.format(int(i/len(attributes) * 100)), end='')
        data[a] = data[a][~z_or_i]
    print('\rFilter data: [Done]')
        
    print('Clip data...', end='')
    extremum_clip(data['z'], clip_treshold)
    extremum_clip(data['intensity'], clip_treshold)
    data['spatial'] = np.array((data['x'], data['y'], data['z'])).T
    attributes.append('spatial')
    print('\rClip data: [Done]')
    
    class TMPLAS(object):
        def __init__(self, d):
            self.__dict__.update(d)
            
    return TMPLAS(data)