#! /usr/bin/env python

import logging, traceback, os, sys
import produtil.log, produtil.setup
from produtil.log import jlogger
import hwrf_wcoss, hwrf_alerts
from hwrf.ensda import reset_ensda_relocate_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():
    """!Run relocation for one ensemble member. The member to run is specified by the
    ENSDA_MEMB environment variable."""
    logger=logging.getLogger('exhwrf_ensda_relocate')
    ENV=os.environ
    ompnum=ENV['OMP_NUM_THREADS']
    pureomp=ENV['PURE_OPENMP_THREADS']
    if ompnum != pureomp:
        os.environ['OMP_NUM_THREADS']=pureomp
    memb=ENV.get('ENSDA_MEMB','NOPE').lower()
    if memb=='nope':
        fail('Aborting: you must specify ENSDA_MEMB')
    imemb=int(memb,10)
    jlogger.info('HWRF ensda member %03d starting'%imemb)

    import hwrf_expt
    hwrf_expt.init_module()
    hwrf_expt.conf.add_fallback_callback(hwrf_alerts.fallback_callback)
    conf=hwrf_expt.conf
    if not hwrf_expt.conf.getbool('config','run_ens_relocation'):
        jlogger.info('Relocation for ensemble is disabled. This job need not be run.')
        sys.exit(0)

    hwrf_expt.init_module(make_ensemble_da=True)
    omemb=hwrf_expt.enkf_prep.member(hwrf_expt.conf.cycle,imemb)
    try:
        fail=conf.getstr('failure','ensda_relocate_mem%03d_failed'%imemb,'none')
        print('ensda_relocate_mem%03d'%imemb, fail)
        if fail=='unexpected_failure':
            raise hwrf.exceptions.UnexpectedFailureTest()
        if fail=='expected_failure':
            raise hwrf.exceptions.ExpectedFailureTest()
        omemb.run_relocate()
    except(SyntaxError,TypeError,ReferenceError,MemoryError,AttributeError,
           AssertionError,NameError,hwrf.exceptions.UnexpectedFailureTest) as ne:
        logger.error('ensda_relocate_mem%03d is aborting due to FATAL ERROR: '
                     %imemb+str(ne),exc_info=True)
        sys.exit(2)
    except Exception as e:
        msg='Could not run ensda relocation for member %s; will not run relocation.' \
            'Unhandled exception: %s'%(imemb,str(e))
        if conf.fallback('ensda_relocate_mem%03d_failed'%imemb,msg+
              '\n\nPython stack information at location of exception:\n\n'
              +traceback.format_exc()):
            logger.error(msg,exc_info=True)
            flag_file='tmpmem%03d.run_ensda_relocate'%imemb
            reset_ensda_relocate_flag_file(conf,flag_file,False,logger)
            set_ecflow_event('canceled',logger)
            return
        raise

    jlogger.info('HWRF ensda member %03d has completed'%imemb)

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