#=============================================================================
#
#    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.
#
#=============================================================================
import os
import shutil
import sys
import unittest

import smtk
import smtk.attribute
import smtk.io
import smtk.operation
import smtk.resource
import smtk.session.vtk
import smtk.testing

from ACE3PFileValidator import ACE3PFileValidator
import util

OPERATION_SUCCEEDED = int(smtk.operation.Operation.SUCCEEDED)  # 3
MODEL_NAME = 'pillbox4'
OUTPUT_FOLDER = 'export_{}'.format(MODEL_NAME)

class Omega3PTest1(smtk.testing.TestCase):

    def setUp(self):
        # Make sure last test result is removed
        export_folder = os.path.join(smtk.testing.TEMP_DIR, OUTPUT_FOLDER)
        if os.path.exists(export_folder):
            shutil.rmtree(export_folder)

        # Initialize smtk managers
        self.res_manager = smtk.resource.Manager.create()
        self.op_manager = smtk.operation.Manager.create()

        smtk.attribute.Registrar.registerTo(self.res_manager)
        smtk.attribute.Registrar.registerTo(self.op_manager)

        smtk.session.vtk.Registrar.registerTo(self.res_manager)
        smtk.session.vtk.Registrar.registerTo(self.op_manager)

        smtk.operation.Registrar.registerTo(self.op_manager)
        self.op_manager.registerResourceManager(self.res_manager)

    def tearDown(self):
        self.res_manager = None
        self.op_manager = None

    def import_resource(self, path):
        import_op = self.op_manager.createOperation(
            'smtk::operation::ImportResource')
        import_op.parameters().find('filename').setValue(path)
        import_result = import_op.operate()
        import_outcome = import_result.findInt('outcome').value(0)
        self.assertEqual(import_outcome, OPERATION_SUCCEEDED)
        resource = import_result.find('resource').value()
        self.assertIsNotNone(resource)
        return resource

    def import_python_op(self, path):
        import_op = self.op_manager.createOperation(
            'smtk::operation::ImportPythonOperation')
        self.assertIsNotNone(import_op)
        import_op.parameters().find('filename').setValue(path)
        import_result = import_op.operate()
        import_outcome = import_result.findInt('outcome').value(0)
        self.assertEqual(import_outcome, OPERATION_SUCCEEDED)

        op_unique_name = import_result.findString("unique_name").value()
        op = self.op_manager.createOperation(op_unique_name)
        return op

    def set_attributes(self, att_resource, model_resource):
        """Populate simulation attributes"""
        # Bounday conditions
        bc_lookup = {
            1: 'Electric',
            2: 'Electric'
        }
        bc_default = 'Exterior'
        uuids = model_resource.entitiesMatchingFlags(smtk.model.FACE, True)
        for uuid in uuids:
            # print('UUID {}'.format(uuid))
            prop_list = model_resource.integerProperty(uuid, 'pedigree id')
            face_id = prop_list[0]
            # print('Face ID {}'.format(face_id))
            bc_type = bc_lookup.get(face_id, bc_default)
            bc_att = att_resource.createAttribute(bc_type)
            self.assertIsNotNone(bc_att)
            self.assertTrue(bc_att.associateEntity(uuid))

        # Material
        uuids = model_resource.entitiesMatchingFlags(smtk.model.VOLUME, True)
        uuid = uuids.pop()
        mat_att = att_resource.createAttribute('Material')

        epsilon_item = mat_att.findDouble('Epsilon')
        self.assertTrue(epsilon_item.setValue(2.2))
        mu_item = mat_att.findDouble('Mu')
        self.assertTrue(mu_item.setValue(0.5))
        self.assertTrue(mat_att.associateEntity(uuid))

        # Analysis items
        finfo_att = att_resource.findAttributes('FrequencyInfo')[0]
        num_item = finfo_att.findInt('NumEigenvalues')
        self.assertTrue(num_item.setValue(3))

        pp_att = att_resource.findAttributes('PostProcess')[0]
        toggle_item = pp_att.findGroup('Toggle')
        prefix_item = toggle_item.find('ModeFilePrefix')
        prefix_item.setIsEnabled(True)

    def check_results(self, path):
        validator = ACE3PFileValidator(path)
        self.assertGreater(validator.get_number_of_lines(), 12)
        sections = [
            'ModelInfo', 'BoundaryCondition', 'FiniteElement', 'EigenSolver', 'PostProcess',
            'SurfaceMaterial', 'Material'
        ]
        for name in sections:
            msg = 'output includes {} section'.format(name)
            self.assertTrue(validator.has_section(name), msg)

    def test_pillbox4(self):
        # Load model file
        gen_filename = '{}.gen'.format(MODEL_NAME)
        gen_path = os.path.join(
            smtk.testing.SOURCE_DIR, 'data', 'model', '3d', 'genesis', gen_filename)
        resource = self.import_resource(gen_path)
        model_resource = smtk.model.Resource.CastTo(resource)
        self.assertIsNotNone(model_resource)

        # Set the model location in lieu of writing it to the file system
        model_filename = '{}.smtk'.format(MODEL_NAME)
        model_location = os.path.join(smtk.testing.TEMP_DIR, model_filename)
        model_resource.setLocation(model_location)

        # Load attribute template
        sbt_path = os.path.join(
            smtk.testing.SOURCE_DIR, 'simulation-workflows', 'ACE3P.sbt')
        self.assertTrue(os.path.exists(sbt_path))
        att_resource = self.import_resource(sbt_path)
        self.assertIsNotNone(att_resource)

        # Associate model to attribute resource
        self.assertTrue(att_resource.associate(model_resource))

        # Create analysis attribute and set to Omega3P
        analysis_att = util.create_analysis_att(att_resource)
        # (Just gotta know how analysis att is generated by smtk)
        # Because top-level is exclusive, analysis item is discrete string
        analysis_item = analysis_att.findString('Analysis')
        analysis_item.setValue('Omega3P')

        # Create attributes for instanced views
        util.create_instanced_atts(att_resource)

        # writer = smtk.io.AttributeWriter()
        # logger = smtk.io.Logger()
        # write_path = os.path.join(smtk.testing.TEMP_DIR, 'test.sbi')
        # write_err = writer.write(att_resource, write_path, logger)
        # print('Write attribute returned iserr?', write_err)

        self.set_attributes(att_resource, model_resource)

        # Get export operator
        py_path = os.path.join(
            smtk.testing.SOURCE_DIR, 'simulation-workflows', 'ACE3P.py')
        self.assertTrue(os.path.exists(py_path))
        export_op = self.import_python_op(py_path)
        self.assertIsNotNone(export_op)

        # Setup export
        params_att = export_op.parameters()
        att_item = params_att.findResource("attributes")
        att_item.setValue(att_resource)

        model_item = params_att.findResource("model")
        model_item.setValue(model_resource)

        # Set MeshFile to the input (genesis) file
        file_item = params_att.findFile('MeshFile')
        file_item.setValue(gen_path)

        # Set the output items
        export_path = os.path.join(smtk.testing.TEMP_DIR, OUTPUT_FOLDER)
        dir_item = params_att.findDirectory('OutputFolder')
        dir_item.setValue(export_path)

        prefix_item = params_att.findString('OutputFilePrefix')
        prefix_item.setValue(MODEL_NAME)

        # Run the operator
        result = export_op.operate()
        outcome = result.findInt('outcome')
        self.assertEqual(outcome.value(), OPERATION_SUCCEEDED)

        # Check that expected file was created
        expected_filename = '{}.omega3p'.format(MODEL_NAME)
        expected_path = os.path.join(export_path, expected_filename)
        self.assertTrue(os.path.exists(expected_path))

        # Validate result
        self.check_results(expected_path)

        # (On success) remove the export folder
        shutil.rmtree(export_path)

if __name__ == '__main__':
    smtk.testing.process_arguments()
    unittest.main()
