# coding: utf-8
#
# Copyright © 2021 MSC.Software Corporation
#  
# Permission to use Scripts
#  
# Permission is granted to MSC.Software customers with valid licenses of the MSC Apex™ software 
# to use script files posted on this page (collectively “Scripts”) solely for customers’ own authorized 
# internal use in conjunction with licensed use of the MSC Apex software.  
# No other use is permitted.  No license to MSC Apex or other software is granted under this Permission.
#  
# SCRIPTS ARE PROVIDED ON AN "AS-IS" BASIS, WITH NO WARRANTIES OF ANY KIND.  USE OF THE SCRIPTS IS AT CUSTOMER'S SOLE RISK.
#  
# IN NO EVENT SHALL MSC.SOFTWARE CORPORATION, ITS AFFILIATES, SUBSIDIARIES, SUPPLIERS, AGENTS, REPRESENTATIVES OR EMPLOYEES BE LIABLE FOR ANY 
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL OR CONSEQUENTIAL DAMAGES HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, TORT, 
# STRICT LIABILITY, OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION USE OF THE SCRIPTS, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.  
#  
# Customer shall indemnify, defend and hold MSC.Software Corporation, its affiliates, subsidiaries, suppliers, agents, 
# employees or representatives (collectively, “Indemnified Parties”) harmless from and against any third-party claims 
# brought against Indemnified Parties rising out of or in connection with Customer’s use of the Scripts.
#
import os
# Timer utilities
from timeit import default_timer as timer
import csv
import lxml.etree as ET
import pandas as pd
import numpy as np
import vtk
from vtk.util import numpy_support

from importlib import reload
# import pydevd_pycharm as rd
# rd = reload(rd)
# rd.settrace('localhost', port=21000, stdoutToServer=True, stderrToServer=True)

