Source code for pycif.plugins.datastreams.fluxes.TNO_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 nearest available files bracketing a given date.

    Lists candidate files in `ref_dir` matching `file_format`, extending
    the search to `ref_dir_next`/`ref_dir_previous` when `dd` is close to a
    month boundary (within 3 hours), sorts them by parsed date, and returns
    the file/date pair nearest to `dd` on-or-before and on-or-after it.

    Args:
        ref_dir (str): Directory to search for candidate files.
        file_format (str): ``strftime``-style filename pattern used both to
            match files and parse their embedded date.
        dd (datetime.datetime): Reference date to bracket.
        ref_dir_next (str): Directory for the following month, searched in
            addition to `ref_dir` when `dd` is near a month boundary.
        ref_dir_previous (str or bool): Directory for the previous month,
            searched in addition to `ref_dir` when `dd` is near a month
            boundary; ``False`` to skip.

    Returns:
        tuple[list[str], list[datetime.datetime]]: ``([file1, file2], [dd1, dd2])``,
        the paths and adjusted dates of the nearest file on-or-before and
        on-or-after `dd` (falling back to the first/last available file, with
        a warning, if none is found strictly before/after `dd`).

    Raises:
        CifError: If no candidate files are found at all.
    """
    # 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]