"""
Script to construct vtp file from MinimalFEM files
"""

import os
import sys

import vtk

class VTPBuilder:
    def __init__(self):
        self.fem_inp_file = None
        self.fem_out_file = None
        self.polydata = None

    def build_polydata(self, fem_inp_file, fem_out_file):
        self.fem_inp_file = fem_inp_file
        self.fem_out_file = fem_out_file

        self.polydata = vtk.vtkPolyData()
        self._build_geometry()
        self._build_attributes()

    def write(self, filename):
        writer = vtk.vtkXMLPolyDataWriter()
        writer.SetFileName(filename)
        writer.SetDataModeToAscii()
        writer.SetInputData(self.polydata)
        ret_code = writer.Write()
        print('Wrote', filename, ', return code', ret_code)
        return ret_code

    def _build_geometry(self):
        """Parse input file to create polydata geometry"""
        points = vtk.vtkPoints()
        polys = vtk.vtkCellArray()

        completed = False
        with open(self.fem_inp_file) as inp:
            # Read points
            inp.readline()
            num_points = int(inp.readline())
            print('num_points', num_points)
            points.SetNumberOfPoints(num_points)
            z = 0
            for i in range(num_points):
                line = inp.readline()
                coords = line.split(' ')
                x = float(coords[0])
                y = float(coords[1])
                points.SetPoint(i, x, y, z)

            # Read faces
            num_faces = int(inp.readline())
            print('num_faces', num_faces)
            for i in range(num_faces):
                line = inp.readline()
                split = line.split(' ')
                ids = [int(s) for s in split]
                polys.InsertNextCell(self._makeVtkIdList(ids))

            self.polydata.SetPoints(points)
            self.polydata.SetPolys(polys)
            # print(self.polydata)

            completed = True

        return completed

    def _build_attributes(self):
        """Parse output file to create displacement and stress attrributes"""
        disp_array = vtk.vtkFloatArray()
        stress_array = vtk.vtkFloatArray()

        completed = False
        with open(self.fem_out_file) as out:
            # Read displacement (point data)
            num_points = self.polydata.GetNumberOfPoints()
            disp_array.SetName('displacement')
            disp_array.SetNumberOfComponents(3)
            disp_array.SetNumberOfTuples(num_points)
            for i in range(num_points):
                line = out.readline()
                x = float(line)
                line = out.readline()
                y = float(line)
                disp_array.InsertTuple3(i, x, y, 0)
            self.polydata.GetPointData().SetVectors(disp_array)

            # Read stress (cell data)
            num_faces = self.polydata.GetNumberOfPolys()
            stress_array.SetName('stress')
            stress_array.SetNumberOfComponents(1)
            stress_array.SetNumberOfTuples(num_faces)
            # for i in range(num_points):
            #     line = out.readline()
            #     stress = float(line)
            #     stress_array.InsertTuple1(i, stress)
            # self.polydata.GetPointData().SetScalars(stress_array)
            for i in range(num_faces):
                line = out.readline()
                stress = float(line)
                stress_array.InsertTuple1(i, stress)
            self.polydata.GetCellData().SetScalars(stress_array)

            completed = True

        return completed


    def _makeVtkIdList(self, it):
        """
        Makes a vtkIdList from a Python iterable.

        :param it: A python iterable.
        :return: A vtkIdList
        """
        vil = vtk.vtkIdList()
        for i in it:
            vil.InsertNextId(int(i))
        return vil

if __name__ == '__main__':
    if len(sys.argv) < 3:
        print('Usage: convertorVTP.py  fem.inp  fem.out  [fem.vtp]')
        sys.exit(-1)

    builder = VTPBuilder()
    builder.build_polydata(sys.argv[1], sys.argv[2])

    if len(sys.argv) > 3:
        out_filename = sys.argv[3]
    else:
        out_filename = 'fem.vtp'
    builder.write(out_filename)

    print('finis')
