# 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,time
import sys
from importlib import reload
# Timer utilities
from timeit import default_timer as timer
import uuid
import numpy as np
import vtk
from vtk.util import numpy_support

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

# this function will return the path to the Common folder that contains common functions used
# across tools it also supports the use of an environment variable in case the folder is not in an
# intersecting path

def getCommonFolder():
    try:
        env_test_path = os.environ["CUSTOM_TOOLS_COMMON_FOLDER"]
        if os.path.exists(env_test_path):
            return env_test_path
    except:
        pass

    path = os.path.dirname(__file__)
    data_path = os.path.join(path, 'Common')
    up = 0
    up_path = '..'
    while (True):
        if os.path.exists(data_path) and up > 2:
            break
        if up == 10:
            return ''
        data_path = os.path.join(path, up_path, 'Common')
        up += 1
        up_path = os.path.join(up_path, '..')
    return data_path

common_path = getCommonFolder()

this_path = os.path.dirname(__file__)
vtk_utils_path = os.path.join(common_path, 'VTK UtilitiesC')
if 'VTK UtilitiesC' not in sys.path:
    sys.path.append(vtk_utils_path)

import apex_to_vtk as a2vtk
a2vtk = reload(a2vtk)

import external_field_reader as crf
crf = reload(crf)

import map_poly_scalar_to_poly_target as mapfield
mapfield = reload(mapfield)

import vtk_display as vtkd
vtkd = reload(vtkd)

import vtk_to_apex as vtk2a
vtk2a = reload(vtk2a)

global current_selection

import apex

def clearDataSet():
    global current_dataset
    current_dataset = {}
    current_dataset['source']={}
    current_dataset['source']['tc'] = time.time()
    current_dataset['source']['fattr'] = None
    current_dataset['source']['filename'] = None
    current_dataset['source']['data'] = {}


    current_dataset['target']={}
    current_dataset['target']['selected_entities']=None
    current_dataset['target']['string'] = None
    current_dataset['target']['data'] = {}

    current_dataset['mapped']={}
    current_dataset['mapped']['option'] = None
    current_dataset['mapped']['tol'] = None
    current_dataset['mapped']['data'] = {}
    return current_dataset

try:
    if len(current_dataset['source']['data'].keys()):
        pass
except:
    current_dataset = clearDataSet()


def reset_current_info(dict={}):
    apex.disableShowOutput()
    current_dataset = clearDataSet()

def get_current_info(dict={}):
    apex.disableShowOutput()
    global current_dataset
    ret_dict = {}
    ret_dict['method'] = 'get_current_info'
    if current_dataset['source']['filename'] != None:
        ret_dict['source_filename']=current_dataset['source']['filename']
    else:
        ret_dict['source_filename']=''
    if current_dataset['target']['string'] != None:
       ret_dict['target'] = current_dataset['target']['string']
    else:
       ret_dict['target']=''
    if current_dataset['mapped']['tol'] != None:
        ret_dict['tol'] = current_dataset['mapped']['tol']
    else:
        ret_dict['tol'] = ''
    try:
        fnames = current_dataset['source']['data']['scalar_field_names']
        fname=list(fnames.keys())[0]
        ret_dict['ftype']=fname
    except:
        ret_dict['ftype'] = ''

    ret_dict['units'] = get_length_pressure_temperature_unit_systems()

    return ret_dict

def display_external_field(dictionary):
    apex.disableShowOutput()
    field_filename = dictionary['field_file_name']
    times = dictionary['times']
    field_name = dictionary['field']
    units_label = dictionary['units']
    units_length = units_label.split('-')[0]
    units_pressure = units_label.split('-')[1]
    units_temperature = units_label.split('-')[2]

    source, scalar_field_names = get_source(field_filename,
                                            units_length = units_length,
                                            units_pressure = units_pressure,
                                            units_temperature =units_temperature)

    if field_name == 'Pressure':
        scalar_bar_units = units_pressure

    elif field_name == 'Temperature':
        scalar_bar_units = units_temperature

    field_topology = scalar_field_names[field_name]

    vtk_mesh_display = vtkd.vtk_display({'units' : scalar_bar_units})
    vtk_mesh_display.display_mesh(source, field_topology, field_name)
    vtk_mesh_display.render()


