"""Data loading utilities for train_xgboost: Excel files and API-based datasets.""" from __future__ import annotations import json import logging import re import urllib.parse import urllib.request import zipfile from datetime import UTC, date, datetime, timedelta from pathlib import Path import pandas as pd from config import CITY_COORDS, parse_iso_date LOGGER = logging.getLogger("static_xgboost") def load_api_test_dataset( city: str, test_year: int, api_start_date: str | None, api_end_date: str | None, ) -> pd.DataFrame: """Load test set from Open-Meteo archive API in canonical schema.""" lat, lon = CITY_COORDS[city] start_date = parse_iso_date(api_start_date, "api-start-date") if api_start_date else date(test_year, 1, 1) end_date = parse_iso_date(api_end_date, "api-end-date") if api_end_date else (datetime.now(UTC).date() - timedelta(days=2)) if end_date < start_date: raise ValueError(f"api-end-date {end_date} is before api-start-date {start_date}.") params = { "latitude": lat, "longitude": lon, "hourly": ( "temperature_2m,relative_humidity_2m,wind_speed_10m,wind_gusts_10m," "surface_pressure,shortwave_radiation,direct_radiation,diffuse_radiation" ), "timezone": "UTC", "start_date": start_date.isoformat(), "end_date": end_date.isoformat(), } url = "https://archive-api.open-meteo.com/v1/archive" query = urllib.parse.urlencode(params) request_url = f"{url}?{query}" with urllib.request.urlopen(request_url, timeout=30) as response: payload = json.loads(response.read().decode("utf-8")) hourly = payload.get("hourly") if not isinstance(hourly, dict): raise RuntimeError("Open-Meteo response missing 'hourly' block.") times = hourly.get("time") or [] if not times: raise RuntimeError("Open-Meteo response contains no hourly time rows.") size = len(times) def col(name: str) -> list[float | None]: values = hourly.get(name) if values is None: return [None] * size return values df = pd.DataFrame( { "timestamp": pd.to_datetime(times, errors="coerce", utc=True), "temperature": col("temperature_2m"), "humidity": col("relative_humidity_2m"), "wind_speed": col("wind_speed_10m"), "wind_gust": col("wind_gusts_10m"), "pressure": col("surface_pressure"), "ghi": col("shortwave_radiation"), "dni": col("direct_radiation"), "dhi": col("diffuse_radiation"), "source_file": f"api:{city}:{start_date}->{end_date}", } ) df = df.dropna(subset=["timestamp", "ghi"]).reset_index(drop=True) LOGGER.info( "Loaded API test dataset for %s: rows=%d range=%s->%s", city, len(df), start_date, end_date, ) return df def _canonical_col_name(raw: str) -> str: """Normalize column name and apply known aliases.""" normalized = re.sub(r"[^a-z0-9]+", "_", raw.strip().lower()).strip("_") aliases = { "temp": "temperature", "winddirection": "wind_direction", "windspeed": "wind_speed", "windgust": "wind_gust", "tair_avg": "temperature", "tamb": "temperature", "rh_avg": "humidity", "rh": "humidity", "bp_cs100_avg": "pressure", "bp": "pressure", "ghi_corr_avg": "ghi", "dni_corr_avg": "dni", "dhi_corr_avg": "dhi", "ws": "wind_speed", "wsgust": "wind_gust", "wd": "wind_direction", "timestamp": "timestamp", "ts": "timestamp", "ws_wvc_1": "wind_speed", } return aliases.get(normalized, normalized) def _read_with_strategies(file_path: Path, nrows: int | None) -> pd.DataFrame | None: """Read a file using format-specific sheet/header strategies.""" name_lower = file_path.name.lower() if name_lower.startswith("pk_isb_") or name_lower.startswith("pk-isb_"): strategies: list[tuple[str | int, int]] = [ ("1h", 18), ("10min", 18), ("day", 18), (0, 0), ] else: strategies = [ (0, 0), ("Sheet1", 0), ("data", 0), ("1h", 18), ] for sheet_name, header_row in strategies: try: return pd.read_excel(file_path, sheet_name=sheet_name, header=header_row, nrows=nrows) except Exception: continue return None def _standardize_frame(df: pd.DataFrame, file_name: str) -> pd.DataFrame: """Normalize columns and produce canonical weather frame.""" work = df.copy() work.columns = [_canonical_col_name(str(col)) for col in work.columns] if "timestamp" in work.columns: work["timestamp"] = pd.to_datetime(work["timestamp"], errors="coerce", utc=True) elif "date" in work.columns and "time" in work.columns: work["timestamp"] = pd.to_datetime( work["date"].astype(str) + " " + work["time"].astype(str), errors="coerce", utc=True, ) elif "date" in work.columns: work["timestamp"] = pd.to_datetime(work["date"], errors="coerce", utc=True) else: return pd.DataFrame() numeric_cols = [ "temperature", "humidity", "wind_direction", "wind_speed", "wind_gust", "pressure", "ghi", "dni", "dhi", ] present_numeric = [col for col in numeric_cols if col in work.columns] for col in present_numeric: work[col] = pd.to_numeric(work[col], errors="coerce") if "ghi" not in work.columns: return pd.DataFrame() keep_cols = ["timestamp", *present_numeric] out = work[keep_cols].copy() out = out.dropna(subset=["timestamp", "ghi"]) out["source_file"] = file_name return out def load_dataset(dataset_dir: Path, max_rows_per_file: int, max_files: int) -> pd.DataFrame: """Load and unify all Excel files from the dataset folder.""" files = sorted(dataset_dir.glob("*.xlsx")) if max_files > 0: files = files[:max_files] if not files: raise FileNotFoundError(f"No .xlsx files found in {dataset_dir}") frames: list[pd.DataFrame] = [] for idx, file_path in enumerate(files, start=1): LOGGER.info("Reading file %d/%d: %s", idx, len(files), file_path.name) nrows = max_rows_per_file if max_rows_per_file > 0 else None try: if not zipfile.is_zipfile(file_path): raise ValueError("File is not a zip file") df = _read_with_strategies(file_path, nrows=nrows) if df is None: LOGGER.warning("Skipping %s (no readable sheet/header strategy)", file_path.name) continue except Exception as exc: LOGGER.warning("Skipping %s (read failed: %s)", file_path.name, exc) continue if df.empty: continue standardized = _standardize_frame(df, file_path.name) if standardized.empty: LOGGER.warning("Skipping %s (required columns missing)", file_path.name) continue frames.append(standardized) if not frames: raise RuntimeError("No usable files found after schema normalization.") merged = pd.concat(frames, ignore_index=True) LOGGER.info( "Loaded %d rows from %d files (max_rows_per_file=%d)", len(merged), len(frames), max_rows_per_file, ) return merged