standalone_inference.py

from __future__ import annotations

import csv
import json
import math
import os
from collections import defaultdict
from datetime import datetime, timedelta, timezone
from itertools import combinations
from pathlib import Path
from typing import Any

import joblib
import numpy as np
import pandas as pd
import requests

REGIONS = ["NSW1", "QLD1", "SA1", "TAS1", "VIC1"]
MARKET_TIME_OFFSET_HOURS = 10
HORIZONS: dict[int, str] = {
    1: "5m",
    3: "15m",
    6: "30m",
}
LOAD_FUELTECHS = {"pumps", "battery_charging"}
RENEWABLE_FUELTECHS = {
    "solar_utility",
    "solar_rooftop",
    "wind",
    "hydro",
    "battery_discharging",
    "bioenergy_biogas",
    "bioenergy_biomass",
}
DEFAULT_BASE_URL = os.environ.get("OPENELECTRICITY_API_URL", "https://api.openelectricity.org.au/v4")
DEFAULT_LOOKBACK_MINUTES = 90
# Ignore the newest API buckets; price/dispatch often lags demand by several minutes.
DEFAULT_SNAPSHOT_LAG_MINUTES = int(os.environ.get("SNAPSHOT_LAG_MINUTES", "10"))
DEFAULT_MODELS_DIR = Path(__file__).resolve().parent / "multi_horizon"
DEFAULT_LATEST_OUTPUT = Path(__file__).resolve().parent / "predictions" / "latest_predictions.json"
DEFAULT_LOG_PATH = Path(__file__).resolve().parent / "predictions" / "prediction_log.csv"

TARGET_SUFFIX_MAP = {
    "price_dollar_per_mwh": "price",
    "demand_mw": "demand",
    "gen_wind_mw": "gen_wind",
    "gen_coal_black_mw": "gen_coal_black",
    "gen_coal_brown_mw": "gen_coal_brown"
}


def strip_timezone(dt: datetime) -> str:
    return dt.replace(tzinfo=None, microsecond=0).isoformat()


def to_float(value: Any) -> float:
    return float("nan") if value is None else float(value)


def is_nan(value: Any) -> bool:
    if value is None:
        return True
    try:
        return bool(math.isnan(float(value)))
    except (TypeError, ValueError):
        return True


def json_safe_number(value: Any) -> float | None:
    if value is None:
        return None
    try:
        x = float(value)
    except (TypeError, ValueError):
        return None
    if math.isnan(x) or math.isinf(x):
        return None
    return x


def parse_interval_utc(timestamp: str) -> datetime:
    """Parse API interval strings (e.g. ...+10:00 or ...Z) to aware UTC."""
    s = str(timestamp).strip()
    if s.endswith("Z"):
        s = s[:-1] + "+00:00"
    ts = datetime.fromisoformat(s)
    if ts.tzinfo is None:
        return ts.replace(tzinfo=timezone.utc)
    return ts.astimezone(timezone.utc)


def mean(values: list[float | None]) -> float:
    filtered = [value for value in values if value is not None and not np.isnan(value)]
    if not filtered:
        return float("nan")
    return float(sum(filtered) / len(filtered))


def sample_std(values: list[float | None]) -> float:
    filtered = np.array([value for value in values if value is not None and not np.isnan(value)], dtype=float)
    if filtered.size < 2:
        return float("nan")
    return float(filtered.std(ddof=1))


def api_request(path: str, params: list[tuple[str, str]], api_key: str, base_url: str = DEFAULT_BASE_URL) -> dict[str, Any]:
    response = requests.get(
        f"{base_url}{path}",
        params=params,
        headers={
            "Authorization": f"Bearer {api_key}",
            "Content-Type": "application/json",
            "Accept": "application/json",
            "User-Agent": "nem-spot-price-predictor/1.0",
        },
        timeout=60,
    )
    response.raise_for_status()
    return response.json()


def network_time_series_to_rows(data: list[dict[str, Any]]) -> list[dict[str, Any]]:
    rows_map: dict[tuple[str, tuple[tuple[str, Any], ...]], dict[str, Any]] = {}

    for series in data:
        metric = series["metric"]
        for result in series.get("results", []):
            columns = result.get("columns", {})
            grouping_items = tuple(sorted(columns.items()))
            for timestamp, value in result.get("data", []):
                row_key = (timestamp, grouping_items)
                row = rows_map.setdefault(
                    row_key,
                    {
                        "interval": timestamp,
                        **columns,
                    },
                )
                row[metric] = value

    return sorted(rows_map.values(), key=lambda row: row["interval"])


