00-download-data.py

#!/usr/bin/env python3
"""
ECMWF SEAS5 and CHIRPS Data Downloader

Downloads SEAS5 seasonal forecast data from ECMWF CDS API and CHIRPS precipitation data.
Run with --help for usage examples and data availability information.
"""

import os
import argparse
import datetime
import subprocess
import cdsapi


# =============================================================================
# SEAS5 DATA AVAILABILITY CONFIGURATION
# =============================================================================
# Historical years have all months available. Current/future years may have
# limited months depending on when the forecasts are released.

# Full historical data is available for these years (all 12 months)
# This should be updated annually to reflect the previous calendar year
SEAS5_FULL_YEARS_END = 2025  # Last year with complete 12 months of data


def get_available_months_for_year(year):
    """
    Get list of available months for a given year.

    Simple logic:
    - Historical years (1981-2025): All 12 months available
    - Current calendar year: Months 1 to current month are available
    - Future years: No data available

    Args:
        year: Year to check (integer)

    Returns:
        list: List of available month numbers (1-12), or empty list if year not available
    """
    year = int(year)
    now = datetime.datetime.now()
    current_year = now.year
    current_month = now.month

    if year < 1981:
        return []
    elif year < current_year:
        # Historical years with complete data
        return list(range(1, 13))  # All months 1-12
    elif year == current_year:
        # Current year - assume months up to current month are available
        return list(range(1, current_month + 1))
    else:
        # Future years
        return []


def validate_year_month_availability(year, months):
    """
    Validate that requested months are available for the specified year.

    Args:
        year: Year to download (integer)
        months: List of month strings (e.g., ["01", "02", "03"])

    Returns:
        tuple: (is_valid, available_months, unavailable_months)
    """
    available = get_available_months_for_year(year)
    available_set = set(available)

    requested_months = [int(m) for m in months]
    unavailable = [m for m in requested_months if m not in available_set]
    valid_months = [m for m in requested_months if m in available_set]

    return (len(unavailable) == 0, valid_months, unavailable)


def print_availability_info(year=None):
    """
    Print information about SEAS5 data availability.

    Args:
        year: Optional specific year to check. If None, prints general info.
    """
    print("\n" + "=" * 70)
    print("SEAS5 DATA AVAILABILITY INFORMATION")
    print("=" * 70)

    now = datetime.datetime.now()
    current_year = now.year
    current_month = now.month

    if year is not None:
        year = int(year)
        available = get_available_months_for_year(year)

        if not available:
            if year < 1981:
                print(f"\nYear {year}: NO DATA AVAILABLE (SEAS5 starts from 1981)")
            else:
                print(f"\nYear {year}: NO DATA AVAILABLE YET (future year)")
        else:
            month_names = ['Jan', 'Feb', 'Mar', 'Apr', 'May', 'Jun',
                          'Jul', 'Aug', 'Sep', 'Oct', 'Nov', 'Dec']
            available_names = [month_names[m-1] for m in available]
            print(f"\nYear {year}: {len(available)} months available")
            print(f"Available months: {', '.join(available_names)}")
            print(f"Month numbers: {', '.join(str(m) for m in available)}")

            if year == current_year and len(available) < 12:
                print(f"\nNote: Current year - automatically assumes months 1-{current_month} are available")
                print(f"      (based on current date: {now.strftime('%Y-%m-%d')})")
    else:
        current_year_available = get_available_months_for_year(current_year)
        print(f"\nHistorical data (1981-{SEAS5_FULL_YEARS_END}): All 12 months available")
        print(f"Current year ({current_year}): Months {current_year_available} available")
        print(f"  (Auto-detected: current month = {current_month})")
        print(f"Future years ({current_year+1}+): No data available yet")

    print("\n" + "-" * 70)
    print("AVAILABILITY LOGIC:")
    print(f"  - Historical years (< {current_year}): All 12 months")
    print(f"  - Current year ({current_year}): Months 1 to {current_month}")
    print("  - Future years: No data")
    print("=" * 70 + "\n")

