""" Represent a triangulated surface using a 3D boolean grid"""
import logging
import numpy as np

from rpl.tools.ray_tracing.bsp_tree_poly import BSP_Element
from rpl.tools.geometry import geom_utils
import data_io


class BSP_Grid(object):
    def __init__(self, node_array, tris, allocate_step=100000):
        """ 
        Store the triangles with an enumeration so that even when they are subdivided their 
        identity is not lost.
        """
        tri_nums = np.arange(len(tris), dtype=np.int32).reshape((len(tris), 1))
        minus_ones = -np.ones((len(tris), 6), dtype=np.int32)

        self.tris = np.hstack((tris, minus_ones, tri_nums))
        
        self.allocate_step = allocate_step
        self.node_array = node_array  # Reference to the full list of nodes
        self._resize()
        self.next_free = len(node_array)

        self.split_cache = np.zeros(len(self.tris), dtype=np.int32)

    def _resize(self):
        """
        Increase node array size by the allocate_step amount.
        """
        self.array_size = len(self.node_array) + self.allocate_step
        self.node_array = np.concatenate((self.node_array, np.zeros((self.allocate_step, 3))))

    def add_node(self, node):
        """
        Adds a new node to the end of the node array (expanding if required). Returns the index of
        the newly added node.
        """
        if self.next_free == self.array_size:
            self._resize()
        self.node_array[self.next_free] = node
        self.next_free += 1
        return self.next_free - 1

    def prepare_add(self, num_add_nodes):
        """
        Make sure that ``num_add_nodes`` can be added later without needing a resize.
        
        Useful if adding nodes from within cython where resizing is tricky.
        """
        if self.next_free + num_add_nodes >= self.array_size:
            self._resize()
        return self.next_free


def make_grid(veh_surfs, settings):
    """
    Make coordinates of voxelated grid based on overall list of vehicle surfaces
    """
    ## Find overall bounding box
    x_min, x_max = 1e30, -1e30
    y_min, y_max = 1e30, -1e30
    z_min, z_max = 1e30, -1e30

    for key, veh_surf in veh_surfs.items():
        x_min, x_max = min(x_min, np.min(veh_surf["x"])), max(x_max, np.max(veh_surf["x"]))
        y_min, y_max = min(y_min, np.min(veh_surf["y"])), max(y_max, np.max(veh_surf["y"]))
        z_min, z_max = min(z_min, np.min(veh_surf["z"])), max(z_max, np.max(veh_surf["z"]))

    x_min, x_max = x_min - settings["voxel_size"],  x_max + settings["voxel_size"]
    y_min, y_max = y_min - settings["voxel_size"],  y_max + settings["voxel_size"]
    z_min, z_max = z_min - settings["voxel_size"],  z_max + settings["voxel_size"]

    ###########################################
    # Create the uniformly spaced grid points
    x_grid = np.arange(x_min, x_max + settings["voxel_size"], settings["voxel_size"])
    y_grid = np.arange(y_min, y_max + settings["voxel_size"], settings["voxel_size"])
    z_grid = np.arange(z_min, z_max + settings["voxel_size"], settings["voxel_size"])

    return x_grid, y_grid, z_grid


def convert_geom(veh_surf, tr_mat):
    """
    Rotate nodes using provided transformation matrix; convert xyz node dict to nodes array
    """
    veh_surf["nodes"] = np.vstack((veh_surf["x"], veh_surf["y"], veh_surf["z"])).T
    veh_surf['nodes'] = np.dot(veh_surf['nodes'], tr_mat[:3, :3])

    veh_surf["x"] = veh_surf['nodes'][:, 0]
    veh_surf["y"] = veh_surf['nodes'][:, 1]
    veh_surf["z"] = veh_surf['nodes'][:, 2]

    return veh_surf


