#! /usr/bin/env python

##@namespace scripts.exhwrf_meanhx
# Run GSI to calcuate ensemble Hx for EnKF
# This job must be run after exhwrf_meanhx job. There is one
# mandatory environment variable that must be set before executing this
# script: $ENSDA_MEMB, which must be set to the integer ENSDA member
# number (generally a number from 1 to 40)
#

import logging, os, sys, produtil.log, produtil.setup
from produtil.log import jlogger
import hwrf_wcoss

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

def main():
    """!Run Hx for one ensemble member. The member to run is specified by the
    ENSDA_MEMB environment variable."""
    logger=logging.getLogger('exhwrf_enshx')
    ENV=os.environ
    memb=ENV.get('ENSDA_MEMB','NOPE').lower()
    if memb=='nope':
            fail('Aborting: you must specify ENSHX_MEMB')
    imemb=int(memb,10)
    jlogger.info('EnKF enshx member %03d starting'%imemb)

    import hwrf_expt
    hwrf_expt.init_module(make_ensemble_da=True)
    if not hwrf_expt.conf.getbool('config','run_ensemble_da',False):
        jlogger.info('enshx is disabled.  This job need not be run.')
        sys.exit(0)

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

    omemb=hwrf_expt.enkf_prep.member(hwrf_expt.conf.cycle,imemb)
    omemb.run_gsihx()
    for prod in omemb.products():
        if not prod.location:
            logger.error('No product: %s'%(prod.did,))
        elif not prod.available:
            logger.error('Product %s not available (location %s)'%(
                    repr(prod.did),repr(prod.location)))
        else:
            dest='%s/%s.ensda_%03d.%s'%(
                hwrf_expt.conf.getdir('com'),
                hwrf_expt.conf.getstr('config','out_prefix'),
                imemb,os.path.basename(prod.location))
            logger.info('%s %s: send to %s'%(
                    str(prod.did),repr(imemb),str(dest)))
            assert(os.path.isabs(dest))
            produtil.fileop.deliver_file(
                prod.location,dest,logger=logger)

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

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