def download_seas5(output_dir="./data", filename_prefix="seas5_precipitation_"):
    """
    Download SEAS5 dataset from ECMWF CDS API
    
    Args:
        output_dir: Directory to save the downloaded data
        filename_prefix: Prefix for the output filename
    
    Returns:
        str: Path to the downloaded file
    """
    print("Downloading SEAS5 data from ECMWF CDS...")
    
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)
    
    # Current date to use in the filename
    current_date = datetime.datetime.now().strftime("%Y%m%d")
    output_file = os.path.join(output_dir, f'{filename_prefix}{current_date}.grib')
    
    # Define the SEAS5 dataset and request parameters
    dataset = "seasonal-monthly-single-levels"
    request = {
        "originating_centre": "ecmwf",
        "system": "51",
        "variable": ["total_precipitation"],
        "year": [
            "1981", "1982", "1983", "1984", "1985", "1986", "1987", "1988", "1989",
            "1990", "1991", "1992", "1993", "1994", "1995", "1996", "1997", "1998",
            "1999", "2000", "2001", "2002", "2003", "2004", "2005", "2006", "2007",
            "2008", "2009", "2010", "2011", "2012", "2013", "2014", "2015", "2016",
            "2017", "2018", "2019", "2020", "2021", "2022", "2023", "2024", "2025"
        ],
        "month": [
            "01", "02", "03", "04", "05", "06", 
            "07", "08", "09", "10", "11", "12"
        ],
        "leadtime_month": ["1", "2", "3", "4", "5", "6"],
        "data_format": "grib",
        "product_type": ["monthly_mean"],
        "area": [23, 21, -12, 53]
    }
    
    try:
        client = cdsapi.Client()
        client.retrieve(dataset, request, output_file)
        print(f"SEAS5 data downloaded successfully to: {output_file}")
        return output_file
    except Exception as e:
        print(f"Error downloading SEAS5 data: {e}")
        return None


def parse_month_input(month_input):
    """
    Parse month input in various formats (single month, comma-separated, or range)

    Args:
        month_input: Input string like "3", "1,2,3,4", "1-6", or "1-12"

    Returns:
        list: List of month strings formatted as two digits (e.g., ["01", "02"])
    """
    months = []

    # Check if input is a range (e.g., "1-6")
    if isinstance(month_input, str) and "-" in month_input:
        start, end = map(int, month_input.split("-"))
        if 1 <= start <= 12 and 1 <= end <= 12:
            months = [f"{m:02d}" for m in range(start, end + 1)]
        else:
            raise ValueError("Month range must be between 1-12")

    # Check if input is comma-separated (e.g., "1,3,5")
    elif isinstance(month_input, str) and "," in month_input:
        for m in month_input.split(","):
            month_num = int(m.strip())
            if 1 <= month_num <= 12:
                months.append(f"{month_num:02d}")
            else:
                raise ValueError(f"Invalid month: {month_num}. Must be between 1-12")

    # Single string that can be converted to integer (e.g., "8")
    elif isinstance(month_input, str):
        try:
            month_num = int(month_input.strip())
            if 1 <= month_num <= 12:
                months.append(f"{month_num:02d}")
            else:
                raise ValueError(f"Invalid month: {month_num}. Must be between 1-12")
        except ValueError:
            raise ValueError(f"Invalid month format: {month_input}. Expected a number, comma-separated numbers, or range (e.g., 1-6)")

    # Single integer or list of integers
    elif isinstance(month_input, int) or isinstance(month_input, list):
        if isinstance(month_input, int):
            month_input = [month_input]
        for m in month_input:
            if 1 <= m <= 12:
                months.append(f"{m:02d}")
            else:
                raise ValueError(f"Invalid month: {m}. Must be between 1-12")

    if not months:
        raise ValueError("No valid months provided")

    return months

