#! /usr/bin/env python

##@namespace scripts.exhwrf_enkf
# Run EnKF
# This job mush be run after exhwrf_enshx job.
#

import os, sys, traceback
import produtil.log, produtil.setup, produtil.cluster
from produtil.log import jlogger
import hwrf_wcoss, hwrf_alerts
from hwrf.enkf import unset_enkfstatus, set_enkfstatus
from hwrf.ensda import reset_ensda_flag_file
import hwrf.exceptions
from produtil.ecflow import set_ecflow_event

def fail(msg):
    """!Logs a message to the produtil.log.jlogger and exits with status 2
    @param msg the error message"""
    jlogger.error(msg)
    sys.exit(2)

def main():
    """!Runs the selected EnKF task."""

    import hwrf_expt
    hwrf_expt.init_module()
    hwrf_expt.conf.add_fallback_callback(hwrf_alerts.fallback_callback)
    conf=hwrf_expt.conf
    logger=conf.log('exhwrf_enkf')
    unset_enkfstatus(conf,logger)

    if conf.getint('config','ensda_opt') == 0:
        jlogger.info('enkf is disabled.  This job need not be run.')
        sys.exit(0)

    logger=conf.log('exhwrf_enkf starts')

    try:
        fail=conf.getstr('failure','enkf_failed','none')
        if fail=='unexpected_failure':
            raise hwrf.exceptions.UnexpectedFailureTest()
        if fail=='expected_failure':
            raise hwrf.exceptions.ExpectedFailureTest()

        if produtil.cluster.name() in ['gyre','tide']:
            hwrf_wcoss.set_vars_for_enkf(logger)
        else:
            logger.info('Not on WCOSS, so not setting WCOSS-specific vars.')

        hwrf_expt.enkf_d02.run()

        set_enkfstatus(conf,logger)

    except(SyntaxError,TypeError,ReferenceError,MemoryError,AttributeError,
           AssertionError,NameError,hwrf.exceptions.UnexpectedFailureTest) as ne:
        logger.error('HWRF GSI is aborting due to FATAL ERROR: '+str(ne),exc_info=True)
        sys.exit(2)
    except Exception as e:
        msg='Could not run enkf for domain d02; Will restart ensemble from GDAS.' \
            'Unhandled exception: %s'%(str(e))
        if conf.fallback('enkf_failed',msg+
              '\n\nPython stack information at location of exception:\n\n'
              +traceback.format_exc()):
            logger.error(msg,exc_info=True)
            flag_file='tmpenkf.run_ensda'
            reset_ensda_flag_file(conf,flag_file,True,False,logger)
            set_ecflow_event('canceled',logger)
            return
        raise

if __name__=='__main__':
    try:
        produtil.setup.setup()
        main()
    except Exception as e:
        jlogger.critical('HWRF enkf is aborting: '+str(e),exc_info=True)
        sys.exit(2)