def find_occupied_voxels(surf, surf_mask, voxel_data):
    """
    Voxels with any triangle from ``surf`` are considered occupied and or'ed with ``group_mask``.
    
    If the supplied ``occupied_voxels`` is None a voxel array is created and returned.
    """
    nodes = surf["nodes"]
    tris = surf["tris"]
    
    x_pts, y_pts, z_pts = [voxel_data[k] for k in ("x_grid", "y_grid", "z_grid")]
    vox_size = voxel_data["vox_size"]
    ## Find the local extents of this part
    min_x, max_x = np.min(surf["x"]) - vox_size, np.max(surf["x"]) + vox_size
    min_y, max_y = np.min(surf["y"]) - vox_size, np.max(surf["y"]) + vox_size
    min_z, max_z = np.min(surf["z"]) - vox_size, np.max(surf["z"]) + vox_size
    
    b_tree = BSP_Grid(nodes, tris)

    # Create BSP tree elements- we're not using a tree, but we are using some of the functions
    b_x_root = BSP_Element(b_tree.tris, b_tree)

    size_i, size_j, size_k = len(x_pts), len(y_pts), len(z_pts)

    ## Create the occupied voxels if none were supplied
    if voxel_data["value"] is None:
        voxel_data["value"] = np.zeros((size_i - 1, size_j - 1, size_k - 1), dtype=np.uint32)

    occupied_voxels = voxel_data["value"]
    
    ## The [1:] is because to make n voxels in a given direction we need n-1 splits
    for i, x_pos in enumerate(x_pts[1:]):
        if x_pos < min_x: continue
        if x_pos > max_x: break
        b_above_x, b_below_x = b_x_root.split_at(0, x_pos)
        b_y_root = b_below_x
        for j, y_pos in enumerate(y_pts[1:]):
            if b_y_root is None:
                break
            if y_pos < min_y: continue
            if y_pos > max_y: break
            b_above_y, b_below_y = b_y_root.split_at(1, y_pos)
            b_z_root = b_below_y
            for k, z_pos in enumerate(z_pts[1:]):
                if b_z_root is None:
                    break
                if z_pos < min_z: continue
                if z_pos > max_z: break
                b_above_z, b_below_z = b_z_root.split_at(2, z_pos)
                
                if not (b_below_z and (len(b_below_z.tris) == 0)):
                    ## There is at least part of triangle here so mark as occupied
                    occupied_voxels[i, j, k] |= surf_mask

                b_z_root = b_above_z
            b_y_root = b_above_y
        b_x_root = b_above_x

    return voxel_data


#############
# Main code
def main(vehicle_comp_coords, tr_mat, voxel_masks, settings):
    """
    Perform voxelization for all vehicle geometries in a list of parts. Combine on a uniform grid.
    """
    for key, veh_surf in vehicle_comp_coords.items():
        # Convert coordinates and find overall best bounding box
        veh_surf = convert_geom(veh_surf, tr_mat)

    x_grid, y_grid, z_grid = make_grid(vehicle_comp_coords, settings)

    voxel_data = {"x_grid": x_grid,
                  "y_grid": y_grid,
                  "z_grid": z_grid,
                  "vox_size": settings["voxel_size"],
                  "csys_trans": tr_mat,
                  "value": None}

    for key, veh_surf in vehicle_comp_coords.items():
        # Build up the voxel_data
        logging.debug("Sampling component: {}".format(key))
        ## Default mask is 1 for anything not in an identified set
        surf_mask = 1
        for mask, geo_set in voxel_masks.items():
            if veh_surf['part_class'] in geo_set:
                surf_mask |= mask

        voxel_data = find_occupied_voxels(veh_surf, surf_mask, voxel_data)

    return voxel_data

if __name__ == "__main__":
    from rpl.tools.api import test_bench_api as tb_api
    SETTINGS = tb_api.load_settings("settings.js")

    DOORS = {'Hatch_Assembly_Rear_Ramp', 'Hatch_Assembly_Personnel_Door'}
    HATCHES = {'Hatch_Assembly_Driver_Commander', 'Hatch_Assembly_Cargo'}
    HULLS = {"Hull_Assembly_Parametric", 'Hull_Assembly_Example_With_Connector'}
    MANIKINS = {"Manikin"}
    # Special labels applied to specific types of voxels
    VOXEL_LABELS = {2: HULLS,
                    4: DOORS,
                    8: HATCHES,
                    16: MANIKINS}

    vehicle_surfs = tb_api.load_geometry(tb_api.get_all_geom_set() - MANIKINS, single_file=False)

    # Modify node coords so object aligns with cartesian axes of occ voxel grid, +z=up
    # Vector to rotate around is cross product of current z axis and sfc normal
    veh_up = np.array([0., 1., 0.])
    rot_around = np.cross(veh_up, np.array([0, 0, 1]))
    rot_ang = -np.arccos(veh_up[2])
    tr_mat = geom_utils.rotation_about_vector(rot_around, rot_ang)

