#! /usr/bin/env python

import logging, os
from hwrf.storminfo import find_tcvitals_for, StormInfoError, parse_tcvitals
from hwrf.numerics import to_datetime_rel, to_datetime
from produtil.log import jlogger

__all__=['BadVitals', 'hmon_vitals', 'vitals_for_storm', 'vitals_for_slot']

class BadVitals(Exception): pass

class hmon_vitals(object):
    def __init__(self,vitals):
        super(hmon_vitals,self).__init__()
        self.vitals=vitals.copy()
        self._split_vitals()
    def _split_vitals(self):
        vitals=self.vitals
        self.STORM_NAME=vitals.stormname     # a str, all uppercase
        self.STORM_LAT=vitals.lat            # a float
        self.STORM_LON=vitals.lon            # a float
        self.STORM_LON_PARENT=self.STORM_LON-10   # a float
        self.CEN_LON=vitals.lon              # a float
        self.CEN_LAT=self.STORM_LAT+5        # a float
        if self.CEN_LAT>44: self.CEN_LAT=45  #
        self.NEST1_LAT=self.STORM_LAT        # a float
        self.NEST1_LON=self.STORM_LON        # a float
        self.NEST2_LAT=self.STORM_LAT        # a float
        self.NEST2_LON=self.STORM_LON        # a float
        self.STORM_NUMBER="%02d"%vitals.stnum # "09"
        self.STORM_BASIN=vitals.basin1        # "L"
        # YYYYMMDDHH date for six hours before the storm:
        self.DATE06 = to_datetime_rel(-6*3600,vitals.when).strftime("%Y%m%d%H")
        self.YMDH = vitals.YMDH
        self.YYYY = vitals.when.strftime("%Y")

def vitals_for_storm(num,basin,date,syndat_dir=None,logger=None):
    if logger is None: logger=jlogger
    if syndat_dir is None: syndat_dir=os.environ['COMINarch']

    if isinstance(num,str): # convert from string to int if needed
        num=int(num,10)
    date=to_datetime(date)
    syndat_file=os.path.join(syndat_dir,date.strftime("syndat_tcvitals.%Y"))

    logger.info('%s: search for %02d%s %s'%(
            syndat_file,num,basin,date.strftime("%Y%m%d%H")))

    vitals=None
    with open(syndat_file,'rt') as f:
        for found_vitals in find_tcvitals_for(
                f,logger,True,date,num,basin):
            vitals=found_vitals

    if not vitals:
        logger.error("No vitals found for %s %s %s"%(
                str(num),str(basin),date.strftime("%Y%m%d%H")))
        return None

    return hmon_vitals(vitals)
    

def vitals_for_slot(islot=None,messagedir=None,logger=None):
    """!Obtains the vitals for storm in slot islot, returns None if no
    vitals exists.

    @returns an hmon_vitals object if vitals are found.  Returns None
    if no vitals are found.

    @var islot the storm slot number, 1 through 8.  Defaults to the
    storm_num environment variable, used by NCO's ecFlow workflow to
    pass the requested storm slot number to the scripts

    @var messagedir the directory with the "nstorms" and "message1"
    through "message8" files.  Defaults to the COMINmsg environment
    variable.

    @var logger a logging object, defaults to the jlogfile"""
    if logger is None: logger=jlogger
    if messagedir is None: messagedir=os.environ['COMINmsg']
    if islot is None: islot=os.environ['storm_num']

    islot=int(islot)
    assert(islot>=1)

    stormlabel=os.environ.get('stormlabel','storm%d'%islot)

    nstorms_file=os.path.join(messagedir,"nstorms")
    with open(nstorms_file,'rt') as f:
        nstorms=int(f.readline())
    if islot>nstorms:
        logger.info("Slot %d is beyond nstorms=%d"%(islot,nstorms))
        return None

    message_file=os.path.join(messagedir,"message%d"%islot)
    if not os.path.exists(message_file):
        logger.info("Slot %d does not exist (no %s)"%(islot,message_file))
        return None

    vitals_list=None
    with open(message_file,'rt') as f:
        vitals_list=parse_tcvitals(f,logger,raise_all=True)
    if not vitals_list:
        logger.info('Did not find any vitals in storm slot %d (%s)'%(
                    islot,message_file))
        return None
    vitals=vitals_list[0]
        
    return hmon_vitals(vitals)