#
# 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.
#
#!/usr/bin/env python
# process: internal
import os
import sys

import numpy as np
import vtk

# Need this because Apex reuses the same Python interpreter between script executions
# So we need to explicitly reload imported modules to see changes.....
from importlib import reload
# PyCharm remote debugger - install in the Apex installation folder
# import pydevd_pycharm as rd
# rd = reload(rd)
# rd.settrace('localhost', port=21000, stdoutToServer=True, stderrToServer=True)

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)


this_path = os.path.dirname(__file__)

# Performance profiler
import cProfile, pstats

# apex_vtk : code that reads the model from Apex
import apex_vtk_dataset_converter as avtkdsc
avtkdsc = reload(avtkdsc)

# model_vtk : code that does is where vtk operations happen on the model
import model_vtk as mvtk
mvtk = reload(mvtk)

# vtk_display : code to display the model from vtk data structures
import vtk_display as vtkdisp
vtkdisp = reload(vtkdisp)

import apex_vtk_fem_field as avff
avff = reload(avff)

# apex scripting API
import apex

# Set to True to export the vtk data as a vtkMultiblockDataset
# You can then open the vtk model in Paraview for example
savemodel=False

# Set to True to print performance profile data for the run
generate_pstats=False

# This method is the entry point that is called from the user defined tool code
# The dictionary contains the dati passed from the ui tool panel
def display_apex_entities_in_vtk_window(dictionary={}):

    print('Setting script unit system to ', apex.getApplicationUnitSystemLabel())
    apex.setScriptUnitSystem(apex.getApplicationUnitSystemLabel())

    # Disable showing output dialog at the end of script execution
    # Uncomment this during debugging sessions
    apex.disableShowOutput()


    # Set up the performance profiler if profiling is requested
    if generate_pstats :
        cp = cProfile.Profile()
        cp.enable()

    # Create the class that holds the vtk representation of the apex model
    # and get the root node of this model
    global model_vtk
    model_vtk = mvtk.model_vtk()

    global model_root
    model_root = model_vtk.get_model_root()

    # Create an instance of the class that converts the apex objects to vtk datasets
    global av
    av=avtkdsc.apex_vtk_dataset_converter()

    # Determine what we're going to process
    # visualization_target_type = dictionary['visualization_target_type']
    model_display_scope = dictionary['model_display_scope']
    save_vtk_file = dictionary["save_vtk_file"]
    
    # DB4L8R - Color options not supported in this version yet
#    color_options = dictionary['color_options']

    # DB4L8R - Clean up this logic when getCurrentSelection() is working properly
    visualization_target = apex.selection.getCurrentSelection()
    visualization_target_count = visualization_target.len()
    target_load = None
    visualization_target_type = None
    if visualization_target_count > 0:
        for target in visualization_target:
        # target_load = visualization_target[0]
            target_load_name = target.name
            if isinstance(target, apex.environment.LoadPressure):
                target_load = apex.environment.getLoadPressure(target_load_name)
                visualization_target_type = 'Pressure'
                break
            elif isinstance(target, apex.environment.LoadTemperature):
                target_load = target
                visualization_target_type = 'Temperature'
                break

    if target_load is None:
        print('No valid loads selected....')
        return

    # DB4L8R - Can change this logic when we're able to read the regiosn nmaes fomr te laods
    #          That will avoid having to read the entre model whenonly the Region is required
    #          For now, we have to read the whole model to fix the pressures so
    if visualization_target_type in ["Pressure", "Temperature"]:
        model, model_map = process_target(target=apex.currentModel(),
                                          mvtk_parent=model_root,
                                          transfer_options='Mesh')

        # Instantiate the display and initialize the rendering pipeline
        vtk_display = vtkdisp.vtk_display(dictionary)
#        vtk_display.initialize_rendering_pipeline()

        # Switch to display whatever vizualization target was requested in the UI
        # Display a pressure field
        if isinstance(target_load, apex.environment.LoadPressure):
            field_name = 'Pressure'
           # DB4L8R - Handle this more gracefully - send an error back to UI
           # Also do this check earlier, before reading the model.
           # Need to keep it here until temp selection is working....
            # DB4L8R Remove this when the pressure laod API is fixed and put in the logic to extract the field from
            # the load property
            # Handle errors in the current API for variable pressure fields
            property = target_load.getLoadPressureProperty()
            pressure_field = None
            try:
                pressure_field = property.getPressureValues()
            except:
                print('Problem.......')
                exit()

