#=============================================================================
#
#  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.
#
#=============================================================================
"""
Export operator for Truchas workflows
"""
from __future__ import print_function
import os
print('loading', os.path.basename(__file__))
import sys
sys.dont_write_bytecode

import smtk
import smtk.attribute
import smtk.io
import smtk.operation
import smtk.session.mesh

# Add the directory containing this file to the python module search list
import inspect
sys.path.insert(0, os.path.dirname(os.path.abspath(inspect.getfile(inspect.currentframe()))))

import internal
from internal import writer
reload(writer)  # for development

class Export(smtk.operation.Operation):

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

    def name(self):
        return "Export Truchas"

    def operateInternal(self):
        try:
            success = ExportCMB(self)
        except:
            print('Error', self.log().convertToString())
            #smtk.ErrorMessage(self.log(), sys.exc_info()[0])
            raise
            return self.createResult(smtk.operation.Operation.Outcome.FAILED)

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

    def createSpecification(self):
        spec = self.createBaseSpecification()
        print('spec:', spec)

        # Load export atts
        source_dir = os.path.abspath(os.path.dirname(__file__))
        print('source_dir:', source_dir)
        sbt_path = os.path.join(source_dir, 'internal', 'truchas-export.sbt')
        print('sbt_path:', sbt_path)
        reader = smtk.io.AttributeReader()
        result = reader.read(spec, sbt_path, self.log())
        print('reader result:', result)

        # Setup result definition
        resultDef = spec.createDefinition('test result', 'result')
        successDef = smtk.attribute.IntItemDefinition.New('success')
        resultDef.addItemDefinition(successDef)

        return spec


# ---------------------------------------------------------------------
#
# Dictionary of formatters for output lines
# Arguments are: (keyword, attribute type, item path, **kwargs)
#
# ---------------------------------------------------------------------

# Predefined strings for "if_condition" arguments
from internal.writer.writer import THERMAL_ANALYSIS, ONLY_THERMAL_ANALYSIS, FLOW_ANALYSIS, \
  VISCOUS_FLOW, INVISCID_FLOW, FLUID_PHASE, MASS_LIMITER, BC_INFLOW, \
  VOID_MATERIAL, ENCLOSURE_RADIATION, MOVING_RADIATION, INDUCTION_HEATING