#    voxel_data = main(vehicle_surfs, tr_mat, VOXEL_LABELS, SETTINGS)
    vox_veh_folder = r"voxelated_models/vehicles/{}/{}".format(SETTINGS["run_id"],
                                                               SETTINGS["voxel_size"])

    vox_veh_file = "voxels_{}_vox{}_hacked".format(SETTINGS["run_id"],
                                                   SETTINGS["voxel_size"])

    try:
        voxel_data = data_io.load_array(vox_veh_folder, vox_veh_file, True)
    except:
        voxel_data = main(vehicle_surfs, tr_mat, VOXEL_LABELS, SETTINGS)
    
    from mayavi import mlab
    xo, yo, zo = np.where(voxel_data["value"] == 1)
    plot_vehicle = mlab.points3d(voxel_data["x_grid"][xo],
                                 voxel_data["y_grid"][yo],
                                 voxel_data["z_grid"][zo],
                                 color=(0.9, 0.9, 0.9),
                                 scale_mode="none", scale_factor=voxel_data["vox_size"],
                                 mode='cube', opacity=1)

    xo, yo, zo = np.where(voxel_data["value"] & 2)
    plot_vehicle = mlab.points3d(voxel_data["x_grid"][xo],
                                 voxel_data["y_grid"][yo],
                                 voxel_data["z_grid"][zo],
                                 color=(1, 1, 1),
                                 scale_mode="none", scale_factor=voxel_data["vox_size"],
                                 mode='cube', opacity=0.05)

    xo, yo, zo = np.where(voxel_data["value"] & 4)
    plot_vehicle = mlab.points3d(voxel_data["x_grid"][xo],
                                 voxel_data["y_grid"][yo],
                                 voxel_data["z_grid"][zo],
                                 color=(1.0, 0.5, 0.5),
                                 scale_mode="none", scale_factor=voxel_data["vox_size"],
                                 mode='cube', opacity=1)

    xo, yo, zo = np.where(voxel_data["value"] & 8)
    plot_vehicle = mlab.points3d(voxel_data["x_grid"][xo],
                                 voxel_data["y_grid"][yo],
                                 voxel_data["z_grid"][zo],
                                 color=(0.6, 0.6, 1.0),
                                 scale_mode="none", scale_factor=voxel_data["vox_size"],
                                 mode='cube', opacity=1)

    # No manikins included, no need to plot them
    # xo, yo, zo = np.where(voxel_data["value"] & 16)
    # plot_vehicle = mlab.points3d(voxel_data["x_grid"][xo],
    #                              voxel_data["y_grid"][yo],
    #                              voxel_data["z_grid"][zo],
    #                              color=(0.5, 1.0, 0.8),
    #                              scale_mode="none", scale_factor=voxel_data["vox_size"],
    #                              mode='cube', opacity=1.0)
    mlab.show()

    # Save the voxelated model of the vehicle (sans door and other excluded parts)
    data_io.save_multi_array(vox_veh_folder, vox_veh_file, voxel_data)
