Source code for pycif.plugins.datastreams.fluxes.CAMSREG_nc.utils

import datetime
import glob
from logging import debug, info, warn
import os
import calendar
from pathlib import Path
import numpy as np
from logging import info
from .....utils.check.errclass import CifError




def find_valid_file(ref_dir, file_format, dd, ref_dir_next, ref_dir_previous=False):
    """Find the pair of files bracketing a reference date.

    Lists the files in ``ref_dir`` (and, when ``dd`` falls near a month
    boundary, in ``ref_dir_next``/``ref_dir_previous``) that match
    ``file_format``, parses their date from the file name, sorts them, and
    returns the nearest available file/date before ``dd`` and the nearest
    one after ``dd`` (used to bracket forecast-style files).

    Args:
        ref_dir (str): directory to search for candidate files.
        file_format (str): strptime/strftime-compatible file name pattern
            used both to parse each candidate file's date and to derive
            ``dd``'s canonical file name.
        dd (datetime.datetime): reference date to bracket.
        ref_dir_next (str): adjacent directory to also search when ``dd`` is
            close to the end of its month.
        ref_dir_previous (str or bool): adjacent directory to also search
            when ``dd`` is close to the start of its month; ``False`` to
            skip this lookup.

    Returns:
        (list[str], list[datetime.datetime]): a 2-element list of file paths
        ``[file_before, file_after]`` and the corresponding 2-element list of
        dates ``[date_before, date_after]`` bracketing ``dd``.

    Raises:
        CifError: if no file matching ``file_format`` is found in
            ``ref_dir`` (and adjacent directories, if searched).
    """
    # Get all files and dates matching the file and format
    list_files_orig = os.listdir(ref_dir)
    # Convert ref date
    ref_date = datetime.datetime.strptime(dd.strftime(file_format), file_format)
    previous_date = ref_date - datetime.timedelta(hours=3)
    if previous_date.month < ref_date.month and ref_dir_previous:
        try:
            list_files_orig += os.listdir(ref_dir_previous)
        except:
            info(f"Did not find any valid files in {ref_dir_previous} with format {file_format}")
    next_date = ref_date + datetime.timedelta(hours=3)
    if next_date.month > ref_date.month:
        try:
            list_files_orig += os.listdir(ref_dir_next)
        except:
            info(f"Did not find any valid files in {ref_dir_previous} with format {file_format}")
    
    list_dates_cur = []
    list_files_cur = []
    for f in list_files_orig:
        try:
            if f.find('idx') < 0:
                list_dates_cur.append(
                    datetime.datetime.strptime(f, file_format))
                list_files_cur.append(f)
        except:
            continue
    
    list_files = np.array(list_files_cur)
    list_dates = np.array(list_dates_cur)
    
    # Sorting along dates
    isort = np.argsort(list_dates)
    list_dates = list_dates[isort]
    list_files = list_files[isort]

    if len(list_files) == 0:
        raise CifError(f"Did not find any valid files in '{ref_dir}' with format '{file_format}'. Please check your yml file")
    
    # Compute deltas
    mask = (list_dates - ref_date) <= datetime.timedelta(0)

    if not np.any(mask):
        warn(f"Did not find any date in '{list_dates}' starting before {ref_date} for the available files found: '{list_files}'."
        )
        warn(f"Assuming the first date {list_dates[0]} is valid already since {ref_date}.")
        # Choose the last date and file
        file_ref1 = Path(ref_dir, list_files[0])
        date_ref1 = list_dates[0]
    else:
        # find nearest previous date
        file_ref1 = Path(ref_dir, list_files[mask][np.argmax(list_dates[mask])])
        date_ref1 = list_dates[mask].max()

    mask = (list_dates - ref_date) >= datetime.timedelta(0)
    if not np.any(mask):
        warn(f"Did not find any date in '{list_dates}' starting after {ref_date} for the available files found: '{list_files}'."
        )
        warn(f"Assuming the last date {list_dates[-1]} is valid until {ref_date}.")
        # Choose the last date and file
        file_ref2 = Path(ref_dir, list_files[-1])
        date_ref2 = list_dates[-1]
    else:
        # find nearest next date
        file_ref2 = Path(ref_dir, list_files[mask][np.argmin(list_dates[mask])])
        date_ref2 = list_dates[mask].min()

    # Reconvert to original date
    dd1 = dd + (date_ref1 - ref_date)
    dd2 = dd + (date_ref2 - ref_date)

    return [str(file_ref1), str(file_ref2)], [dd1, dd2]