#=============================================================================
#
#  Copyright (c) Kitware, Inc.
#  All rights reserved.
#  See LICENSE.txt for details.
#
#  This software is distributed WITHOUT ANY WARRANTY; without even
#  the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
#  PURPOSE.  See the above copyright notice for more information.
#
#=============================================================================

""" generate_mesh.py:

Generate a mesh from an RGG model using MeshKit + Cubit.

"""
import sys

from . import generate_mesh_xml

import json
import os
import rggsession
import shutil
import smtk
import smtk.attribute
import smtk.io
import smtk.mesh
import smtk.model
import smtk.operation
import subprocess
import sys
import tempfile
import threading


def _meshAssembly(assembly, assygen, cubit, dirpath):
    # read <assembly>.inp and create <assembly>.py, <assembly>.jou and <assembly>.template.jou
    executable_args = [assygen, assembly]
    assygen_task = subprocess.Popen(executable_args, cwd=dirpath, stdout=subprocess.PIPE)
    assygen_task.wait()

    # read <assembly>.py and create <assembly>.sat
    executable_args = [cubit, '-nographics', '-batch', assembly + '.py']
    cubit_py_task = subprocess.Popen(executable_args, cwd=dirpath, stdout=subprocess.PIPE)
    cubit_py_task.wait()

    # read <assembly>.jou (which references <assembly>.sat) and create <assembly>.cub
    executable_args = [cubit, '-nographics', '-batch', assembly + '.jou']
    cubit_jou_task = subprocess.Popen(executable_args, cwd=dirpath, stdout=subprocess.PIPE)
    cubit_jou_task.wait()


