Create app.py

app.py

@@ -0,0 +1,587 @@
+import streamlit as st
+import pandas as pd
+import numpy as np
+import yfinance as yf
+import plotly.graph_objects as go
+from plotly.subplots import make_subplots 
+import plotly.express as px
+from datetime import datetime, timedelta
+from statsmodels.tsa.statespace.sarimax import SARIMAX
+from prophet import Prophet
+from sklearn.ensemble import RandomForestRegressor
+import xgboost as xgb
+from sklearn.preprocessing import StandardScaler, MinMaxScaler
+from keras.models import Sequential
+from keras.optimizers import Adam
+from keras.layers import Dense, LSTM
+from sklearn.model_selection import train_test_split, TimeSeriesSplit
+from sklearn.metrics import mean_squared_error, mean_absolute_percentage_error
+import requests
+from bs4 import BeautifulSoup
+import base64
+import warnings
+from ta.trend import SMAIndicator, EMAIndicator
+from ta.momentum import RSIIndicator
+from ta.volatility import BollingerBands
+from pmdarima import auto_arima
+warnings.filterwarnings('ignore')
+
+# List of companies (display name, ticker symbol)
+COMPANIES = [
+    ("Apple", "AAPL"), ("Microsoft", "MSFT"), ("Amazon", "AMZN"), ("Google", "GOOGL"),
+    ("Facebook", "FB"), ("Tesla", "TSLA"), ("NVIDIA", "NVDA"), ("JPMorgan Chase", "JPM"),
+    ("Johnson & Johnson", "JNJ"), ("Visa", "V"), ("Procter & Gamble", "PG"), ("UnitedHealth", "UNH"),
+    ("Home Depot", "HD"), ("Mastercard", "MA"), ("Bank of America", "BAC"), ("Disney", "DIS"),
+    ("Netflix", "NFLX"), ("Coca-Cola", "KO"), ("Pepsi", "PEP"), ("Adobe", "ADBE")
+]
+
+class StockPredictor:
+    def __init__(self, data, model_type):
+        self.data = data
+        self.model_type = model_type
+        self.model = None
+        self.scaler = None
+        self.lstm_scaler = None
+
+    def preprocess_data(self):
+        self.data['Date'] = pd.to_datetime(self.data.index)
+        self.data = self.data.reset_index(drop=True)
+
+        # Enhanced Feature Engineering
+        self.data['DayOfWeek'] = self.data['Date'].dt.dayofweek
+        self.data['Month'] = self.data['Date'].dt.month
+        self.data['Year'] = self.data['Date'].dt.year
+        self.data['IsMonthEnd'] = self.data['Date'].dt.is_month_end.astype(int)
+
+        # Technical Indicators
+        self.data['SMA_20'] = SMAIndicator(close=self.data['Close'], window=20).sma_indicator()
+        self.data['EMA_20'] = EMAIndicator(close=self.data['Close'], window=20).ema_indicator()
+        self.data['RSI'] = RSIIndicator(close=self.data['Close']).rsi()
+        bb = BollingerBands(close=self.data['Close'], window=20, window_dev=2)
+        self.data['BB_High'] = bb.bollinger_hband()
+        self.data['BB_Low'] = bb.bollinger_lband()
+        
+        # Log returns
+        self.data['LogReturn'] = np.log(self.data['Close'] / self.data['Close'].shift(1))
+
+        # Handle NaN values 
+        self.data.dropna(inplace=True)
+
+        # Define features for the model
+        self.features = ['Open', 'High', 'Low', 'Close', 'Volume', 'SMA_20', 'EMA_20', 'RSI', 'BB_High', 'BB_Low', 'LogReturn', 'DayOfWeek', 'Month', 'Year', 'IsMonthEnd']
+
+        # Apply scaling for XGBoost and RandomForest
+        if self.model_type in ['XGBoost', 'RandomForest']:
+            self.scaler = StandardScaler()
+            self.data[self.features] = self.scaler.fit_transform(self.data[self.features])
+
+        # Additional preprocessing for LSTM
+        if self.model_type == 'LSTM':
+            self.lstm_scaler = MinMaxScaler(feature_range=(0, 1))
+            self.data['Scaled_Close'] = self.lstm_scaler.fit_transform(self.data[['Close']])
+
+    def create_lstm_dataset(self, look_back=60):
+        scaled_data = self.data['Scaled_Close'].values
+        x, y = [], []
+        for i in range(look_back, len(scaled_data)):
+            x.append(scaled_data[i-look_back:i])
+            y.append(scaled_data[i])
+        return np.array(x), np.array(y)
+
+    def train_model(self):
+        try:
+            if self.model_type == 'LSTM':
+                x, y = self.create_lstm_dataset()
+                x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, shuffle=False)
+
+                model = Sequential([
+                    LSTM(50, return_sequences=True, input_shape=(x_train.shape[1], 1)),
+                    LSTM(50, return_sequences=False),
+                    Dense(25),
+                    Dense(1)
+                ])
+
+                model.compile(optimizer=Adam(learning_rate=0.001), loss='mean_squared_error')
+                model.fit(x_train, y_train, epochs=50, batch_size=32, validation_data=(x_test, y_test), verbose=0)
+
+                self.model = model
+
+            elif self.model_type == 'SARIMA':
+                train_data = self.data['Close']
+                # Use auto_arima to find optimal parameters
+                from pmdarima import auto_arima
+                auto_model = auto_arima(train_data, start_p=1, start_q=1, max_p=3, max_q=3, m=12,
+                                        start_P=0, seasonal=True, d=1, D=1, trace=True,
+                                        error_action='ignore', suppress_warnings=True, stepwise=True)
+                
+                self.model = SARIMAX(train_data, order=auto_model.order, seasonal_order=auto_model.seasonal_order)
+                self.model = self.model.fit(disp=False)
+
+            elif self.model_type == 'Prophet':
+                df = self.data[['Date', 'Close']].rename(columns={'Date': 'ds', 'Close': 'y'})
+                self.model = Prophet(
+                    changepoint_prior_scale=0.05,
+                    seasonality_prior_scale=10,
+                    holidays_prior_scale=10,
+                    daily_seasonality=True,
+                    weekly_seasonality=True,
+                    yearly_seasonality=True
+                )
+                for feature in ['SMA_20', 'EMA_20', 'RSI', 'BB_High', 'BB_Low']:
+                    self.model.add_regressor(feature)
+                    df[feature] = self.data[feature]
+                self.model.fit(df)
+
+            elif self.model_type == 'XGBoost':
+                X = self.data[self.features]
+                y = self.data['Close']
+                X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=False)
+                
+                param_grid = {
+                    'max_depth': [3, 5],
+                    'learning_rate': [0.01, 0.1],
+                    'n_estimators': [100, 200]
+                }
+                model = xgb.XGBRegressor(objective='reg:squarederror')
+                grid_search = GridSearchCV(estimator=model, param_grid=param_grid, cv=3, n_jobs=-1, verbose=0)
+                grid_search.fit(X_train, y_train)
+                
+                self.model = grid_search.best_estimator_
+
+            elif self.model_type == 'RandomForest':
+                X = self.data[self.features]
+                y = self.data['Close']
+                X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, shuffle=False)
+                
+                param_grid = {
+                    'n_estimators': [100, 200],
+                    'max_depth': [10, 20]
+                }
+                model = RandomForestRegressor(random_state=42)
+                grid_search = GridSearchCV(estimator=model, param_grid=param_grid, cv=3, n_jobs=-1, verbose=0)
+                grid_search.fit(X_train, y_train)
+                
+                self.model = grid_search.best_estimator_
+
+            return True
+
+        except Exception as e:
+            print(f"Error training {self.model_type} model: {str(e)}")
+            return False
+
+    def predict(self, days=30):
+        try:
+            if self.model_type == 'LSTM':
+                last_sequence = self.data['Scaled_Close'].values[-60:].reshape(1, 60, 1)
+                predictions = []
+                for _ in range(days):
+                    pred = self.model.predict(last_sequence)
+                    predictions.append(pred[0, 0])
+                    last_sequence = np.roll(last_sequence, -1, axis=1)
+                    last_sequence[0, -1, 0] = pred[0, 0]
+                return self.lstm_scaler.inverse_transform(np.array(predictions).reshape(-1, 1)).flatten()
+
+            elif self.model_type == 'SARIMA':
+                forecast = self.model.get_forecast(steps=days)
+                return forecast.predicted_mean
+
+            elif self.model_type == 'Prophet':
+                future = self.model.make_future_dataframe(periods=days)
+                for feature in ['SMA_20', 'EMA_20', 'RSI', 'BB_High', 'BB_Low']:
+                    future[feature] = self.data[feature].iloc[-1]  # Use last known value
+                forecast = self.model.predict(future)
+                return forecast['yhat'][-days:].values
+
+            elif self.model_type in ['XGBoost', 'RandomForest']:
+                last_data = self.data[self.features].iloc[-1:].values
+                predictions = []
+                for _ in range(days):
+                    pred = self.model.predict(last_data)
+                    predictions.append(pred[0])
+                    # Update last_data for next prediction
+                    last_data = np.roll(last_data, -1, axis=1)
+                    last_data[0, -5] = pred[0]  # Assuming 'Close' is the 5th from last feature
+                return predictions
+
+        except Exception as e:
+            print(f"Error predicting with {self.model_type} model: {str(e)}")
+            return None
+
+    def evaluate_model(self, test_data):
+        predictions = self.predict(len(test_data))
+        mse = mean_squared_error(test_data['Close'], predictions)
+        mape = mean_absolute_percentage_error(test_data['Close'], predictions)
+        rmse = np.sqrt(mse)
+        return mse, mape, rmse
+
+def fetch_stock_data(ticker):
+    try:
+        end_date = datetime.now()
+        start_date = datetime(2000, 1, 1) 
+        data = yf.download(ticker, start=start_date, end=end_date)
+        return data
+    except Exception as e:
+        st.error(f"Error fetching data for {ticker}: {str(e)}")
+        return None
+
+def create_test_plot(train_data, test_data, predicted_data, company_name):
+    fig = go.Figure()
+
+    fig.add_trace(go.Scatter(
+        x=train_data.index,
+        y=train_data['Close'],
+        mode='lines',
+        name='Training Data',
+        line=dict(color='blue')
+    ))
+
+    fig.add_trace(go.Scatter(
+        x=test_data.index,
+        y=test_data['Close'],
+        mode='lines',
+        name='Actual (Test) Data',
+        line=dict(color='green')
+    ))
+
+    if predicted_data is not None:
+        fig.add_trace(go.Scatter(
+            x=test_data.index,  # Align predicted data with test data
+            y=predicted_data['yhat'][-len(test_data):],
+            mode='lines',
+            name='Predicted Data',
+            line=dict(color='red', dash='dash')
+        ))
+
+    fig.update_layout(
+        title=f'{company_name} Stock Price Prediction (Test Model)',
+        xaxis_title='Date',
+        yaxis_title='Close Price',
+        template='plotly_dark',
+        hovermode='x unified',
+        xaxis_rangeslider_visible=True,
+        legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1)
+    )
+    return fig
+
+def create_prediction_plot(data, predicted_data, company_name):
+    fig = go.Figure()
+
+    fig.add_trace(go.Scatter(
+        x=data.index,
+        y=data['Close'],
+        mode='lines',
+        name='Historical Data',
+        line=dict(color='cyan')
+    ))
+
+    if predicted_data is not None:
+        future_dates = pd.date_range(start=data.index[-1] + pd.Timedelta(days=1), periods=len(predicted_data))
+        fig.add_trace(go.Scatter(
+            x=future_dates,
+            y=predicted_data['yhat'],
+            mode='lines',
+            name='Predicted Data',
+            line=dict(color='yellow')
+        ))
+
+    fig.update_layout(
+        title=f'{company_name} Stock Price Prediction',
+        xaxis_title='Date',
+        yaxis_title='Close Price',
+        template='plotly_dark',
+        hovermode='x unified',
+        xaxis_rangeslider_visible=True,
+        legend=dict(orientation="h", yanchor="bottom", y=1.02, xanchor="right", x=1)
+    )
+    return fig
+
+def create_candlestick_plot(data, company_name):
+    fig = go.Figure(data=[go.Candlestick(x=data.index,
+                                         open=data['Open'],
+                                         high=data['High'],
+                                         low=data['Low'],
+                                         close=data['Close'])])
+    fig.update_layout(
+        title=f'{company_name} Stock Price Candlestick Chart',
+        xaxis_title='Date',
+        yaxis_title='Price',
+        template='plotly_dark',
+        xaxis_rangeslider_visible=True
+    )
+    return fig
+
+def fetch_news(company_name):
+    try:
+        url = f"https://news.google.com/rss/search?q={company_name}+stock&hl=en-US&gl=US&ceid=US:en"
+        response = requests.get(url)
+        soup = BeautifulSoup(response.content, features='xml')
+        news_items = soup.findAll('item')
+        
+        news = []
+        for item in news_items[:5]:
+            news.append({
+                'title': item.title.text,
+                'link': item.link.text,
+                'pubDate': item.pubDate.text
+            })
+        
+        return news
+    except Exception as e:
+        st.error(f"Error fetching news: {str(e)}")
+        return []
+
+def get_table_download_link(df):
+    csv = df.to_csv(index=False)
+    b64 = base64.b64encode(csv.encode()).decode()
+    href = f'<a href="data:file/csv;base64,{b64}" download="stock_data.csv">Download CSV File</a>'
+    return href
+
+def main():
+    st.set_page_config(page_title="Stock Price Predictor", layout="wide")
+    st.title("Advanced Stock Price Predictor using Prophet")
+
+    st.sidebar.title("Options")
+    app_mode = st.sidebar.selectbox("Choose the app mode", ["Test Model", "Predict Stock Prices"])
+
+    if app_mode == "Test Model":
+        test_model()
+    else:
+        predict_stock_prices()
+
+def test_model():
+    st.header("Test Prophet Model")
+
+    col1, col2 = st.columns(2)
+
+    with col1:
+        company = st.selectbox("Select Company", [company for company, _ in COMPANIES])
+        test_split = st.slider("Test Data Split", 0.1, 0.5, 0.2, 0.05)
+
+    if st.button("Train and Test Model"):
+        with st.spinner("Fetching data and training model..."):
+            company_name, ticker = next((name, symbol) for name, symbol in COMPANIES if name == company)
+
+            data = fetch_stock_data(ticker)
+
+            if data is not None:
+                st.subheader("Stock Data Information")
+                st.write(data.info())
+                st.write(data.describe())
+                st.dataframe(data.head())
+
+                st.markdown(get_table_download_link(data), unsafe_allow_html=True)
+
+                split_index = int(len(data) * (1 - test_split))
+                train_data = data.iloc[:split_index]
+                test_data = data.iloc[split_index:]
+
+                predictor = StockPredictor(train_data)
+                predictor.preprocess_data()
+                if predictor.train_model():
+                    test_pred = predictor.predict(days=len(test_data))
+
+                    if test_pred is not None:
+                        mse, mape, rmse = predictor.evaluate_model(test_data)
+                        accuracy = 100 - mape * 100
+
+                        st.subheader("Model Performance")
+                        st.metric("Prediction Accuracy", f"{accuracy:.2f}%")
+                        st.metric("Mean Squared Error", f"{mse:.4f}")
+                        st.metric("Root Mean Squared Error", f"{rmse:.4f}")
+
+                        plot = create_test_plot(predictor.data, test_data, test_pred, company_name)
+                        st.plotly_chart(plot, use_container_width=True)
+                else:
+                    st.error("Failed to train the Prophet model. Please try a different dataset.")
+
+def predict_stock_prices():
+    st.header("Predict Stock Prices")
+
+    col1, col2 = st.columns(2)
+
+    with col1:
+        company = st.selectbox("Select Company", [company for company, _ in COMPANIES])
+        days_to_predict = st.slider("Days to Predict", 1, 365, 30)
+
+    if st.button("Predict Stock Prices"):
+        with st.spinner("Fetching data and making predictions..."):
+            company_name, ticker = next((name, symbol) for name, symbol in COMPANIES if name == company)
+
+            data = fetch_stock_data(ticker)
+
+            if data is not None:
+                st.subheader("Stock Data Information")
+                st.write(data.info())
+                st.write(data.describe())
+                st.dataframe(data.head())
+
+                st.markdown(get_table_download_link(data), unsafe_allow_html=True)
+
+                predictor = StockPredictor(data)
+                predictor.preprocess_data()
+                if predictor.train_model():
+                    predictions = predictor.predict(days=days_to_predict)
+
+                    if predictions is not None:
+                        plot = create_prediction_plot(data, predictions, company_name)
+                        st.plotly_chart(plot, use_container_width=True)
+
+                        candlestick_plot = create_candlestick_plot(data, company_name)
+                        st.plotly_chart(candlestick_plot, use_container_width=True)
+
+                        st.subheader("Predicted Prices")
+                        pred_df = predictions[['ds', 'yhat', 'yhat_lower', 'yhat_upper']].tail(days_to_predict)
+                        pred_df.columns = ['Date', 'Predicted Price', 'Lower Bound', 'Upper Bound']
+                        st.dataframe(pred_df)
+
+                        news = fetch_news(company_name)
+                        st.subheader("Latest News")
+                        for item in news:
+                            st.markdown(f"[{item['title']}]({item['link']}) ({item['pubDate']})")
+                else:
+                    st.error("Failed to train the Prophet model. Please try a different dataset.")
+
+def explore_data():
+    st.header("Explore Stock Data")
+
+    col1, col2 = st.columns(2)
+    
+    with col1:
+        company = st.selectbox("Select Company", [company for company, _ in COMPANIES])
+    
+    with col2:
+        period = st.selectbox("Select Time Period", ["1mo", "3mo", "6mo", "1y", "2y", "5y", "max"])
+
+    company_name, ticker = next((name, symbol) for name, symbol in COMPANIES if name == company)
+
+    if st.button("Explore Data"):
+        with st.spinner("Fetching and analyzing data..."):
+            data = yf.download(ticker, period=period)
+
+            if data is not None and not data.empty:
+                st.subheader(f"{company_name} Stock Data")
+                
+                # Create tabs for different visualizations
+                tab1, tab2, tab3, tab4, tab5 = st.tabs(["Price History", "OHLC", "Technical Indicators", "Volume & Turnover", "Statistics"])
+                
+                with tab1:
+                    # Stock Price History
+                    fig = go.Figure()
+                    fig.add_trace(go.Scatter(x=data.index, y=data['Open'], mode='lines', name='Open'))
+                    fig.add_trace(go.Scatter(x=data.index, y=data['High'], mode='lines', name='High'))
+                    fig.add_trace(go.Scatter(x=data.index, y=data['Low'], mode='lines', name='Low'))
+                    fig.add_trace(go.Scatter(x=data.index, y=data['Close'], mode='lines', name='Close'))
+                    
+                    # Add rolling mean and standard deviation
+                    data['Rolling_Mean'] = data['Close'].rolling(window=20).mean()
+                    data['Rolling_Std'] = data['Close'].rolling(window=20).std()
+                    fig.add_trace(go.Scatter(x=data.index, y=data['Rolling_Mean'], mode='lines', name='20-day Rolling Mean', line=dict(dash='dash')))
+                    fig.add_trace(go.Scatter(x=data.index, y=data['Rolling_Std'], mode='lines', name='20-day Rolling Std', line=dict(dash='dot')))
+                    
+                    fig.update_layout(title=f"{company_name} Stock Price History",
+                                      xaxis_title="Date",
+                                      yaxis_title="Price",
+                                      hovermode="x unified",
+                                      template="plotly_dark")
+                    st.plotly_chart(fig, use_container_width=True)
+
+                with tab2:
+                    # OHLC Chart
+                    ohlc_fig = go.Figure(data=[go.Candlestick(x=data.index,
+                                                              open=data['Open'],
+                                                              high=data['High'],
+                                                              low=data['Low'],
+                                                              close=data['Close'])])
+                    ohlc_fig.update_layout(title=f"{company_name} OHLC Chart",
+                                           xaxis_title="Date",
+                                           yaxis_title="Price",
+                                           template="plotly_dark",
+                                           xaxis_rangeslider_visible=False)
+                    st.plotly_chart(ohlc_fig, use_container_width=True)
+
+                with tab3:
+                    # Technical Indicators
+                    data['SMA_20'] = SMAIndicator(close=data['Close'], window=20).sma_indicator()
+                    data['EMA_20'] = EMAIndicator(close=data['Close'], window=20).ema_indicator()
+                    bb = BollingerBands(close=data['Close'], window=20, window_dev=2)
+                    data['BB_High'] = bb.bollinger_hband()
+                    data['BB_Low'] = bb.bollinger_lband()
+                    data['RSI'] = RSIIndicator(close=data['Close']).rsi()
+
+                    fig = make_subplots(rows=2, cols=1, shared_xaxes=True, 
+                                        vertical_spacing=0.03, 
+                                        subplot_titles=("Price and Indicators", "RSI"),
+                                        row_heights=[0.7, 0.3])
+
+                    fig.add_trace(go.Scatter(x=data.index, y=data['Close'], mode='lines', name='Close'), row=1, col=1)
+                    fig.add_trace(go.Scatter(x=data.index, y=data['SMA_20'], mode='lines', name='SMA 20'), row=1, col=1)
+                    fig.add_trace(go.Scatter(x=data.index, y=data['EMA_20'], mode='lines', name='EMA 20'), row=1, col=1)
+                    fig.add_trace(go.Scatter(x=data.index, y=data['BB_High'], mode='lines', name='BB High'), row=1, col=1)
+                    fig.add_trace(go.Scatter(x=data.index, y=data['BB_Low'], mode='lines', name='BB Low'), row=1, col=1)
+
+                    fig.add_trace(go.Scatter(x=data.index, y=data['RSI'], mode='lines', name='RSI'), row=2, col=1)
+                    fig.add_hline(y=70, line_dash="dash", line_color="red", row=2, col=1)
+                    fig.add_hline(y=30, line_dash="dash", line_color="green", row=2, col=1)
+
+                    fig.update_layout(height=800, title_text=f"{company_name} Technical Indicators",
+                                      hovermode="x unified", template="plotly_dark")
+                    fig.update_xaxes(rangeslider_visible=False, row=2, col=1)
+                    fig.update_yaxes(title_text="Price", row=1, col=1)
+                    fig.update_yaxes(title_text="RSI", row=2, col=1)
+
+                    st.plotly_chart(fig, use_container_width=True)
+
+                with tab4:
+                    # Volume and Turnover
+                    fig = make_subplots(rows=2, cols=1, shared_xaxes=True, 
+                                        vertical_spacing=0.03, 
+                                        subplot_titles=("Volume", "Turnover (if available)"),
+                                        row_heights=[0.5, 0.5])
+
+                    fig.add_trace(go.Bar(x=data.index, y=data['Volume'], name='Volume'), row=1, col=1)
+                    
+                    if 'Turnover' in data.columns:
+                        fig.add_trace(go.Bar(x=data.index, y=data['Turnover'], name='Turnover'), row=2, col=1)
+                    else:
+                        fig.add_annotation(text="Turnover data not available", xref="paper", yref="paper", x=0.5, y=0.25, showarrow=False)
+
+                    fig.update_layout(height=600, title_text=f"{company_name} Volume and Turnover",
+                                      hovermode="x unified", template="plotly_dark")
+                    fig.update_xaxes(rangeslider_visible=False, row=2, col=1)
+                    fig.update_yaxes(title_text="Volume", row=1, col=1)
+                    fig.update_yaxes(title_text="Turnover", row=2, col=1)
+
+                    st.plotly_chart(fig, use_container_width=True)
+
+                with tab5:
+                    # Display key statistics
+                    st.subheader("Key Statistics")
+                    col1, col2, col3 = st.columns(3)
+                    with col1:
+                        st.metric("Current Price", f"${data['Close'].iloc[-1]:.2f}")
+                        st.metric("52 Week High", f"${data['High'].max():.2f}")
+                    with col2:
+                        st.metric("Volume", f"{data['Volume'].iloc[-1]:,}")
+                        st.metric("52 Week Low", f"${data['Low'].min():.2f}")
+                    with col3:
+                        returns = (data['Close'].pct_change() * 100).dropna()
+                        st.metric("Avg Daily Return", f"{returns.mean():.2f}%")
+                        st.metric("Return Volatility", f"{returns.std():.2f}%")
+
+                    # Correlation Heatmap
+                    correlation = data[['Open', 'High', 'Low', 'Close', 'Volume']].corr()
+                    heatmap_fig = px.imshow(correlation, text_auto=True, aspect="auto", color_continuous_scale='Viridis')
+                    heatmap_fig.update_layout(title="Correlation Heatmap", template="plotly_dark")
+                    st.plotly_chart(heatmap_fig, use_container_width=True)
+
+                # Display news
+                st.subheader("Latest News")
+                news = fetch_news(company_name)
+                for item in news:
+                    st.markdown(f"[{item['title']}]({item['link']}) ({item['pubDate']})")
+
+            else:
+                st.error("Failed to fetch data. Please try again.")
+
+if __name__ == "__main__":
+    main()