def download_current_month_seas5(output_dir="./data", filename_prefix="seas5_precipitation_",
                                  month_input=None, year=None, year_start=None, year_end=None,
                                  validate_availability=False, skip_unavailable=False):
    """
    Download SEAS5 dataset from ECMWF CDS API for specific months and year(s)

    Args:
        output_dir: Directory to save the downloaded data
        filename_prefix: Prefix for the output filename
        month_input: Month(s) to download data for - can be:
                    - Single integer (1-12)
                    - String range like "1-6"
                    - Comma-separated string like "1,3,5"
                    - List of integers [1, 3, 5]
        year: Single year to download (4-digit integer, defaults to current year if None)
        year_start: Start year for year range (use with year_end for multi-year download)
        year_end: End year for year range (use with year_start for multi-year download)
        validate_availability: If True, validate months against known availability before download
        skip_unavailable: If True, download available months for current year instead of failing

    Returns:
        str or list: Path(s) to the downloaded file(s), or None if download failed.
                    Returns a list when multiple files are downloaded (e.g., historical + current year)

    Raises:
        ValueError: If requested months are not available and skip_unavailable is False

    Note:
        SEAS5 data availability (automatically determined):
        - Historical years (< current year): All months (1-12) available
        - Current year: Months 1 to current month available
        - Future years: No data available

        When --skip-unavailable is used with a year range including current year:
        - Historical years are downloaded with all requested months
        - Current year is downloaded separately with only available months
        - This results in two separate GRIB files that can be merged in processing

        Use --check-availability --year YYYY to verify before downloading.

    Examples:
        # Single year download
        download_current_month_seas5(month_input="1-12", year=2025)

        # Year range download (1981-2025)
        download_current_month_seas5(month_input="1-12", year_start=1981, year_end=2025)

        # Historical + current year (downloads 2 files)
        download_current_month_seas5(month_input="1-12", year_start=2025, year_end=2026, skip_unavailable=True)
    """
    # Determine years to download
    if year_start is not None and year_end is not None:
        # Year range mode
        if year_start > year_end:
            print(f"ERROR: --year-start ({year_start}) must be <= --year-end ({year_end})")
            return None
        years = list(range(year_start, year_end + 1))
        year_display = f"{year_start}-{year_end}"
    elif year is not None:
        # Single year mode
        years = [year]
        year_display = str(year)
    else:
        # Default to current year
        year = datetime.datetime.now().year
        years = [year]
        year_display = str(year)

    # Parse month input to get list of months in proper format
    months = parse_month_input(month_input)
    requested_month_ints = [int(m) for m in months]

    # Get current year for validation logic
    current_year = datetime.datetime.now().year

    # Validate availability for each year if requested or if any year >= current_year
    needs_validation = validate_availability or any(y >= current_year for y in years)

    # Separate years into groups: historical (full months) and current year (partial months)
    historical_years = []
    current_year_download = None  # Will hold (year, available_months) if current year is in range

    if needs_validation:
        for y in years:
            available = get_available_months_for_year(y)

            if y < current_year:
                # Historical year - all months should be available
                historical_years.append(y)
            elif y == current_year:
                # Current year - check which requested months are available
                available_requested = [m for m in requested_month_ints if m in available]

                if not available_requested:
                    if skip_unavailable:
                        print(f"WARNING: Year {y} - no requested months available. Skipping year.")
                    else:
                        print(f"\nERROR: None of the requested months are available for year {y}.")
                        print(f"Available months for {y}: {', '.join(str(m) for m in available) if available else 'None'}")
                        print(f"\nTo check availability: python 00-download-data.py --check-availability --year {y}")
                        return None
                else:
                    unavailable = [m for m in requested_month_ints if m not in available]
                    if unavailable:
                        if skip_unavailable:
                            print(f"INFO: Year {y} - only months {', '.join(str(m) for m in available_requested)} available.")
                            print(f"      Will download available months separately.")
                            current_year_download = (y, available_requested)
                        else:
                            print(f"\nERROR: Requested months {', '.join(str(m) for m in unavailable)} are not available for year {y}.")
                            print(f"Available months for {y}: {', '.join(str(m) for m in available)}")
                            print(f"\nTo download available months only: add --skip-unavailable flag")
                            return None
                    else:
                        # All requested months are available for current year
                        current_year_download = (y, available_requested)
            else:
                # Future year
                if skip_unavailable:
                    print(f"WARNING: Year {y} - future year, no data available. Skipping.")
                else:
                    print(f"\nERROR: Year {y} is a future year with no data available.")
                    return None

        if not historical_years and not current_year_download:
            print("ERROR: No valid years to download after validation.")
            return None
    else:
        # No validation needed - all years are historical
        historical_years = years

    # Create output directory if needed
    os.makedirs(output_dir, exist_ok=True)

    downloaded_files = []
    current_date = datetime.datetime.now().strftime("%Y%m%d")

    # Download historical years (if any)
    if historical_years:
        year_str_list = [str(y) for y in historical_years]
        if len(historical_years) == 1:
            year_info = f"year{historical_years[0]}"
            year_display_hist = str(historical_years[0])
        else:
            year_info = f"years{min(historical_years)}-{max(historical_years)}"
            year_display_hist = f"{min(historical_years)}-{max(historical_years)}"

        month_info = "months_" + "_".join(months) if len(months) <= 4 else f"months_{len(months)}_months"
        output_file = os.path.join(output_dir, f'{filename_prefix}{current_date}_{year_info}_{month_info}.grib')

        print(f"\nDownloading SEAS5 data for years {year_display_hist}, months {', '.join(months)}...")
        print(f"Total years: {len(historical_years)}, Total months per year: {len(months)}")

        dataset = "seasonal-monthly-single-levels"
        request = {
            "originating_centre": "ecmwf",
            "system": "51",
            "variable": ["total_precipitation"],
            "year": year_str_list,
            "month": months,
            "leadtime_month": ["1", "2", "3", "4", "5", "6"],
            "data_format": "grib",
            "product_type": ["monthly_mean"],
            "area": [23, 21, -12, 53]
        }

        try:
            client = cdsapi.Client()
            client.retrieve(dataset, request, output_file)
            print(f"Historical data downloaded successfully to: {output_file}")
            downloaded_files.append(output_file)
        except Exception as e:
            print(f"Error downloading historical SEAS5 data: {e}")
            return None

    # Download current year separately (if needed)
    if current_year_download:
        cy_year, cy_months = current_year_download
        cy_months_str = [f"{m:02d}" for m in cy_months]

        year_info = f"year{cy_year}"
        month_info = "months_" + "_".join(cy_months_str) if len(cy_months_str) <= 4 else f"months_{len(cy_months_str)}_months"
        output_file = os.path.join(output_dir, f'{filename_prefix}{current_date}_{year_info}_{month_info}.grib')

        print(f"\nDownloading SEAS5 data for year {cy_year}, months {', '.join(cy_months_str)}...")

        dataset = "seasonal-monthly-single-levels"
        request = {
            "originating_centre": "ecmwf",
            "system": "51",
            "variable": ["total_precipitation"],
            "year": [str(cy_year)],
            "month": cy_months_str,
            "leadtime_month": ["1", "2", "3", "4", "5", "6"],
            "data_format": "grib",
            "product_type": ["monthly_mean"],
            "area": [23, 21, -12, 53]
        }

        try:
            client = cdsapi.Client()
            client.retrieve(dataset, request, output_file)
            print(f"Current year data downloaded successfully to: {output_file}")
            downloaded_files.append(output_file)
        except Exception as e:
            print(f"Error downloading current year SEAS5 data: {e}")
            # Continue if we at least have historical data
            if not downloaded_files:
                return None

    # Return result
    if len(downloaded_files) == 1:
        print(f"\nSEAS5 data downloaded successfully to: {downloaded_files[0]}")
        return downloaded_files[0]
    elif len(downloaded_files) > 1:
        print(f"\n{'='*70}")
        print(f"Downloaded {len(downloaded_files)} files:")
        for f in downloaded_files:
            print(f"  - {f}")
        print(f"\nTo use with 01-run-process-spi.py, specify:")
        print(f"  --seas51-main-file {downloaded_files[0]} \\")
        print(f"  --seas51-additional-files {' '.join(downloaded_files[1:])}")
        print(f"{'='*70}")
        return downloaded_files
    else:
        return None


