# author: Zhiyuan Liu
# date: 2018.6.1 Happy Children's Day
# this file transform legacy srep to new srep not only about the file format, but divide crest spoke and inner spokes
from __future__ import print_function
from xml.etree.ElementTree import Element, SubElement
from ElementTree_pretty import prettify
# import xml.etree.cElementTree as etree
import vtk
import srep
import os
import sys

if len(sys.argv) < 3:
    print('Usage: ' + sys.argv[0] + '[input legacy m3d file name] [output prefix] [optional: epsilon]')
    exit(-1)

epsilon = 0.02 # movement of crest spoke away from its nearby interior point
if len(sys.argv) == 4:
    epsilon = float(sys.argv[3])

m3d_filename = sys.argv[1]
outPrefix = sys.argv[2]
outputUp = outPrefix + '/up.vtp'
outputDown = outPrefix + '/down.vtp'
outputCrest = outPrefix + '/crest.vtp'
"""
The purpose of this part of code is to play around m3d to fit into the new s-rep format framework
"""

up_spokes_polydata = vtk.vtkPolyData()
down_spokes_polydata = vtk.vtkPolyData()
crest_spokes_polydata = vtk.vtkPolyData()

# read in m3d file
s = srep.srep()
s.readSrepFromM3D(m3d_filename)

writer = vtk.vtkXMLPolyDataWriter()
writer.SetDataModeToAscii()

"""
Given an s-rep,
1. read in medial points
2. set spoke vectors as point data
3. set spoke vector length as point data as well
4. create a header that contains meta data (e.g., number of rows, number of columns, mesh type)
5. create xml header
"""
nRows = s.fig.numRows
nCols = s.fig.numCols
# create
root = Element('s-rep')
nRowsXMLElement = SubElement(root, 'nRows')
nRowsXMLElement.text = str(nRows)

nColsXMLElement = SubElement(root, 'nCols')
nColsXMLElement.text = str(nCols)

meshTypeXMLElement = SubElement(root, 'meshType')
meshTypeXMLElement.text = 'Quad'

colorXMLElement = SubElement(root, 'color')

redXMLElement = SubElement(colorXMLElement, 'red')
redXMLElement.text = str(0)

greenXMLElement = SubElement(colorXMLElement, 'green')
greenXMLElement.text = str(0.5)

blueXMLElement = SubElement(colorXMLElement, 'blue')
blueXMLElement.text = str(0)

isMeanFlagXMLElement = SubElement(root, 'isMean')
isMeanFlagXMLElement.text = 'False'

meanStatPathXMLElement = SubElement(root, 'meanStatPath')
meanStatPathXMLElement.text = ''
# later this can be done as a parameter
upSpokeXMLElement = SubElement(root, 'upSpoke')
upSpokeXMLElement.text = os.path.join(outputUp)

downSpokeXMLElement = SubElement(root, 'downSpoke')
downSpokeXMLElement.text = os.path.join(outputDown)

crestSpokeXMLElement = SubElement(root, 'crestSpoke')
crestSpokeXMLElement.text = os.path.join(outputCrest)

file_handle = open(outPrefix + '/header.xml','w')
file_handle.write(prettify(root))
file_handle.close()

# medial points and polygons first
medial_points = vtk.vtkPoints()
medial_points.SetDataTypeToDouble() # important, this fix the bug that new srep has different precision with legacy one, you can find it by runningapplyTps2NewSrep program
medial_polys = vtk.vtkCellArray()

# This will be curves
crest_points = vtk.vtkPoints()
crest_polys = vtk.vtkCellArray()
#
up_spoke_directions = vtk.vtkDoubleArray()
up_spoke_directions.SetNumberOfComponents(3)
up_spoke_directions.SetName("spokeDirection")

up_spoke_lengths = vtk.vtkDoubleArray()
up_spoke_lengths.SetNumberOfComponents(1)
up_spoke_lengths.SetName("spokeLength")

down_spoke_directions = vtk.vtkDoubleArray()
down_spoke_directions.SetNumberOfComponents(3)
down_spoke_directions.SetName("spokeDirection")

