Update stock_analysis.py

stock_analysis.py

@@ -1,26 +1,27 @@
 import pandas as pd
+import numpy as np
 import plotly.graph_objects as go
 from datetime import timedelta
 from statsmodels.tsa.arima.model import ARIMA
-from config import FORECAST_PERIOD, ticker_dict
+from config import FORECAST_PERIOD, ticker_dict, CONFIDENCE_INTERVAL
 from data_fetcher import get_stock_data, get_company_info
 
 def is_business_day(a_date):
     return a_date.weekday() < 5
 
 def forecast_series(series, model="ARIMA", forecast_period=FORECAST_PERIOD):
-    predictions = list()
+    predictions = []
+    confidence_intervals = []
+    
     if series.shape[1] > 1:
         series = series['Close'].values.tolist()
     
     if model == "ARIMA":
-        for _ in range(forecast_period):
-            model = ARIMA(series, order=(5, 1, 0))
-            model_fit = model.fit()
-            output = model_fit.forecast()
-            yhat = output[0]
-            predictions.append(yhat)
-            series.append(yhat)
+        model = ARIMA(series, order=(5, 1, 0))
+        model_fit = model.fit()
+        forecast = model_fit.forecast(steps=forecast_period, alpha=(1 - CONFIDENCE_INTERVAL))
+        predictions = forecast.predicted_mean
+        confidence_intervals = forecast.conf_int()
     elif model == "Prophet":
         # Implement Prophet forecasting method
         pass
@@ -28,9 +29,9 @@ def forecast_series(series, model="ARIMA", forecast_period=FORECAST_PERIOD):
         # Implement LSTM forecasting method
         pass
 
-    return predictions
+    return predictions, confidence_intervals
 
-def get_stock_graph_and_info(idx, stock, interval, graph_type, forecast_method):
+def get_stock_graph_and_info(idx, stock, interval, graph_type, forecast_method, start_date, end_date):
     stock_name, ticker_name = stock.split(":")
     
     if ticker_dict[idx] == 'FTSE 100':
@@ -38,27 +39,33 @@ def get_stock_graph_and_info(idx, stock, interval, graph_type, forecast_method):
     elif ticker_dict[idx] == 'CAC 40':
         ticker_name += '.PA'
 
-    series = get_stock_data(ticker_name, interval)
-    predictions = forecast_series(series, model=forecast_method)
+    series = get_stock_data(ticker_name, interval, start_date, end_date)
+    predictions, confidence_intervals = forecast_series(series, model=forecast_method)
 
     last_date = pd.to_datetime(series['Date'].values[-1])
-    forecast_week = []
-    i = 1
-    while len(forecast_week) < FORECAST_PERIOD:
-        next_date = last_date + timedelta(days=i)
-        if is_business_day(next_date):
-            forecast_week.append(next_date)
-        i += 1
-
-    predictions = predictions[:len(forecast_week)]
-    forecast_week = forecast_week[:len(predictions)]
+    forecast_dates = pd.date_range(start=last_date + timedelta(days=1), periods=FORECAST_PERIOD)
+    forecast_dates = [date for date in forecast_dates if is_business_day(date)]
 
-    forecast = pd.DataFrame({"Date": forecast_week, "Forecast": predictions})
+    forecast = pd.DataFrame({
+        "Date": forecast_dates,
+        "Forecast": predictions,
+        "Lower_CI": confidence_intervals.iloc[:, 0],
+        "Upper_CI": confidence_intervals.iloc[:, 1]
+    })
 
     if graph_type == 'Line Graph':
         fig = go.Figure()
         fig.add_trace(go.Scatter(x=series['Date'], y=series['Close'], mode='lines', name='Historical'))
         fig.add_trace(go.Scatter(x=forecast['Date'], y=forecast['Forecast'], mode='lines', name='Forecast'))
+        fig.add_trace(go.Scatter(
+            x=forecast['Date'].tolist() + forecast['Date'].tolist()[::-1],
+            y=forecast['Upper_CI'].tolist() + forecast['Lower_CI'].tolist()[::-1],
+            fill='toself',
+            fillcolor='rgba(0,100,80,0.2)',
+            line=dict(color='rgba(255,255,255,0)'),
+            hoverinfo="skip",
+            showlegend=False
+        ))
     else:  # Candlestick Graph
         fig = go.Figure(data=[go.Candlestick(x=series['Date'],
                                              open=series['Open'],
@@ -67,6 +74,15 @@ def get_stock_graph_and_info(idx, stock, interval, graph_type, forecast_method):
                                              close=series['Close'],
                                              name='Historical')])
         fig.add_trace(go.Scatter(x=forecast['Date'], y=forecast['Forecast'], mode='lines', name='Forecast'))
+        fig.add_trace(go.Scatter(
+            x=forecast['Date'].tolist() + forecast['Date'].tolist()[::-1],
+            y=forecast['Upper_CI'].tolist() + forecast['Lower_CI'].tolist()[::-1],
+            fill='toself',
+            fillcolor='rgba(0,100,80,0.2)',
+            line=dict(color='rgba(255,255,255,0)'),
+            hoverinfo="skip",
+            showlegend=False
+        ))
 
     fig.update_layout(title=f"Stock Price of {stock_name}",
                       xaxis_title="Date",