#=============================================================================
#
#  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 script for Ground Water applications
"""

import sys
import smtk

# so devs don't complain about pyc files in source
sys.dont_write_bytecode = True

# Check if adhcommon already loaded, and either import or reload
adh = sys.modules.get('adhcommon')
if adh is None:
  import adhcommon as adh
else:
  reload(adh)

# ---------------------------------------------------------------------
#
# Define placeholder/passthrough functions for custom writer functions
# Actual implementations, prefixed by "_", are further below
#
# ---------------------------------------------------------------------
def write_outputinterval(scope, item, card_format, context_id):
  return _write_outputinterval(scope, item, card_format, context_id)


# ---------------------------------------------------------------------
#
# Dictionary of formatters for each output card
#
# There are 3 card format types: val, idval, bc
# Arguments are (item name, opcode, comment=None, subitems=None)
# These might (should?) be migrated to adhcommon.py
#
# ---------------------------------------------------------------------
fmt = adh.CardFormat
format_table = {
  'Solvers': [
    fmt.val('MaxNonLinearIters', 'IP NIT', comment='! Maximum number of non-linear iterations'),
    fmt.val('NonLinearTolMaxNorm', 'IP NTL', comment='! Non-Linear absolute tolerance'),
    fmt.val('NonLinearTolMaxChange', 'IP ITL', comment='! Non-Linear incremental tolerance'),
    fmt.val('MaxLinearIters', 'IP MIT', comment='! Maximum number of linear iterations'),
    fmt.val('MemoryIncrementBlockSize', 'OP INC', comment='! Incremental memory size'),
    fmt.val('PreconditioningBlocks', 'OP BLK', comment='! Number of blocks per processor for pre-conditioner'),
    fmt.val('PreconditionerType', 'OP PRE', comment='! Preconditioner type'),
  ],

  'Time': [
    fmt.val('JulianDay', 'TC JUL', comment='! Julian day the simulation starts on, used for the MET file'),

    fmt.val('StartTime', 'TC T0', comment='! Starting time of the simulation', \
      subitem_names=['Value1', 'Value2']),
    fmt.val('EndTime', 'TC TF', comment='! Final time of the simulation', \
      subitem_names=['Value1', 'Value2']),
    fmt.val('TimestepSize', 'TC IDT', comment='! The XY Series that will control the time step size'),
    fmt.val('AdaptiveTimeStep', 'TC NDP  ! Non-Adaptive time steps: forces adaptive time stepping to be turned of', comment='! Non-Adaptive time steps: forces adaptive time stepping to be turned off'),
    fmt.val('OutputInterval', 'OC ', comment='! Output time interval', custom_writer=write_outputinterval),
    fmt.val('PrintAdaptedMeshes', 'OC ADP', comment='! Print out the adapted mesh'),
  ],

  'Globals': [
    fmt.multival(['RayToADHSocket', 'ADHToRaySocket'], 'OP SOC', \
      comment='! Socket numbers for adh/raycaster communication (in/out for adh)'),
    fmt.val('Gravity', 'MP G', comment='! Gravity, (m)/(hr^2)'),
    fmt.val('WaterSpecificHeat', 'MP SHW', comment='! Specific heat of water, Units = (W-hr)/(g K)'),
    fmt.val('AirSpecificHeat', 'MP SHG', comment='!  Specific heat of gas, Units = (W-hr)/(g K)'),
    fmt.val('WaterSpecificGravity', 'MP SGW', comment='! Specific gravity of water'),
    fmt.val('GasSpecificGravity', 'MP SGG', comment='! Specific gravity of gas'),
    fmt.val('WaterThermalConductivity', 'MP TKW', comment='! Thermal conductivity of water, Units = (W)/(m K)'),
    fmt.val('AirThermalConductivity', 'MP TKG', comment='! Thermal conductivity of gas, Units = (W)/(m K)'),
    fmt.val('ReferenceDensity', 'MP RHO', comment='! Reference density, g/m^3'),
    fmt.val('ReferenceViscosity', 'MP VIS', comment='! Reference viscosity, Units = ?'),
  ],

  'Material': [
    fmt.idval('Porosity', 'MP POR', comment='! Porosity'),
    fmt.idval('HydraulicConductivity', 'MP K', comment='! Hydraulic conductivity', \
      subitem_names=['Value1', 'Value2', 'Value3', 'Value4', 'Value5', 'Value6']),
    fmt.idval('PresSatCurveIndex', 'MP SAT', comment='! Pressure saturation curve, XY series number'),
    fmt.idval('PresRelCondCurveIndex', 'MP KR', comment='! Pressure relative conductivity curve, XY series number'),
    fmt.idval('SpecificStorage', 'MP SS', comment='! Specific storage'),
    fmt.idval('SolidSpecificHeat', 'MP SHS', comment='! Specific heat of solids'),
    fmt.idval('SolidSpecificGravity', 'MP SGS', comment='! Specific gravity of solids'),
    fmt.idval('Albedo', 'MP ALB', comment='! Albedo'),
    fmt.idval('BulkEmissivity', 'MP EMS', comment='! Bulk emissivity'),
    fmt.idval('DryThermalConductivity', 'MP TKD', comment='! Dry/Solid Thermal Conductivity'),
    fmt.idval('SatThermalConductivity', 'MP TKS', comment='! Sat Thermal Conductivity'),
    fmt.idval('AnisotropyTensor', 'MP TKA', comment='! Anisotropy tensor for thermal conductivity', \
      subitem_names=['Value1', 'Value2', 'Value3', 'Value4', 'Value5', 'Value6']),
    fmt.idval('FractionSandGravel', 'MP FSG', comment='! Mass fraction coarse grain (sand and gravel)'),
    fmt.idval('ResidualSaturation', 'MP RSD', comment='! Residual saturation by gravity drainage'),
    fmt.idval('QuartzFraction', 'MP QTZ', comment='! Mass fraction quartz'),
    fmt.idval('VanGenuchtenAlpha', 'MP VGA', comment='! van Genuchten curve alpha (1/L)'),
    fmt.idval('VanGenuchtenN', 'MP VGN', comment='! van Genuchten curve exponent'),
    fmt.idval('VanGenuchtenMaxCP', 'MP VGP', comment='! van Genuchten curve max capillary pressure head'),
    fmt.idval('VanGenuchtenNumXY', 'MP VGX', comment='! number of van Genuchten curve entries'),
    fmt.idval('Tortuosity', 'MP TOR', comment='! Tortuosity'),
    fmt.idval('LongitudinalDispersivity', 'MP DPL', comment='! Longitudinal dispersivity'),
    fmt.idval('TransverseDispersivity', 'MP DPT', comment='! Transverse dispersivity'),
    fmt.idval('MaxRefineLevels', 'MP ML', comment='! Refinement levels'),
    fmt.idval('FlowRefineTol', 'MP FRT', comment='! Refinement Tolerance'),
    fmt.idval('MaterialRGB', 'RGB', comment='! MaterialRGB'),
    fmt.idval('MaterialTran', 'TRAN', comment='! MaterialTran'),
    fmt.idval('MaterialEmit', 'EMIT', comment='! MaterialEmit'),
  ],

  # Boundary Conditions
  'SpecifiedHead': fmt.bc('Value', 'DB FLW', comment='! Dirichlet boundary condition for flow'),
  'SpecifiedFlux': fmt.bc('Value', 'NB FLW', comment='! Neumann boundary condition for flow'),
  'FlowInjectionWell': fmt.bc('Value', 'WL FLW', comment='! FlowInjectionWell'),
  'METData': fmt.bc('Value', 'NB MET'),
  'GroundSurfaceHeatFlux': fmt.bc('Value', 'NB HFX', comment='! Heat Flux boundary condition for temperature'),
  'RayCaster': fmt.bc('Value', 'NB RAY', comment='! RayCaster'),
  'BottomBoundaryTemp': fmt.bc('Value', 'DB TMP', comment='! Dirichlet boundary condition for temperature'),
}


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

    Returns boolean indicating success
    Parameters
    ----------
    spec: Top-level object passed in from CMB
    '''
    #print 'Enter ExportCMB()'

    # Initialize scope instance to store spec values and other info
    scope = adh.init_scope(spec)
    if scope.logger.hasErrors():
      print scope.logger.convertToString()
      print 'FILES NOT WRITTEN because of errors'
      return False
    scope.format_table = format_table

    print 'Analysis types:', scope.analysis_types
    if not scope.analysis_types:
      msg = 'No analysis types selected'
      print 'WARNING:', msg
      scope.logger.addWarning(msg)
    else:
      print 'Categories:', sorted(list(scope.categories))

    if not scope.output_filename:
      msg = 'Cannot export -- output filename missing'
      print 'ERROR:', msg
      scope.logger.addError(msg)
      return False

    # Open output file and start exporting content
    completed = False
    with open(scope.output_filename, 'w') as scope.output:
      if 'Heat' in scope.categories:
        scope.output.write('OP HT\n')
      if 'Flow' in scope.categories:
        scope.output.write('OP GW\n')
      n = len(scope.constituent_dict)
      scope.output.write('OP TRN %d\n' % n)

      # Call write-content functions in specified top-level order
      att_type_list = [
        'Solvers', 'Time', 'Material', 'BoundaryCondition','Globals'
      ]
      for att_type in att_type_list:
        ok = adh.write_section(scope, att_type)
        #if not ok:
        #    break

      # Write material id cards
      adh.write_MID_cards(scope)

      # Write function attributes
      adh.write_functions(scope)

      # Write MTS cards (material id for each model domain)
      adh.write_MTS_cards(scope)

      # Write NDS & EGS cards for boundary conditions
      adh.write_bc_sets(scope)

      # Last line
      scope.output.write('END\n')
      print 'Wrote', scope.output_filename
      completed = True

    if not completed:
      print 'WARNING: Export terminated unexpectedly -- output might be invalid.'

    return completed


# ---------------------------------------------------------------------
def _write_outputinterval(scope, item, card_format, context_id):
  '''Writes OutputInterval card

  Can be be specified either of 2 ways:
    "OP INC" for fixed time interval
    "OP SRS" for interval specified by function
  '''
  if not item.isDiscrete:
    msg = 'Expected OutputInterval to be discrete item - ignoring'
    print 'ERROR:', msg
    scope.logger.addError(msg)
    return False

  output_list = list()

  # Fixed interval is index 0
  if 0 == item.discreteIndex(0):
    output_list.append('OC INT')

    # Time value
    subitem = item.activeChildItem(0)
    if 'FixedInterval' != subitem.name():
      msg = 'Unexpected subitem type \"%s\"' % subitem.name()
      print 'ERROR:', msg
      scope.logger.addError(msg)
    value_item = smtk.attribute.to_concrete(subitem)
    output_list += adh.get_values_as_strings(scope, value_item)

    # Time units
    subitem = item.activeChildItem(1)
    if 'FixedIntervalUnits' != subitem.name():
      msg = 'Unexpected subitem type \"%s\"' % subitem.name()
      print 'ERROR:', msg
      scope.logger.addError(msg)
    units_item = smtk.attribute.to_concrete(subitem)
    index = units_item.discreteIndex(0)
    output_list.append(str(index))

  # Function-specified interval is index 1
  elif 1 == item.discreteIndex(0):
    output_list.append('OC SRS')

    # Time function
    subitem = item.activeChildItem(0)
    if subitem.name() is not 'IntervalFunction':
      msg = 'Unexpected subitem type \"%s\"' % subitem.name()
      print 'ERROR:', msg
      scope.logger.addError(msg)
    value_item = smtk.attribute.to_concrete(subitem)
    output_list += adh.get_values_as_strings(scope, value_item)
  else:
    msg = 'Unexpected discrete value %d' % item.discreteIndex(0)
    print 'ERROR:', msg
    scope.logger.addError(msg)
    return False

  # Comment
  if card_format.comment is not None:
    output_list.append(card_format.comment)

  # Join output_list into one string
  output_string = ' '.join(output_list)
  scope.output.write(output_string)
  scope.output.write('\n')
  return True
