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

""" generate_cross_sections.py:

Generate a cross sections file using mcc3.

"""
import sys

from . import generate_cross_sections_xml

import datetime
import json
import os
import rggsession
import shutil
import smtk
import smtk.attribute
import smtk.io
import smtk.mesh
import smtk.model
import smtk.operation
import subprocess
import sys
import tempfile

class GenerateCrossSections(smtk.operation.Operation):

    def __init__(self):
        smtk.operation.Operation.__init__(self)
        self.endf_version = 7

    def name(self):
        return "generate cross sections"

    def is_exe(self, fpath):
        return os.path.isfile(fpath) and os.access(fpath, os.X_OK)

    def rm_dir(self, dpath):
        try:
            shutil.rmtree(dpath)
        except OSError as e:
            smtk.WarningMessage(self.log(), 'Could not remove directory \'%s\'' % dpath)

    def which(self, executable):
        """ Find an executable using the following queues (in order):
        1) check user-defined inputs
        2) check relative to current executable
        3) check relative to PATH
        """
        # 1) check user-defined inputs
        file_att = self.parameters().findFile(executable)
        if file_att and file_att.isEnabled():
            exe_file = file_att.value()
            if self.is_exe(exe_file):
                return exe_file

        # 2) check relative to current executable
        current_exe = sys.executable
        fpath, fname = os.path.split(current_exe)

        exe_file = os.path.join(fpath, '../bin', executable)
        if self.is_exe(exe_file):
            return exe_file

        # 3) check relative to PATH
        for path in os.environ["PATH"].split(os.pathsep):
            exe_file = os.path.join(path, executable)
            if self.is_exe(exe_file):
                return exe_file

        return None

    def ableToOperate(self):
        if not smtk.operation.Operation.ableToOperate(self):
            return False

        # Ensure that the necessary executables can be found
        return self.which('mcc3.x') != None or self.is_exe(self.pyarc.user_object.mcc3_executable)

    def operateInternal(self):
        smtk.InfoMessage(self.log(), 'generate cross sections')

        # Access the input model
        model = smtk.model.Entity.CastTo(self.parameters().associations().value(0))

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

        # Access assygen, coregen and cubit
        mcc3 = self.which('mcc3.x')

        mcclibdir = None
        pwlibdir = None

        if mcc3:
            smtk.InfoMessage(self.log(), 'mcc3 executable is %s' % mcc3)
            mcclibdir = os.path.join(os.path.dirname(mcc3), 'lib/lib.mcc.e70')
            pwlibdir = os.path.join(os.path.dirname(mcc3), 'lib/lib.pw.200.e70')

        if not os.path.exists(mcclibdir):
            smtk.WarningMessage(self.log(), 'could not locate mcc lib dir')
            return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        if not os.path.exists(pwlibdir):
            smtk.WarningMessage(self.log(), 'could not locate pw lib dir')
            return self.createResult(smtk.operation.Operation.Outcome.FAILED)

        # Access the output file name
        filename = self.parameters().find('filename').value()

        # Access the MCC3 input attributes
        mcc3Att = smtk.attribute.Attribute.CastTo(
            self.parameters().find('mcc3_atts').value(0))

        # Construct a temporary scratch space for generating input files
        temp_dir = tempfile.mkdtemp()

        temp_inputfile = os.path.join(temp_dir, "input")
        success = False
        with open(temp_inputfile, 'w') as inp:
            inp.write('! Generated by CMB {}\n\n'.format(
                datetime.datetime.now().strftime('%d-%b-%Y  %H:%M')))

            # library block
            inp.write('$library\n')
            inp.write('c_mcclibdir\t="{}"\n'.format(mcclibdir))
            inp.write('c_pwlibdir\t="{}","."\n'.format(pwlibdir))
            inp.write('/\n')

            # control block
            inp.write('$control\n')
            inp.write('c_group_structure\t={}\n'.format(mcc3Att.find('egroupname').value()))
            inp.write('i_number_region\t={}\n'.format(len(materials)))
            inp.write('c_geometry_type\t=mixture\n')
            inp.write('i_scattering_order\t={}\n'.format(mcc3Att.find('scattering_order').value()))
            inp.write('c_inelastic_treatment\t={}\n'.format(
              mcc3Att.find('inelastic_treatment').value()))
            inp.write('l_buckling_search\t=F\n')
            inp.write('/\n')

            # material block
            inp.write('$material\n')
            material_index = 1
            isotope_names = {}
            for material in materials:
                inp.write('c_composition_name({})= {}\n'.format(
                  material_index, material.name().upper()[:8]))
                inp.write('t_composition(:,{})=\n'.format(material_index))
                for i in range(material.numberOfComponents()):
                    name = material.component(i)[:6]
                    isotope_names[name] = isotope_names.get(name, 0) + 1
                    if isotope_names[name] == 1:
                      unique_name = material.component(i)[:6]
                    else:
                      unique_name = name[:5] + str(isotope_names[name])
                    inp.write('{}\t{}\t{}\t{}\n'.format(
                      material.component(i).upper().ljust(5, '_') + str(self.endf_version),
                      unique_name, material.numberDensity(i), material.temperature()))
                material_index = material_index + 1
            inp.write('/\n')

            # output block
            inp.write('$output\n')
            inp.write('l_edit_leakage\t=T\n')
            inp.write('/\n')
            success = True

        if not success:
            smtk.WarningMessage(self.log(), 'could not generate mcc3 input file')
            self.rm_dir(temp_dir)
            return self.createResult(smtk.operation.Operation.Outcome.FAILED)
        else:
            smtk.InfoMessage(self.log(), 'generated mcc3 input file')

        task = subprocess.Popen([mcc3], cwd=temp_dir, stdout=subprocess.PIPE,
                                  universal_newlines=True)
        for stdout_line in iter(task.stdout.readline, ""):
            print(stdout_line)
            smtk.InfoMessage(self.log(), stdout_line)
        task.stdout.close()
        task.wait()

        # Copy the results to the user-specified output file
        from shutil import copyfile
        copyfile(os.path.join(temp_dir, "ISOTXS.merged"), filename)

        # Clean up the temporary work directory
        self.rm_dir(temp_dir)

        # Return with success
        result = self.createResult(smtk.operation.Operation.Outcome.SUCCEEDED)
        return result

    def createSpecification(self):
        spec = smtk.attribute.Resource.create()
        reader = smtk.io.AttributeReader()
        reader.readContents(spec, generate_cross_sections_xml.description, self.log())
        return spec
