FitToHeightMap

#!/usr/bin/env python3

from dataclasses import dataclass
from pathlib import Path

# noinspection PyUnresolvedReferences
import vtkmodules.vtkInteractionStyle
# noinspection PyUnresolvedReferences
import vtkmodules.vtkRenderingOpenGL2
from vtkmodules.vtkCommonColor import vtkNamedColors
from vtkmodules.vtkCommonCore import vtkLookupTable
from vtkmodules.vtkFiltersCore import vtkProbeFilter
from vtkmodules.vtkFiltersGeneral import vtkWarpScalar
from vtkmodules.vtkFiltersGeometry import vtkImageDataGeometryFilter
from vtkmodules.vtkFiltersModeling import vtkFitToHeightMapFilter
from vtkmodules.vtkFiltersSources import vtkPlaneSource
from vtkmodules.vtkIOImage import vtkDEMReader
from vtkmodules.vtkInteractionWidgets import (
    vtkTextRepresentation,
    vtkTextWidget)
from vtkmodules.vtkRenderingCore import (
    vtkActor,
    vtkPolyDataMapper,
    vtkRenderWindow,
    vtkRenderWindowInteractor,
    vtkRenderer,
    vtkTextActor,
    vtkTextProperty
)


def get_program_parameters():
    import argparse
    description = 'Cover a height map with vtkPolyData generated by a vtkPlaneSource.'
    epilogue = '''
    '''
    parser = argparse.ArgumentParser(description=description, epilog=epilogue,
                                     formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument('filename', help='SainteHelens.dem.')
    args = parser.parse_args()
    return args.filename


def main():
    file_name = get_program_parameters()

    fn = Path(file_name)
    if not fn.is_file():
        print(f'{fn} not found.')
        return

    # Create the RenderWindow, Renderer
    #
    colors = vtkNamedColors()

    renderers = [
        vtkRenderer(viewport=(0, 0, 1.0 / 3.0, 1), background=colors.GetColor3d('Wheat')),
        vtkRenderer(viewport=(1.0 / 3.0, 0, 2.0 / 3.0, 1), background=colors.GetColor3d('BurlyWood')),
        vtkRenderer(viewport=(2.0 / 3.0, 0, 1, 1), background=colors.GetColor3d('Tan')),
    ]

    render_window = vtkRenderWindow(size=(1200, 400), window_name='FitToHeightMap')
    for i in range(0, 3):
        render_window.AddRenderer(renderers[i])

    interactor = vtkRenderWindowInteractor()
    interactor.render_window = render_window

    # Create the pipelines. Load terrain data.
    lut = vtkLookupTable(hue_range=(0.6, 0), saturation_range=(1.0, 0), value_range=(0.5, 1.0))

    # Read the data: a height field results.
    dem_reader = vtkDEMReader(file_name=file_name)
    dem_reader.update()

    surface = vtkImageDataGeometryFilter()

    # Warp the surface in the vertical direction.
    warp = vtkWarpScalar(scale_factor=1, use_normal=True, normal=(0, 0, 1))

    # Show the terrain
    scalar_range = dem_reader.output.scalar_range

    dem_mapper = vtkPolyDataMapper(scalar_visibility=True, scalar_range=scalar_range, lookup_table=lut)
    dem_reader >> surface >> warp >> dem_mapper

    dem_actor = vtkActor(mapper=dem_mapper)

    # Create polygon(s) to fit. z-values should be height of DEM
    # to accommodate ProbeFilter.
    z_level = dem_reader.GetOutput().GetBounds()[5]
    dem_bounds = dem_reader.output.bounds

    plane = vtkPlaneSource(origin=(dem_bounds[0], dem_bounds[2], z_level),
                           point1=(dem_bounds[1], dem_bounds[2], z_level),
                           point2=(dem_bounds[0], dem_bounds[3], z_level),
                           resolution=(128, 128))

    # # Get the scalars from the DEM
    probe_dem = vtkProbeFilter(source_data=dem_reader.output)
    plane >> probe_dem

    # Fit polygons to surface using two strategies.
    # Both will share the same lookup tables.
    point_fit = vtkFitToHeightMapFilter()
    probe_dem >> point_fit
    point_fit.SetHeightMapConnection(dem_reader.GetOutputPort())
    point_fit.SetFittingStrategyToPointProjection()
    point_fit.UseHeightMapOffsetOn()

    point_mapper = vtkPolyDataMapper(scalar_visibility=True, scalar_range=scalar_range, lookup_table=lut)
    point_fit >> point_mapper

    point_actor = vtkActor(mapper=point_mapper)

    # Fit polygons to surface (cell strategy).
    cell_fit = vtkFitToHeightMapFilter()
    probe_dem >> cell_fit
    cell_fit.SetHeightMapConnection(dem_reader.GetOutputPort())
    cell_fit.SetFittingStrategyToCellAverageHeight()
    cell_fit.UseHeightMapOffsetOn()

    cell_mapper = vtkPolyDataMapper(scalar_visibility=True, scalar_range=scalar_range, lookup_table=lut)
    cell_fit >> cell_mapper

    cell_actor = vtkActor(mapper=cell_mapper)

    # Render the three representations.
    renderers[0].AddActor(dem_actor)
    renderers[1].AddActor(point_actor)
    renderers[2].AddActor(cell_actor)

    # Now add some text.
    text = {0: 'Original Terrain', 1: 'Point Fit', 2: 'Cell Fit'}
    # Create the TextActors.
    text_actors = list()
    text_representations = list()
    text_widgets = list()
    text_property = vtkTextProperty(color=colors.GetColor3d('DarkSlateGray'), bold=True, italic=False, shadow=False,
                                    font_size=12, font_family_as_string='Courier',
                                    justification=TextProperty.Justification.VTK_TEXT_CENTERED,
                                    vertical_justification=TextProperty.VerticalJustification.VTK_TEXT_CENTERED)

    text_positions = get_text_positions(list(text.values()), justification=TextProperty.Justification.VTK_TEXT_CENTERED,
                                        vertical_justification=TextProperty.VerticalJustification.VTK_TEXT_BOTTOM
                                        )

    for k, v in text.items():
        text_actors.append(
            vtkTextActor(input=v, text_scale_mode=vtkTextActor.TEXT_SCALE_MODE_NONE, text_property=text_property))

        # Create the text representation. Used for positioning the text actor.
        text_representations.append(vtkTextRepresentation(enforce_normalized_viewport_bounds=True))
        text_representations[k].GetPositionCoordinate().value = text_positions[v]['p']
        text_representations[k].GetPosition2Coordinate().value = text_positions[v]['p2']

        # Create the TextWidgets.
        text_widgets.append(
            vtkTextWidget(representation=text_representations[k], text_actor=text_actors[k],
                          default_renderer=renderers[k], interactor=interactor, selectable=False))

    # Look down the x-axis.
    renderers[0].active_camera.SetPosition(1, 0, 0)
    renderers[0].active_camera.SetFocalPoint(0, 1, 0)
    renderers[0].active_camera.SetViewUp(0, 0, 1)
    renderers[0].ResetCamera()

    # Rotate to an oblique view.
    renderers[0].active_camera.Azimuth(30.0)
    renderers[0].active_camera.Elevation(60.0)

    # Share the cameras.
    renderers[1].active_camera = renderers[0].active_camera
    renderers[2].active_camera = renderers[0].active_camera

    render_window.Render()

    for k in text.keys():
        text_widgets[k].On()

    interactor.Start()


def get_text_positions(names, justification=0, vertical_justification=0, width=0.96, height=0.1):
    """
    Get viewport positioning information for a list of names.

    :param names: The list of names.
    :param justification: Horizontal justification of the text, default is left.
    :param vertical_justification: Vertical justification of the text, default is bottom.
    :param width: Width of the bounding_box of the text in screen coordinates.
    :param height: Height of the bounding_box of the text in screen coordinates.
    :return: A list of positioning information.
    """
    # The gap between the left or right edge of the screen and the text.
    dx = 0.02
    width = abs(width)
    if width > 0.96:
        width = 0.96

    y0 = 0.01
    height = abs(height)
    if height > 0.9:
        height = 0.9
    dy = height
    if vertical_justification == TextProperty.VerticalJustification.VTK_TEXT_TOP:
        y0 = 1.0 - (dy + y0)
        dy = height
    if vertical_justification == TextProperty.VerticalJustification.VTK_TEXT_CENTERED:
        y0 = 0.5 - (dy / 2.0 + y0)
        dy = height

    name_len_min = 0
    name_len_max = 0
    first = True
    for k in names:
        sz = len(k)
        if first:
            name_len_min = name_len_max = sz
            first = False
        else:
            name_len_min = min(name_len_min, sz)
            name_len_max = max(name_len_max, sz)
    text_positions = dict()
    for k in names:
        sz = len(k)
        delta_sz = width * sz / name_len_max
        if delta_sz > width:
            delta_sz = width

        if justification == TextProperty.Justification.VTK_TEXT_CENTERED:
            x0 = 0.5 - delta_sz / 2.0
        elif justification == TextProperty.Justification.VTK_TEXT_RIGHT:
            x0 = 1.0 - dx - delta_sz
        else:
            # Default is left justification.
            x0 = dx

        # For debugging!
        # print(
        #     f'{k:16s}: (x0, y0) = ({x0:3.2f}, {y0:3.2f}), (x1, y1) = ({x0 + delta_sz:3.2f}, {y0 + dy:3.2f})'
        #     f', width={delta_sz:3.2f}, height={dy:3.2f}')
        text_positions[k] = {'p': [x0, y0, 0], 'p2': [delta_sz, dy, 0]}

    return text_positions


@dataclass(frozen=True)
class TextProperty:
    @dataclass(frozen=True)
    class Justification:
        VTK_TEXT_LEFT: int = 0
        VTK_TEXT_CENTERED: int = 1
        VTK_TEXT_RIGHT: int = 2

    @dataclass(frozen=True)
    class VerticalJustification:
        VTK_TEXT_BOTTOM: int = 0
        VTK_TEXT_CENTERED: int = 1
        VTK_TEXT_TOP: int = 2


if __name__ == '__main__':
    main()