Component.py

"""Component class definition file."""

import shutil
import subprocess
import time

import numpy as np
import vtk
from parapy.core import Attribute, Input, Part, child
from parapy.core.validate import IsInstance, OneOf, Range
from parapy.exchange import STEPWriter
from parapy.geom import (
    CloseSurface,
    GeomBase,
    IntersectedShapes,
    Line,
    Plane,
    Point,
    RotatedShape,
    SewnShell,
    Solid,
    Vector,
)
from vtkmodules.util.numpy_support import vtk_to_numpy

from model.auxiliary import dedupe_pts
from model.paths import (
    AS_PATH,
    ROT_POINT_CLOUD_DIR,
    ROT_SHAPE_DIR,
    TCL_SCRIPT_DIR,
)


class Component(GeomBase):
    """Represent the Component class.

    Each component that is part of the environment is represented by an instance of this class.
    """

    # CONSTANTS
    NORMAL_VECTOR_DICT: dict = {"xy": [0, 0, 1],
                                "xz": [0, 1, 0],
                                "yx": [0, 0, 1],
                                "yz": [1, 0, 0],
                                "zx": [0, 1, 0],
                                "zy": [1, 0, 0]}

    AXES_COLORS: list[str] = ["red", "green", "blue"]

    # Required Inputs
    name: str = Input(validator=IsInstance(str)) # Name of the component
    indexes: list[int] = Input() # Indexes of oml step-array
    orientation: str = Input(validator=OneOf(("xy", "xz", "yx", "yz", "zx", "zy"))) # Orientation plane

    # Optional Inputs: orientation deduction
    orientation_vector_interpolation_percentage: float = Input(0.99, validator=Range(0,1)) # In case of wing: percentage of chord on which orientation vector is based: 0.75 is quarter-chord line, 1 is LE.
    intersect_plane_percentage_min: float = Input(0.1, validator=Range(0,1)) # Position of first intersection-plane as percentage of half-length
    intersect_plane_percentage_max: float = Input(0.1, validator=Range(0,1)) # Position of second intersection-plane as percentage of half-length

    # Optional Inputs: ASitus cell creation
    min_cell_size: int = Input(1000, validator=Range(50, 1500)) # Minimum cell size of voxelization
    max_cell_Size: int = Input(1500, validator=Range(100, 3000)) # Maximum cell size of voxelization
    cube_cells: bool = Input(True, validator=IsInstance(bool)) # Bool to set voxels as cubes
    local_dedupe_thresh: float = Input(50)

    # Optional Inputs: Self-Made Cell Generation Method
    self_made_point_cloud: bool = Input(False)
    distance_per_point: int = Input(500) # distance between consecutive points

    # Optional Inputs: tolerance on sewing surfaces together
    fuse_tolerance: float = Input(50) # tolerance for connecting edges, such as the trailing edge of wings

    # Inputs for displaying the component
    transparency: float = Input(0.8)
    color: float = Input("yellow")

    @indexes.validator
    def indexes(self, indexes):
        """Validate the indexes input."""
        if not isinstance(indexes, list):
            return False, "Indexes must be a list of integers."
        msgs = ""
        valid = True
        for index in indexes:
            if index >= len(self.parent.oml_children) or index < 0:
                msg = f"\nIndex {index} of {self.name} component out of range. Max index = {len(self.parent.oml_children) - 1}."
                msgs += msg
                valid = False
        if valid:
            return valid
        else:
            return valid, msgs

    @local_dedupe_thresh.validator
    def local_dedupe_thresh(self, value):
        """Validate the dedupe threshold."""
        if value >= self.min_cell_size:
            msg = f"Local dedupe threshold is set too high with respect to minimum cell size: {value} >= {self.min_cell_size}."
            return False, msg
        else:
            return True

    @Part
    def step_parts(self):
         return Solid(quantify=len(self.indexes),
                      built_from=self.parent.oml_children[self.indexes[child.index]],
                      label=str(child.index))

    @Attribute
    def fused_shell(self):
        """Fuse shells if the component is made of multiple shells, such as for wings."""
        fused = self.parent.oml_children[self.indexes[0]]
        if len(self.indexes) > 1:
            for i in self.indexes[1:]:
                fused = SewnShell(
                    built_from=[fused, self.parent.oml_children[i]],
                    tolerance=self.fuse_tolerance,
                    hidden=True
                )
        return fused

    @Part
    def closed_solid(self):
        """Close any open edges of the component."""
        return CloseSurface(built_from=self.fused_shell,
                            transparency=self.transparency,
                            color=self.color)

    @Attribute
    def point_cloud(self) -> list[Point]:
        """Create a point cloud contained within the closed_solid.

        The point cloud is created based on the chosen creation method:
        either self-made (Dennis-algorithm), or through Analysis Situs.

        Return: A list of points inside the closed_solid
        """
        if self.self_made_point_cloud:
            return self.create_self_made_point_cloud()
        else:
            return self.create_asitus_point_cloud()

    def create_asitus_point_cloud(self):
        """Create points attribute from .vtu file.

        .vtu file is generated using analysis situs from the voxelized STEP file.
        Only outputs points, not the connections between them.
        """
        """Read the VTU File."""
        reader = vtk.vtkXMLUnstructuredGridReader()
        reader.SetFileName(str(ROT_POINT_CLOUD_DIR / f"rotated_{self.name}.vtu"))
        reader.Update()
        data = reader.GetOutput()

        print(f"Updating the point-cloud for {self.name}")
        """Extract the list of points from the data."""
        points_array = vtk_to_numpy(data.GetPoints().GetData())
        points = []
        for point in points_array:
            points.append(Point(point[0], point[1], point[2]).rotate(self.rotation_vector_and_angle[0],
                                                                     -self.rotation_vector_and_angle[1]))
        points = dedupe_pts(points, self.local_dedupe_thresh)

        return points

    ## Parts to deduce orientation
    @Attribute
    def min_max_planes(self):
        bbox = self.closed_solid.bbox

        max_point = bbox.position.location.translate(x=bbox.width / 2 * self.intersect_plane_percentage_max, y=bbox.length / 2 * self.intersect_plane_percentage_max, z=bbox.height / 2 * self.intersect_plane_percentage_max)
        min_point = bbox.position.location.translate(x=-bbox.width / 2 * self.intersect_plane_percentage_min, y=-bbox.length / 2 * self.intersect_plane_percentage_min, z=-bbox.height / 2 * self.intersect_plane_percentage_min)

        return [min_point, max_point]

    @Part
    def intersect_planes(self):
        return Plane(quantify=2,
                     normal=self.NORMAL_VECTOR_DICT[self.orientation],
                     reference=self.min_max_planes[child.index],
                     v_dim=5000,
                     hidden=True)

    @Part
    def intersection_curves_min(self):
        return IntersectedShapes(shape_in=self.closed_solid,
                                 tool=self.intersect_planes[0])

    @Part
    def intersection_curves_max(self):
        return IntersectedShapes(shape_in=self.closed_solid,
                                 tool=self.intersect_planes[1])

    @Attribute
    def orientation_vectors(self):

        sample_points_max = np.concatenate(
            [edge.sample_points for edge in self.intersection_curves_max.edges],
            axis=0
        )

        # Considering the minimum plane
        sample_points_min = np.concatenate(
            [edge.sample_points for edge in self.intersection_curves_min.edges],
            axis=0
        )

        maxplane_max_arg = np.argmax(sample_points_max, axis=0)
        maxplane_min_arg = np.argmin(sample_points_max, axis=0)
        minplane_max_arg = np.argmax(sample_points_min, axis=0)
        minplane_min_arg = np.argmin(sample_points_min, axis=0)

        if self.orientation[0] == "x":
            vector_point_max = Point(*map(float, sample_points_max[maxplane_max_arg[0]])).interpolate(Point(*map(float, sample_points_max[maxplane_min_arg[0]])), self.orientation_vector_interpolation_percentage)
            vector_point_min = Point(*map(float, sample_points_min[minplane_max_arg[0]])).interpolate(Point(*map(float, sample_points_min[minplane_min_arg[0]])), self.orientation_vector_interpolation_percentage)
            vector = vector_point_min.vector_to(vector_point_max)
        elif self.orientation[0] == "y":
            vector_point_max = Point(*map(float, sample_points_max[maxplane_max_arg[1]])).interpolate(Point(*map(float, sample_points_max[maxplane_min_arg[1]])), self.orientation_vector_interpolation_percentage)
            vector_point_min = Point(*map(float, sample_points_min[minplane_max_arg[1]])).interpolate(Point(*map(float, sample_points_min[minplane_min_arg[1]])), self.orientation_vector_interpolation_percentage)
            vector = vector_point_min.vector_to(vector_point_max)
        elif self.orientation[0] == "z":
            vector_point_max = Point(*map(float, sample_points_max[maxplane_max_arg[2]])).interpolate(Point(*map(float, sample_points_max[maxplane_min_arg[2]])), self.orientation_vector_interpolation_percentage)
            vector_point_min = Point(*map(float, sample_points_min[minplane_max_arg[2]])).interpolate(Point(*map(float, sample_points_min[minplane_min_arg[2]])), self.orientation_vector_interpolation_percentage)
            vector = vector_point_min.vector_to(vector_point_max)
        else:
            vector = Vector(0,0,0)
            vector_point_min = Point(0,0,0)
            vector_point_max = Point(0,0,0)

        return [vector, vector_point_min, vector_point_max]

    @Attribute
    def rotation_vector_and_angle(self):
        vector = self.orientation_vectors[0]
        if "x" not in self.orientation:
            angle, rotation_vector = vector.angle_and_axis(Vector(1, 0, 0))
        elif "y" not in self.orientation:
            angle, rotation_vector = vector.angle_and_axis(Vector(0,1,0))
        elif "z" not in self.orientation:
            angle, rotation_vector = vector.angle_and_axis(Vector(0,0,1))
        else:
            angle, rotation_vector = [0, Vector(0,0,0)]

        return rotation_vector, angle

    @Part
    def rotated_solid(self):
        return Solid(built_from=self.closed_solid.rotated(rotation_axis=self.rotation_vector_and_angle[0],
                                                                            angle=self.rotation_vector_and_angle[1]),
                     hidden=True)

    @Part
    def orientation_vector_visual(self):
        return Line(reference=self.orientation_vectors[1],
                    direction=self.orientation_vectors[0],
                    line_thickness=3.0,
                    color=self.AXES_COLORS[0],
                    v_dim=10000)

    @Part
    def step_writer(self):
        return STEPWriter(nodes=self.rotated_solid,
                          filename=str(ROT_SHAPE_DIR / f"rotated_{self.name}.stp"))

    ## Functions to Write to ASitus
    def write_tcl_script(self):
        """Write the TCL script based on the settings provided."""
        filepath = TCL_SCRIPT_DIR / f"{self.name}_script.tcl"

        tcl_script = f"""load-step rotated_{self.name}.stp
generate-facets -lin 1 -ang 0.5
ddf-build-svo -min {self.min_cell_size} -max {self.max_cell_Size} -prec 1{' -cube' if self.cube_cells else ''} -owner 0:2:3:1
set-param 0:2:23:1 26 3
set-param 0:2:23:1 27 1
ddf-dump-vtu 0:2:23:1 rotated_{self.name}.vtu
"""

        with filepath.open("w") as f:
            f.write(tcl_script)

    def call_analysis_situs(self):
        """Call Analysis Situs using the command line."""
        # Generate the required TCL script
        self.write_tcl_script()

        # Copy the TCL script from local directory to AS directory
        src = str(TCL_SCRIPT_DIR / f"{self.name}_script.tcl")
        dest = str(AS_PATH)
        shutil.copy(src, dest)

        # Copy the rotated shape from local directory to AS directory
        src = str(ROT_SHAPE_DIR / f"rotated_{self.name}.stp")
        dest = str(AS_PATH)
        shutil.copy(src, dest)

        # Execute the script
        cmd = f'asiExe.exe /runscript="{self.name}_script.tcl" > log_{self.name}.txt'
        subprocess.run(cmd, cwd=r"C:\Program Files\ASitus", shell=True)

        # Copy the result from AS directory to local directory
        src = str(AS_PATH / f"rotated_{self.name}.vtu")
        dest = str(ROT_POINT_CLOUD_DIR / f"rotated_{self.name}.vtu")
        shutil.copy(src, dest)

        delattr(self, 'point_cloud')

    ## --------------------------------------------
    ## Self-made Point Creation Algorithm
    def create_self_made_point_cloud(self) -> list[Point]:
        """Create point cloud."""
        points = []

        aligned_part, rotations = self.component_aligned_to_axes
        rotation_axes = list(rotations)

        print(f"Creating point cloud for {self.name}")
        start = time.time()

        # get dimensions and position of the component's bounding box
        center = aligned_part.oriented_bbox.center
        width = aligned_part.oriented_bbox.width
        length = aligned_part.oriented_bbox.length
        height = aligned_part.oriented_bbox.height

        # number of points in each axis
        n_x = int((width / 2) / self.distance_per_point)
        n_y = int((length / 2) / self.distance_per_point)
        n_z = int((height / 2) / self.distance_per_point)

        # divide the bounding box into equally spaced points and check if each is inside the component
        # points are produced outwards from the center. Produces better results, especially for wings
        # was_inside = False
        for i in [center[0] + q * self.distance_per_point for q in range(-n_x, n_x + 1)]:
            for j in [center[1] + q * self.distance_per_point for q in range(-n_y, n_y + 1)]:
                for k in [center[2] + q * self.distance_per_point for q in range(-n_z, n_z + 1)]:
                    point = Point(float(i), float(j), float(k))

                    is_inside = aligned_part.solids[0].is_point_inside(point, include_boundary=True)
                    # if point is inside the aligned_solid, rotate the point and add it to the list
                    if is_inside:
                        # rotate points back to the original parts position. Apply rotations in reverse order
                        point = point.rotate(rotation_axes[1], -rotations[rotation_axes[1]],
                                             rotation_axes[0], -rotations[rotation_axes[0]]
                                             )
                        points.append(point)
                        # was_inside = True
                    # # if the point is not inside the aligned_solid, and the previous point was, break out of inner loop
                    # # as following points will also be outside the shape. (TRUE FOR CONVEX SHAPES ONLY)
                    # elif not is_inside and was_inside:
                    #     was_inside = False
                    #     break

        end = time.time()
        print(f"Elapsed time creating {self.name} point cloud: {end - start:.4f} seconds.")

        return points

    @Attribute
    def component_aligned_to_axes(self):
        """Create a copy of the closed_solid that is aligned with the reference axes.

        Return: The rotated part and the rotation operations applied
        """
        rotations = {}
        # align to x-axis
        bbox_vector_x = np.array(self.closed_solid.oriented_bbox.orientation.x) / np.linalg.norm(np.array(self.closed_solid.oriented_bbox.orientation.x))
        x_axis_vector = np.array([1, 0, 0])
        x_rotation_axis = np.cross(bbox_vector_x, x_axis_vector)
        x_rotation_axis_vector = Vector(
            float(x_rotation_axis[0]),
            float(x_rotation_axis[1]),
            float(x_rotation_axis[2]),
        )
        rotations[x_rotation_axis_vector] = np.arccos(
            np.dot(bbox_vector_x, x_axis_vector)
        )
        # rotate part
        aligned_part = RotatedShape(
            shape_in=self.closed_solid,
            rotation_point=Point(0, 0, 0),
            vector=x_rotation_axis_vector,
            angle=rotations[x_rotation_axis_vector]
        )

        # align to z-axis
        bbox_vector_z = np.array(aligned_part.oriented_bbox.orientation.z) / np.linalg.norm(np.array(aligned_part.oriented_bbox.orientation.z))
        z_axis_vector = np.array([0, 0, 1])
        z_rotation_axis = np.cross(bbox_vector_z, z_axis_vector)
        z_rotation_axis_vector = Vector(
            float(z_rotation_axis[0]),
            float(z_rotation_axis[1]),
            float(z_rotation_axis[2]),
        )
        rotations[z_rotation_axis_vector] = np.arccos(
            np.dot(bbox_vector_z, z_axis_vector)
        )
        # rotate part
        aligned_part = RotatedShape(
            shape_in=aligned_part,
            rotation_point=Point(0, 0, 0),
            vector=z_rotation_axis_vector,
            angle=rotations[z_rotation_axis_vector],
        )

        return aligned_part, rotations