# Please order the namelists alphabetically in this table
card = writer.CardFormat
format_table = {
  'ALTMESH': [
    card('altmesh_file', item_path='model', is_custom=True),
    card('altmesh_coordinate_scale_factor', item_path='coordinate-scale-factor'),
  ],
  'BC': [
    card('bc_name', use_name_for_value=True),
    card('bc_variable', item_path='variable'),
    card('bc_type', is_custom=True),
    card('surface_name',literal_value='from mesh file'),
    card('mesh_surface', use_model_entities_for_value=True),
    card('bc_value', is_custom=True, expression_keyword='bc_table'),
    card('inflow_material', item_path='variable/inflow/inflow-material',
      if_condition=BC_INFLOW),
    card('inflow_temperature', item_path='variable/inflow/inflow-temperature',
      if_condition=BC_INFLOW)
  ],
  'BODY': [
    card('surface_name', literal_value='from mesh file'),
    card('mesh_material_number'),
    card('material_number', is_custom=True),
    card('temperature', att_type='TemperatureInitialCondition', item_path='temperature',
      expression_keyword='temperature_function'),
    card('velocity', att_type='VelocityInitialCondition', item_path='velocity')
  ],
  # 'CHAPARRAL': [
  #   card('blocking_enclosure', item_path='blocking-enclosure'),
  #   card(None, item_path='partial-enclosure-area',
  #     set_condition='partial-enclosure'),
  #   card('partial_enclosure', item_path='partial-enclosure-area',
  #     use_condition_for_boolean='partial-enclosure'),
  #   card('partial_area', item_path='partial-enclosure-area',
  #     if_condition='partial-enclosure'),
  #   card('bsp_max_tree_depth', item_path='bsp-max-tree-depth'),
  #   card('bsp_min_leaf_length', item_path='bsp-min-leaf-length'),
  #   card('spatial_tolerance', item_path='spatial-tolerance'),
  #   card('hemicube_resolution', item_path='hemicube-resolution'),
  #   card('min_separation', item_path='min-separation'),
  #   card('max_subdivisions', item_path='max-subdivisions'),
  #   card('smoothing_tolerance', item_path='smoothing-tolerance'),
  #   card('smoothing_max_iter', item_path='smoothing-max-iter'),
  #   card('smoothing_weigth', item_path='smoothing-weight'),
  #   card('verbosity_level', item_path='verbosity-level')
  # ],
  'DIFFUSION_SOLVER': [
    card('stepping_method', literal_value='Adaptive BDF2',
      if_condition=ONLY_THERMAL_ANALYSIS),
    card('stepping_method',literal_value='Non-adaptive BDF1',
      if_condition=FLOW_ANALYSIS),

    card('abs_enthalpy_tol', item_path='abs-enthalpy-tol',
      if_condition=ONLY_THERMAL_ANALYSIS),
    card('rel_enthalpy_tol', item_path='rel-enthalpy-tol',
      if_condition=ONLY_THERMAL_ANALYSIS),
    card('abs_temp_tol', item_path='abs-temperature-tol',
      if_condition=ONLY_THERMAL_ANALYSIS),
    card('rel_temp_tol', item_path='rel-temperature-tol',
      if_condition=ONLY_THERMAL_ANALYSIS),

    card('residual_atol', item_path='residual-atol',
      if_condition=FLOW_ANALYSIS),
    card('residual_rtol', item_path='residual-rtol',
      if_condition=FLOW_ANALYSIS),
    card('nlk_preconditioner', item_path='nlk-preconditioner',
      literal_value='Hypre_AMG'),
    card('max_nlk_itr', item_path='max-nlk-itr'),
    card('max_nlk_vec', item_path='max-nlk-vec'),
    card('nlk_tol', item_path='nlk-tol',
        if_condition=ONLY_THERMAL_ANALYSIS),
    card('nlk_vec_tol', item_path='nlk-vec-tol'),
    card('verbose_stepping', item_path='verbose-stepping'),
    card('pc_amg_cycles', item_path='pc-amg-cycles'),
    card('cond_vfrac_threshold', item_path='cond-vfrac-threshold',
      if_condition=FLOW_ANALYSIS)
  ],
  'DS_BOUNDARY_CONDIATION': [],
  'DS_INTERFACE_CONDITION': [],
  'DS_SOURCE': [
    card('equation', literal_value='temperature'),
    card('cell_set_ids', use_model_entities_for_value=True),
    card('source_constant', item_path='source',
      expression_keyword='source_function')
  ],
  'ELECTROMAGNETICS': [
    card('em_domain_type', item_path='em-domain-type'),
    card('symmetry_axis', item_path='symmetry-axis'),
    card('source_frequency', att_type='em-source', item_path='source-type/constant-source-frequency'),
    card('uniform_source', att_type='em-source', item_path='source-type/constant-uniform-source'),
    card('cg_stopping_tolerance', item_path='cg-stopping-tolerance'),
    card('maximum_cg_iterations', item_path='maximum-cg-iterations'),
    card('num_etasq', item_path='num-etasq'),
    card('steps_per_cycle', item_path='steps-per-cycle'),
    card('ss_stopping_tolerance', item_path='ss-stopping-tolerance'),
    card('maximum_source_cycles', item_path='maximum-source-cycles'),
    card('output_level', item_path='output-level')
  ],
  # 'ENCLOSURE': [
  #   card('name', item_path='enclosure-name'),
  #   card('mesh_file', is_custom=True),
  #   card('coord_scale_factor', att_type="mesh", item_path='coordinate-scale-factor'),
  #   card('symmetries', item_path='symmetries', is_custom=True),
  #   card('side_set_ids', is_custom=True),
  #   card('moving_enclosure', use_condition_for_boolean=MOVING_RADIATION),
  #   card('displacement_set_ids', is_custom=True),
  #   card('displacement_sequence', item_path='moving-radiation/displacement-sequence',
  #     is_custom=True),
  #   card('enclosure_file_sequence_prefix',
  #     item_path='moving-radiation/enclosure-file-prefix',
  #     if_condition=MOVING_RADIATION),
  #   card('ignore_block_ids',
  #     item_path='ignore-block-ids',
  #     use_model_entities_for_value=True)
  # ],
  # 'ENCLOSURE_RADIATION': [
  #   card('name', item_path='enclosure-name'),
  #   card('ambient_constant', item_path='ambient-temperature',
  #     expression_keyword='ambient_function'),
  #   card('error_tolerance', item_path='error-tolerance'),
  #   card('moving_enclosure', use_condition_for_boolean=MOVING_RADIATION),
  #   card('linear_interpolation', item_path='moving-radiation/linear-interpolation',
  #     if_condition=MOVING_RADIATION),
  #   card('skip_geometry_check', use_condition_for_boolean=MOVING_RADIATION,
  #     if_condition=MOVING_RADIATION),
  #   card('enclosure_file_sequence_prefix',
  #     item_path='moving-radiation/enclosure-file-prefix',
  #     if_condition=MOVING_RADIATION),
  #   card('time_sequence', item_path='moving-radiation/displacement-sequence',
  #     is_custom=True)
  # ],
  # 'ENCLOSURE_SURFACE': [
  #   card('name', use_name_for_value=True),
  #   card('enclosure_name', att_type='enclosure-radiation',
  #     item_path='enclosure/enclosure-name'),
  #   card('face_block_ids', use_model_entities_for_value=True),
  #   card('emissivity', item_path='emissivity')
  # ],
  'FUNCTION': [
    card('name', use_name_for_value=True),
    card('type', item_path='type')
  ],
  'INDUCTION_COIL': [
    card('nturns', item_path='type/nturns'),
    card('center', item_path='type/center'),
    card('radius', item_path='type/radius'),
    card('length', item_path='type/length'),
    card('current', item_path='current')
  ],
  'LINEAR_SOLVER': [
    card('name', item_path='', use_name_for_value=True),
    card('method', item_path='method'),
    card('preconditioning_method', item_path='preconditioning-method'),
    card('preconditioning_steps', item_path='preconditioning-steps'),
    card('relaxation_parameter', item_path='relaxation-parameter'),
    card('stopping_criterion', item_path='stopping-criterion'),
    card('convergence_criterion', item_path='convergence-criterion'),
    card('krylov_vectors', item_path='krylov-vectors'),
    card('output_mode', item_path='output-mode')
  ],
  'MATERIAL': [
    # card('density', item_path='material-type/density'),
    card('void_temperature', item_path='void-temperature',
      if_condition=VOID_MATERIAL),
    card('sound_speed', item_path='sound-speed',
      if_condition=VOID_MATERIAL)
  ],
  'MATERIAL_SYSTEM': [
    card('transition_temps_low',
      item_path='material-type/solid-transition-temperature'),
    card('transition_temps_high',
      item_path='material-type/liquid-transition-temperature'),
    card('latent_heat', item_path='material-type/thermal-two-phase/latent-heat')
  ],
  'MESH': [
    card('mesh_file', is_custom=True),
    card('coordinate_scale_factor', item_path='coordinate-scale-factor'),
    card('exodus_block_modulus', item_path='exodus-block-modulus'),
    card('interface_side_sets', is_custom=True)
  ],
  'NUMERICS': [
    card('dt_init', att_type='numerics', item_path='dt_init'),
    card('dt_grow', att_type='numerics', item_path='dt_grow'),
    card('dt_max', att_type='numerics', item_path='dt_max'),
    card('dt_min', att_type='numerics', item_path='dt_min'),
    card('t', att_type='outputs', item_path='start-time'),
    card('discrete_ops_type', item_path='analysis/fluid/flow-numerics/discrete-ops-type',
      if_condition=FLOW_ANALYSIS),
    card('projection_linear_solution', item_path='analysis/fluid/projection-linear-solver',
      use_name_for_value=True, if_condition=FLOW_ANALYSIS),
    card('viscous_implictness', item_path='analysis/fluid/viscous-flow-model/viscous-implicitness',
      if_condition=VISCOUS_FLOW),
    card('viscous_linear_solution',
      item_path='analysis/fluid/viscous-flow-model/viscous-linear-solver',
      use_name_for_value=True, if_condition=VISCOUS_FLOW),
    card('viscous_number', item_path='analysis/fluid/flow-numerics/viscous-number',
      if_condition=FLOW_ANALYSIS),
    card('volume_track_subcycles', item_path='analysis/fluid/flow-numerics/volume-track-subcycles',
      if_condition=FLOW_ANALYSIS),
    card('body_force_implicitness', item_path='analysis/fluid/flow-numerics/body-force-implicitness',
      if_condition=FLOW_ANALYSIS),
    card('mass_limiter', item_path='analysis/fluid/flow-numerics/mass-limiter',
      if_condition=FLOW_ANALYSIS, as_boolean=True, set_condition=MASS_LIMITER),
    # Note that mass_limiter_cutoff must come AFTER mass_limiter
    card('mass_limiter_cutoff', item_path='analysis/fluid/flow-numerics/mass-limiter/mass-limiter-cutoff',
      if_condition=MASS_LIMITER),
    card('mechanical_energy_bound', item_path='analysis/fluid/flow-numerics/mechanical-energy-bound',
      if_condition=FLOW_ANALYSIS),
    card('momentum_solidify_implicitness', item_path='analysis/fluid/flow-numerics/momentum-solidify-implicitness',
      if_condition=FLOW_ANALYSIS)
  ],
  'OUTPUTS': [
    # Writer has custom code for output time lists
    card('output_t', item_path='output-times/time'),
    card('output_dt', item_path='output-times/time')
  ],
  'PHASE': [
    card('density', item_path='density', as_property=True),
    card('density-deviation', item_path='density-deviation', as_property=True),
    card('specific heat', item_path='specific-heat', as_property=True),
    card('conductivity', item_path='conductivity', as_property=True),
    card('viscosity', item_path='fluid/viscosity', if_condition=FLUID_PHASE, as_property=True),
    card('electrical conductivity', item_path='electrical-properties/electrical-conductivity',
      if_condition=INDUCTION_HEATING,as_property=True),
    card('electrical susceptibility', item_path='electrical-properties/electrical-susceptibility',
      if_condition=INDUCTION_HEATING,as_property=True),
    card('magnetic susceptibility', item_path='electrical-properties/magnetic-susceptibility',
      if_condition=INDUCTION_HEATING,as_property=True)
  ],
  'PHYSICAL_CONSTANTS': [
    card('absolute_zero', item_path='absolute-zero'),
    card('stefan_boltzmann', item_path='stefan-boltzmann')
  ],
  'PHYSICS': [
    card('flow', use_condition_for_boolean=FLOW_ANALYSIS),
    card('heat_transport', use_condition_for_boolean=THERMAL_ANALYSIS),
    card('electromagnetics', use_condition_for_boolean=INDUCTION_HEATING),
    card('body_force', att_type='physics', item_path='fluid-body-force',
      if_condition=FLOW_ANALYSIS)
    # card('inviscid', if_condition=FLOW_ANALYSIS,
    #   use_condition_for_boolean=INVISCID_FLOW)
  ],
  'PROBE': [
    card('probe_name', use_name_for_value=True),
    card('probe_description', item_path='description'),
    card('probe_coords', item_path='coords')
  ],
  'SIMULATION_CONTROL': [
    card('phase_start_times', item_path='simulation-control/phase-start-times'),
    card('phase_init_dt_factor', item_path='simulation-control/phase-init-dt-factor')
  ]
}