#            adjustPressureValues(pressure_field)

            load_region, retained_region = extract_load_and_nonload_regions_from_model(model,
                                                                                       model_map,
                                                                                       pressure_field)

            # Generate an mbds of polydatas with the fem field added as a cell/point field
            mbds, field_topology = avff.add_fem_field_to_mbds(load_region, pressure_field, field_name)

            # Save the vtk data if requested
            if save_vtk_file != '':
                # Export the mesh and field data
                comp_filter = vtk.vtkCompositeDataGeometryFilter()
                comp_filter.SetInputData(mbds)
                writer = vtk.vtkXMLPolyDataWriter()
                writer.SetInputConnection(comp_filter.GetOutputPort())
                writer.SetFileName(save_vtk_file)
                writer.SetDataModeToBinary()
                writer.Write()

            # Set up the display of the meshes that have field data
            vtk_display.display_mesh(mbds, field_topology, field_name)

            if retained_region is not None and model_display_scope == 'All':
                vtk_display.display_mesh(retained_region)

            vtk_display.render()

        elif isinstance(target_load, apex.environment.LoadTemperature):
            field_name = 'Temperature'
            load_temperature_property = target_load.temperatureProperty
            temperature_values = load_temperature_property.temperatureValues
#            temperature_values = fix_temperature_load(temperature_values)

            load_region, retained_region = extract_load_and_nonload_regions_from_model(model,
                                                                                       model_map,
                                                                                       temperature_values)

            mbds, field_topology = avff.add_fem_field_to_mbds(load_region,
                                                                          temperature_values,
                                                                          field_name)

            # Save the vtk data if requested
            if save_vtk_file != '':
                # Export the mesh and field data
                comp_filter = vtk.vtkCompositeDataGeometryFilter()
                comp_filter.SetInputData(mbds)
                writer = vtk.vtkXMLPolyDataWriter()
                writer.SetInputConnection(comp_filter.GetOutputPort())
                writer.SetFileName(save_vtk_file)
                writer.SetDataModeToBinary()
                writer.Write()

            vtk_display.display_mesh(mbds, field_topology, field_name)

            if retained_region is not None and model_display_scope == 'All':
                vtk_display.display_mesh(retained_region)

#                vtk_display.display_nodal_scalar_field(temperature_values, model, model_map)
            vtk_display.render()



def process_target(target, mvtk_parent, transfer_options = "Both"):
    if isinstance(target, apex.Model) or isinstance(target, apex.Assembly):
        process_assembly(apex_assembly=target, mvtk_parent=mvtk_parent, transfer_options=transfer_options)
    elif isinstance(target, apex.Part):
        process_part(target, mvtk_parent, transfer_options)
    elif isinstance(target, apex.geometry.GeometryBody) or isinstance(target, apex.mesh.MeshBody):
        process_body(target, mvtk_parent, transfer_options)
    else:
        print("Not a valid selection....exiting")
    return model_vtk.get_model_root(), model_vtk.get_model_map_with_block_id()

# Recursively process Assemblies and Parts
def process_assembly(apex_assembly, mvtk_parent, transfer_options = "Both"):
    #Add the assembly to the vtk model
    mvtk_parent = model_vtk.create_assembly(mvtk_parent, apex_assembly.pathName)
    for apex_part in apex_assembly.parts:
        process_part(apex_part, mvtk_parent, transfer_options)
    for apex_assembly in apex_assembly.assemblies:
        process_assembly(apex_assembly, mvtk_parent, transfer_options)