def check_grib_file(grib_file_path):
    """
    Check a GRIB file to display available time information
    
    Args:
        grib_file_path: Path to the GRIB file to check
    """
    print(f"Checking GRIB file: {grib_file_path}")
    
    try:
        import xarray as xr
        
        # Open the dataset with specific backend settings
        ds = xr.open_dataset(grib_file_path, engine='cfgrib',
                            backend_kwargs=dict(time_dims=('forecastMonth', 'time')))
        
        # Extract and display time information
        print("\nGRIB file contents summary:")
        print("-------------------------")
        
        # Display dataset dimensions
        print(f"Dimensions: {dict(ds.dims)}")
        
        # Print time-related coordinates
        for time_dim in ['time', 'forecastMonth', 'valid_time']:
            if time_dim in ds.coords:
                time_values = ds[time_dim].values
                if len(time_values) > 0:
                    print(f"\n{time_dim}:")
                    print(f"  - Start: {time_values[0]}")
                    print(f"  - End: {time_values[-1]}")
                    print(f"  - Total values: {len(time_values)}")
        
        # Show variables in the dataset
        print("\nVariables in dataset:")
        for var_name, var in ds.variables.items():
            if var_name not in ds.dims and var_name not in ds.coords:
                print(f"  - {var_name}: {var.dims}")
        
        ds.close()
        
    except Exception as e:
        print(f"Error checking GRIB file: {e}")