def network_time_series_to_rows_by_region(data: list[dict[str, Any]]) -> list[dict[str, Any]]:
    """Merge market metrics on (interval, region); v4 uses column `region` not `network_region`."""
    merged: dict[tuple[str, str], dict[str, Any]] = {}
    for series in data:
        metric = series["metric"]
        for result in series.get("results", []):
            columns = result.get("columns", {}) or {}
            region = columns.get("network_region") or columns.get("region")
            if not region:
                continue
            region_s = str(region)
            for timestamp, value in result.get("data", []):
                key = (timestamp, region_s)
                row = merged.setdefault(
                    key,
                    {
                        "interval": timestamp,
                        "network_region": region_s,
                        "region": region_s,
                    },
                )
                row[metric] = value
    return sorted(merged.values(), key=lambda row: (row["interval"], row.get("network_region", row.get("region", ""))))


def filter_timestamps_by_lag(timestamps: list[str], lag_minutes: int) -> list[str]:
    """Only intervals at or before (now UTC - lag): settlement fields are usually complete."""
    if lag_minutes <= 0:
        return list(timestamps)
    cutoff = datetime.now(timezone.utc) - timedelta(minutes=lag_minutes)
    return [t for t in timestamps if parse_interval_utc(t) <= cutoff]


def trim_trailing_missing_prices(snapshots: list[dict[str, Any]], min_keep: int = 13) -> None:
    while len(snapshots) > min_keep and any(
        is_nan(snapshots[-1].get(f"{region}_price_dollar_per_mwh")) for region in REGIONS
    ):
        snapshots.pop()


def forward_fill_region_prices(snapshots: list[dict[str, Any]]) -> None:
    last_good: dict[str, float] = {}
    for snap in snapshots:
        for region in REGIONS:
            key = f"{region}_price_dollar_per_mwh"
            val = to_float(snap.get(key))
            if is_nan(val):
                prev = last_good.get(region)
                if prev is not None and not is_nan(prev):
                    snap[key] = prev
            else:
                last_good[region] = val


def get_market_time_parts(timestamp: str) -> dict[str, int]:
    ts = datetime.fromisoformat(timestamp.replace("Z", "+00:00"))
    market_ts = ts + timedelta(hours=MARKET_TIME_OFFSET_HOURS)
    js_day_of_week = (market_ts.weekday() + 1) % 7
    is_weekend = 1 if js_day_of_week in {0, 6} else 0
    is_business_hour = 1 if is_weekend == 0 and 9 <= market_ts.hour < 17 else 0
    return {
        "hour": market_ts.hour,
        "day_of_week": js_day_of_week,
        "month": market_ts.month,
        "is_weekend": is_weekend,
        "is_business_hour": is_business_hour,
    }


