"""
inference.py — SARIMAX Groundwater Level Forecasting
=====================================================
Usage
-----
    from inference import load_model, forecast
    model = load_model()
    predictions, conf_int = forecast(model, steps=12, exog_future=X_future)
"""

import json
import joblib
import numpy as np
import pandas as pd
from pathlib import Path

MODEL_PATH  = Path(__file__).parent / "sarimax_model.pkl"
CONFIG_PATH = Path(__file__).parent / "model_config.json"


def load_model():
    """Load the fitted SARIMAX model."""
    return joblib.load(MODEL_PATH)


def load_config():
    """Load model configuration and metrics."""
    with open(CONFIG_PATH) as f:
        return json.load(f)


def forecast(model, steps: int, exog_future: pd.DataFrame):
    """
    Generate a forecast.

    Parameters
    ----------
    model       : fitted SARIMAX result (from load_model)
    steps       : number of months ahead to forecast
    exog_future : DataFrame with columns [temperature, precipitation, wind_speed]
                  Must have a DatetimeIndex with freq='MS' and exactly `steps` rows.

    Returns
    -------
    predictions : pd.Series  — point forecasts
    conf_int    : pd.DataFrame — 95% confidence interval
    """
    cfg           = load_config()
    required_cols = cfg['exog_cols']

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

    fc          = model.get_forecast(steps=steps, exog=exog_future[required_cols])
    predictions = fc.predicted_mean
    conf_int    = fc.conf_int()
    return predictions, conf_int


if __name__ == "__main__":
    model = load_model()
    cfg   = load_config()
    print(f"Model loaded: SARIMAX{cfg['order']}x{cfg['seasonal_order']}")
    print(f"Test RMSE: {cfg['test_metrics']['RMSE']}")

    # Dummy forecast — replace with real future met data
    idx   = pd.date_range(start="2024-01-01", periods=12, freq="MS")
    X_fut = pd.DataFrame({
        "temperature"  : np.random.uniform(3, 15, 12),
        "precipitation": np.random.uniform(20, 120, 12),
        "wind_speed"   : np.random.uniform(10, 25, 12),
    }, index=idx)

    preds, ci = forecast(model, steps=12, exog_future=X_fut)
    print(pd.DataFrame({"predicted": preds, "lower": ci.iloc[:,0], "upper": ci.iloc[:,1]}))