def handle_target_selection(dictionary):
    global current_selection
    apex.disableShowOutput()
    dict = {}
    selected_entities = apex.selection.getCurrentSelection()
    if not selected_entities: return;
    current_selection = current_dataset['target']['selected_entities']
    if selected_entities.len() == 0:
        if len(current_selection) > 0:
            selected_entities=current_selection
            if selected_entities == current_dataset['target']['selected_entities']:
                pass
            else:
                current_dataset['target']['selection']=None
    else:
        current_dataset['target']['selection']=None
        current_selection = selected_entities

    dict['selected_entities'] = ''

    if selected_entities.len() == 1:
        if isinstance(selected_entities[0], apex.IName):
            dict['selected_entities'] = selected_entities[0].name
        else:
            dict['selected_entities'] = str(type(selected_entities[0]))
    else:
        dict['selected_entities'] = str(selected_entities.len()) + ' Entities'
    current_dataset['target']['string'] = dict['selected_entities']
    return dict

def map_field_to_selected_entities(mapping_inputs):

    # Disable Apex script completion message
    #apex.disableShowOutput()

    #DB4L8R - Return a message to the status bar or to a popup
    if apex.selection.getCurrentSelection().len() == 0:
        return

    apex.setScriptUnitSystem(apex.getApplicationUnitSystemLabel())

    global current_selection

    # Unpack the UI options
    field_source = mapping_inputs['source']
    field_name = mapping_inputs['field']
    source_units = mapping_inputs['units']
    source_units_length = source_units.split('-')[0]
    source_units_pressure = source_units.split('-')[1]
    source_units_temperature = source_units.split('-')[2]

    # Get the mapping data from the UI
    mapping_algorithm = mapping_inputs['algorithm']

    boundary_only = False
    if mapping_inputs['outerOnly'] == 'True':
        boundary_only = True

    preview = True
    if mapping_inputs['method'] == 'Create':
        preview=False
        load_name=mapping_inputs['load_name']

    tolerance = None
    if mapping_inputs['tolerance'] != 'None':
        tolerance = float(mapping_inputs['tolerance'])
    else:
        tolerance=None

    # Get the mapped mesh
    #   Internally gets the source from the cache of generates a new one if necessary
    #   Internally gets the target from the cache or generate a new one if necessary
    #   Will return a cached mapped mesh if one exists and matches the inputs, otherwise it will generate a new one
    mapped_mesh = get_mapped_mesh(field_source, field_name, mapping_algorithm, boundary_only, tolerance, source_units)

    # Preview the mapped mesh or create the load
    if preview:
        if field_name == 'Pressure':
            scalar_bar_units = source_units_pressure
        elif field_name == 'Temperature':
            scalar_bar_units = source_units_temperature

        vtk_display = vtkd.vtk_display({'units': scalar_bar_units})

        if tolerance is not None:
            vtk_display.display_mesh(filter_by_distance(mapped_mesh, tolerance),
                                     'Point',
                                     field_name)
#            export_vtk_file(mapped_mesh, r'C:\Scratch\mapped_after_threshold_before_display')
        else:
           vtk_display.display_mesh(mapped_mesh, 'Point', field_name)
#            export_vtk_file(mapped_mesh, r'C:\Scratch\mapped_before_display')

        vtk_display.render()
    else:
        if tolerance is not None:
            vtk2a.create_apex_load_from_vtk_field(filter_by_distance(mapped_mesh, tolerance),
                                                  field_name,
                                                  tolerance)
        else:
            vtk2a.create_apex_load_from_vtk_field(mapped_mesh, field_name, tolerance)


def get_source(source_filename,
               units_length = None,
               units_pressure = None,
               units_temperature = None):

    # Check if we have a source field in the cache
    source_field = None
    try:
        source_field = current_dataset['source']['data']['poly']
    except:
        source_field = None

    # If we have source in the cache, validate that it matches the source settings from the UI
    # If it does, return the cached sourced
    # If not, generate a new one by reading the source file
    # DB3L8R - This code is messy - clean it up
    if source_field is not None:
        cached_source_filename = source_field.GetFieldData().GetAbstractArray('source_filename').GetValue(0)
        cached_scalar_fields_array = source_field.GetFieldData().GetAbstractArray('scalar_fields')
        cached_scalar_fields = {}
        for i in range(0, cached_scalar_fields_array.GetNumberOfValues(), 2):
            field_name = cached_scalar_fields_array.GetValue(i)
            field_location = cached_scalar_fields_array.GetValue(i+1)
            cached_scalar_fields[field_name] = field_location
        cached_original_units_length = source_field.GetFieldData().GetAbstractArray('original_units_length').GetValue(0)
        cached_units_pressure = source_field.GetFieldData().GetAbstractArray('units_pressure').GetValue(0)
        cached_units_temperature = source_field.GetFieldData().GetAbstractArray('units_temperature').GetValue(0)
        if cached_source_filename == source_filename:
            if cached_original_units_length != units_length:
                cached_current_units_length = source_field.GetFieldData().GetAbstractArray('current_units_length').GetValue(0)
                if cached_current_units_length == cached_original_units_length:
                    source_field.GetFieldData().GetAbstractArray('units_length').SetValue(0, units_length)
                else:
