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

"""BlockMesh operation"""

import json
import os
import shutil
import sys
import time

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 FoamMixin


class BlockMesh(smtk.operation.Operation, FoamMixin):
    """Configures and optionally runs blockMesh to generate background mesh.

    Generates controlDict and blockMeshDict files.
    Runs blockMesh either sync or async mode.
    Uses foam_operation.sbt for template
    """

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

        # Do NOT store any smtk resources as member data (causes memory leak)

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

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

    def operateInternal(self):
        """"""
        project = self._get_project()
        if project is None:
            return self.createResult(smtk.operation.Operation.Outcome.UNABLE_TO_OPERATE)

        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)

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

        # Check attributes
        att_names = [
            'controlDict', 'BackgroundGeometry', 'BlockMeshBoundaryConditions', 'BlockMeshSize']
        if not self._check_attributes(att_resource, att_names):
            return self.createResult(smtk.operation.Operation.Outcome.UNABLE_TO_OPERATE)

        # Create case directory if needed
        project_dir = os.path.abspath(os.path.dirname(project.location()))
        case_dir = os.path.join(project_dir, 'foam/blockMesh')
        constant_dir = os.path.join(case_dir, 'constant')
        system_dir = os.path.join(case_dir, 'system')
        for path in [constant_dir, system_dir]:
            if not os.path.exists(path):
                os.makedirs(path)
        self.case_dir = case_dir

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

        # Copy blockMeshDict file
        from_path = os.path.join(workflows_folder, 'internal/foam/system/blockMeshDict')
        shutil.copy2(from_path, system_dir)

        # Append boundary info to blockMeshDict file
        bmd_file = os.path.join(system_dir, 'blockMeshDict')
        if not self._append_boundary(bmd_file, att_resource):
            return self.createResult(smtk.operation.Operation.Outcome.FAILED)

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

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

        if not self._write_allrun_file(case_dir, '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('blockMesh', case_dir, run_item):
                return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        result = self._create_result()
        return result

    def _append_boundary(self, bmd_file: str, att_resource: smtk.attribute.Resource) -> bool:
        mesh_att = att_resource.findAttribute('BlockMeshBoundaryConditions')

        # These are the face definitions that make sense to me
        # left='0374', right='1562', bottom='0123', top='4765', back='0451', front='3267')
        side_dict = dict(
            # These are the face definitions that actually "work"
            left='0374', right='1562', bottom='0451', top='3267', back='0123', front='4765')
        side_list = ['left', 'right', 'bottom', 'top', 'back', 'front']

        # Generate boundary-type dictionary from mesh_att
        bc_dict = dict()

        frontback_item = mesh_att.findString('FrontBackSides')
        if frontback_item is not None:
            frontback_value = frontback_item.value()
            bc_dict['front'] = frontback_value
            bc_dict['back'] = frontback_value

        bottom_item = mesh_att.findString('BottomSide')
        if bottom_item is not None:
            bc_dict['bottom'] = bottom_item.value()

        indent = ' ' * 4
        complete = False
        with open(bmd_file, 'a') as fp:
            fp.write('\nboundary\n')
            fp.write('(\n')

            for side_name in side_list:
                bc_type = bc_dict.get(side_name, 'patch')

                fp.write('{}{}\n'.format(indent, side_name))
                fp.write('{}{{\n'.format(indent))
                fp.write('{}type {};\n'.format(indent * 2, bc_type))
                fp.write('{}faces\n'.format(indent * 2))
                fp.write('{}(\n'.format(indent * 2))

                face_strings = side_dict.get(side_name)
                face_list = list(face_strings)
                face_line = '{}({})\n'.format(indent * 3, ' '.join(face_list))
                fp.write(face_line)

                fp.write('{});\n'.format(indent * 2))
                fp.write('{}}}\n'.format(indent))

            fp.write(');\n')

            fp.write('\n')
            fp.write('// ************************************************************************* //\n')

            complete = True

        return complete

    def _generate_blockMeshConfig(self, case_dir: str, att_resource: smtk.attribute.Resource) -> bool:
        """Writes constant/include/blockMeshConfig file. Returns True on success."""
        # Build (key,value) list
        kv_list = list()
        complete = False

        # scale_att = att_resource.findAttribute('Scale')
        # scale_item = scale_att.findDouble('scale')
        # kv_list.append(('scale', scale_item.value()))
        kv_list.append(('scale', 1.0))

        geom_att = att_resource.findAttribute('BackgroundGeometry')

        box_item = geom_att.itemAtPath('box')
        names = ['Xmin', 'Xmax', 'Ymin', 'Ymax', 'Zmin', 'Zmax']
        for i, name in enumerate(names):
            kv_list.append((name, box_item.value(i)))

        mesh_att = att_resource.findAttribute('BlockMeshSize')
        size_item = mesh_att.findString('MeshSize')
        cell_counts = self._get_cell_counts(size_item, box_item)
        if cell_counts is None:
            return False
        names = ['xCells', 'yCells', 'zCells']
        for t in zip(names, cell_counts):
            kv_list.append(t)

        grading_item = mesh_att.findInt('simpleGrading')
        names = ['xGrading', 'yGrading', 'zGrading']
        for i, name in enumerate(names):
            kv_list.append((name, grading_item.value(i)))

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

        path = os.path.join(include_dir, 'blockMeshConfig')
        with open(path, 'w') as fp:
            self.cd_writer.write_kvlist(fp, kv_list, double_space=False)
            complete = True

        return complete

    def _get_cell_counts(self, meshsize_item: smtk.attribute.StringItem, box_item: smtk.attribute.DoubleItem) -> bool:
        """Parses options in MeshSize item."""
        # Each option has one child item, so find the active one
        numcells_item = meshsize_item.findChild(
            'numcells-eachdir', smtk.attribute.SearchStyle.IMMEDIATE_ACTIVE)
        if numcells_item is not None:
            cell_counts = [None] * 3
            for i in range(3):
                cell_counts[i] = numcells_item.value(i)
            return cell_counts

        # Remaining cases need the lengths of each direction and the max length
        lengths = [None] * 3
        lengths[0] = box_item.value(1) - box_item.value(0)
        lengths[1] = box_item.value(3) - box_item.value(2)
        lengths[2] = box_item.value(5) - box_item.value(4)
        max_length = max(lengths)

        # Determine target cell size
        numcells = [None] * 3
        cell_size = None

        maxdir_child = meshsize_item.findChild(
            'numcells-maxdir', smtk.attribute.SearchStyle.IMMEDIATE_ACTIVE)
        if maxdir_child is not None:
            cell_size = max_length / maxdir_child.value()

        relative_child = meshsize_item.findChild(
            'relative-meshsize', smtk.attribute.SearchStyle.IMMEDIATE_ACTIVE)
        if relative_child is not None:
            cell_size = max_length * relative_child.value()

        absolute_child = meshsize_item.findChild(
            'absolute-meshsize', smtk.attribute.SearchStyle.IMMEDIATE_ACTIVE)
        if absolute_child is not None:
            cell_size = absolute_child.value()

        if cell_size is None:
            self.log().addError('Error finding mesh size')
            return None

        for i in range(3):
            numcells[i] = int(lengths[i] / cell_size)

        return numcells