class GenerateMesh(smtk.operation.Operation):

    def __init__(self):
        smtk.operation.Operation.__init__(self)

    def name(self):
        return "generate mesh"

    def which(self, executable):
        """ Find an executable using the following queues (in order):
        1) check user-defined inputs
        2) check relative to current executable
        3) check relative to PATH
        """
        def is_exe(fpath):
            return os.path.isfile(fpath) and os.access(fpath, os.X_OK)

        # 1) check user-defined inputs
        file_att = self.parameters().findFile(executable)
        if file_att and file_att.isEnabled():
            exe_file = file_att.value()
            if is_exe(exe_file):
                return exe_file

        # 2) check relative to current executable
        current_exe = sys.executable
        fpath, fname = os.path.split(current_exe)

        exe_file = os.path.join(fpath, '../bin', executable)
        if is_exe(exe_file):
            return exe_file

        # 3) check relative to PATH
        for path in os.environ["PATH"].split(os.pathsep):
            exe_file = os.path.join(path, executable)
            if is_exe(exe_file):
                return exe_file

        return None

    def ableToOperate(self):
        if not smtk.operation.Operation.ableToOperate(self):
            return False

        # Ensure that the necessary executables can be found
        return self.which('coregen') and self.which('assygen') and \
          (self.which('cubit') != None or self.which('Cubit') != None)

    def operateInternal(self):
        smtk.InfoMessage(self.log(), 'generate mesh')

        # Access the input model
        model = smtk.model.Entity.CastTo(self.parameters().associations().value(0))

        # Access assygen, coregen and cubit
        assygen = self.which('assygen')
        coregen = self.which('coregen')
        cubit = self.which('cubit')
        if cubit is None:
            cubit = self.which('Cubit')

        smtk.InfoMessage(self.log(), 'assygen executable is %s' % assygen)
        smtk.InfoMessage(self.log(), 'coregen executable is %s' % coregen)
        smtk.InfoMessage(self.log(), 'cubit executable is %s' % cubit)

        # Construct a temporary scratch space for generating input files
        temp_dir = tempfile.mkdtemp()

        exportInp = smtk.session.rgg.ExportInp.create()
        exportInp.parameters().associate(model)
        exportInp.parameters().find('directory').setValue(temp_dir)

        exportInpResult = exportInp.operate()

        # Fetch rgg core and its assemblies
        # TODO: the access keys used to get these items should be python wrapped
        coreGroup = model.resource().findEntitiesByProperty('rggType', '_rgg_core')[0]
        assemblyGroups = model.resource().findEntitiesByProperty('rggType', '_rgg_assembly')

        # Construct mesh files for each assembly
        assembly_threads = []
        for assembly in assemblyGroups:
            assembly_threads.append(
                threading.Thread(target=_meshAssembly,
                                 args=(assembly.name(), assygen, cubit, temp_dir)))
        map(lambda x: x.start(), assembly_threads)
        map(lambda x: x.join(), assembly_threads)

        # Construct mesh file for the assembled core
        executable_args = [coregen, coreGroup.name()]
        coregen_task = subprocess.Popen(executable_args, cwd=temp_dir, stdout=subprocess.PIPE)
        coregen_task.wait()

        # Access the generated mesh file

        # First, get its name
        core = json.loads(coreGroup.stringProperty('core_descriptions')[0])
        mesh_file = os.path.join(temp_dir, core['exportParameters']['outputFileName'])
        smtk.InfoMessage(self.log(), 'mesh file: %s' % mesh_file)

        # Then, import the file as a new mesh
        mesh_resource = smtk.mesh.Resource.create()
        smtk.io.importMesh(mesh_file, mesh_resource)

        # Merge contact points to make the mesh conformal
        mesh_resource.meshes().mergeCoincidentContactPoints()

        # If there are no sidesets, create one for the outer shell
        if not mesh_resource.neumanns():
            shell = mesh_resource.meshes().extractShell()
            mesh_resource.setNeumannOnMeshes(shell, smtk.mesh.Neumann(99))

        # TODO: at least one of these calls is necessary. This should be cleaner
        mesh_resource.classifyTo(model.resource())
        mesh_resource.modelResource = model.resource()
        mesh_resource.associateToModel(model.id())

        # Classify each mesh element block onto its corresponding model entity.

        # Cubit constructs blocks based on material type, and the generated
        # input for Cubit artifically assigns materials by traversing the core
        # in the following order (see ExportInp.cxx:71):
        #  - assemblies
        #    - ducts
        #      - segments
        #        - layers
        #    - pins
        #      - pieces
        #        - layers
        #      - cell material
        #  - background mesh
        core_group = smtk.model.Group(model.resource().findEntitiesByProperty(
            'rggType', smtk.session.rgg.Core.typeDescription)[0])
        core = smtk.session.rgg.Core(core_group)
        assembly_groups = model.resource().findEntitiesByProperty(
            'rggType', smtk.session.rgg.Assembly.typeDescription)
        assemblies = [smtk.session.rgg.Assembly(smtk.model.Group(assembly_group))
         for assembly_group in assembly_groups]

        # A generator that traverses the model entities of the core in the same
        # way as they were traversed upon export to Cubit
        def owningComponents():
            for assembly in assemblies:
                duct_aux = smtk.model.AuxiliaryGeometry(model.resource(), assembly.associatedDuct())
                duct = smtk.session.rgg.Duct(duct_aux)
                for segment in duct.segments():
                    for layer in segment.layers:
                        if layer[0] != 0:
                            yield duct_aux
                pins_aux = [smtk.model.AuxiliaryGeometry(model.resource(), pinId)
                             for pinId in assembly.layout()]
                for pin_aux in pins_aux:
                    pin = smtk.session.rgg.Pin(pin_aux)
                    for layer_material in pin.layerMaterials():
                        if layer_material.subMaterialIndex != 0:
                            yield pin_aux
                    if pin.cellMaterialIndex() != 0:
                        yield pin_aux
            yield core_group

        meshsets = [mesh_resource.meshes(domain) for domain in mesh_resource.domains()]
        for (entity, meshset) in zip(owningComponents(), meshsets):
            meshset.setModelEntity(entity)

        # Clean up the temporary work directory
        try:
            shutil.rmtree(temp_dir)
        except OSError as e:
            smtk.WarningMessage(self.log(),
                                'Could not remove temporary directory \'%s\'' % temp_dir)

        # Return with success
        result = self.createResult(smtk.operation.Operation.Outcome.SUCCEEDED)
        result.find('resource').setValue(mesh_resource)
        return result

    def createSpecification(self):
        spec = smtk.attribute.Resource.create()
        reader = smtk.io.AttributeReader()
        reader.readContents(spec, generate_mesh_xml.description, self.log())
        return spec