down_spoke_lengths = vtk.vtkDoubleArray()
down_spoke_lengths.SetNumberOfComponents(1)
down_spoke_lengths.SetName("spokeLength")

crest_spoke_directions = vtk.vtkDoubleArray()
crest_spoke_directions.SetNumberOfComponents(3)
crest_spoke_directions.SetName("spokeDirection")

crest_spoke_lengths = vtk.vtkDoubleArray()
crest_spoke_lengths.SetNumberOfComponents(1)
crest_spoke_lengths.SetName("spokeLength")

"""
Read in
  1. medial points,
  2. up spokeLength, up spokeDirection
  3. down spokeLength, down spokeDirection
"""
for r in range(nRows):
    for c in range(nCols):
        current_atom = s.fig.atoms[r, c]
        current_point = current_atom.hub.P
        current_id = medial_points.InsertNextPoint(current_point)

        if r < nRows - 1 and c < nCols - 1:
            quad = vtk.vtkQuad()
            quad.GetPointIds().SetId(0, current_id)
            quad.GetPointIds().SetId(1, current_id + nCols)
            quad.GetPointIds().SetId(2, current_id + nCols + 1)
            quad.GetPointIds().SetId(3, current_id + 1)
            medial_polys.InsertNextCell(quad)

        current_up_spoke = current_atom.topSpoke
        current_up_spoke_direction = current_up_spoke.U
        up_spoke_directions.InsertNextTuple(current_up_spoke_direction)
        current_up_spoke_length = current_up_spoke.r
        up_spoke_lengths.InsertNextTuple1(current_up_spoke_length)

        current_down_spoke = current_atom.botSpoke
        current_down_spoke_direction = current_down_spoke.U
        down_spoke_directions.InsertNextTuple(current_down_spoke_direction)
        current_down_spoke_length = current_down_spoke.r
        down_spoke_lengths.InsertNextTuple1(current_down_spoke_length)

"""
Construct crest vtp
  1. read points, direction, and length in clockwise
  2. write it as a curve  
"""
# \TODO: There is a lot repeating of the same code. Need to think if I can abstract the repetition into a function
# first row
r = 0
for c in range(nCols - 1):
    current_crest_atom = s.fig.atoms[r, c]
    current_crest_spoke = current_crest_atom.crestSpoke
    current_crest_spoke_direction = current_crest_spoke.U

    current_crest_point = current_crest_atom.hub.P + epsilon * current_crest_spoke_direction
    current_crest_point_id = crest_points.InsertNextPoint(current_crest_point)
    line = vtk.vtkLine()
    if current_crest_point_id > 0:
        line = vtk.vtkLine()
        line.GetPointIds().SetId(0, current_crest_point_id - 1)
        line.GetPointIds().SetId(1, current_crest_point_id)
        crest_polys.InsertNextCell(line)
    crest_spoke_directions.InsertNextTuple(current_crest_spoke_direction)
    current_crest_spoke_length = current_crest_spoke.r
    crest_spoke_lengths.InsertNextTuple1(current_crest_spoke_length)

# last column
c = nCols - 1
for r in range(nRows - 1):
    current_crest_atom = s.fig.atoms[r, c]
    current_crest_spoke = current_crest_atom.crestSpoke
    current_crest_spoke_direction = current_crest_spoke.U

    current_crest_point = current_crest_atom.hub.P + epsilon * current_crest_spoke_direction
    current_crest_point_id = crest_points.InsertNextPoint(current_crest_point)
    line = vtk.vtkLine()
    if current_crest_point_id > 0:
        line = vtk.vtkLine()
        line.GetPointIds().SetId(0, current_crest_point_id - 1)
        line.GetPointIds().SetId(1, current_crest_point_id)
        crest_polys.InsertNextCell(line)
    crest_spoke_directions.InsertNextTuple(current_crest_spoke_direction)
    current_crest_spoke_length = current_crest_spoke.r
    crest_spoke_lengths.InsertNextTuple1(current_crest_spoke_length)

