#=============================================================================
#
#  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.
#
#=============================================================================
"""
Root writer script for Truchas workflows
"""
import os
print 'loading', os.path.basename(__file__)
import sys

import smtk
if 'pybind11' == smtk.wrappingProtocol():
    #print 'Using pybind11 bindings'
    import smtk.simulation

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

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

# Please order the namelists alphabetically in this table
card = writer.CardFormat
format_table = {
  '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', use_model_entities_for_value=True),
    card('material_number', is_custom=True),
    card('temperature', item_path='temperature',
      expression_keyword='temperature_function'),
    card('velocity', 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')
  ],
  '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='enable/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', literal_value='tabular')
  ],
  '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='material-type/void-temperature',
      if_condition=VOID_MATERIAL),
    card('sound_speed', item_path='material-type/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='material-type/density', as_property=True),
    card('specific heat', item_path='material-type/thermal/specific-heat',
      if_condition='fluid', as_property=True),
    card('specific heat', item_path='material-type/thermal/specific-heat',
      if_condition='solid', as_property=True),
    card('specific heat',
      item_path='material-type/thermal-two-phase/specific-heat',
      if_condition='two-phase',
      as_property=True),
    card('conductivity', item_path='material-type/conductivity', as_property=True),
    card('viscosity', item_path='material-type/viscosity', \
      if_condition=[VISCOUS_FLOW, FLUID_PHASE], as_property=True),
    card('density deviation', item_path='material-type/density-deviation', \
      if_condition=FLUID_PHASE, as_property=True)
  ],
  'PHYSICAL_CONSTANTS': [
    card('absolute_zero', item_path='absolute-zero'),
    card('stefan_boltzmann', item_path='stefan-boltzmann')
  ],
  'PHYSICS': [
    card('fluid_flow', use_condition_for_boolean=FLOW_ANALYSIS),
    card('heat_transfer', use_condition_for_boolean=THERMAL_ANALYSIS),
    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='enable/enclosure',
    if_condition=ENCLOSURE_RADIATION,
    custom_method='_write_enclosure'),
  namelist('ENCLOSURE_RADIATION',
    att_type='enclosure-radiation',
    base_item_path='enable/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', att_type='solver'),
  # Note: internally, the materials writer includes PHASE and MATERIAL_SYSTEM
  namelist('MATERIAL', custom_method='_write_materials'),
  namelist('DIFFUSION_SOLVER',
    if_condition=ONLY_THERMAL_ANALYSIS,
    att_type='solver',
    base_item_path='analysis/thermal'),
  namelist('DIFFUSION_SOLVER',
    if_condition=[THERMAL_ANALYSIS, FLOW_ANALYSIS],
    att_type='solver',
    base_item_path='analysis/thermal-plus-fluid/thermal-solver'),
  namelist('CHAPARRAL', att_type='enclosure-radiation',
    base_item_path='enable/chaparral',
    if_condition=ENCLOSURE_RADIATION),
  namelist('NUMERICS', att_type='solver'),
  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(
    'DS_BOUNDARY_CONDITION',
    att_type='thermal-surface-condition',
    base_item_path='type/ds-boundary-condition',
    custom_method='_write_thermal_condition'),
  namelist(
    'DS_INTERFACE_CONDITION',
    att_type='thermal-surface-condition',
    base_item_path='type/ds-interface-condition',
    custom_method='_write_thermal_condition'),
  namelist('BC',
    att_type='boundary-condition', custom_method='_write_bc'),
  namelist('BODY', att_type='body', custom_method='_write_body'),
  namelist('DS_SOURCE', att_type='ds-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')
]


# ---------------------------------------------------------------------
def ExportCMB(spec):
  '''
  Entry function, called by CMB to write export file
  '''
  logger = spec.getLogger()

  # Get export attributes
  export_spec_att = None
  export_atts = spec.getExportAttributes()
  if export_atts is not None:
    att_list = export_atts.findAttributes('ExportSpec')
    if att_list:
      export_spec_att = att_list[0]

  if export_spec_att is None:
    msg = 'No ExportSpec instance -- cannot export'
    print 'WARNING:', msg
    logger.addError(msg)
    return False

  # Get project name
  project_name_item = export_spec_att.findString('ProjectName')
  project_name = project_name_item.value(0)
  if not project_name:
    project_name = 'truchas'
    msg = 'No Project Name specified; using \"%s\"' % project_name
    print 'WARNING:', msg
    logger.addWarning(msg)
  print 'project_name', project_name
  project_name = project_name

  # Initialize project path (folder)
  project_path = None
  item = export_spec_att.find('ProjectDir')
  if item is not None:
    if 'shiboken' == smtk.wrappingProtocol():
      dir_item = smtk.attribute.to_concrete(item)
    else:
      dir_item = item
    project_path = dir_item.value(0)
    print 'project_path', project_path

  if not project_path:
    project_path = os.getcwd()
    msg = 'No project_path specified -- using', project_path
    print 'ERROR:', msg
    logger.addWarning(msg)
    return False

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

  # Initialize writer object
  output_filename = '%s.inp' % project_name
  output_path = os.path.join(project_path, output_filename)
  truchas_writer = writer.Writer(spec)
  completed = truchas_writer.write(output_path, namelist_sequence, format_table)
  print 'Writer completion status %s' % completed
  sys.stdout.flush()
  return completed