# This list sets the order that namelists are written to the Truchas input file
namelist = writer.Namelist
namelist_sequence = [
  namelist('MESH', att_type='mesh', custom_method='_write_mesh'),
  namelist('OUTPUTS', att_type='outputs', custom_method='_write_outputs'),
  # namelist('ENCLOSURE', att_type='enclosure-radiation',
  #   base_item_path='enclosure',
  #   if_condition=ENCLOSURE_RADIATION,
  #   custom_method='_write_enclosure'),
  # namelist('ENCLOSURE_RADIATION',
  #   att_type='enclosure-radiation',
  #   base_item_path='enclosure',
  #   if_condition=ENCLOSURE_RADIATION,
  #   custom_method='_write_enclosure_radiation'),
  # namelist('ENCLOSURE_SURFACE',
  #   att_type='thermal-surface-condition',
  #   base_item_path='type',
  #   if_condition=ENCLOSURE_RADIATION,
  #   custom_method='_write_thermal_condition'),
  namelist('PHYSICS', custom_method='_write_physics'),
  # Note: internally, the materials writer includes PHASE and MATERIAL_SYSTEM
  namelist('MATERIAL', custom_method='_write_materials', separator='MATERIALS'),
  namelist('FUNCTION', att_type='fn.material', custom_method='_write_function'),
  namelist('DIFFUSION_SOLVER',
    if_condition=ONLY_THERMAL_ANALYSIS,
    att_type='ht.solver',
    base_item_path='analysis/thermal'),
  namelist('DIFFUSION_SOLVER',
    if_condition=[THERMAL_ANALYSIS, FLOW_ANALYSIS],
    att_type='ht.solver',
    base_item_path='analysis/thermal-plus-fluid/thermal-solver'),
  # namelist('CHAPARRAL', att_type='enclosure-radiation',
  #   base_item_path='chaparral',
  #   if_condition=ENCLOSURE_RADIATION),
  namelist('NUMERICS', att_type='numerics'),
  # namelist('LINEAR_SOLVER',
  #   if_condition=FLOW_ANALYSIS,
  #   att_type='solver',
  #   base_item_path='analysis/fluid/projection-linear-solver'),
  # namelist('LINEAR_SOLVER',
  #   if_condition=VISCOUS_FLOW,
  #   att_type='solver',
  #   base_item_path='analysis/fluid/viscous-flow-model/viscous-linear-solver'),
  namelist('ALTMESH', if_condition=INDUCTION_HEATING, att_type='electromagnetics',
    custom_method='_write_mesh', separator='INDUCTION HEATING'),
  namelist('ELECTROMAGNETICS', if_condition=INDUCTION_HEATING, att_type='electromagnetics'),
  namelist('INDUCTION_COIL', if_condition=INDUCTION_HEATING, att_type='induction-coil'),
  namelist(
    'DS_BOUNDARY_CONDITION',
    att_type='ht.boundary',
    custom_method='_write_thermal_condition',
    separator='BOUNDARY CONDITIONS'),
  namelist(
    'DS_INTERFACE_CONDITION',
    att_type='ht.interface',
    custom_method='_write_thermal_condition',
    separator='INTERFACE CONDITIONS'),
  namelist('BC',
    att_type='boundary-condition', custom_method='_write_bc'),
  namelist('BODY', att_type='phase', custom_method='_write_body', separator='BODIES'),
  namelist('FUNCTION', att_type='fn.initial-condition', custom_method='_write_function'),
  namelist('DS_SOURCE', att_type='ht.source'),
  namelist('PROBE', att_type='probe'),
  namelist(
    'SIMULATION_CONTROL',
    att_type='simulation-control',
    custom_method='_write_simcontrol'),
  # namelist('FUNCTION', att_type='tabular-function', custom_method='_write_function'),
  namelist('PHYSICAL_CONSTANTS', att_type='physical-constants')
]

