#! /usr/bin/env python

##@namespace scripts.exhwrf_init
# Runs the spectral, initial and boundary processing and generates
# WRF input files for use by other scripts.  Which work is done is
# selected by several environment variables:
#
# * $INIT_MODEL=GFS --- analysis time initialization, which does NOT
#   have to be from the GFS model.  
# * $INIT_MODEL=GDAS1 --- initialize from prior cycle of a global model,
#   which does NOT have to be from the GDAS model.  
# * $INIT_FHR=N --- for an integer N, the forecast hour of the prior
#   GDAS cycle.  This is generally a number from 3 to 9.
# * $INIT_PARTS=ANL --- only run what is needed for the initial time of
#   the forecast, when relocation and GSI are disabled
# * $INIT_PARTS=3DVAR --- all inputs needed for the relocation and GSI
# * $INIT_PARTS=BDY --- boundary conditions only
# * $INIT_PARTS=ALL --- run all parts of the initialization
#
# @note Models other than GDAS and GFS can be used as initial, boundary
# and previous cycle forecast conditions by modifying the hwrf config
# files.  However, the $INIT_MODEL variable must still be set to GFS or
# GDAS.  For example, the GEFS-based HWRF ensemble runs a 20 member
# forecast ensemble using the GEFS 20 member forecast, but it still sets
# $INIT_MODEL=GFS when running its exhwrf_init scripts with GEFS data.

import os, sys, produtil.log, produtil.setup
from produtil.log import jlogger
from produtil.ecflow import set_ecflow_label

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

def main():
    """!Run the selected parts of the initialization."""
    ENV=os.environ
    init_model=ENV['INIT_MODEL'].lower()
    init_fhr=int(ENV.get('INIT_FHR','0'))
    init_parts=ENV['INIT_PARTS'].lower()
    if init_model!='gfs' and init_model!='gdas1':
        fail('Aborting: init_model="%s" must be "gfs" or "gdas1"'%(init_model,))
    if init_model=='gdas1' and init_fhr<1:
        fail('Aborting: when init_model=gdas1, init_fhr must be > 1 (init_fhr=%d)'%(init_fhr,))
    if init_model=='gfs': init_fhr=0
    
    import hwrf_expt
    hwrf_expt.init_module()
    os.chdir(hwrf_expt.conf.getdir('WORKhwrf'))
    if init_model=='gfs':
        init=hwrf_expt.gfs_init
    elif not hwrf_expt.conf.getbool('config','run_gsi'):
        jlogger.info('GSI is disabled.  This job need not be run.')
        sys.exit(0)
    else:
        init=None
        logger=hwrf_expt.fgat_init.log()
        logger.info('search for fgat hour %d'%(init_fhr,))
        for fhr,init in hwrf_expt.fgat_init.fhr_and_init():
            if abs(fhr-init_fhr)<0.01:
                logger.info('fhr %d is init_fhr %d'%(fhr,init_fhr))
                #init.run()
                break
            else:
                logger.info('fhr %d is not init_fhr %d'%(fhr,init_fhr))
        assert(init is not None)

    if init_parts == 'parent':
        init.run_through_anl()
    elif init_parts == '3dvar':
        init.run_through_anl()
        init.run_init_after_anl()
    elif init_parts == 'bdy':
        init.run_real_bdy()
    elif init_parts == 'all':
        init.run_through_anl()
        init.run_init_after_anl()
        init.run_real_bdy()
    else:
        fail('Aborting: invalid value of INIT_PARTS: "%s" (must be "parent," "3dvar" or "bdy")'%(init_parts,))
    set_ecflow_label('status','success')
    

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