# Process the Part and the bodies it composes
def process_part(apex_part, mvtk_parent, transfer_options):
    mvtk_parent = model_vtk.create_part(mvtk_parent, apex_part.pathName)
    if transfer_options == "Both" or transfer_options == "Mesh":
        for meshbody in apex_part.meshes:
            if isinstance(meshbody, apex.mesh.PointMesh):
                print("Skipping PointMesh")
            else:
                dataset = av.ugrid_from_mesh(meshbody)
                model_vtk.add_body_to_part(mvtk_parent, dataset, meshbody.pathName)
    if transfer_options == "Both" or transfer_options == "Geometry":
        for geombody in apex_part.geometryBodies:
            dataset = av.ugrid_from_geombody(geombody)
            model_vtk.add_body_to_part(mvtk_parent, dataset, geombody.name)

# Process the Part and bodies it composes
# Ignores the Mesh/Geometry options - if the user picks a body they likely want to process it
def process_body(body, mvtk_parent, transfer_options):
    #Create a Part to hold the body
    mvtk_parent = model_vtk.create_part(mvtk_parent, "Part 1")
    if isinstance(body, apex.mesh.PointMesh):
        print("Skipping PointMesh")
    else:
        if isinstance(body, apex.mesh.MeshBody):
            dataset = av.ugrid_from_mesh(body)
            model_vtk.add_body_to_part(mvtk_parent, dataset, body.pathName)
        elif  isinstance(body, apex.geometry.GeometryBody):
            dataset = av.ugrid_from_geombody(body)
            model_vtk.add_body_to_part(mvtk_parent, dataset, body.pathName)


def extract_load_and_nonload_regions_from_model(model, model_map, load):

    # Add all of the blocks that do not have field data to the block filter
    all_blocks = []
    omitted_blocks = []

    for block_index in model_map.values():
        all_blocks.append(block_index[1])

    blocks_with_fields_filter = vtk.vtkExtractBlock()
    blocks_with_fields_filter.SetInputData(model)

    # Add the indices of all blocks that are referenced in the field
    for mesh_body_name in load.keys():
        index = model_map.get(mesh_body_name, None)[1]
        blocks_with_fields_filter.AddIndex(index)
        omitted_blocks.append(index)

    blocks_with_fields_filter.Update()

    # Remove the blocks that have fields from the full set of blocks in the model
    retained_blocks_filter = None
    if len(all_blocks) > len(omitted_blocks):
        retained_blocks = np.setdiff1d(all_blocks,
                                       omitted_blocks,
                                       assume_unique=True)

        retained_blocks_filter = vtk.vtkExtractBlock()
        retained_blocks_filter.SetInputData(model)

        for retained_block in retained_blocks:
            retained_blocks_filter.AddIndex(int(retained_block))

        retained_blocks_filter.Update()

        return blocks_with_fields_filter.GetOutput(), retained_blocks_filter.GetOutput()
    else:
        return blocks_with_fields_filter.GetOutput(), None


def initialize_load_vizualization(vizualization_target_type, vizualization_target, model_display_scope, color_options):
    if model_display_scope == 'All':
        process_target(target=apex.currentModel(), mvtk_parent=model_root, transfer_options = 'Mesh')
        if vizualization_target_type == "Pressure":
            pressure_name = vizualization_target.name
            pressure_fixed = fix_pressure_definition(model_root, pressure_name)
            print(pressure_fixed)

def initialize_model(model_display_scope, visualization_target):
    if model_display_scope == 'All':
        process_target(target=apex.currentModel(), mvtk_parent=model_root, transfer_options = 'Mesh')
    elif model_display_scope == 'Region':
        if isinstance(visualization_target, [apex.environment.VariablePressure]):
            # Iterate over the bodies and then process that set of bodies
            print('Not yet supported')

# Prints performance data by modules of interest
def output_pstats(profiler):
    sortby = 'tottime'
    ps = pstats.Stats(profiler)
    ps.sort_stats(sortby)
    ps.print_stats()
    ps.print_stats('apex_vtk_dataset_converter|model_vtk|display_apex_model_in_vtk_window')
    ps.print_stats('vtk_display')
    ps.print_stats('vtkC')
    ps.print_stats('numpy_support')
    ps.print_stats('mesh')
    ps.print_stats('extensions')
    ps.print_stats('leaf')
    ps.print_stats('enum')
    ps.print_stats('site_packages/numpy')