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

"""simpleFoam solver operation"""

import glob
import io
import os
import shutil
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 FoamMixin, DictTemplate


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

    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 OpenFOAM solver"

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

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

        Steps include
          * Copy mesh file from snappyHexMesh
          * Copy 0.orig directory to 0
          * Generate 0.orig files from templates
          * Copy system files fvSchemes, fvSolution to system directory
          * Generate 0/include/initialConditions file
        """
        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()))

        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)

        solver = self.cd_writer.get_application_name(att_resource)
        self.case_dir = os.path.join(self.project_dir, 'foam', solver)

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

        self._reset_directory()

        # Copy mesh from snappyHexMesh case
        from_dir = os.path.join(self.case_dir, os.pardir, 'snappyHexMesh/constant/polyMesh')
        if not os.path.exists(from_dir):
            self.log().addError('No mesh at snappyMesh/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)

        # Create 0.orig/include folder
        include_folder = os.path.join(self.case_dir, '0.orig/include')
        if not os.path.exists(include_folder):
            os.makedirs(include_folder)

        start_dir = os.path.join(self.workflows_folder, 'internal/foam')

        # Generate dictionary for rendering 0.orig files
        field_dict = self._generate_fielddict(att_resource)
        for fieldname in ['epsilon', 'k', 'nut', 'omega', 'p', 'U']:
            source_file = os.path.join(start_dir, 'templates', fieldname)
            dest_file = os.path.join(self.case_dir, '0.orig', fieldname)
            if not self._generate_fieldfile(source_file, dest_file, field_dict):
                self.log().addError('Failed to generate {}'.format(dest_file))
                return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        # Copy transportProperties file
        source_file = os.path.join(start_dir, 'constant/transportProperties')
        dest_file = os.path.join(self.case_dir, 'constant/transportProperties')
        shutil.copyfile(source_file, dest_file)

        # Generate config files
        self._generate_initialConditions(att_resource)
        self._generate_transportPropertiesConfig(att_resource)
        self._generate_turbulenceProperties(self.case_dir, att_resource, start_dir)

        # Generate turbulence model files
        turb_att = att_resource.findAttribute('TurbulenceModel')
        turb_item = turb_att.findString('TurbulenceModel')
        # We are presently hard-coded to RAS models
        # ras_item = turb_item.findChild('RASModel', smtk.attribute.SearchStyle.IMMEDIATE_ACTIVE)
        # ras_model = ras_item.value()

        # Check for 0 directory
        dest_dir = os.path.join(self.case_dir, '0')
        if not os.path.exists(dest_dir):
            source_dir = os.path.join(self.case_dir, '0.orig')
            shutil.copytree(source_dir, dest_dir)

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

        if not self._write_allrun_file(
                self.case_dir, solver, copy_casename='snappyHexMesh', restore0=True):
            return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        # Check "Run" option
        run_item = self.parameters().findString('run')
        if run_item.isEnabled():
            app = self.cd_writer.get_application_name(att_resource)
            if not self._run_openfoam(app, self.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')

        # 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 _reset_directory(self):
        """Delete any content from previous jobs."""

        rel_paths = ['0']
        for rel_path in rel_paths:
            path = os.path.join(self.case_dir, rel_path)
            if os.path.exists(path):
                shutil.rmtree(path)

        pathname = '{}/processor*'.format(self.case_dir)
        proc_dirs = glob.glob(pathname)
        for proc_dir in proc_dirs:
            shutil.rmtree(proc_dir)

    def _generate_initialConditions(self, att_resource: smtk.attribute.Resource):
        """Generates initial conditions to include/0.orig directory."""
        kv_list = list()

        pressure_att = att_resource.findAttribute('PressureBoundaryCondition')
        pressure_item = pressure_att.findDouble('Pressure')
        t = ('Pressure', pressure_item.value())
        kv_list.append(t)

        velocity_att = att_resource.findAttribute('VelocityBoundaryCondition')
        velocity_item = velocity_att.findDouble('Velocity')
        if velocity_item.numberOfValues() == 3:  # for backward compatibility
            data = [velocity_item.value(i) for i in range(3)]
        else:
            data = [velocity_item.value(), 0.0, 0.0]
        t = ('FlowVelocity', data)
        kv_list.append(t)

        turb_att = att_resource.findAttribute('TurbulenceModel')
        turb_item = turb_att.findString('TurbulenceModel')
        turb_value = turb_item.value()
        epsilon = k = omega = 0.0
        if turb_value == 'RAS':
            IM_ACTIVE = smtk.attribute.SearchStyle.IMMEDIATE_ACTIVE  # save typing
            ras_item = turb_item.findChild('RASModel', IM_ACTIVE)

            epsilon_item = ras_item.findChild('turbulentEpsilon', IM_ACTIVE)
            if epsilon_item is not None:
                epsilon = epsilon_item.value()
            k_item = ras_item.findChild('turbulentKE', IM_ACTIVE)
            if k_item is not None:
                k = k_item.value()
            omega_item = ras_item.findChild('turbulentOmega', IM_ACTIVE)
            if omega_item is not None:
                omega = omega_item.value()

        kv_list.append(('turbulentEpsilon', epsilon))
        kv_list.append(('turbulentKE', k))
        kv_list.append(('turbulentOmega', omega))

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

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

    def _generate_transportPropertiesConfig(self, att_resource: smtk.attribute.Resource):
        """"""
        properties_att = att_resource.findAttribute('PhysicalProperties')
        kviscosity_item = properties_att.findDouble('KinematicViscosity')
        t = ('KinematicViscosity', kviscosity_item.value())
        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, 'transportPropertiesConfig')
        with open(path, 'w') as fp:
            self.cd_writer.write_kvlist(fp, kv_list)

    def _generate_fielddict(self, att_resource: smtk.attribute.Resource) -> dict:
        """Uses boundary settings in block mesh to set bottom and front/back data."""
        field_dict = dict()  # return value
        mesh_att = att_resource.findAttribute('BlockMeshBoundaryConditions')

        # FrontBackSides is one of empty, zeroGradient, symmetryPlane
        front_back_item = mesh_att.findString('FrontBackSides')
        front_back_value = front_back_item.value()
        for field in ['epsilon', 'k', 'nut', 'omega', 'p', 'U']:
            key = '{}_front_back'.format(field)
            field_dict[key] = front_back_value

        # Special cases
        if front_back_value == 'zeroGradient':
            field_dict['k_front_back'] = 'slip'
            field_dict['nut_front_back'] = 'calculated'

        # Set lookup tables for turbulence params
        epsilon_dict = dict(wall='epsilonWallFunction', zeroGradient='zeroGradient')
        k_dict = dict(wall='kqRWallFunction', zeroGradient='slip')
        nut_dict = dict(wall='nutkWallFunction', zeroGradient='calculated')
        omega_dict = dict(wall='omegaWallFunction', zeroGradient='slip')
        U_dict = dict(wall='noSlip', zeroGradient='zeroGradient')

        lookup_dict = dict(
            epsilon=epsilon_dict, k=k_dict, nut=nut_dict, omega=omega_dict, U=U_dict)

        def set_field_items(field_dict, side_name, item_value, lookup_dict):
            """Populate field_dict for given side and value"""
            for field_name in ['epsilon', 'k', 'nut', 'omega', 'U']:
                d = lookup_dict.get(field_name, {})
                dict_value = d.get(item_value)
                key = '{}_{}'.format(field_name, side_name)
                field_dict[key] = dict_value

        # BottomSide is either zeroGradient or Wall
        bottom_item = mesh_att.findString('BottomSide')
        bottom_value = bottom_item.value()
        field_dict['p_bottom'] = bottom_value
        set_field_items(field_dict, 'bottom', bottom_value, lookup_dict)

        # Same thing for TopSide
        top_item = mesh_att.findString('TopSide')
        top_value = top_item.value()
        field_dict['p_top'] = top_value
        set_field_items(field_dict, 'top', top_value, lookup_dict)

        turb_att = att_resource.findAttribute('TurbulenceModel')
        turb_item = turb_att.findString('TurbulenceModel')
        turb_value = turb_item.value()

        return field_dict

    def _generate_fieldfile(self, template_path: str, dest_path: str, field_dict: dict) -> bool:
        """"""
        # Load template file
        template_string = None
        with open(template_path) as fin:
            template_string = fin.read()
        if template_string is None:
            self.log().addError('Failed to find/load template file {}'.format(template_path))
            return False
        dict_template = DictTemplate(template_string)

        # Generate output string and write to dest_path
        output_string = dict_template.substitute(field_dict)
        if dest_path.endswith('U'):
            print('dest_path', dest_path)
            print('field_dict', field_dict)
            print('U_top', field_dict.get('U_top'))
        with open(dest_path, 'w') as fout:
            fout.write(output_string)

        return True

    def _generate_turbulenceProperties(
            self, case_dir: str, att_resource: smtk.attribute.Resource, start_dir: str) -> bool:
        """"""
        turb_dict = None
        turb_att = att_resource.findAttribute('TurbulenceModel')
        turb_item = turb_att.findString('TurbulenceModel')
        turb_value = turb_item.value()
        if turb_value == 'laminar':
            turb_dict = dict(simulation_type='laminar', ras_model='kEpsilon')
        elif turb_value == 'RAS':
            ras_item = turb_item.findChild('RASModel', smtk.attribute.SearchStyle.IMMEDIATE_ACTIVE)
            ras_model = ras_item.value()
            turb_dict = dict(simulation_type='RAS', ras_model=ras_model)
        else:
            self.log().addError('unrecognized turbulend value {}'.format(turb_value))
            return False

        # Load template file
        template_string = None
        template_path = os.path.join(start_dir, 'templates/turbulenceProperties')
        with open(template_path) as fin:
            template_string = fin.read()
        if template_string is None:
            self.log().addError('Failed to find/load template file {}'.format(template_path))
            return False
        dict_template = DictTemplate(template_string)

        # Generate output string and write to dest_path
        dest_path = os.path.join(case_dir, 'constant/turbulenceProperties')
        output_string = dict_template.substitute(turb_dict)
        with open(dest_path, 'w') as fout:
            fout.write(output_string)

        return True

    def _run_process(self, case_dir: str, run_item: smtk.attribute.IntItem, app: str = 'simpleFoam') -> bool:
        """"""
        logfilename = '{}.log'.format(app)
        self.logfile = os.path.join(case_dir, logfilename)

        args = []
        if hasattr(smtk, 'use_openfoam_docker') and smtk.use_openfoam_docker:
            args = self._docker_run_commands(case_dir)

        run_mode = run_item.value()
        if run_mode == 'sync':  # run and wait for completion
            args.append(app)
            with open(self.logfile, 'w') as fp:
                completed_proc = subprocess.run(
                    args, stdout=fp, stderr=fp, cwd=case_dir, universal_newlines=True)
                if completed_proc.returncode == 0:
                    self.status = ProcessStatus.Completed
                    success = True

                    # Touch .foam file for visualization
                    foam_filename = '{}.foam'.format(app)
                    foamfile = os.path.join(case_dir, foam_filename)
                    with open(foamfile, 'a'):
                        os.utime(foamfile, None)

                else:
                    self.log().addError('{} returned code {}'.format(app, completed_proc.returncode))
                    self.status = ProcessStatus.Error
                    success = False
            return success

        elif run_mode == 'async':  # launch and return
            command = '{0} > {1}'.format(app, self.logfile)
            args.append(command)
            proc = subprocess.Popen(args, shell=True, cwd=case_dir, universal_newlines=True)
            self.log().addRecord(smtk.io.Logger.Info, 'Started process {}'.format(proc.pid))
            self.status = ProcessStatus.Started
            self.pid = proc.pid
            return True

        else:
            self.log().addError('internal error - unrecognized run_mode {}'.format(run_mode))
            return False
