nvidia_agent_core.py

from __future__ import annotations

from pathlib import Path
from typing import Any

import joblib
import pandas as pd


MODEL_PATH = Path("nvidia_price_model.pkl")
CSV_PATH = Path("nvda_2014_to_2026.csv")

CORRELATION_KEYWORDS = [
    "news correlation",
    "news impact",
    "price impact",
    "why did stock",
    "why did nvda",
    "sell off",
    "selloff",
    "event impact",
    "correlate",
    "correlation",
    "geopolitical shock",
    "taiwan",
    "invasion",
    "invade",
    "blockade",
    "war",
    "sanction",
]

ML_KEYWORDS = [
    "predict",
    "stock price",
    "share price",
    "closing price",
    "tomorrow",
    "forecast",
    "next day",
    "ml model",
    "7 day",
    "7-day",
]

RESULTS_STRATEGY_KEYWORDS = [
    "revenue",
    "profit",
    "financial",
    "gross",
    "earnings",
    "balance sheet",
    "annual report",
    "report",
    "risk",
    "blackwell",
    "strategy",
    "future plans",
    "roadmap",
]

OUTLOOK_KEYWORDS = ["news", "latest", "outlook", "market", "analyst", "recent"]


def route_query(query: str) -> str:
    """Deterministic route used by notebook, Streamlit, Gradio, and Telegram."""
    q = query.lower()
    if any(keyword in q for keyword in CORRELATION_KEYWORDS):
        return "correlation_agent"
    if any(keyword in q for keyword in RESULTS_STRATEGY_KEYWORDS):
        return "results_strategy_agent"
    if any(keyword in q for keyword in ML_KEYWORDS):
        return "ML_agent"
    if any(keyword in q for keyword in OUTLOOK_KEYWORDS):
        return "outlook_agent"
    return "general_agent"


def load_price_checkpoint(model_path: str | Path = MODEL_PATH) -> dict[str, Any]:
    model_path = Path(model_path)
    if not model_path.exists():
        raise FileNotFoundError(f"{model_path} not found. Run Cell-8 or train_nvidia_ml_model.py first.")
    return joblib.load(model_path)


def build_future_from_checkpoint(model, checkpoint: dict[str, Any], periods: int = 7) -> pd.DataFrame:
    future = model.make_future_dataframe(periods=periods, freq="B")
    regressor_columns = checkpoint.get("regressor_columns", [])
    latest_regressors = checkpoint.get("latest_regressors", {})

    missing = [col for col in regressor_columns if col not in latest_regressors]
    if missing:
        raise ValueError(f"Checkpoint is missing latest regressor values for: {missing}")

    for col in regressor_columns:
        future[col] = latest_regressors[col]
    return future


def apply_residual_ensemble(
    forecast: pd.DataFrame,
    future: pd.DataFrame,
    checkpoint: dict[str, Any],
) -> pd.DataFrame:
    residual_model = checkpoint.get("residual_model")
    if residual_model is None:
        return forecast

    regressor_columns = checkpoint.get("regressor_columns", [])
    residual_feature_columns = checkpoint.get(
        "residual_feature_columns",
        regressor_columns + ["prophet_yhat"],
    )

    features = future[regressor_columns].copy()
    features["prophet_yhat"] = forecast["yhat"].values
    features = features[residual_feature_columns]

    correction = residual_model.predict(features)
    adjusted = forecast.copy()
    adjusted["prophet_yhat"] = adjusted["yhat"]
    adjusted["residual_correction"] = correction
    adjusted["yhat"] = adjusted["prophet_yhat"] + correction
    adjusted["yhat_lower"] = adjusted["yhat_lower"] + correction
    adjusted["yhat_upper"] = adjusted["yhat_upper"] + correction
    return adjusted


def format_forecast_table(forecast: pd.DataFrame, days: int = 7) -> pd.DataFrame:
    pred = forecast.tail(days)[["ds", "yhat", "yhat_lower", "yhat_upper"]].copy()
    pred = pred.rename(
        columns={
            "ds": "Date",
            "yhat": "Predicted_Close",
            "yhat_lower": "Lower_Bound",
            "yhat_upper": "Upper_Bound",
        }
    )
    pred["Date"] = pd.to_datetime(pred["Date"]).dt.strftime("%Y-%m-%d")
    for col in ["Predicted_Close", "Lower_Bound", "Upper_Bound"]:
        pred[col] = pred[col].round(2)
    return pred[["Date", "Predicted_Close", "Lower_Bound", "Upper_Bound"]]


def run_ensemble_forecast(
    periods: int = 7,
    model_path: str | Path = MODEL_PATH,
) -> tuple[dict[str, Any], pd.DataFrame, pd.DataFrame]:
    checkpoint = load_price_checkpoint(model_path)
    model = checkpoint["prophet_model"]
    future = build_future_from_checkpoint(model, checkpoint, periods=periods)
    prophet_forecast = model.predict(future)
    ensemble_forecast = apply_residual_ensemble(prophet_forecast, future, checkpoint)
    pred_df = format_forecast_table(ensemble_forecast, days=min(7, periods))
    return checkpoint, ensemble_forecast, pred_df


def predict_nvidia_stock_payload(periods: int = 7) -> dict[str, Any]:
    checkpoint, _forecast, pred_df = run_ensemble_forecast(periods=periods)
    next_day_prediction = float(pred_df.iloc[0]["Predicted_Close"])
    final_prediction = float(pred_df.iloc[-1]["Predicted_Close"])
    last_close = float(checkpoint.get("last_close", 0))
    expected_move_pct = ((final_prediction / last_close) - 1) * 100 if last_close else None

    return {
        "prediction": next_day_prediction,
        "day_7_prediction": final_prediction,
        "expected_7day_move_pct": round(expected_move_pct, 2) if expected_move_pct is not None else None,
        "forecast_table": pred_df.to_dict(orient="records"),
        "model_version": checkpoint.get("model_version", "unknown"),
        "uses_residual_model": checkpoint.get("residual_model") is not None,
        "regressors": checkpoint.get("regressor_columns", []),
        "residual_features": checkpoint.get("residual_feature_columns", []),
        "backtest_mape": round(float(checkpoint.get("backtest_mape", 0)), 3),
        "directional_accuracy": round(float(checkpoint.get("directional_accuracy", 0)), 2),
        "last_close": last_close,
    }


def forecast_markdown(payload: dict[str, Any]) -> str:
    pred_df = pd.DataFrame(payload.get("forecast_table", []))
    table_md = pred_df.to_markdown(index=False) if not pred_df.empty else "No forecast table returned."
    expected_move = payload.get("expected_7day_move_pct")
    move_text = f"{expected_move:+.2f}%" if isinstance(expected_move, (int, float)) else "n/a"

    return f"""**NVIDIA 7-Business-Day Stock Forecast (Prophet + Residual ML Ensemble)**

**Next Trading Day Close:** **${payload.get('prediction', 0):,.2f}**
**Day-7 Expected Close:** **${payload.get('day_7_prediction', 0):,.2f}** ({move_text} vs latest close)

**7-Day Outlook:**
{table_md}

**Model:** `{payload.get('model_version', 'unknown')}`
**Uses residual ML correction:** `{payload.get('uses_residual_model')}`
**Backtested MAPE:** `{payload.get('backtest_mape')}%`
**Directional Accuracy:** `{payload.get('directional_accuracy')}%`
"""