Update stock_analysis.py

stock_analysis.py

@@ -1,59 +1,93 @@
-import gradio as gr
-from datetime import datetime, date, timedelta
-from config import index_options, time_intervals, START_DATE, END_DATE, FORECAST_PERIOD
-from data_fetcher import get_stocks_from_index
-from stock_analysis import get_stock_graph_and_info
-
-def validate_date(date_string):
-    try:
-        return datetime.strptime(date_string, "%Y-%m-%d").date()
-    except ValueError:
-        return None
-
-demo = gr.Blocks()
-
-with demo:
-    d1 = gr.Dropdown(index_options, label='Please select Index...', info='Will be adding more indices later on', interactive=True)
-    d2 = gr.Dropdown(label='Please Select Stock from your selected index', interactive=True)
-    d3 = gr.Dropdown(time_intervals, label='Select Time Interval', value='1d', interactive=True)
-    d4 = gr.Radio(['Line Graph', 'Candlestick Graph'], label='Select Graph Type', value='Line Graph', interactive=True)
-    d5 = gr.Dropdown(['ARIMA', 'Prophet', 'LSTM'], label='Select Forecasting Method', value='ARIMA', interactive=True)
+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, 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 = []
+    confidence_intervals = []
     
-    # New date inputs using Textbox
-    date_start = gr.Textbox(label="Start Date (YYYY-MM-DD)", value=START_DATE.strftime("%Y-%m-%d"))
-    date_end = gr.Textbox(label="End Date (YYYY-MM-DD)", value=END_DATE.strftime("%Y-%m-%d"))
-
-    out_graph = gr.Plot()
-    out_fundamentals = gr.DataFrame()
-
-    inputs = [d1, d2, d3, d4, d5, date_start, date_end]
-    outputs = [out_graph, out_fundamentals]
-
-    def update_stock_options(index):
-        stocks = get_stocks_from_index(index)
-        return gr.Dropdown(choices=stocks)
-
-    def process_inputs(*args):
-        idx, stock, interval, graph_type, forecast_method, start_date, end_date = args
-        
-        start = validate_date(start_date)
-        end = validate_date(end_date)
-        
-        if start is None or end is None:
-            return gr.Textbox(value="Invalid date format. Please use YYYY-MM-DD."), None
-        
-        if start > end:
-            return gr.Textbox(value="Start date must be before end date."), None
-        
-        return get_stock_graph_and_info(idx, stock, interval, graph_type, forecast_method, start, end)
-
-    d1.change(update_stock_options, d1, d2)
-    d2.change(process_inputs, inputs, outputs)
-    d3.change(process_inputs, inputs, outputs)
-    d4.change(process_inputs, inputs, outputs)
-    d5.change(process_inputs, inputs, outputs)
-    date_start.change(process_inputs, inputs, outputs)
-    date_end.change(process_inputs, inputs, outputs)
-
-if __name__ == "__main__":
-    demo.launch()
+    if series.shape[1] > 1:
+        series = series['Close'].values.tolist()
+    
+    if model == "ARIMA":
+        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
+    elif model == "LSTM":
+        # Implement LSTM forecasting method
+        pass
+
+    return predictions, confidence_intervals
+
+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':
+        ticker_name += '.L' if ticker_name[-1] != '.' else 'L'
+    elif ticker_dict[idx] == 'CAC 40':
+        ticker_name += '.PA'
+
+    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_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_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'],
+                                             high=series['High'],
+                                             low=series['Low'],
+                                             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",
+                      yaxis_title="Price")
+
+    fundamentals = get_company_info(ticker_name)
+
+    return fig, fundamentals