test_sequence = [
  namelist('MATERIAL', custom_method='_write_materials', separator='MATERIALS'),
  namelist('FUNCTION', att_type='fn.material', custom_method='_write_function')
]



# ---------------------------------------------------------------------
def ExportCMB(export_op):
  '''
  Entry function, called by CMB to write export file
  '''

  operator_spec = export_op.parameters()
  logger = export_op.log()

  # Get output filename
  output_file_item = operator_spec.findFile('output-file')
  output_file = output_file_item.value(0)
  output_dir = os.path.dirname(output_file)

  # Create output folder if needed
  if not os.path.exists(output_dir):
    os.makedirs(output_dir)

  # Get mesh filename
  mesh_file_item = operator_spec.findFile('mesh-file')
  mesh_file = mesh_file_item.value(0)

  altmesh_file = 'Not Specified'
  altmesh_file_item = operator_spec.findFile('alt-mesh-file')
  if altmesh_file_item and altmesh_file_item.isEnabled():
    altmesh_file = altmesh_file_item.value(0)

  # Initialize writer
  truchas_writer = writer.Writer(operator_spec, logger,
    mesh_filename=mesh_file, altmesh_filename=altmesh_file)
  completed = truchas_writer.write(output_file, namelist_sequence, format_table)
  # completed = truchas_writer.write(output_file, test_sequence, format_table)
  print('Writer completion status %s' % completed)
  sys.stdout.flush()
  return completed
