#! /usr/bin/env python

##@namespace scripts.exhwrf_check_init
#Checks the initialization jobs' output and database entries to see
#if the initialization appears to have succeeded.  This is not a
#comprehensive check, but it tries to be as thorough as possible
#without taking more than a half minute to run.

import os, sys
import produtil.setup, produtil.log
import hwrf.mpipomtc
import hwrf_expt
from produtil.log import jlogger

def check_ocean_init(logger):
    """!Checks to see if the ocean initialization succeeded for the chosen ocean model.
    @param logger a logging.Logger for log messages
    @returns True if the ocean init ran, False otherwise"""
    conf=hwrf_expt.conf
    run_ocean=conf.getbool('config','run_ocean',True)
    ocean_model=conf.getstr('config','ocean_model','')

    if not run_ocean:
        logger.info('Ocean is disabled via run_ocean option.  Skipping ocean checks.')
        return True

    if not hwrf_expt.ocstatus.get(logger):
        logger.warning('The ocean status file says ocean is disabled.')
        return False

    if ocean_model=='HYCOM':
        if not hwrf_expt.hycominit.is_completed():
            logger.warning('The hycominit completion flag is off.  '
                           'Ocean init failed.')
            return False
    elif ocean_model=='POM':
        if not hwrf_expt.pominit.is_completed():
            logger.warning('The pominit completion flag is off.  '
                           'Ocean init failed.')
            return False
    else:
        logger.error('Invalid value %s for ocean_model.'%(repr(ocean_model),))
        return False

    return True

def check_wave_init(logger):
    conf=hwrf_expt.conf
    run_wave=conf.getbool('config','run_wave',True)
    wave_model=conf.getstr('config','wave_model','')

    if not run_wave:
        logger.info('Wave is disabled via run_wave option.  Skipping wave checks.')
        return True

    if not hwrf_expt.wvstatus.get(logger):
        logger.warning('The wave status file says wave is disabled.')
        return False

    if wave_model=='WW3':
        if not hwrf_expt.ww3init.is_completed():
            logger.warning('The ww3init completion flag is off.  '
                           'Wave init failed.')
            return False
    else:
        logger.error('Invalid value %s for wave_model.'%(repr(wave_model),))
        return False

    return True

def check_gsi(logger):
    """!Checks to see if the GSI data assimilation system ran
    correctly for all domains for which it should have run.
    @param logger a logging.Logger for log messages
    @returns True if the GSI ran for all expected domains, False otherwise"""
    okay=True
    conf=hwrf_expt.conf
    if hwrf.gsi.get_gsistatus(conf,'gsi_d02',logger) and not hwrf_expt.gsi_d02.is_completed():
        logger.warning(
            'GSI gsi_d02 Task state is not COMPLETED.  Domain 2 GSI failed.')
        okay=False

    if hwrf.gsi.get_gsistatus(conf,'gsi_d03',logger) and not hwrf_expt.gsi_d03.is_completed():
        logger.warning(
            'GSI gsi_d03 Task state is not COMPLETED.  Domain 3 GSI failed.')
        okay=False
    return okay

def check_enkf(logger):
    """!Checks to see if the ENKF data assimilation system ran
    correctly for all domains for which it should have run.
    @param logger a logging.Logger for log messages
    @returns True if the ENKF ran for all expected domains, False otherwise"""
    okay=True
    conf=hwrf_expt.conf
    if hwrf.enkf.get_enkfstatus(conf,'enkf_d02',logger) and not hwrf_expt.enkf_d02.is_completed():
        logger.warning(
            'EnKF enkf_d02 Task state is not COMPLETED.  Domain 2 EnKF failed.')
        okay=False
    return okay

def check_relocate(conf,who,init,relocate,logger,when):
    """!Checks to see if the vortex relocation system ran correctly.
    @param conf an hwrf.config.HWRFConfig for configuration info
    @param who a string name for this relocation job, for messages
    @param init an hwrf.init.HWRFInit or subclass, the grids
      input to the relocation
    @param relocate an hwrf.relocate.RelocationTask to check
    @param logger a logging.Logger for messages
    @param when the expected forecast time of this relocation"
    @returns True if the relocation job succeeded, or False otherwise"""
    i=init
    r=relocate
    okay=True
    if not r.completed:
        logger.warning('%s: FAIL: relocation job state is not COMPLETE (state %s)'
                       %(who,r.strstate))
        okay=False
    for domain in i.wrfghost:
        # Get the product for this domain and time:
        if domain.is_moad():
            p=r.wrfinput_at_time(when,domain)
        else:
            p=r.get_ghost(domain)
            
        # Is the product available, and does it have a valid location?
        if not p.available:
            logger.warning('%s: FAIL: relocation output product %s '
                           'not available'%(who,p.did))
            okay=False
        if p.location is None or p.location=='':
            logger.warning('%s: FAIL: relocation output product %s has '
                           'no location'%(who,p.did))
            okay=False
        elif not produtil.fileop.isnonempty(p.location):
            logger.warning(
                '%s: FAIL: relocation output %s is empty or does not exist at '
                '%s'%(who,p.did,p.location))
            okay=False

        # Get the cold/warm start info and such:
        try:
            rinfo=r.rinfo
            logger.info('%s: relocation info: %s'%(who,rinfo))
            logger.info('%s:   from file: %s'%(who,rinfo.from_file))
            if rinfo.warm_cold_flag is None:
                logger.warning(
                    '%s: FAIL: warm/cold state is unknown (None).'%(who,))
                okay=False
            elif rinfo.warm_cold_flag is not hwrf.relocate.WARM:
                if rinfo.cold_ok:
                    logger.info('%s: cold start, but cold_ok=True'%(who,))
                elif conf.getbool('config','expect_cold_start',False):
                    logger.warning('%s: cold start: no prior cycle found'%(who,))
                else:
                    logger.error('%s: FAIL: unexpected cold start'%(who,))
                    okay=False
            else:
                logger.info('%s: warm start'%(who,))
        except Exception as e:
            logger.error('%s: error determining cold vs. warm start: %s'
                         %(who,str(e)),exc_info=True)
            okay=False
    return okay