# bottom row
r = nRows - 1
for c in range(nCols-1, 0, -1):
    current_crest_atom = s.fig.atoms[r, c]
    current_crest_spoke = current_crest_atom.crestSpoke
    current_crest_spoke_direction = current_crest_spoke.U

    current_crest_point = current_crest_atom.hub.P + epsilon * current_crest_spoke_direction
    current_crest_point_id = crest_points.InsertNextPoint(current_crest_point)
    line = vtk.vtkLine()
    if current_crest_point_id > 0:
        line = vtk.vtkLine()
        line.GetPointIds().SetId(0, current_crest_point_id - 1)
        line.GetPointIds().SetId(1, current_crest_point_id)
        crest_polys.InsertNextCell(line)
    crest_spoke_directions.InsertNextTuple(current_crest_spoke_direction)
    current_crest_spoke_length = current_crest_spoke.r
    crest_spoke_lengths.InsertNextTuple1(current_crest_spoke_length)

# first column
c = 0
for r in range(nRows-1, 0, -1):
    current_crest_atom = s.fig.atoms[r, c]
    current_crest_spoke = current_crest_atom.crestSpoke
    current_crest_spoke_direction = current_crest_spoke.U

    current_crest_point = current_crest_atom.hub.P + epsilon * current_crest_spoke_direction
    current_crest_point_id = crest_points.InsertNextPoint(current_crest_point)
    line = vtk.vtkLine()
    if current_crest_point_id > 0:
        line = vtk.vtkLine()
        line.GetPointIds().SetId(0, current_crest_point_id - 1)
        line.GetPointIds().SetId(1, current_crest_point_id)
        crest_polys.InsertNextCell(line)
    crest_spoke_directions.InsertNextTuple(current_crest_spoke_direction)
    current_crest_spoke_length = current_crest_spoke.r
    crest_spoke_lengths.InsertNextTuple1(current_crest_spoke_length)
# to complete loop
line = vtk.vtkLine()
line.GetPointIds().SetId(0, current_crest_point_id)
line.GetPointIds().SetId(1, 0)
crest_polys.InsertNextCell(line)

polyLine = vtk.vtkPolyLine()
polyLine.GetPointIds().SetNumberOfIds(crest_points.GetNumberOfPoints()+1)
for i in range(crest_points.GetNumberOfPoints()):
    # polyLine->GetPointIds()->SetId(i, i);
    polyLine.GetPointIds().SetId(i, i)

polyLine.GetPointIds().SetId(crest_points.GetNumberOfPoints(),0)
cells = vtk.vtkCellArray()
cells.InsertNextCell(polyLine)

crest_spokes_polydata.SetPoints(crest_points)
crest_spokes_polydata.SetLines(cells)
crest_spokes_polydata.GetPointData().AddArray(crest_spoke_directions)
crest_spokes_polydata.GetPointData().SetActiveVectors("spokeDirection")
crest_spokes_polydata.GetPointData().AddArray(crest_spoke_lengths)
crest_spokes_polydata.GetPointData().SetActiveScalars("spokeLength")
writer.SetFileName(outputCrest)
writer.SetInputData(crest_spokes_polydata)
writer.Update()

up_spokes_polydata.SetPoints(medial_points)
up_spokes_polydata.SetPolys(medial_polys)
up_spokes_polydata.GetPointData().AddArray(up_spoke_directions)
up_spokes_polydata.GetPointData().SetActiveVectors("spokeDirection")
up_spokes_polydata.GetPointData().AddArray(up_spoke_lengths)
up_spokes_polydata.GetPointData().SetActiveScalars("spokeLength")
writer.SetFileName(outputUp)
writer.SetInputData(up_spokes_polydata)
writer.Update()

down_spokes_polydata.SetPoints(medial_points)
down_spokes_polydata.SetPolys(medial_polys)
down_spokes_polydata.GetPointData().AddArray(down_spoke_directions)
down_spokes_polydata.GetPointData().SetActiveVectors("spokeDirection")
down_spokes_polydata.GetPointData().AddArray(down_spoke_lengths)
down_spokes_polydata.GetPointData().SetActiveScalars("spokeLength")
writer.SetFileName(outputDown)
writer.SetInputData(down_spokes_polydata)
writer.Update()


print('Finished transformation and save to files:')
print(outputUp)
print(outputDown)
print(outputCrest)