def fetch_recent_region_series(
    api_key: str,
    lookback_minutes: int = DEFAULT_LOOKBACK_MINUTES,
    base_url: str = DEFAULT_BASE_URL,
    snapshot_lag_minutes: int = DEFAULT_SNAPSHOT_LAG_MINUTES,
) -> list[dict[str, Any]]:
    now = datetime.now(timezone.utc)
    date_start = strip_timezone(now - timedelta(minutes=lookback_minutes))

    generation_response = api_request(
        "/data/network/NEM",
        [
            ("metrics", "power"),
            ("metrics", "energy"),
            ("metrics", "market_value"),
            ("interval", "5m"),
            ("date_start", date_start),
            ("primary_grouping", "network_region"),
            ("secondary_grouping", "fueltech"),
        ],
        api_key,
        base_url,
    )
    market_response = api_request(
        "/market/network/NEM",
        [
            ("metrics", "price"),
            ("metrics", "demand"),
            ("metrics", "curtailment_solar_utility"),
            ("metrics", "curtailment_wind"),
            ("interval", "5m"),
            ("date_start", date_start),
            ("primary_grouping", "network_region"),
        ],
        api_key,
        base_url,
    )

    generation_rows = network_time_series_to_rows(generation_response.get("data", []))
    market_rows = network_time_series_to_rows_by_region(market_response.get("data", []))

    generation_by_ts_region: dict[tuple[str, str], list[dict[str, Any]]] = defaultdict(list)
    for row in generation_rows:
        region = row.get("network_region") or row.get("region")
        if region in REGIONS:
            generation_by_ts_region[(row["interval"], region)].append(row)

    market_by_ts_region: dict[tuple[str, str], dict[str, Any]] = {}
    for row in market_rows:
        region = row.get("network_region") or row.get("region")
        if region in REGIONS:
            market_by_ts_region[(row["interval"], region)] = row

    timestamps = sorted({row["interval"] for row in generation_rows} | {row["interval"] for row in market_rows})
    lag = max(0, int(snapshot_lag_minutes))
    timestamps_lagged = filter_timestamps_by_lag(timestamps, lag)
    if len(timestamps_lagged) < 13:
        timestamps_lagged = filter_timestamps_by_lag(timestamps, max(0, lag - 5))
    if len(timestamps_lagged) < 13:
        timestamps_lagged = filter_timestamps_by_lag(timestamps, 0)
    if len(timestamps_lagged) < 13:
        raise ValueError(
            f"After SNAPSHOT_LAG_MINUTES={lag}, only {len(timestamps_lagged)} intervals remain; "
            f"need >= 13. Increase lookback_minutes (got {lookback_minutes}) or lower lag."
        )

    region_snapshots: list[dict[str, Any]] = []

    for timestamp in timestamps_lagged:
        row: dict[str, Any] = {"timestamp": timestamp}
        for region in REGIONS:
            generation_items = generation_by_ts_region.get((timestamp, region), [])
            market_item = market_by_ts_region.get((timestamp, region), {})
            current_price = to_float(market_item.get("price"))
            current_demand = to_float(market_item.get("demand"))
            wind_mw = 0.0
            coal_black_mw = 0.0
            coal_brown_mw = 0.0
            total_net_power = 0.0
            renewables_mw = 0.0

            for item in generation_items:
                fueltech = str(item.get("fueltech", ""))
                power = to_float(item.get("power"))
                if fueltech in LOAD_FUELTECHS:
                    continue
                total_net_power += max(0.0, power)
                if fueltech in RENEWABLE_FUELTECHS:
                    renewables_mw += power
                if fueltech == "wind":
                    wind_mw = power
                elif fueltech == "coal_black":
                    coal_black_mw = power
                elif fueltech == "coal_brown":
                    coal_brown_mw = power

            renewables_pct = (renewables_mw / total_net_power * 100.0) if total_net_power > 0 else float("nan")

            row[f"{region}_price_dollar_per_mwh"] = current_price
            row[f"{region}_demand_mw"] = current_demand
            row[f"{region}_renewables_pct"] = renewables_pct
            row[f"{region}_gen_wind_mw"] = wind_mw
            row[f"{region}_gen_coal_black_mw"] = coal_black_mw
            row[f"{region}_gen_coal_brown_mw"] = coal_brown_mw

        region_snapshots.append(row)

    trim_trailing_missing_prices(region_snapshots, min_keep=13)
    forward_fill_region_prices(region_snapshots)

    return region_snapshots