def check_fgat_relocate(conf,logger):
    """!Checks all steps of the FGAT relocation.
    @param conf the hwrf.config.HWRFConfig for configuration info
    @param logger a logging.Logger for log messages
    @returns True if the relocation jobs ALL succeeded, or False otherwise"""
    okay=True
    reloc=hwrf_expt.fgat_init.get_relocates()
    for t,r in reloc.items():
        i=hwrf_expt.fgat_init.init_at_time(t)
        who="FGAT time "+t.strftime('%Y%m%d%H')
        if not check_relocate(conf,who,i,r,logger,t):
            okay=False
            logger.error(who+': failed.')
    return okay

def check_gfs_relocate(conf,logger):
    """!Calls check_relocate() to check the GFS relocation
    @param conf an hwrf.config.HWRFConfig for configuration info
    @param logger a logging.Logger for log messages
    @returns True if the GFS relocation ran correctly and False otherwise"""
    if not check_relocate(conf,'GFS relocate',hwrf_expt.gfs_init,
                          hwrf_expt.gfs_init.rstage3,logger,
                          hwrf_expt.conf.cycle):
        logger.error('GFS relocate failed.')
        return False
    return True

def main():
    """!Checks the entire HWRF initialization to see if all needed
    inputs to the forecast are present and all init tasks ran
    correctly."""
    produtil.setup.setup()
    hwrf_expt.init_module()
    conf=hwrf_expt.conf
    logger=conf.log("check_init")

    ocean_flag=conf.getbool('config','run_ocean')
    wave_flag=conf.getbool('config','run_wave')
    gsi_flag=conf.getbool('config','run_gsi') 
    reloc_flag=conf.getbool('config','run_relocation')
    enkf_flag=run_enkf=conf.getbool('config','run_ensemble_da') and conf.getint('config','ensda_opt')>0

    okay=True

    if ocean_flag:
        if check_ocean_init(logger):
            logger.info('Ocean init succeeded.')
        elif conf.syndat.basin1 in hwrf_expt.non_ocean_basins:
            logger.info('Ocean init aborted, but basin is not supported.')
            hwrf_expt.runwrf.remove_ocean()
        else:
            logger.error('Ocean init failed.')
            okay=False
            hwrf_expt.runwrf.remove_ocean()
    else:
        logger.info('Ocean is disabled.  Skipping ocean checks.')

    if wave_flag:
        if check_wave_init(logger):
            logger.info('Wave init succeeded.')
        elif conf.syndat.basin1 in hwrf_expt.non_wave_basins:
            logger.info('Wave init aborted, but basin is not supported.')
            hwrf_expt.runwrf.remove_wave()
        else:
            logger.error('Wave init failed.')
            okay=False
            hwrf_expt.runwrf.remove_wave()
    else:
        logger.info('Wave is disabled.  Skipping wave checks.')

    if gsi_flag:
        if not check_gsi(logger):
            logger.error('GSI failed.')
            okay=False
        if not check_fgat_relocate(conf,logger):
            logger.error('FGAT relocate failed.')
            okay=False
    elif reloc_flag:
        logger.info('GSI is disabled.  Skipping GSI and FGAT '
                    'relocation checks.')
        if not check_gfs_relocate(conf,logger):
            logger.error('GFS relocate failed.')
    else:
        logger.info('Relocation and GSI are disabled.  '
                    'Skipping relocation checks.')

    if enkf_flag:
        if not check_enkf(logger):
            logger.error('EnKF failed.')
            okay=False

    logger.info('Asking the forecast object to check if all input files '
                'are available.')
    print(type(hwrf_expt.runwrf).__name__)
    have_input=hwrf_expt.runwrf.check_all_inputs()
    if not have_input:
        okay=False
        logger.error('FAILURE: WRF or POM inputs are missing')
    if not have_input:
        logger.critical('FORECAST INPUTS ARE MISSING!!')
        sys.exit(1)
    elif not okay:
        logger.critical('INIT JOBS DID NOT SUCCEED!!')
        sys.exit(1)

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