# Class to display a VTKMultiBlockDataSet in a render window
class CradleFieldReader():
    def __init__(self):
        self.field_filename = None
        self.poly = None
        self.scalar_field_names = {}

    def set_file_name(self, filename):
        self.field_filename = filename

    def get_output(self):
        return self.poly, self.scalar_field_names

    def update(self):
        filename, file_extension = os.path.splitext(self.field_filename)

        if file_extension == '.xml' or file_extension == '.XML':
            point_coordinates, num_polygons, polygon_topologies, scalar_fields = self.__read_xml()
            self.__create_poly(point_coordinates, num_polygons, polygon_topologies, scalar_fields)
        elif file_extension == '.csv' or file_extension == '.txt':
            point_coordinates, scalar_fields = self.__read_delimited()
            self.__create_poly(point_coordinates, None, None, scalar_fields)
        elif file_extension == '.vtp':
            self.__create_poly(None, None, None, None)
            return

        for fieldname in scalar_fields.keys():
            self.scalar_field_names[fieldname] = 'Point'

    def __preview_xml(self):

        dict={}
        dict['temp'] = False
        dict['pres'] = False
        dict['locs'] = True

        with open(self.field_filename,'r',encoding='utf-8') as xmlfile:
            head = [next(xmlfile) for x in range(10)]
        for h in head:
            if 'PRES' in h:
                dict['pres'] = True
            elif 'TEMP' in h:
                dict['temp'] = True

        return dict

    def __preview_vtp(self):

        dict={}
        dict['temp'] = False
        dict['pres'] = False
        dict['locs'] = True

        reader = vtk.vtkXMLPolyDataReader()
        reader.SetFileName(self.field_filename)
        reader.Update()
        vtk_poly = reader.GetOutput()
        if vtk_poly.GetCellData().GetArray('Pressure'):
            dict['pres'] = True
        elif vtk_poly.GetPointData().GetArray('Pressure'):
            dict['pres'] = True
        elif vtk_poly.GetPointData().GetArray('Temperature'):
            dict['temp'] = True
        return dict

    def __read_xml(self):

        tree = ET.parse(self.field_filename)

        available_fields = []
        for comment in tree.xpath('(//variable_output/comment())'):
            if 'PRES' in str(comment):
                available_fields.append('PRES')
            elif 'TEMP' in str(comment):
                available_fields.append('TEMP')

        position_elements = tree.xpath('(//Position)')

        point_coordinates = np.empty((len(position_elements),3))
        i=0
        for position_element in position_elements:
            point_coordinates[i,0] = position_element.attrib.get('x')
            point_coordinates[i,1] = position_element.attrib.get('y')
            point_coordinates[i,2] = position_element.attrib.get('z')
            i+=1

        polygon_elements = tree.xpath('(//POLYGON)')

        polygon_vertex_counts = np.empty(len(polygon_elements), dtype=np.int32)
        i=0
        for polygon_element in polygon_elements:
            polygon_vertex_counts[i] = int(polygon_element.xpath('count(Vertex)'))
            i+=1

        scalar_elements = tree.xpath('(//Scalar)')

        scalar_fields = {}
        for field in available_fields:
            field_array = np.empty(int(len(scalar_elements)))
            i=0
            for scalar_element in scalar_elements:
                field_array[i] = scalar_element.attrib.get(field)
                i+=1
            if field == 'PRES':
                scalar_fields['Pressure'] = field_array
            elif field == 'TEMP':
                scalar_fields['Temperature'] = field_array

        polygon_topologies = np.arange(np.sum(polygon_vertex_counts))
        insertion_indices = np.cumsum(polygon_vertex_counts)
        insertion_indices = np.insert(insertion_indices, 0, 0)[:-1]
        polygon_topologies = np.insert(polygon_topologies,
                                       insertion_indices,
                                       polygon_vertex_counts,
                                       axis=0)

        return point_coordinates, len(polygon_elements), polygon_topologies, scalar_fields

    def __create_poly(self, point_coordinates, num_polys, poly_topologies, scalar_fields):

        if point_coordinates is None and scalar_fields is None:
            reader = vtk.vtkXMLPolyDataReader()
            reader.SetFileName(self.field_filename)
            reader.Update()
            vtk_poly=reader.GetOutput()
            self.poly=vtk_poly
            if vtk_poly.GetCellData().GetArray('Pressure'):
                c2p = vtk.vtkCellDataToPointData()
                c2p.PassCellDataOn()
                c2p.SetInputDataObject(vtk_poly)
                c2p.Update()
                self.poly = c2p.GetOutput()
                self.scalar_field_names['Pressure'] = 'Point'
            elif vtk_poly.GetPointData().GetArray('Pressure'):
                self.scalar_field_names['Pressure'] = 'Point'
            elif vtk_poly.GetPointData().GetArray('Temperature'):
                self.scalar_field_names['Temperature'] = 'Point'
            return

        vtk_poly=vtk.vtkPolyData()
        vtk_points = vtk.vtkPoints()
        vtk_points.SetData(numpy_support.numpy_to_vtk(point_coordinates, deep=False, array_type=vtk.VTK_FLOAT))

        # Add the Points to the vtkPolyData
        vtk_poly.SetPoints(vtk_points)

        for field_name in scalar_fields.keys():
        # Add the original Node ID's to the vtkPoints as a scalar array
            vtk_field_values = numpy_support.numpy_to_vtk(scalar_fields[field_name], deep=False, array_type=vtk.VTK_FLOAT)
            vtk_field_values.SetNumberOfComponents(1)
            vtk_field_values.SetName(field_name)
            vtk_poly.GetPointData().AddArray(vtk_field_values)


        if num_polys is not None:
            vtk_polygons = vtk.vtkCellArray()
            # Add the cell definitions to ths vtkCell data structure
            vtk_polygons.SetCells(num_polys,
                                  numpy_support.numpy_to_vtk(poly_topologies,
                                                             deep=False,
                                                             array_type=vtk.vtkIdTypeArray().GetDataType()))
            vtk_poly.SetPolys(vtk_polygons)

        self.poly = vtk_poly


    def __read_delimited(self):

        # Gather up inputs for the pandas read
        has_header = False
        sample_bytes = 2048
        dialect = None
        with open(self.field_filename, 'r') as csvfile:
            sniffer = csv.Sniffer()
            has_header = sniffer.has_header(csvfile.read(sample_bytes))
            csvfile.seek(0)
            dialect = sniffer.sniff(csvfile.read(sample_bytes))
            csvfile.close()

        if has_header:
            has_header_pd = 0
        else:
            has_header_pd = None

        # Read the file
        field = pd.read_csv(self.field_filename, sep=dialect.delimiter, header=has_header_pd)

        # Figure out what columns we have
        x_index = None
        y_index = None
        z_index = None
        press_index = None
        temp_index = None
        field_type = None

        num_columns = len(field.columns)
        if has_header and len(field.columns) >= 4:
        # Simple test for Cradle .csv files
            if field.columns[0] == 'NO.':
                for idx, header in enumerate(field.columns):
                    if ('X' == header):
                        x_index = idx
                    elif ('Y' == header):
                        y_index = idx
                    elif ('Z' == header):
                        z_index = idx
                    elif ('PRES' == header):
                        press_index = idx
                        field_type = 'Pressure'
                    elif ('TEMP' == header):
                        temp_index = idx
                        field_type = 'Temperature'
            else:
                for idx, header in enumerate(field.columns):
                    if 'x' in str(header).lower():
                        x_index = idx
                    elif 'y' in str(header).lower():
                        y_index = idx
                    elif 'z' in str(header).lower():
                        z_index = idx
                    elif 'pr' in str(header).lower():
                        press_index = idx
                    elif 'te' in str(header).lower():
                        temp_index = idx
        elif len(field.columns) >= 4:
            x_index = 0
            y_index = 1
            z_index = 2
            field_type = 'Temperature'
            temp_index = 3

        # Check if we have minimum viable data
        # Coordinates and pressures or temperature
        if None in [x_index, y_index, z_index] or (press_index is None and temp_index is None):
            return

        x_coords = field.iloc[:, x_index].values
        y_coords = field.iloc[:, y_index].values
        z_coords = field.iloc[:, z_index].values

        point_coordinates = np.stack([x_coords, y_coords, z_coords], axis=1)

        scalar_fields = {}

        if press_index is not None:
            pressures = field.iloc[:, press_index].values
            scalar_fields['Pressure'] = pressures

        if temp_index is not None:
            temperatures = field.iloc[:, temp_index].values
            scalar_fields['Temperature'] = temperatures

        return point_coordinates, scalar_fields


    def preview_file(self):
        filename, file_extension = os.path.splitext(self.field_filename)

        dict={}
        dict['filename']=self.field_filename
        dict['method']='preview_file'
        dict['file_extension'] = file_extension
        if file_extension == '.xml' or file_extension == '.XML':
            preview_data = self.__preview_xml()
        elif file_extension == '.vtp':
            preview_data = self.__preview_vtp()
        elif file_extension == '.csv' or file_extension == '.txt':
            preview_data = self.__preview_delimited()
        dict['preview_data']=preview_data
        return dict


    def __preview_delimited(self):
        dict={}

        # Gather up inputs for the pandas read
        has_header = False
        sample_bytes = 2048
        dialect = None
        headers = None
        with open(self.field_filename, 'r') as csvfile:
            sniffer = csv.Sniffer()
            has_header = sniffer.has_header(csvfile.read(sample_bytes))
            csvfile.seek(0)
            dialect = sniffer.sniff(csvfile.read(sample_bytes))
            csvfile.seek(0)
            reader = csv.reader(csvfile)
            if has_header:
                headers = reader.__next__()
            else:
                first_row = reader.__next__()
            csvfile.close()

        x_index = None
        y_index = None
        z_index = None
        temp_index = None
        press_index = None

        if has_header:
            index = 0
            for header in headers:
                if 'x' in str(header).lower():
                    x_index = index
                elif 'y' in str(header).lower():
                    y_index = index
                elif 'z' in str(header).lower():
                    z_index = index
                elif 'pr' in str(header).lower():
                    press_index = index
                elif 'te' in str(header).lower():
                    temp_index = index

            dict['headers']=headers
            dict['locs'] = True
            dict['temp'] = False
            dict['pres'] = False

            if None in [x_index, y_index, z_index]:
                dict['locs']=False

            if temp_index is not None:
                dict['temp'] = True


            if press_index is not None:
                dict['pres'] = True

        else:
            dict['headers']=first_row

        return dict