from modules.data import Module, debug_decorator
import numpy as np
import os


class SaveLas(Module):
    create_folder = False

    @debug_decorator
    def __call__(self,  las_data=None, las_info_dict=None, overwrite=False,
                 filename=None,
                 default=None, folder='SaveLas',
                 call_number=None, call_total=None,
                 **kwargs):
        '''
        Save las file
        '''
        # Possibly set folder from 'default'.
        folder = default if default is not None else folder
        # Use folder
        basename = os.path.basename(self.save_dir)
        if basename != folder:
            dirname = os.path.dirname(self.save_dir)
            self.save_dir = os.path.join(dirname, folder)

        # Create folder if needed
        if not os.path.isdir(self.save_dir):
            os.makedirs(self.save_dir)

        file_ending = '.laz'
        if call_total > 1:
            file_ending = f'_{call_number:05d}{file_ending}'

        if filename is None:
            filename = 'las' + file_ending

        filename = os.path.join(self.save_dir, filename)
        if os.path.exists(filename) and not overwrite:
            # Skip if file already exists
            return False

        las_data.write(filename)


class MakeLas(Module):
    ''' '''
    create_folder = False

    @debug_decorator
    def __call__(self,
                 terrain_temp=[], terrain_prim=[],
                 points_per_meter=50,
                 segmentation_array=None,
                 default=None,
                 **_):
        # Get latest terrain from temp/heap
        from utils.utils import get_terrain
        terrain = get_terrain(terrain_temp, terrain_prim, remove=False)

        # Construct an 'interpolator'
        from utils.interpolator import Interpolator
        interpolator = Interpolator(terrain.array, terrain.extent)

        # Possibly set points-per-meter from 'default'
        points_per_meter = default if default is not None else points_per_meter

        import laspy
        las = laspy.create(point_format=0)

        extent = terrain.extent
        xlim = extent[:2]
        ylim = extent[2:]

        num_points = int(points_per_meter*(xlim[1]-xlim[0])*(ylim[1]-ylim[0]))

        # Generate random x and y coordinates
        x = np.random.uniform(*xlim, num_points)
        y = np.random.uniform(*ylim, num_points)

        # Fetch height at x and y coordinates
        z = interpolator.interpolator(np.array([x, y]).T)

        # Possibly assign a classification, according to a 'segmentation_array'
        # generated by 'RenderSegmentation' module
        if segmentation_array is not None:
            segmentation_interpolator = Interpolator(
                segmentation_array, terrain.extent)
            value = segmentation_interpolator.interpolator(np.array([x, y]).T)

            # Set initial classification
            las.classification = np.ones_like(x)

            pick_ground = (value < 0.1)
            las.classification[pick_ground] = 2

            pick_rock = (value >= 0.1)
            las.classification[pick_rock] = 10

        # Setup points in las object
        las.x = x
        las.y = y
        las.z = z

        return {
            'las_data': las,
            }