def build_feature_row(region_series: list[dict[str, Any]]) -> dict[str, float | int | str]:
    if len(region_series) < 13:
        raise ValueError(f"Need at least 13 recent 5-minute intervals, got {len(region_series)}.")

    snapshots = region_series[-13:]
    feature_row: dict[str, float | int | str] = {"timestamp": snapshots[-1]["timestamp"]}
    feature_row.update(get_market_time_parts(str(snapshots[-1]["timestamp"])))

    for region in REGIONS:
        prices = [to_float(snapshot.get(f"{region}_price_dollar_per_mwh")) for snapshot in snapshots]
        demands = [to_float(snapshot.get(f"{region}_demand_mw")) for snapshot in snapshots]
        renewables = [to_float(snapshot.get(f"{region}_renewables_pct")) for snapshot in snapshots]
        winds = [to_float(snapshot.get(f"{region}_gen_wind_mw")) for snapshot in snapshots]

        current_price = prices[-1]
        current_demand = demands[-1]
        current_renewables = renewables[-1]
        current_wind = winds[-1]
        current_coal_black = to_float(snapshots[-1].get(f"{region}_gen_coal_black_mw", 0.0))
        current_coal_brown = to_float(snapshots[-1].get(f"{region}_gen_coal_brown_mw", 0.0))

        feature_row[f"{region}_price_dollar_per_mwh"] = current_price
        feature_row[f"{region}_demand_mw"] = current_demand
        feature_row[f"{region}_renewables_pct"] = current_renewables
        feature_row[f"{region}_gen_wind_mw"] = current_wind
        feature_row[f"{region}_gen_coal_black_mw"] = current_coal_black
        feature_row[f"{region}_gen_coal_brown_mw"] = current_coal_brown
        feature_row[f"{region}_price_lag_1"] = prices[-2]
        feature_row[f"{region}_price_lag_2"] = prices[-3]
        feature_row[f"{region}_price_lag_6"] = prices[-7]
        feature_row[f"{region}_price_lag_12"] = prices[-13]
        feature_row[f"{region}_demand_rollmean_6"] = mean(demands[-6:])
        feature_row[f"{region}_demand_rollmean_12"] = mean(demands[-12:])
        feature_row[f"{region}_wind_rollmean_6"] = mean(winds[-6:])
        feature_row[f"{region}_wind_rollmean_12"] = mean(winds[-12:])
        feature_row[f"{region}_gen_wind_mw"] = current_wind

        feature_row[f"{region}_price_ramp_1"] = current_price - prices[-2]
        feature_row[f"{region}_price_ramp_2"] = current_price - prices[-3]
        feature_row[f"{region}_price_ramp_6"] = current_price - prices[-7]
        feature_row[f"{region}_price_ramp_12"] = current_price - prices[-13]
        feature_row[f"{region}_price_ramp_1_to_6"] = prices[-2] - prices[-7]
        feature_row[f"{region}_price_volatility"] = sample_std([prices[-1], prices[-2], prices[-3], prices[-7], prices[-13]])
        feature_row[f"{region}_demand_ramp_6"] = current_demand - feature_row[f"{region}_demand_rollmean_6"]
        feature_row[f"{region}_demand_ramp_12"] = current_demand - feature_row[f"{region}_demand_rollmean_12"]
        feature_row[f"{region}_wind_ramp_6"] = current_wind - feature_row[f"{region}_wind_rollmean_6"]
        feature_row[f"{region}_wind_ramp_12"] = current_wind - feature_row[f"{region}_wind_rollmean_12"]
        feature_row[f"{region}_wind_share_change"] = current_wind * current_renewables / 100.0

    feature_row["hour_sin"] = float(np.sin(2 * np.pi * float(feature_row["hour"]) / 24))
    feature_row["hour_cos"] = float(np.cos(2 * np.pi * float(feature_row["hour"]) / 24))
    feature_row["dow_sin"] = float(np.sin(2 * np.pi * float(feature_row["day_of_week"]) / 7))
    feature_row["dow_cos"] = float(np.cos(2 * np.pi * float(feature_row["day_of_week"]) / 7))

    for left, right in combinations(REGIONS, 2):
        left_price = float(feature_row[f"{left}_price_dollar_per_mwh"])
        right_price = float(feature_row[f"{right}_price_dollar_per_mwh"])
        spread = left_price - right_price
        feature_row[f"spread_{left}_{right}"] = spread
        feature_row[f"abs_spread_{left}_{right}"] = abs(spread)

    return feature_row


def load_metadata(models_dir: Path, region: str, clean_suffix: str, horizon_steps: int) -> dict[str, Any] | None:
    metadata_path = models_dir / f"{region.lower()}_{clean_suffix}_lightgbm_residual_blend_tplus{horizon_steps}_metadata.json"
    if not metadata_path.exists():
        return None
    return json.loads(metadata_path.read_text(encoding="utf-8"))



def predict_residual(model: Any, model_input: pd.DataFrame) -> float:
    booster = getattr(model, "booster_", None) or getattr(model, "_Booster", None)
    if booster is not None:
        return float(booster.predict(model_input)[0])
    return float(model.predict(model_input)[0])


def append_prediction_log(log_path: Path, rows: list[dict[str, Any]]) -> None:
    log_path.parent.mkdir(parents=True, exist_ok=True)
    fieldnames = [
        "prediction_generated_at",
        "source_snapshot_at",
        "region",
        "target_metric",
        "horizon_label",
        "horizon_minutes",
        "model_name",
        "alpha",
        "current_value",
        "predicted_residual",
        "predicted_value",
    ]
    expected_header = ",".join(fieldnames)
    file_mode = "a"
    write_header = not log_path.exists()

    if log_path.exists():
        with log_path.open("r", encoding="utf-8", newline="") as handle:
            first_line = handle.readline().strip()
        if first_line != expected_header:
            file_mode = "w"
            write_header = True

    with log_path.open(file_mode, encoding="utf-8", newline="") as handle:
        writer = csv.DictWriter(handle, fieldnames=fieldnames)
        if write_header:
            writer.writeheader()
        writer.writerows(rows)


