#!/usr/bin/env python
#------------------------------------------------------------------------------#
#                                   Imports                                    #
#------------------------------------------------------------------------------#
from math import pi, cos, sin, sqrt
import vtk

LEVEL = 6

#------------------------------------------------------------------------------#
#                        Koch Snowflake as vtkPolyLine                         #
#------------------------------------------------------------------------------#
def as_polyline(points, level):
  # Use the points from the previous iteration to create the points of the next
  # level. There is an assumption on my part that the curve is traversed in a
  # counterclockwise fashion. If the initial triangle above is written to
  # describe clockwise motion, the points will face inward instead of outward.
  for i in range(level):
    temp = vtk.vtkPoints()
    # The first point of the previous vtkPoints is the first point of the next vtkPoints.
    temp.InsertNextPoint(*points.GetPoint(0))

    # Iterate over "edges" in the vtkPoints
    for i in range(1, points.GetNumberOfPoints()):
      x0, y0, z0 = points.GetPoint(i - 1)
      x1, y1, z1 = points.GetPoint(i)
      t = sqrt((x1 - x0)**2 + (y1 - y0)**2)
      nx = (x1 - x0)/t # x-component of edge unit tangent
      ny = (y1 - y0)/t # y-component of edge unit tangent

      # the points describing the Koch snowflake edge
      temp.InsertNextPoint(x0 + nx*t/3, y0 + ny*t/3, 0.)
      temp.InsertNextPoint(x0 + nx*t/2 + ny*t*sqrt(3)/6, y0 + ny*t/2 - nx*t*sqrt(3)/6, 0.)
      temp.InsertNextPoint(x0 + nx*2*t/3, y0 + ny*2*t/3, 0.)
      temp.InsertNextPoint(x0 + nx*t, y0 + ny*t, 0.)

    points = temp

  # draw the outline
  lines = vtk.vtkCellArray()
  pl = vtk.vtkPolyLine()
  pl.GetPointIds().SetNumberOfIds(points.GetNumberOfPoints())
  for i in range(points.GetNumberOfPoints()):
    pl.GetPointIds().SetId(i, i)
  lines.InsertNextCell(pl)

  # complete the polydata
  polydata = vtk.vtkPolyData()
  polydata.SetLines(lines)
  polydata.SetPoints(points)

  return polydata

#------------------------------------------------------------------------------#
#                 Koch Snowflake as collection of vtkTriangles                 #
#------------------------------------------------------------------------------#
def as_triangles(indices, cellarray, level, data):
  if len(indices) >= 3:
    stride = len(indices)//4
    indices.append(indices[-1] + 1)

    triangle = vtk.vtkTriangle()
    triangle.GetPointIds().SetId(0, indices[ stride])
    triangle.GetPointIds().SetId(1, indices[2*stride])
    triangle.GetPointIds().SetId(2, indices[3*stride])

    cellarray.InsertNextCell(triangle)
    data.InsertNextValue(level)

    as_triangles(indices[       0:   stride], cellarray, level + 1, data)
    as_triangles(indices[  stride: 2*stride], cellarray, level + 1, data)
    as_triangles(indices[2*stride: 3*stride], cellarray, level + 1, data)
    as_triangles(indices[3*stride:       -1], cellarray, level + 1, data)

#------------------------------------------------------------------------------#
#                                 Main Method                                  #
#------------------------------------------------------------------------------#
if __name__ == "__main__":
  # Initially, set up the points to be an equilateral triangle. Note that the
  # first point is the same as the last point to make this a closed curve when
  # I create the vtkPolyLine.
  points = vtk.vtkPoints()
  for i in range(4):
    points.InsertNextPoint(cos(2.*pi*i/3), sin(2*pi*i/3.), 0.)

  outline_pd  = as_polyline(points, LEVEL)
  # You have already gone through the trouble of putting the points in the
  # right places - so "all" you need todo now is to create polygons from the
  # points that are in the vtkPoints.

  # The points that are passed in, have an overlap of the beginning and the
  # end. For this next trick, I will need a list of the indices in the
  # vtkPoints. They're consecutive, so thats pretty straightforward.

  indices = [i for i in range(outline_pd.GetPoints().GetNumberOfPoints() + 1)]
  triangles   = vtk.vtkCellArray()

  # Set this up for each of the initial sides, then call the recursive function.
  stride = (len(indices) - 1)//3

  # The cell data will allow us to color the triangles based on the level of
  # the iteration of the Koch snowflake.
  data = vtk.vtkIntArray()
  data.SetNumberOfComponents(0)
  data.SetName("Iteration Level")

  # This is the starting triangle.
  t = vtk.vtkTriangle()
  t.GetPointIds().SetId(0,        0)
  t.GetPointIds().SetId(1,   stride)
  t.GetPointIds().SetId(2, 2*stride)
  triangles.InsertNextCell(t)
  data.InsertNextValue(0)

  as_triangles(indices[       0:   stride + 1], triangles, 1, data)
  as_triangles(indices[  stride: 2*stride + 1], triangles, 1, data)
  as_triangles(indices[2*stride:           -1], triangles, 1, data)

  triangle_pd = vtk.vtkPolyData()
  triangle_pd.SetPoints(outline_pd.GetPoints())
  triangle_pd.SetPolys(triangles)
  triangle_pd.GetCellData().SetScalars(data)

  #-----------------#
  # rendering stuff #
  #-----------------#
  outline_mapper = vtk.vtkPolyDataMapper()
  outline_mapper.SetInputData(outline_pd)

  lut = vtk.vtkLookupTable()
  lut.SetNumberOfTableValues(256)
  lut.SetHueRange(0.6, 0.6)
  lut.SetSaturationRange(0., 1.)
  lut.Build()

  triangle_mapper = vtk.vtkPolyDataMapper()
  triangle_mapper.SetInputData(triangle_pd)
  triangle_mapper.SetScalarRange(0.0, LEVEL)
  triangle_mapper.SetLookupTable(lut)

  outline_actor = vtk.vtkActor()
  outline_actor.SetMapper(outline_mapper)

  triangle_actor = vtk.vtkActor()
  triangle_actor.SetMapper(triangle_mapper)

  outline_ren = vtk.vtkRenderer()
  outline_ren.AddActor(outline_actor)
  outline_ren.SetViewport(0.0, 0.0, 0.5, 1.0)

  triangle_ren = vtk.vtkRenderer()
  triangle_ren.AddActor(triangle_actor)
  triangle_ren.SetViewport(0.5, 0.0, 1.0, 1.0)
  triangle_ren.SetActiveCamera(outline_ren.GetActiveCamera())

  renw = vtk.vtkRenderWindow()
  renw.AddRenderer(outline_ren)
  renw.AddRenderer(triangle_ren)
  renw.SetSize(800, 400)

  iren = vtk.vtkRenderWindowInteractor()
  iren.SetRenderWindow(renw)

  outline_ren.ResetCamera()
  renw.Render()
  iren.Start()