def download_chirps(output_dir="./data", filename="chirps-v3.0.monthly.nc"):
    """
    Download CHIRPS precipitation data using wget
    
    Args:
        output_dir: Directory to save the downloaded data
        filename: Name for the output file
    
    Returns:
        str: Path to the downloaded file
    """
    print("Downloading CHIRPS precipitation data...")
    
    # Create output directory if it doesn't exist
    os.makedirs(output_dir, exist_ok=True)
    
    output_file = os.path.join(output_dir, filename)
    url = "https://data.chc.ucsb.edu/products/CHIRPS-2.0/global_monthly/netcdf/chirps-v2.0.monthly.nc"
    
    try:
        # Check if file already exists
        if os.path.exists(output_file):
            print(f"CHIRPS data already exists at: {output_file}")
            return output_file
            
        # Use wget to download the file
        subprocess.run(["wget", url, "-O", output_file], check=True)
        print(f"CHIRPS data downloaded successfully to: {output_file}")
        return output_file
    except subprocess.CalledProcessError as e:
        print(f"Error downloading CHIRPS data: {e}")
        return None
    except Exception as e:
        print(f"Unexpected error: {e}")
        return None

def main():
    # Create epilog with data availability information
    epilog_text = """
DATA AVAILABILITY NOTES:
------------------------
SEAS5 data availability varies by year:
  - Historical years (1981-2025): All months (1-12) are available
  - Current year (2026): Only months with released forecasts are available
    (ECMWF releases new forecasts around the 13th of each month)
  - Future years: No data available

IMPORTANT: When downloading data for the current year (2026), you must only
request months that have been released. Requesting unavailable months will
cause the CDS API to fail.

To check which months are available:
  python 00-download-data.py --check-availability --year 2026

To automatically validate and skip unavailable months:
  python 00-download-data.py --months 1-6 --year 2026 --validate-availability --skip-unavailable

EXAMPLES:
---------
# Download single year historical data (all months available):
  python 00-download-data.py --output-dir ../data --months 1-12 --year 2025

# Download full historical dataset (1981-2025, all months):
  python 00-download-data.py --output-dir ../data --months 1-12 --year-start 1981 --year-end 2025

# Download partial historical range:
  python 00-download-data.py --output-dir ../data --months 1-12 --year-start 2000 --year-end 2025

# Download current year (only available months):
  python 00-download-data.py --output-dir ../data --months 1 --year 2026

# Download historical + current year (skip unavailable):
  python 00-download-data.py --output-dir ../data --months 1-12 --year-start 2025 --year-end 2026 --skip-unavailable

# Check data availability:
  python 00-download-data.py --check-availability --year 2026

# Validate before downloading:
  python 00-download-data.py --months 1-3 --year 2026 --validate-availability
"""

    parser = argparse.ArgumentParser(
        description="Download ECMWF SEAS5 and CHIRPS data",
        epilog=epilog_text,
        formatter_class=argparse.RawDescriptionHelpFormatter
    )

    # Output options
    parser.add_argument("--output-dir", default="./data",
                        help="Directory to save downloaded data (default: ./data)")

    # Data source selection
    source_group = parser.add_argument_group('Data Source Selection')
    source_group.add_argument("--seas5-only", action="store_true",
                              help="Download only SEAS5 data (no CHIRPS)")
    source_group.add_argument("--chirps-only", action="store_true",
                              help="Download only CHIRPS data (no SEAS5)")

    # SEAS5 specific options
    seas5_group = parser.add_argument_group('SEAS5 Options')
    seas5_group.add_argument("--months", type=str, metavar="MONTHS",
                             help="Months to download. Formats: single (3), "
                                  "comma-separated (1,2,3), or range (1-12)")
    seas5_group.add_argument("--year", type=int,
                             help="Single year for SEAS5 data. IMPORTANT: For 2026+, only released "
                                  "months are available. Use --check-availability to verify. "
                                  "Cannot be used with --year-start/--year-end.")
    seas5_group.add_argument("--year-start", type=int, metavar="YEAR",
                             help="Start year for multi-year download (e.g., 1981). "
                                  "Use with --year-end for downloading year ranges.")
    seas5_group.add_argument("--year-end", type=int, metavar="YEAR",
                             help="End year for multi-year download (e.g., 2025). "
                                  "Use with --year-start for downloading year ranges.")

    # Data availability options
    avail_group = parser.add_argument_group('Data Availability',
                                            'Options for checking and validating data availability')
    avail_group.add_argument("--check-availability", action="store_true",
                             help="Check and display SEAS5 data availability for a year "
                                  "(use with --year to check a specific year)")
    avail_group.add_argument("--validate-availability", action="store_true",
                             help="Validate requested months against known availability before "
                                  "downloading (automatically enabled for year >= 2026)")
    avail_group.add_argument("--skip-unavailable", action="store_true",
                             help="Skip unavailable months/years instead of failing "
                                  "(use with --validate-availability or year ranges including 2026+)")

    # Utility options
    util_group = parser.add_argument_group('Utility Options')
    util_group.add_argument("--check-available-grib", type=str, metavar="GRIB_FILE_PATH",
                            help="Check and display time information for a specified GRIB file")

    args = parser.parse_args()

    # If check-availability option is provided, show availability info and exit
    if args.check_availability:
        print_availability_info(args.year)
        return

    # If check-available-grib option is provided, do that and exit
    if args.check_available_grib:
        check_grib_file(args.check_available_grib)
        return

    # Handle mutually exclusive options
    if sum([args.seas5_only, args.chirps_only, args.months is not None]) > 1:
        print("Error: Cannot specify multiple download options together")
        print("Use --seas5-only OR --chirps-only OR --months")
        return

    # Validate year arguments
    has_year = args.year is not None
    has_year_range = args.year_start is not None or args.year_end is not None

    if has_year and has_year_range:
        print("Error: Cannot use --year together with --year-start/--year-end")
        print("Use either --year for single year OR --year-start and --year-end for year range")
        return

    if (args.year_start is not None) != (args.year_end is not None):
        print("Error: --year-start and --year-end must be used together")
        return

    if args.year_start is not None and args.year_end is not None:
        if args.year_start > args.year_end:
            print(f"Error: --year-start ({args.year_start}) must be <= --year-end ({args.year_end})")
            return
        if args.year_start < 1981:
            print(f"Error: --year-start ({args.year_start}) cannot be before 1981 (SEAS5 data starts from 1981)")
            return

    # Validate that year options are only used with appropriate download options
    if (has_year or has_year_range) and args.months is None and not args.check_availability:
        print("Error: --year/--year-start/--year-end require --months to specify which months to download")
        return

    # Validate --skip-unavailable requires --validate-availability
    if args.skip_unavailable and not args.validate_availability:
        print("Warning: --skip-unavailable requires --validate-availability. Enabling validation.")
        args.validate_availability = True

    # Download SEAS5 data for specific month if requested
    if args.months is not None:
        seas5_file = download_current_month_seas5(
            args.output_dir,
            month_input=args.months,
            year=args.year,
            year_start=args.year_start,
            year_end=args.year_end,
            validate_availability=args.validate_availability,
            skip_unavailable=args.skip_unavailable
        )
        if seas5_file is None:
            print("\nDownload failed. Please check the error messages above.")
            return
    # Download full SEAS5 data if requested or if no specific option is specified
    elif args.seas5_only or not args.chirps_only:
        seas5_file = download_seas5(args.output_dir)
    
    # Download CHIRPS data if requested or if neither option is specified
    if args.chirps_only or (not args.seas5_only and args.months is None):
        chirps_file = download_chirps(args.output_dir)
    
    print("Data download process complete")

if __name__ == "__main__":
    main()