#+--------------------------------------------------------------------------------------------+
# Model for Stock Price Prediction via Linear Regression
# Using today vs tomorrow analysis
# app.py for deployment on Hugging Face
# 
# Written by: Prakash R. Kota
# Location: East Greenbush, NY
#
#
# Based on: 2a_HF_linear_model_app.ipynb
# and Based on: 1a_HF_linear_model.ipynb
#
# Written on: 25 Mar 2025
# Last update: 25 Mar 2025
#+--------------------------------------------------------------------------------------------+

import os
import joblib
import yfinance as yf
import pandas as pd
import numpy as np
from datetime import datetime
from sklearn.metrics import mean_absolute_percentage_error
from sklearn.linear_model import LinearRegression
import gradio as gr

# --- Global Variables --- #
Stock = "NVDA"
model_path = f"./model/{Stock}_lr_model.pkl"
start_date = "2020-01-01"
test_start_date = "2025-01-01"
today = datetime.today().strftime('%Y-%m-%d')

# --- Load Model --- #
if not os.path.exists(model_path):
    raise FileNotFoundError(f"Model not found at {model_path}. Please train and save it first.")

model = joblib.load(model_path)

def run_prediction():
    # --- Download Data --- #
    data = yf.download(Stock, start=start_date, end=today)[["Close"]]
    data.dropna(inplace=True)
    data.index = pd.to_datetime(data.index)

    # --- Create X and y for prediction --- #
    xactual = data[:-1]["Close"].values.reshape(-1, 1)  # Today's close
    ytrue = data[1:]["Close"]                           # Tomorrow's close (Series)
    dates = data[1:].index

    # Keep only test range (2025+)
    mask = dates >= test_start_date
    xactual = xactual[mask]
    ytrue = ytrue[mask]
    dates = dates[mask]

    # Predict
    typred = model.predict(xactual)

    # --- Build DataFrame --- #
    pred_df = pd.DataFrame({
        "Date": dates,
        "Actual Close": ytrue.squeeze().values,
        "Predicted Close": typred.flatten()
    })

    # --- Calculate Metrics --- #
    pred_df["% Error Raw"] = ((pred_df["Predicted Close"] - pred_df["Actual Close"]) / pred_df["Actual Close"]) * 100

    # Compute MAPE and SAPE first
    mape = np.mean(np.abs(pred_df["% Error Raw"]))
    sape = np.std(np.abs(pred_df["% Error Raw"]))

    # Format the columns
    pred_df["Actual Close"] = pred_df["Actual Close"].round(2)
    pred_df["Predicted Close"] = pred_df["Predicted Close"].round(2)
    pred_df["% Error"] = pred_df["% Error Raw"].apply(lambda x: f"$ {x:+.2f}")
    pred_df.drop(columns=["% Error Raw"], inplace=True)

    # Add MAPE Range per row to table
    pred_df["±MAPE Range"] = pred_df["Predicted Close"].apply(
        lambda x: f"${x * (1 - mape/100):.2f} to ${x * (1 + mape/100):.2f}"
    )


    # --- Next Day Prediction --- #
    latest_close = float(data["Close"].iloc[-1])
    latest_date = data.index[-1].strftime("%Y-%m-%d")
    next_pred = float(model.predict(np.array([[latest_close]]))[0])
    next_date = (data.index[-1] + pd.tseries.offsets.BDay(1)).strftime("%Y-%m-%d")
    now = datetime.now().strftime("%Y-%m-%d %H:%M:%S")

    # Expected Ranges
    mape_range = (next_pred * (1 - mape/100), next_pred * (1 + mape/100))
    sape_range = (next_pred * (1 - sape/100), next_pred * (1 + sape/100))

    summary = f"""
Prediction for {Stock} made at {now}:
Last available date: {latest_date}, Close = ${latest_close:.2f}
Predicted closing price for next trading day ({next_date}): ${next_pred:.2f}
Expected range (±MAPE): ${mape_range[0]:.2f} to ${mape_range[1]:.2f}
Expected range (±SAPE):  ${sape_range[0]:.2f} to ${sape_range[1]:.2f}
"""

    # Sort and format
    pred_df = pred_df.sort_values("Date", ascending=False)
    pred_df["Date"] = pred_df["Date"].dt.strftime("%Y-%m-%d")

    return summary, pred_df

# --- Gradio Interface --- #
description = f"""<h3>Linear Regression Stock Prediction</h3>
<p>This app loads a trained linear regression model for <b>{Stock}</b> and predicts the next trading day's close based on the last available price. 
It also shows historical prediction accuracy from 2025 onward.</p>"""

demo = gr.Interface(
    fn=run_prediction,
    inputs=[],
    outputs=[
        gr.Textbox(label="Prediction Summary", lines=6),
        gr.Dataframe(label="Prediction Table (2025+)", wrap=True)
    ],
    title="Stock Prediction using Linear Regression",
    description=description,
    allow_flagging="never"
)

if __name__ == "__main__":
    demo.launch()