def generate_prediction_payload(
    api_key: str,
    models_dir: Path = DEFAULT_MODELS_DIR,
    latest_output: Path | None = DEFAULT_LATEST_OUTPUT,
    log_path: Path | None = DEFAULT_LOG_PATH,
    lookback_minutes: int = DEFAULT_LOOKBACK_MINUTES,
    base_url: str = DEFAULT_BASE_URL,
    snapshot_lag_minutes: int = DEFAULT_SNAPSHOT_LAG_MINUTES,
) -> dict[str, Any]:
    region_series = fetch_recent_region_series(
        api_key,
        lookback_minutes=lookback_minutes,
        base_url=base_url,
        snapshot_lag_minutes=snapshot_lag_minutes,
    )
    feature_row = build_feature_row(region_series)
    source_snapshot_at = str(feature_row["timestamp"])
    prediction_generated_at = datetime.now(timezone.utc).replace(microsecond=0).isoformat().replace("+00:00", "Z")

    region_predictions: list[dict[str, Any]] = []
    prediction_log_rows: list[dict[str, Any]] = []

    for region in REGIONS:
        # Define base dictionary for the region's current values
        region_data = {
            "prediction_generated_at": prediction_generated_at,
            "source_snapshot_at": source_snapshot_at,
            "region": region,
            "current_values": {},
            "forecasts": {horizon: {} for horizon in HORIZONS.values()}
        }

        for raw_suffix, clean_suffix in TARGET_SUFFIX_MAP.items():
            raw_current = feature_row.get(f"{region}_{raw_suffix}")
            current_val = to_float(raw_current) if raw_current is not None else float("nan")
            if clean_suffix in {"gen_wind", "gen_coal_black", "gen_coal_brown"} and is_nan(current_val):
                current_val = 0.0
            region_data["current_values"][clean_suffix] = json_safe_number(current_val)

            for horizon_steps, horizon_label in HORIZONS.items():
                metadata = load_metadata(models_dir, region, clean_suffix, horizon_steps)
                
                # E.g. TAS doesn't have black coal models, so we skip if metadata is None
                if not metadata:
                    continue
                    
                model_path = models_dir / f"{region.lower()}_{clean_suffix}_lightgbm_residual_blend_tplus{horizon_steps}.pkl"
                model = joblib.load(model_path)
                
                feature_columns = metadata["feature_columns"]
                alpha = float(metadata["alpha"])
                model_input = pd.DataFrame([{column: feature_row.get(column, np.nan) for column in feature_columns}])
                
                predicted_residual = predict_residual(model, model_input)
                if is_nan(current_val):
                    predicted_level = float("nan")
                else:
                    predicted_level = float(current_val) + alpha * float(predicted_residual)

                region_data["forecasts"][horizon_label][clean_suffix] = {
                    "model_name": metadata["model_name"],
                    "alpha": alpha,
                    "horizon_steps": horizon_steps,
                    "horizon_minutes": int(metadata.get("horizon_minutes", horizon_steps * 5)),
                    "predicted_residual": predicted_residual,
                    "predicted_value": json_safe_number(predicted_level),
                }

                prediction_log_rows.append(
                    {
                        "prediction_generated_at": prediction_generated_at,
                        "source_snapshot_at": source_snapshot_at,
                        "region": region,
                        "target_metric": clean_suffix,
                        "horizon_label": horizon_label,
                        "horizon_minutes": horizon_steps * 5,
                        "model_name": metadata["model_name"],
                        "alpha": alpha,
                        "current_value": json_safe_number(current_val),
                        "predicted_residual": predicted_residual,
                        "predicted_value": json_safe_number(predicted_level),
                    }
                )

        region_predictions.append(region_data)

    latest_payload = {
        "prediction_generated_at": prediction_generated_at,
        "source_snapshot_at": source_snapshot_at,
        "prediction_horizons_minutes": [steps * 5 for steps in HORIZONS.keys()],
        "regions": region_predictions,
    }

    if log_path is not None:
        append_prediction_log(log_path, prediction_log_rows)
    if latest_output is not None:
        latest_output.parent.mkdir(parents=True, exist_ok=True)
        latest_output.write_text(
            f"{json.dumps(latest_payload, indent=2, allow_nan=False)}\n",
            encoding="utf-8",
        )

    return latest_payload