inference.py

"""
inference.py — TCN Groundwater Level Forecasting
=================================================
Usage
-----
    from inference import load_model, forecast
    model, scaler_features, scaler_target = load_model()
    prediction = forecast(model, scaler_features, scaler_target, X_window)
"""

import json
import joblib
import numpy as np
import pandas as pd
from pathlib import Path
from tensorflow.keras.models import load_model as keras_load

MODEL_PATH          = Path(__file__).parent / "tcn_model.keras"
SCALER_FEATURES_PATH = Path(__file__).parent / "scaler_features.pkl"
SCALER_TARGET_PATH  = Path(__file__).parent / "scaler_target.pkl"
CONFIG_PATH         = Path(__file__).parent / "model_config.json"


def load_model():
    """Load the TCN model and both scalers."""
    model           = keras_load(MODEL_PATH)
    scaler_features = joblib.load(SCALER_FEATURES_PATH)
    scaler_target   = joblib.load(SCALER_TARGET_PATH)
    print("TCN model and scalers loaded.")
    return model, scaler_features, scaler_target


def load_config():
    with open(CONFIG_PATH) as f:
        return json.load(f)


def forecast(model, scaler_features, scaler_target, X_window: pd.DataFrame):
    """
    Predict the next month's groundwater level.

    Parameters
    ----------
    model           : loaded Keras TCN model
    scaler_features : fitted MinMaxScaler for all 4 input features
    scaler_target   : fitted MinMaxScaler for water_level target
    X_window        : DataFrame with columns [water_level, temperature,
                      precipitation, wind_speed] and exactly 24 rows
                      (the lookback window)

    Returns
    -------
    prediction : float — forecasted water level in original units (m)
    """
    cfg      = load_config()
    required = cfg['features']
    lookback = cfg['lookback_months']

    missing = [c for c in required if c not in X_window.columns]
    if missing:
        raise ValueError(f"X_window is missing columns: {missing}")
    if len(X_window) != lookback:
        raise ValueError(f"X_window must have {lookback} rows, got {len(X_window)}")

    X_scaled   = scaler_features.transform(X_window[required])
    X_input    = X_scaled.reshape(1, lookback, len(required))

    y_scaled   = model.predict(X_input, verbose=0).flatten()
    prediction = scaler_target.inverse_transform(y_scaled.reshape(-1, 1)).flatten()[0]

    return float(prediction)


if __name__ == "__main__":
    model, scaler_features, scaler_target = load_model()
    cfg = load_config()
    print(f"Model: {cfg['architecture']}")
    print(f"Test RMSE: {cfg['test_metrics']['RMSE']} m")

    # Dummy window — replace with real data
    dummy = pd.DataFrame({
        'water_level'  : np.random.uniform(60, 75, 24),
        'temperature'  : np.random.uniform(3, 15, 24),
        'precipitation': np.random.uniform(20, 120, 24),
        'wind_speed'   : np.random.uniform(10, 25, 24),
    })
    pred = forecast(model, scaler_features, scaler_target, dummy)
    print(f"\nForecast (next month): {pred:.4f} m")