#! /usr/bin/env python

##@namespace scripts.exhwrf_wrfout_archive
# Compresses native wrfout files for later archiving by exhwrf_wrfout_archive
#
# To enable this script, you must set the wrfout option in the [archive]
# configuration section:
# @code[.conf]
# [archive]
# wrfout=hpss:/NCEPDEV/emc-hwrf/2year/Jane.Doe/wrfout/{out_prefix}.tar
# @endcode
#
# The extension must be .tar, and the only archiving method supported is
# hpss:/ (which uses htar)

import os, sys
import produtil.setup, produtil.log, produtil.run, produtil.cd
import hwrf.numerics
import hwrf_expt
import produtil.workpool

from produtil.log import jlogger
from produtil.run import *

def archive_one(ncks_cmd,logger,native,compressed):
    logger.info('STARTING: %s => %s'%(native,compressed))
    def copier(infile,tempname,temp):
        checkrun(ncks_cmd[infile,tempname],logger=logger)
    produtil.fileop.deliver_file(native,compressed,
                                 keep=True,copier=copier,logger=logger)
    logger.info('FINISHED: %s => %s'%(native,compressed))

def main():
    ENV=os.environ
    ompnum=ENV['OMP_NUM_THREADS']
    pureomp=ENV['PURE_OPENMP_THREADS']
    if ompnum != pureomp:
        os.environ['OMP_NUM_THREADS']=pureomp
    hwrf_expt.init_module()
    conf=hwrf_expt.conf
    if not conf.getstr('archive','wrfout',''):
        jlogger.info('No wrfout option in [archive] section.  Will not make wrfout archive.')
        sys.exit(0)

    logger=conf.log('wrfout_compress')

    if not os.environ.get('WRFOUT_COMPRESS_STARTED_SELF',''):
        # Top-level script apruns itself:
        checkrun(openmp(exe(os.path.realpath(__file__))).env(
            WRFOUT_COMPRESS_STARTED_SELF='yes'),threads=8,logger=logger)
        return
        
    files=list()
    dt=hwrf.numerics.to_timedelta('6:00:00')
    t0=conf.cycle
    wrf=hwrf_expt.runwrf.wrf()
    ncks_path=hwrf_expt.conf.getexe('ncks','')
    if not ncks_path:
        ncks_path=produtil.fileop.find_exe('ncks')
    ncks_cmd=alias(batchexe(ncks_path)['-4','-L','6','-O'])

    with produtil.cd.NamedDir(hwrf_expt.runwrf.location):
        for i in range(22):
            for dom in wrf:
                t=t0+dt*i
                out=dom.get_output('auxhist3',t)
                if out is None:
                    out=dom.get_output('history',t)
                if out is None:
                    out=dom.get_output('auxhist2',t)
                if out is None:
                    logger.error('%s: could not determine wrfout for '
                                 'domain %s'%(t.strftime('%Y%m%d%H'),
                                              str(dom)))
                if out is not None:
                    if not os.path.exists(out.path()):
                        logger.error('%s: does not exist'%(out.path(),))
                    if not produtil.fileop.isnonempty(out.path(),):
                        logger.error('%s: is empty'%(out.path(),))
                    files.append(out.path())

        with produtil.workpool.WorkPool(8,logger=logger) as workpool:
            for native in files:
                compressed=native+'.nc4'
                workpool.add_work(archive_one,[ncks_cmd,logger,native,compressed])
if __name__=='__main__':
    try:
        thread_logger=bool(os.environ.get('WRFOUT_COMPRESS_STARTED_SELF',''))
        produtil.setup.setup(thread_logger=thread_logger)
        jlogger.info("HWRF wrfout_compress job starting")
        main()
        jlogger.info("HWRF wrfout_compress job completed")
    except Exception as e:
        jlogger.critical('HWRF rundir archive is aborting: '+str(e),
                         exc_info=True)
        sys.exit(2)