#                    print('Generated new source')
                    return read_source_from_file(source_filename,
                                                 units_length,
                                                 units_pressure,
                                                 units_temperature)
            if cached_units_pressure != units_pressure:
                source_field.GetFieldData().GetAbstractArray('units_pressure').SetValue(0, units_pressure)
            if cached_units_temperature != units_temperature:
                source_field.GetFieldData().GetAbstractArray('units_temperature').SetValue(0, units_temperature)
#            print('Using cached source')
            return source_field, cached_scalar_fields
        else:
#            print('Generated new source')
            return read_source_from_file(source_filename,
                                         units_length,
                                         units_pressure,
                                         units_temperature)
    else:
#        print('Generated new source')
        return read_source_from_file(source_filename,
                                     units_length,
                                     units_pressure,
                                     units_temperature)

    # If we don't have a field in te cache or if the cached version is from a different file read from the input file
def read_source_from_file(source_filename,
                          units_length=None,
                          units_pressure=None,
                          units_temperature=None):

    reader = crf.CradleFieldReader()
    reader.set_file_name(source_filename)
    reader.update()
    source_field, scalar_field_names = reader.get_output()

    if source_filename is not None and source_field is not None:  # Add the unit information to the poly
        filename_vtk = vtk.vtkStringArray()
        filename_vtk.SetNumberOfValues(1)
        filename_vtk.SetName('source_filename')
        filename_vtk.SetValue(0, source_filename)
        source_field.GetFieldData().AddArray(filename_vtk)

        num_scalar_fields = len(scalar_field_names.keys())

        scalar_fields = vtk.vtkStringArray()
        scalar_fields.SetNumberOfValues(num_scalar_fields * 2)
        scalar_fields.SetName('scalar_fields')
        i = 0
        for scalar_field_name, scalar_field_location in scalar_field_names.items():
            scalar_fields.SetValue(i, scalar_field_name)
            scalar_fields.SetValue(i+1, scalar_field_location)
            i+=2
        source_field.GetFieldData().AddArray(scalar_fields)

        # Add the unit information to the poly
        if units_length is not None:  # Add the unit information to the poly
            original_length_units = vtk.vtkStringArray()
            original_length_units.SetNumberOfValues(1)
            original_length_units.SetName('original_units_length')
            original_length_units.SetValue(0, units_length)
            source_field.GetFieldData().AddArray(original_length_units)

            current_length_units = vtk.vtkStringArray()
            current_length_units.SetNumberOfValues(1)
            current_length_units.SetName('current_units_length')
            current_length_units.SetValue(0, units_length)
            source_field.GetFieldData().AddArray(current_length_units)

        if units_pressure is not None:  # Add the unit information to the poly
            pressure_units = vtk.vtkStringArray()
            pressure_units.SetNumberOfValues(1)
            pressure_units.SetName('units_pressure')
            pressure_units.SetValue(0, units_pressure)
            source_field.GetFieldData().AddArray(pressure_units)

        if units_temperature is not None:  # Add the unit information to the poly
            temperature_units = vtk.vtkStringArray()
            temperature_units.SetNumberOfValues(1)
            temperature_units.SetName('units_temperature')
            temperature_units.SetValue(0, units_temperature)
            source_field.GetFieldData().AddArray(temperature_units)

            source_guid = uuid.uuid4()
            guid_vtk = vtk.vtkStringArray()
            guid_vtk.SetNumberOfValues(1)
            guid_vtk.SetName('guid')
            guid_vtk.SetValue(0, str(source_guid))
            source_field.GetFieldData().AddArray(guid_vtk)

        global current_dataset
        current_dataset['source']['data']['poly'] = source_field
        current_dataset['source']['filename'] = source_filename

    return source_field, scalar_field_names

