"""
inference.py — LSTM Groundwater Level Forecasting
==================================================
Usage
-----
    from inference import load_model, forecast
    model, scaler_X, scaler_y = load_model()
    prediction = forecast(model, scaler_X, scaler_y, 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 / "lstm_model.keras"
SCALER_X_PATH = Path(__file__).parent / "scaler_X.pkl"
SCALER_Y_PATH = Path(__file__).parent / "scaler_y.pkl"
CONFIG_PATH   = Path(__file__).parent / "model_config.json"


def load_model():
    """Load the LSTM model and both scalers."""
    model    = keras_load(MODEL_PATH)
    scaler_X = joblib.load(SCALER_X_PATH)
    scaler_y = joblib.load(SCALER_Y_PATH)
    print("LSTM model and scalers loaded.")
    return model, scaler_X, scaler_y


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


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

    Parameters
    ----------
    model     : loaded Keras model
    scaler_X  : fitted MinMaxScaler for features
    scaler_y  : fitted MinMaxScaler for 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_X.transform(X_window[required])
    X_input  = X_scaled.reshape(1, lookback, len(required))

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

    return float(prediction)


if __name__ == "__main__":
    model, scaler_X, scaler_y = 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
    import numpy as np
    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_X, scaler_y, dummy)
    print(f"\nForecast (next month): {pred:.4f} m")