[((2930, 3002), 'numpy.arange', 'np.arange', (['x_min', "(x_max + settings['voxel_size'])", "settings['voxel_size']"], {}), "(x_min, x_max + settings['voxel_size'], settings['voxel_size'])\n", (2939, 3002), True, 'import numpy as np\n'), ((3017, 3089), 'numpy.arange', 'np.arange', (['y_min', "(y_max + settings['voxel_size'])", "settings['voxel_size']"], {}), "(y_min, y_max + settings['voxel_size'], settings['voxel_size'])\n", (3026, 3089), True, 'import numpy as np\n'), ((3104, 3176), 'numpy.arange', 'np.arange', (['z_min', "(z_max + settings['voxel_size'])", "settings['voxel_size']"], {}), "(z_min, z_max + settings['voxel_size'], settings['voxel_size'])\n", (3113, 3176), True, 'import numpy as np\n'), ((3475, 3516), 'numpy.dot', 'np.dot', (["veh_surf['nodes']", 'tr_mat[:3, :3]'], {}), "(veh_surf['nodes'], tr_mat[:3, :3])\n", (3481, 3516), True, 'import numpy as np\n'), ((4565, 4597), 'rpl.tools.ray_tracing.bsp_tree_poly.BSP_Element', 'BSP_Element', (['b_tree.tris', 'b_tree'], {}), '(b_tree.tris, b_tree)\n', (4576, 4597), False, 'from rpl.tools.ray_tracing.bsp_tree_poly import BSP_Element\n'), ((7387, 7422), 'rpl.tools.api.test_bench_api.load_settings', 'tb_api.load_settings', (['"""settings.js"""'], {}), "('settings.js')\n", (7407, 7422), True, 'from rpl.tools.api import test_bench_api as tb_api\n'), ((8159, 8184), 'numpy.array', 'np.array', (['[0.0, 1.0, 0.0]'], {}), '([0.0, 1.0, 0.0])\n', (8167, 8184), True, 'import numpy as np\n'), ((8289, 8342), 'rpl.tools.geometry.geom_utils.rotation_about_vector', 'geom_utils.rotation_about_vector', (['rot_around', 'rot_ang'], {}), '(rot_around, rot_ang)\n', (8321, 8342), False, 'from rpl.tools.geometry import geom_utils\n'), ((8967, 9001), 'numpy.where', 'np.where', (["(voxel_data['value'] == 1)"], {}), "(voxel_data['value'] == 1)\n", (8975, 9001), True, 'import numpy as np\n'), ((9022, 9224), 'mayavi.mlab.points3d', 'mlab.points3d', (["voxel_data['x_grid'][xo]", "voxel_data['y_grid'][yo]", "voxel_data['z_grid'][zo]"], {'color': '(0.9, 0.9, 0.9)', 'scale_mode': '"""none"""', 'scale_factor': "voxel_data['vox_size']", 'mode': '"""cube"""', 'opacity': '(1)'}), "(voxel_data['x_grid'][xo], voxel_data['y_grid'][yo],\n    voxel_data['z_grid'][zo], color=(0.9, 0.9, 0.9), scale_mode='none',\n    scale_factor=voxel_data['vox_size'], mode='cube', opacity=1)\n", (9035, 9224), False, 'from mayavi import mlab\n'), ((9407, 9440), 'numpy.where', 'np.where', (["(voxel_data['value'] & 2)"], {}), "(voxel_data['value'] & 2)\n", (9415, 9440), True, 'import numpy as np\n'), ((9461, 9660), 'mayavi.mlab.points3d', 'mlab.points3d', (["voxel_data['x_grid'][xo]", "voxel_data['y_grid'][yo]", "voxel_data['z_grid'][zo]"], {'color': '(1, 1, 1)', 'scale_mode': '"""none"""', 'scale_factor': "voxel_data['vox_size']", 'mode': '"""cube"""', 'opacity': '(0.05)'}), "(voxel_data['x_grid'][xo], voxel_data['y_grid'][yo],\n    voxel_data['z_grid'][zo], color=(1, 1, 1), scale_mode='none',\n    scale_factor=voxel_data['vox_size'], mode='cube', opacity=0.05)\n", (9474, 9660), False, 'from mayavi import mlab\n'), ((9843, 9876), 'numpy.where', 'np.where', (["(voxel_data['value'] & 4)"], {}), "(voxel_data['value'] & 4)\n", (9851, 9876), True, 'import numpy as np\n'), ((9897, 10099), 'mayavi.mlab.points3d', 'mlab.points3d', (["voxel_data['x_grid'][xo]", "voxel_data['y_grid'][yo]", "voxel_data['z_grid'][zo]"], {'color': '(1.0, 0.5, 0.5)', 'scale_mode': '"""none"""', 'scale_factor': "voxel_data['vox_size']", 'mode': '"""cube"""', 'opacity': '(1)'}), "(voxel_data['x_grid'][xo], voxel_data['y_grid'][yo],\n    voxel_data['z_grid'][zo], color=(1.0, 0.5, 0.5), scale_mode='none',\n    scale_factor=voxel_data['vox_size'], mode='cube', opacity=1)\n", (9910, 10099), False, 'from mayavi import mlab\n'), ((10282, 10315), 'numpy.where', 'np.where', (["(voxel_data['value'] & 8)"], {}), "(voxel_data['value'] & 8)\n", (10290, 10315), True, 'import numpy as np\n'), ((10336, 10538), 'mayavi.mlab.points3d', 'mlab.points3d', (["voxel_data['x_grid'][xo]", "voxel_data['y_grid'][yo]", "voxel_data['z_grid'][zo]"], {'color': '(0.6, 0.6, 1.0)', 'scale_mode': '"""none"""', 'scale_factor': "voxel_data['vox_size']", 'mode': '"""cube"""', 'opacity': '(1)'}), "(voxel_data['x_grid'][xo], voxel_data['y_grid'][yo],\n    voxel_data['z_grid'][zo], color=(0.6, 0.6, 1.0), scale_mode='none',\n    scale_factor=voxel_data['vox_size'], mode='cube', opacity=1)\n", (10349, 10538), False, 'from mayavi import mlab\n'), ((11212, 11223), 'mayavi.mlab.show', 'mlab.show', ([], {}), '()\n', (11221, 11223), False, 'from mayavi import mlab\n'), ((11315, 11381), 'data_io.save_multi_array', 'data_io.save_multi_array', (['vox_veh_folder', 'vox_veh_file', 'voxel_data'], {}), '(vox_veh_folder, vox_veh_file, voxel_data)\n', (11339, 11381), False, 'import data_io\n'), ((635, 674), 'numpy.hstack', 'np.hstack', (['(tris, minus_ones, tri_nums)'], {}), '((tris, minus_ones, tri_nums))\n', (644, 674), True, 'import numpy as np\n'), ((3391, 3447), 'numpy.vstack', 'np.vstack', (["(veh_surf['x'], veh_surf['y'], veh_surf['z'])"], {}), "((veh_surf['x'], veh_surf['y'], veh_surf['z']))\n", (3400, 3447), True, 'import numpy as np\n'), ((4792, 4855), 'numpy.zeros', 'np.zeros', (['(size_i - 1, size_j - 1, size_k - 1)'], {'dtype': 'np.uint32'}), '((size_i - 1, size_j - 1, size_k - 1), dtype=np.uint32)\n', (4800, 4855), True, 'import numpy as np\n'), ((8217, 8236), 'numpy.array', 'np.array', (['[0, 0, 1]'], {}), '([0, 0, 1])\n', (8225, 8236), True, 'import numpy as np\n'), ((8254, 8274), 'numpy.arccos', 'np.arccos', (['veh_up[2]'], {}), '(veh_up[2])\n', (8263, 8274), True, 'import numpy as np\n'), ((8772, 8826), 'data_io.load_array', 'data_io.load_array', (['vox_veh_folder', 'vox_veh_file', '(True)'], {}), '(vox_veh_folder, vox_veh_file, True)\n', (8790, 8826), False, 'import data_io\n'), ((4192, 4209), 'numpy.min', 'np.min', (["surf['x']"], {}), "(surf['x'])\n", (4198, 4209), True, 'import numpy as np\n'), ((4222, 4239), 'numpy.max', 'np.max', (["surf['x']"], {}), "(surf['x'])\n", (4228, 4239), True, 'import numpy as np\n'), ((4271, 4288), 'numpy.min', 'np.min', (["surf['y']"], {}), "(surf['y'])\n", (4277, 4288), True, 'import numpy as np\n'), ((4301, 4318), 'numpy.max', 'np.max', (["surf['y']"], {}), "(surf['y'])\n", (4307, 4318), True, 'import numpy as np\n'), ((4350, 4367), 'numpy.min', 'np.min', (["surf['z']"], {}), "(surf['z'])\n", (4356, 4367), True, 'import numpy as np\n'), ((4380, 4397), 'numpy.max', 'np.max', (["surf['z']"], {}), "(surf['z'])\n", (4386, 4397), True, 'import numpy as np\n'), ((7917, 7942), 'rpl.tools.api.test_bench_api.get_all_geom_set', 'tb_api.get_all_geom_set', ([], {}), '()\n', (7940, 7942), True, 'from rpl.tools.api import test_bench_api as tb_api\n'), ((1187, 1220), 'numpy.zeros', 'np.zeros', (['(self.allocate_step, 3)'], {}), '((self.allocate_step, 3))\n', (1195, 1220), True, 'import numpy as np\n'), ((2320, 2341), 'numpy.min', 'np.min', (["veh_surf['x']"], {}), "(veh_surf['x'])\n", (2326, 2341), True, 'import numpy as np\n'), ((2355, 2376), 'numpy.max', 'np.max', (["veh_surf['x']"], {}), "(veh_surf['x'])\n", (2361, 2376), True, 'import numpy as np\n'), ((2413, 2434), 'numpy.min', 'np.min', (["veh_surf['y']"], {}), "(veh_surf['y'])\n", (2419, 2434), True, 'import numpy as np\n'), ((2448, 2469), 'numpy.max', 'np.max', (["veh_surf['y']"], {}), "(veh_surf['y'])\n", (2454, 2469), True, 'import numpy as np\n'), ((2506, 2527), 'numpy.min', 'np.min', (["veh_surf['z']"], {}), "(veh_surf['z'])\n", (2512, 2527), True, 'import numpy as np\n'), ((2541, 2562), 'numpy.max', 'np.max', (["veh_surf['z']"], {}), "(veh_surf['z'])\n", (2547, 2562), True, 'import numpy as np\n')]