# =============================================================================
#
#  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.
#
# =============================================================================

"""SnappyHexMesh operation"""

import glob
import io
import os
import shutil
import subprocess
import sys

import smtk
import smtk.attribute
import smtk.io
import smtk.operation
import smtk.project

OP_SUCCEEDED = int(smtk.operation.Operation.Outcome.SUCCEEDED)

# Make sure __file__ is set when using modelbuilder
import inspect
source_file = os.path.abspath(inspect.getfile(inspect.currentframe()))
__file__ = source_file

# Make sure this folder is in sys.path
source_dir = os.path.dirname(__file__)
if source_dir not in sys.path:
    sys.path.insert(0, source_dir)

from foam_mixin import DictTemplate, FoamMixin


class SnappyHexMesh(smtk.operation.Operation, FoamMixin):
    """Generates inputs files for project analysis mesh and runs snappyHexMesh."""

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

        # Do NOT store any smtk resources as member data (causes memory leak)
        self.project_dir = None
        self.case_dir = None
        self.workflows_folder = None

    def name(self):
        return "Setup and run snappyHexMesh"

    def createSpecification(self):
        spec = self._create_specification(app='snappyHexMesh')
        return spec

    def operateInternal(self):
        """Generate the OpenFOAM input files for snappyHexMesh

        Steps include
          * Clear current polyMesh and copy from background case
          * Generate system/controlDict
          * Copy system files fvSchemes, fvSolution, meshQualityDict to system
          * Generate system/snappyHexMeshDict (inserting refinement regions)
          * Generate constant/include/snappyHexMeshConfig
        """
        project = self._get_project()
        if project is None:
            return self.createResult(smtk.operation.Operation.Outcome.UNABLE_TO_OPERATE)
        self.project_dir = os.path.abspath(os.path.dirname(project.location()))
        self.case_dir = os.path.join(self.project_dir, 'foam/snappyHexMesh')

        att_resource = self._get_attribute_resource(project)
        # print('att_resource:', att_resource)
        if att_resource is None:
            return self.createResult(smtk.operation.Operation.Outcome.UNABLE_TO_OPERATE)

        self.workflows_folder = self._get_workflows_folder()
        if self.workflows_folder is None:
            return self.createResult(smtk.operation.Operation.Outcome.UNABLE_TO_OPERATE)

        # Create case directories if needed
        constant_dir = os.path.join(self.case_dir, 'constant')
        system_dir = os.path.join(self.case_dir, 'system')
        logs_dirs = os.path.join(self.case_dir, 'logs')

        for path in [constant_dir, system_dir, logs_dirs]:
            if not os.path.exists(path):
                os.makedirs(path)

        # Copy mesh from blockMesh case
        from_dir = os.path.join(self.case_dir, os.pardir, 'blockMesh/constant/polyMesh')
        if not os.path.exists(from_dir):
            self.log().addError('No block mesh at blockMesh/constant/polyMesh')
            return self.createResult(smtk.operation.Operation.Outcome.UNABLE_TO_OPERATE)
        to_dir = os.path.join(self.case_dir, 'constant/polyMesh')
        if os.path.exists(to_dir):
            shutil.rmtree(to_dir)
        shutil.copytree(from_dir, to_dir)

        self._copy_system_files(att_resource)

        # Generate controlDict file
        if not self.cd_writer.generate_controlDict(self.case_dir, att_resource):
            return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        # Note that we generate 2 snappyHexMesh files
        # 1. snappyHexMeshDict: we insert refinement regions into a template
        # 2. snappyHexMeshDictConfig: we write OpenFOAM variables here
        self._generate_snappyHexMeshDict(att_resource)

        # Generate config files
        self._generate_snappyHexMeshConfig(att_resource)
        self._generate_meshQualityConfig(att_resource)

        # Generate Allrun & Allclean files
        if not self._write_allclean_file(self.case_dir):
            return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        foam_apps = ['surfaceFeatureExtract', 'snappyHexMesh -overwrite']
        if not self._write_allrun_file(self.case_dir, foam_apps, copy_casename='blockMesh'):
            return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        # Check "Run" option
        run_item = self.parameters().findString('run')
        if run_item.isEnabled():
            if not self._run_openfoam('snappyHexMesh', case_dir, run_item):
                return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        result = self._create_result()
        return result

    def _copy_system_files(self, att_resource):
        """"""
        start_dir = os.path.join(self.workflows_folder, 'internal/foam')
        source_dir = os.path.join(start_dir, 'system')
        dest_dir = os.path.join(self.case_dir, 'system')
        shutil.copytree(source_dir, dest_dir, dirs_exist_ok=True)

        # Copy fvSchemes and fvSolution for this application
        app = self.cd_writer.get_application_name(att_resource)
        filenames = ['fvSchemes', 'fvSolution']
        for filename in filenames:
            source_file = os.path.join(start_dir, app, filename)
            if not os.path.exists(source_file):
                self.log().addError('Unable to find {}'.format(source_file))
                continue
            dest_file = os.path.join(self.case_dir, 'system', filename)
            shutil.copyfile(source_file, dest_file)

    def _generate_meshQualityConfig(self, att_resource: smtk.attribute.Resource):
        """Generates config file to include in meshQualityDict."""
        minweight = 0.2
        meshquality_att = att_resource.findAttribute('MeshQualityControls')
        if meshquality_att is not None:
            minweight_item = meshquality_att.findDouble('minFaceWeight')
            minweight = minweight_item.value()
        t = ('MinFaceWeight', minweight)
        kv_list = [t]

        # Create include directory if needed
        include_dir = os.path.join(self.case_dir, 'constant/include')
        if not os.path.exists(include_dir):
            os.makedirs(include_dir)

        path = os.path.join(include_dir, 'meshQualityDictConfig')
        with open(path, 'w') as fp:
            self.cd_writer.write_kvlist(fp, kv_list)

    def _generate_snappyHexMeshDict(self, att_resource: smtk.attribute.Resource):
        """Generates snappyHexMeshDict file from template."""
        geom_lines = ['']  # for //@RefinementGeometry
        rr_lines = ['']   # for //@RefinementRegions
        indent4 = ' ' * 4
        indent8 = ' ' * 8
        indent12 = ' ' * 12
        indent16 = ' ' * 16

        att_list = att_resource.findAttributes('RefinementRegion')
        for att in att_list:
            # //@RefinementGeometry
            box_item = att.findDouble('box')
            geom_lines.append(indent4 + att.name())
            geom_lines.append(indent4 + '{')
            geom_lines.append(indent8 + 'type   box;')

            min_string = 'min    ({} {} {});'.format(
                box_item.value(0), box_item.value(2), box_item.value(4))
            geom_lines.append(indent8 + min_string)
            max_string = 'max    ({} {} {});'.format(
                box_item.value(1), box_item.value(3), box_item.value(5))
            geom_lines.append(indent8 + max_string)

            geom_lines.append(indent4 + '}')

            # //@RefinementRegions
            # For some reason, extra indent must be added
            mode_item = att.findString('mode')
            mode = mode_item.value()
            rr_lines.append(indent8 + att.name())
            rr_lines.append(indent8 + '{')
            rr_lines.append(indent12 + 'mode   {};'.format(mode))
            if mode in ['inside', 'outside']:
                level_item = mode_item.findChild(
                    'level', smtk.attribute.SearchStyle.IMMEDIATE_ACTIVE)
                level = level_item.value()
                rr_lines.append(indent12 + 'levels ((1.0 {}));'.format(level))

                increment_item = mode_item.findChild(
                    'levelIncrement', smtk.attribute.SearchStyle.IMMEDIATE_ACTIVE)
                if increment_item and increment_item.isEnabled():
                    range_item = increment_item.find('range')
                    split_item = increment_item.find('split')
                    increment_string = '({} {} ({} {} {}))'.format(
                        range_item.value(0), range_item.value(1),
                        split_item.value(0), split_item.value(1), split_item.value(2))
                    rr_lines.append(indent12 + 'levelIncrement {};'.format(increment_string))
            elif mode == 'distance':
                levels_group = mode_item.findChild(
                    'distance-levels', smtk.attribute.SearchStyle.IMMEDIATE_ACTIVE)
                rr_lines.append(indent12 + 'levels')
                rr_lines.append(indent12 + '(')
                for i in range(levels_group.numberOfGroups()):
                    distance = levels_group.find('distance').value()
                    level = levels_group.find('level').value()
                    rr_lines.append(indent16 + '({} {})'.format(distance, level))
                rr_lines.append(indent12 + ');')
            rr_lines.append(indent8 + '}')

        # Read template file
        start_dir = os.path.join(self.workflows_folder, 'internal/foam')

        template_path = os.path.join(start_dir, 'templates/snappyHexMeshDict')
        with open(template_path) as fin:
            template_string = fin.read()
        dict_template = DictTemplate(template_string)

        # Apply refinement regions text
        geom_string = '\n'.join(geom_lines)
        rr_string = '\n'.join(rr_lines)
        output_string = dict_template.substitute(
            RefinementGeometry=geom_string,
            RefinementRegions=rr_string)

        # Write to case directory
        target_path = os.path.join(self.case_dir, 'system/snappyHexMeshDict')
        with open(target_path, 'w') as fout:
            fout.write(output_string)

    def _generate_snappyHexMeshConfig(self, att_resource: smtk.attribute.Resource):
        """Generates config file to include in snappyHexMeshDict."""
        kv_list = list()
        target_att = att_resource.findAttribute('Target')

        filename_item = target_att.findString('Filename')
        filename = filename_item.value()
        kv_list.append(('ModelFileName', filename))

        root, ext = os.path.splitext(filename)
        kv_list.append(('ModelName', root))

        emesh_filename = '{}.eMesh'.format(root)
        kv_list.append(('EMeshFileName', emesh_filename))

        model_point_item = target_att.findDouble('InsidePoint')
        coords = [model_point_item.value(i) for i in range(3)]
        kv_list.append(('LocationInModel', coords))

        # Castellation settings
        castell_att = att_resource.findAttribute('Castellation')
        castell_group = castell_att.findGroup('Castellation')
        # Settings list is (item-name, configfile-name)
        castell_settings = [
            ('InsidePoint', 'LocationInMesh'),
            ('maxGlobalCells', 'MaxGlobalCells'),
            ('minRefinementCells', 'MinRefinementCells'),
            ('nCellsBetweenLevels', 'NCellsBetweenLevels'),
            ('EdgeRefinementLevel', 'EdgeRefinementLevel'),
            ('SurfaceRefinementLevels', 'SurfaceRefinementLevels'),
            ('resolveFeatureAngle', 'ResolveFeatureAngle'),
        ]
        castell_list = self._build_kvlist(castell_group, castell_settings)

        # Main snap controls
        snap_att = att_resource.findAttribute('SnapControls')
        snap_group = snap_att.findGroup('SnapControls')
        snap_settings = [
            ('nSmoothPatch', 'NSmoothPatch'),
            ('tolerance', 'Tolerance'),
            ('nSolveIter', 'NSolveIter'),
            ('nRelaxIter', 'NRelaxIter'),
        ]
        snap_list = self._build_kvlist(snap_group, snap_settings)

        # Feature snap controls
        fs_group = snap_group.find('FeatureSnapping')
        fs_settings = [
            ('nFeatureSnapIter', 'NFeatureSnapIter'),
            ('implicitFeatureSnap', 'ImplicitFeatureSnap'),
            ('explicitFeatureSnap', 'ExplicitFeatureSnap'),
            ('multiRegionFeatureSnap', 'MultiRegionFeatureSnap'),
        ]
        snap_list += self._build_kvlist(fs_group, fs_settings)

        # Create include directory if needed
        include_dir = os.path.join(self.case_dir, 'constant/include')
        if not os.path.exists(include_dir):
            os.makedirs(include_dir)

        path = os.path.join(include_dir, 'snappyHexMeshConfig')
        with open(path, 'w') as fp:
            self.cd_writer.write_kvlist(fp, kv_list)
            self.cd_writer.write_kvsection(fp, castell_list, 'castell')
            self.cd_writer.write_kvsection(fp, snap_list, 'snapping')

    def _build_kvlist(self, group_item: smtk.attribute.GroupItem, settings: list) -> list:
        """"""
        kv_list = list()
        for item_name, key in settings:
            item = group_item.find(item_name)
            if item.type() == smtk.attribute.Item.VoidType:
                value = 'on' if item.isEnabled() else 'off'
            elif item.numberOfValues() == 1:
                value = item.value()
            else:
                value = [item.value(i) for i in range(item.numberOfValues())]
            t = (key, value)
            kv_list.append(t)
        return kv_list
