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

"""Top level writer for PyARC son files
"""

import datetime
import os
from shutil import copyfile
import string

import rggsession

import smtk
import smtk.model

from .nemesh_writer import *

class ProteusSNWriter:
    '''Top level writer for Proteus-SN (.inp) files
    '''
    def __init__(self):
        self.scope = None

    def write(self, scope):
        ''''''
        self.scope = scope
        self.tab = '    '

        self._write_driver()
        self._write_material_assignment()
        self._write_nemesh()
        return True

    def _write_driver(self):
        """Write driver input file.

        The driver is represented by a single/instanced attribute
        """
        driver_att_type = 'driver_input'
        att_list = self.scope.sim_atts.findAttributes(driver_att_type)
        if not att_list:
            raise RuntimeError('No {} attribute'.format(driver_att_type))
        elif len(att_list) > 1:
            tpl = 'Only 1 {} attribute allowed, found {}'
            raise RuntimeError(tpl.format(driver_att_type, len(att_list)))
        driver_att = att_list[0]

        # Copy cross section file into output directory
        crosssection_file = driver_att.find('SOURCEFILE_XS').value()
        if os.path.isfile(crosssection_file):
            local_crosssection_file = os.path.join(self.scope.output_path, 'sn2nd.isotxs')
            print('copying %s to %s' % (crosssection_file, local_crosssection_file))
            copyfile(crosssection_file, local_crosssection_file)
            driver_att.find('SOURCEFILE_XS').setValue(local_crosssection_file)

        # Open output file
        driver_filename = 'sn2nd.inp'
        output_file = os.path.join(self.scope.output_path, driver_filename)
        complete = False
        with open(output_file, 'w') as out:
            out.write('! Generated by CMB {}\n\n'.format(
                datetime.datetime.now().strftime('%d-%b-%Y  %H:%M')))

            # Write items in sequence
            num_items = driver_att.numberOfItems()
            for i in range(num_items):
                item = driver_att.item(i)
                self._write_driver_item(item, out)

            complete = True

        return complete

    def _write_driver_item(self, item, output_stream):
        if item.type() == smtk.attribute.Item.GroupType:
            # Note that all groups have only 1 subgroup
            output_stream.write('\n')
            num_items = item.numberOfItemsPerGroup()
            for i in range(num_items):
                child_item = item.item(i)
                self._write_driver_item(child_item, output_stream)
            return

        value = None
        if item.type() == smtk.attribute.Item.VoidType:
            value = 'YES' if item.isEnabled() else 'NO'
        elif item.name() == 'MESHSIMULATE':   # Special case
            value = item.value() if item.isEnabled() else 0
        elif item.type() == smtk.attribute.Item.FileType and item.isSet():
            value = os.path.relpath(item.value(), self.scope.output_path)
        elif hasattr(item, 'value') and item.isSet():
            value = item.value()
        if value is None:
            output_stream.write('! {}  not specified\n'.format(item.name()))
        else:
            tab = 35
            text_formatter = '{:<%s}{:}\n' % tab
            if item.type() == smtk.attribute.Item.DoubleType:
                text_formatter = '{:<%s}{:e}\n' % tab
            line = text_formatter.format(item.name(), value)
            output_stream.write(line)

    def _write_material_assignment(self):
        """Write material assignment input file.

        The material assignment is translated from the material descriptions of
        the model entities onto which the mesh element blocks are classified
        """
        core_group = smtk.model.Group(self.scope.model_resource.findEntitiesByProperty(
            'rggType', smtk.session.rgg.Core.typeDescription)[0])
        core = smtk.session.rgg.Core(core_group)
        assembly_groups = self.scope.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]

        materials_str = smtk.model.Model(self.scope.model_entity).stringProperty(
            smtk.session.rgg.Material.label)
        materials = [smtk.session.rgg.Material(material_str) for material_str in materials_str]

        materials_by_name = {}
        for material in materials:
            materials_by_name[material.name()] = material

        material_names = smtk.model.Model(self.scope.model_entity).stringProperty("materials")

        # Traverse the model in the same order as in ExportInp.cxx:71 to ensure
        # that mesh element blocks appear in the same order as materials.
        #
        # TODO: Should each model element with a uniquely defined material have
        #       its own AuxiliaryGeometry representation onto which we can
        #       assign the material? That way, we could avoid the need to
        #       preserve traversal order (and it would be much easier to
        #       generalize the export process to other classified meshes).
        def blockMaterials():
            for assembly in assemblies:
                duct_aux = smtk.model.AuxiliaryGeometry(
                  self.scope.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 material_names[layer[0]]
                pins_aux = [smtk.model.AuxiliaryGeometry(self.scope.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 material_names[layer_material.subMaterialIndex]
                    if pin.cellMaterialIndex() != 0:
                        yield material_names[pin.cellMaterialIndex()]
            # TODO: yield appropriate material for background mesh
            yield "Mat_Coolant"


        # Open output file
        material_filename = 'sn2nd.assignment'
        output_file = os.path.join(self.scope.output_path, material_filename)
        complete = False
        with open(output_file, 'w') as out:
            out.write('! Generated by CMB {}\n\n'.format(
                datetime.datetime.now().strftime('%d-%b-%Y  %H:%M')))

            for material in materials:
                out.write('! {} definition\n'.format(material.name()))
                for i in range(material.numberOfComponents()):
                    out.write('MATERIAL_DEF\t{}\t{}\t{}\n'.format(
                      material.name(), material.component(i), material.numberDensity(i)))
                out.write('\n')

            densities = ''

            out.write('! region assignments\n')
            meshsets = [self.scope.mesh_resource.meshes(domain) \
                          for domain in self.scope.mesh_resource.domains()]
            for (material, domain, meshset) in zip(
              blockMaterials(), self.scope.mesh_resource.domains(), meshsets):
                entity_id = meshset.modelEntityIds()[0]
                entity = self.scope.model_resource.find(entity_id)
                out.write('REGION_ALIAS\tREGION_{:09d}\t{:}\n'.format(
                  domain.value(), material))

                # Access material by name, but fall back to a valid material
                # (a workaround for background mesh material that may not be
                # defined; see above TODO)
                m = materials_by_name.get(material, materials_by_name.keys()[0])
                densities += 'REGION_PROPERTY\tREGION_{:09d}\tDensity(g/cc)\t{}\n'.format(
                  domain.value(), m.massDensity())


            out.write('\n! region properties\n' + densities)

            complete = True

        return complete

    def _write_nemesh(self):
        """Write mesh input file.
        """
        # Open output file
        mesh_filename = 'sn2nd.nemesh'
        output_file = os.path.join(self.scope.output_path, mesh_filename)
        write_mesh = nemesh_writer()
        return write_mesh.write(output_file, self.scope.mesh_resource, self.scope.model_entity)