def get_target(selected_entities):

    target_mesh = None
    # Check if we have a target mesh in the cache
    try:
        target_mesh = current_dataset['target']['dataset']
    except:
        target_mesh = None
    message = None
    # if we have a target mesh in the cache, validate that it matches the UI settings
    if target_mesh is not None:
        # Check the units
        target_mesh_units = target_mesh.GetFieldData().GetAbstractArray('unit_system_label').GetValue(0)
        if target_mesh_units == apex.getApplicationUnitSystemLabel():
            # Check the selected entities
            if current_dataset['target'].get('selected_entities', None) is not None:
                cached_selections = current_dataset['target']['selected_entities']
                if len(cached_selections) == len(selected_entities):
                    cached_names = [entity.name for entity in cached_selections]
                    current_names = [entity.name for entity in selected_entities]
                    selection_changed = False
                    for cached_name, selected_name in zip(cached_names, current_names):
                        if cached_name !=  selected_name:
                            target_mesh = None
                            message = "Using cached target"
                            break

    # If we do not have a target in the cache or if the cached version does not match the current inputs
    # Create a new one form the selected entities
    message = "Using cached target"
    if target_mesh is None:
        message = 'Generated new target'
        target_mesh = get_target_from_apex(selected_entities)

#    print(message)

    return target_mesh

def get_target_from_apex(selected_entities):
    # If we need a new target, fetch from Apex using the selected entities
    target_mesh = a2vtk.create_vtk_dataset_from_apex_entities(selected_entities)

    # Add a guid
    target_guid = uuid.uuid4()
    guid_vtk = vtk.vtkStringArray()
    guid_vtk.SetNumberOfValues(1)
    guid_vtk.SetName('guid')
    guid_vtk.SetValue(0, str(target_guid))
    target_mesh.GetFieldData().AddArray(guid_vtk)

    # Add the target and the selected entities to the cache
    current_dataset['target']['dataset'] = target_mesh
    current_dataset['target']['selected_entities'] = selected_entities

    return target_mesh

def get_mapped_mesh(field_source,
                    field_name,
                    mapping_algorithm,
                    boundary_only,
                    tolerance,
                    source_units):

    units_length = source_units.split('-')[0]
    units_pressure = source_units.split('-')[1]
    units_temperature = source_units.split('-')[2]

    global current_dataset

    mapped_mesh = None
    # Check if we have a mapped mesh in the cache
    try:
        mapped_mesh = current_dataset['mapped']['dataset']
    except:
        mapped_mesh = None

    # Get the source mesh
    source_mesh, scalar_field_names = get_source(field_source,
                                                 units_length=units_length,
                                                 units_pressure=units_pressure,
                                                 units_temperature=units_temperature)

    source_guid_vtk = source_mesh.GetFieldData().GetAbstractArray('guid')
    source_guid = source_guid_vtk.GetValue(0)

#    export_vtk_file(source_mesh, r'C:\Scratch\Source')

    # Get the target mesh
    target_entities = apex.selection.getCurrentSelection()
    target_mesh = None
    if len(target_entities) > 0:
        target_mesh = get_target(target_entities)

    target_guid_vtk = target_mesh.GetFieldData().GetAbstractArray('guid')
    target_guid = target_guid_vtk.GetValue(0)

#    export_vtk_file(target_mesh, r'C:\Scratch\Target')

    # Validate the mapped mesh
    if mapped_mesh is not None:
        cached_field_name = current_dataset['mapped']['field_name']
        cached_mapping_algortihm = current_dataset['mapped']['mapping_algorithm']
        cached_boundary_only = current_dataset['mapped']['boundary_only']
        cached_tolerance_ready = current_dataset['mapped']['tolerance_ready']
        cached_tolerance = current_dataset['mapped']['tolerance']
        cached_source_guid = current_dataset['mapped']['source_guid']
        cached_target_guid = current_dataset['mapped']['target_guid']

        mapped_mesh_is_up_to_date = True
        if cached_source_guid != source_guid: mapped_mesh_is_up_to_date = False
        if cached_target_guid != target_guid: mapped_mesh_is_up_to_date = False
        if cached_field_name != field_name: mapped_mesh_is_up_to_date = False
        if cached_mapping_algortihm != mapping_algorithm: mapped_mesh_is_up_to_date = False
        if cached_boundary_only != boundary_only: mapped_mesh_is_up_to_date = False
        if tolerance is not None and cached_tolerance_ready is None: mapped_mesh_is_up_to_date = False

        if mapped_mesh_is_up_to_date:
#            print('Using cached mapped mesh')
            return mapped_mesh

    mapped_mesh = mapfield.map_scalar_field_to_mesh(field_name,
                                                mapping_algorithm,
                                                source_mesh,
                                                target_mesh,
                                                boundary_only,
                                                tolerance)

    current_dataset['mapped']['dataset'] = mapped_mesh
    current_dataset['mapped']['field_name'] = field_name
    current_dataset['mapped']['mapping_algorithm'] = mapping_algorithm
    current_dataset['mapped']['boundary_only'] = boundary_only
    current_dataset['mapped']['tolerance'] = tolerance
    current_dataset['mapped']['source_guid'] = source_guid
    current_dataset['mapped']['target_guid'] = target_guid
    current_dataset['mapped']['tolerance_ready'] = None
    if tolerance is not None:
        current_dataset['mapped']['tolerance_ready'] = True

    print('Generated new mapped mesh')
    return mapped_mesh


def preview_external_file(dictionary):
    apex.disableShowOutput()
    # reset_current_info()
    field_source = dictionary['filename']
    reader = crf.CradleFieldReader()
    reader.set_file_name(field_source)
    dictionary=reader.preview_file()
    return dictionary

def export_vtk_file(dataset, filename, binary = True):
    writer = None
    if isinstance(dataset, vtk.vtkMultiBlockDataSet):
        writer = vtk.vtkXMLMultiBlockDataWriter()
    elif isinstance(dataset, vtk.vtkUnstructuredGrid):
        writer = vtk.vtkXMLUnstrucuredGridWriter()
    elif isinstance(dataset, vtk.vtkPolyData):
        writer = vtk.vtkXMLPolyDataWriter()
    elif isinstance(dataset, vtk.vtkImageData):
        writer = vtk.vtkXMLImageDataWriter()
    else:
        print('No writer configure for ', type(dataset))
        return

    extension = writer.GetDefaultFileExtension()

    if binary == True:
        writer.SetDataModeToBinary()
    else:
        writer.SetDataModeToBinary()

    writer.SetInputData(dataset)

    writer.SetFileName(filename + '.' + extension)

    writer.Write()

def filter_by_distance(mesh, tolerance):

    nmds_iter = mesh.NewIterator()
    nmds_iter.VisitOnlyLeavesOn()
    nmds_iter.SkipEmptyNodesOn()
    nmds_iter.InitTraversal()
    nmds_iter.GoToFirstItem()

    out_mesh = vtk.vtkMultiBlockDataSet()

    block_id = 1
    while (nmds_iter.IsDoneWithTraversal() == 0):
        dataset = nmds_iter.GetCurrentDataObject()

        # Set up a threshold filter
        threshold_filter = vtk.vtkThreshold()
        threshold_filter.SetInputDataObject(dataset)

        # Apply the tolerance
        threshold_filter.ThresholdByLower(tolerance)

        # For now, hard code to thresholding the Distance field
        threshold_filter.SetInputArrayToProcess(0, 0, 0, vtk.vtkDataObject.FIELD_ASSOCIATION_POINTS, "Distance")
        threshold_filter.Update()

#        mesh.SetDataSet(nmds_iter, threshold_filter.GetOutput())
        out_mesh.SetBlock(block_id, threshold_filter.GetOutput())
        block_id += 1

        nmds_iter.GoToNextItem()

#    return mesh
    return out_mesh


def get_length_pressure_temperature_unit_systems():
    unit_system_labels = ['m-Pa-K',
                          'm-Pa-C',
                          'in-psi-F',
                          'ft-psf-F',
                          'cm-bar-K',
                          'cm-kbar-K',
                          'cm-Gbar-K',
                          'cm-hPa-K',
                          'cm-dyn/cm2-K',
                          'cm-Mbar-K',
                          'mm-Pa-K',
                          'mm-kPa-K',
                          'mm-MPa-K',
                          'mm-GPa-K',
                          'mm-kgf/mm2-K']

    return unit_system_labels