Update app.py

app.py

@@ -1,156 +1,193 @@
 import gradio as gr
-import yfinance as yf
+import matplotlib.pyplot as plt
+from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
+import mplfinance as mpf
 import pandas as pd
+import yfinance as yf
+from datetime import datetime, timedelta
 import numpy as np
-from sklearn.model_selection import train_test_split
 from sklearn.preprocessing import MinMaxScaler
-from sklearn.linear_model import LinearRegression
-from sklearn.metrics import mean_squared_error, r2_score
-import mplfinance as mpf
-import matplotlib.pyplot as plt
-
-
-def get_stock_data(symbol, timeframe):
-    """Fetches stock data from Yahoo Finance."""
-    ticker = yf.Ticker(symbol)
-
-    # Calculate period based on timeframe
-    if timeframe in ['1m', '5m', '15m', '30m']:
-        period = "1d"
-    elif timeframe in ['1h']:
-        period = "5d"
-    else:
-        period = "60d"
+from tensorflow.keras.models import Sequential
+from tensorflow.keras.layers import LSTM, Dense
+from tensorflow.keras.callbacks import ModelCheckpoint, EarlyStopping
+import os
+
+
+class CandlestickApp:
+    def __init__(self):
+        self.current_symbol = None
+        self.data = None
+        self.prediction_data = None
+        self.model = None
+        self.model_path = "models/"  # Specify the directory to save models
+
+    def get_stock_data(self, symbol, timeframe, start_date, end_date):
+        try:
+            ticker = yf.Ticker(symbol)
+            data = ticker.history(start=start_date, end=end_date, interval=timeframe)  # Now with start/end dates
+            return data
+        except Exception as e:
+            return None
+
+    def calculate_indicators(self, data):
+        # Calculate RSI, SMA20, SMA50 (as before)
+        delta = data['Close'].diff()
+        gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
+        loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
+        rs = gain / loss
+        data['RSI'] = 100 - (100 / (1 + rs))
+
+        data['SMA20'] = data['Close'].rolling(window=20).mean()
+        data['SMA50'] = data['Close'].rolling(window=50).mean()
+
+        # Add more indicators as needed
+        # ...
+        return data
+
+    def plot_candlestick_chart(self, data, symbol, timeframe):
+        # Chart plotting logic (remains the same)
+        fig = plt.figure(figsize=(12, 6))
+        ax1 = fig.add_subplot(211)
+        ax2 = fig.add_subplot(212, sharex=ax1)
+
+        mpf.plot(data,
+                type='candle',
+                style='charles',
+                ax=ax1,
+                volume=ax2,
+                show_nontrading=True)
+
+        # Add moving averages
+        if len(data) >= 20:
+            ax1.plot(data.index, data['SMA20'], label='SMA20', color='blue', alpha=0.7)
+        if len(data) >= 50:
+            ax1.plot(data.index, data['SMA50'], label='SMA50', color='red', alpha=0.7)
+
+        # Add Prediction (if available)
+        if self.prediction_data is not None:
+            ax1.scatter(self.prediction_data['timestamp'],
+                           self.prediction_data['price'],
+                           color='purple',
+                           marker='*',
+                           s=100,
+                           label='Prediction')
+        ax1.legend()
+        ax1.set_title(f"{symbol} - {timeframe}")
+        ax1.tick_params(axis='x', rotation=45)
+        fig.tight_layout()
+
+        canvas = FigureCanvas(fig)
+        image = np.frombuffer(canvas.tostring_rgb(), dtype='uint8')
+        image = image.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+        return image
+
+
+    def predict_next_movement(self, data):
+        # Improved prediction using LSTM and potentially larger model
+
+         if self.model is None:
+            self.create_lstm_model(data)
+
+        # Prepare data for prediction
+        dataset = data['Close'].values.reshape(-1, 1)
+        scaler = MinMaxScaler(feature_range=(0, 1))
+        dataset = scaler.fit_transform(dataset)
+
+        # Prepare the input sequence 
+        look_back = 20 
+        X_test = []
+        X_test.append(dataset[-look_back:, 0])
+        X_test = np.array(X_test)
+        X_test = np.reshape(X_test, (X_test.shape[0], X_test.shape[1], 1))
+
+        # Make prediction
+        predicted_price = self.model.predict(X_test)
+        predicted_price = scaler.inverse_transform(predicted_price)[0][0]
+
+        # Store prediction data
+        self.prediction_data = {
+            'timestamp': data.index[-1] + pd.Timedelta(self.timeframe_var.get()),
+            'price': predicted_price
+        }
+
+        return predicted_price
     
-    data = ticker.history(period=period, interval=timeframe)
-
-    if data.empty:
-        raise ValueError(f"No data found for symbol '{symbol}' with timeframe '{timeframe}'.")
-    
-    return data
-
-
-
-def calculate_indicators(data):
-    """Calculates technical indicators."""
-    data['SMA20'] = data['Close'].rolling(window=20).mean()  # Simple Moving Average (20 days)
-    data['EMA20'] = data['Close'].ewm(span=20, adjust=False).mean()  # Exponential Moving Average (20 days)
-    data['RSI'] = calculate_rsi(data['Close'])  # Relative Strength Index
-    data['MACD'], data['MACD_Signal'], _ = calculate_macd(data['Close'])  # Moving Average Convergence Divergence
-    data['Stochastic_K'], data['Stochastic_D'] = calculate_stochastic(data['High'], data['Low'], data['Close'])  # Stochastic Oscillator
-    return data
-
-
-
-def calculate_rsi(close_prices, period=14):
-    """Calculates the Relative Strength Index (RSI)."""
-    delta = close_prices.diff()
-    gain = (delta.where(delta > 0, 0)).rolling(window=period).mean()
-    loss = (-delta.where(delta < 0, 0)).rolling(window=period).mean()
-    rs = gain / loss
-    rsi = 100 - (100 / (1 + rs))
-    return rsi
-
-def calculate_macd(close_prices, fast_period=12, slow_period=26, signal_period=9):
-    """Calculates the Moving Average Convergence Divergence (MACD)."""
-    fast_ema = close_prices.ewm(span=fast_period, adjust=False).mean()
-    slow_ema = close_prices.ewm(span=slow_period, adjust=False).mean()
-    macd = fast_ema - slow_ema
-    macd_signal = macd.ewm(span=signal_period, adjust=False).mean()
-    macd_histogram = macd - macd_signal
-    return macd, macd_signal, macd_histogram
-
-def calculate_stochastic(high_prices, low_prices, close_prices, period=14):
-    """Calculates the Stochastic Oscillator."""
-    lowest_low = low_prices.rolling(window=period).min()
-    highest_high = high_prices.rolling(window=period).max()
-    k = ((close_prices - lowest_low) / (highest_high - lowest_low)) * 100
-    d = k.rolling(window=3).mean()
-    return k, d
-
-def predict_next_day(symbol, timeframe):
-    """Predicts the next day's closing price."""
-    data = get_stock_data(symbol, timeframe)
-    data = calculate_indicators(data)
-
-    # Prepare data for training
-    data = data.dropna()
-    X = data[['SMA20', 'EMA20', 'RSI', 'MACD', 'MACD_Signal', 'Stochastic_K', 'Stochastic_D']]
-    y = data['Close']
-
-    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
-
-    # Scale the data
-    scaler = MinMaxScaler()
-    # Create a linear regression model
-    model = LinearRegression()
-    model.fit(X_train, y_train)
-
-    # Predict next day's closing price
-    last_data_point = data.iloc[-1]
-    last_data_point = last_data_point[['SMA20', 'EMA20', 'RSI', 'MACD', 'MACD_Signal', 'Stochastic_K', 'Stochastic_D']]
-    predicted_price = model.predict([last_data_point.values])[0]
-
-    # Calculate model evaluation metrics
-    y_pred = model.predict(X_test)
-    mse = mean_squared_error(y_test, y_pred)
-    rmse = np.sqrt(mse)
-    r2 = r2_score(y_test, y_pred)
-
-    print(f"Mean Squared Error: {mse:.2f}")
-    print(f"Root Mean Squared Error: {rmse:.2f}")
-    print(f"R-squared: {r2:.2f}")
-
-    return predicted_price
-
-def plot_candlestick(data, symbol, timeframe, predicted_price=None):
-    """Plots the candlestick chart with technical indicators."""
-    if data.empty:
-        raise ValueError("No valid data to plot. Please check your inputs.")
-
-    fig, ax = plt.subplots(figsize=(12, 6))
-    mpf.plot(data, type='candle', style='charles', ax=ax, volume=True, show_nontrading=True)
-
-    # Add moving averages
-    ax.plot(data.index, data['SMA20'], label='SMA20', color='blue', alpha=0.7)
-    ax.plot(data.index, data['EMA20'], label='EMA20', color='red', alpha=0.7)
-
-    # Add prediction
-    if predicted_price is not None:
-        last_timestamp = data.index[-1] + pd.Timedelta(timeframe)
-        ax.scatter(last_timestamp, predicted_price, color='green', marker='*', s=100, label='Prediction')
-
-    ax.legend()
-    ax.set_title(f"{symbol} - {timeframe}")
-    ax.tick_params(axis='x', rotation=45)
-    fig.tight_layout()
-    return fig
-
+    def create_lstm_model(self, data):
+        # Enhanced Model Training (larger model, more features, callbacks)
+
+        # Standardize the data for the model
+        dataset = data['Close'].values.reshape(-1, 1)
+        scaler = MinMaxScaler(feature_range=(0, 1))
+        dataset = scaler.fit_transform(dataset)
+        train_size = int(len(dataset) * 0.8)
+        test_size = len(dataset) - train_size
+        train_data, test_data = dataset[0:train_size,:], dataset[train_size:len(dataset),:]
+
+        # Create dataset for LSTM with possible modifications
+        def create_dataset(dataset, look_back=1):
+            X, Y = [], []
+            for i in range(len(dataset)-look_back-1):
+                a = dataset[i:(i+look_back), 0]
+                X.append(a)
+                Y.append(dataset[i + look_back, 0])
+            return np.array(X), np.array(Y)
+
+        look_back = 100
+        X_train, Y_train = create_dataset(train_data, look_back)
+        X_train = np.reshape(X_train, (X_train.shape[0], X_train.shape[1], 1))
+
+        # Create and fit the LSTM network; potentially with more layers
+        model = Sequential()
+        model.add(LSTM(units=256, return_sequences=True, input_shape=(X_train.shape[1], 1)))
+        model.add(LSTM(units=128, return_sequences=True)) # Add more LSTM layers if needed
+        model.add(LSTM(units=64))
+        model.add(Dense(1))
+        model.compile(loss='mean_squared_error', optimizer='adam')
+
+        os.makedirs(self.model_path, exist_ok=True)
+        filepath = os.path.join(self.model_path,"stock_prediction_model.h5") 
+        checkpoint = ModelCheckpoint(filepath, monitor='loss', verbose=1, save_best_only=True, mode='min') # Save best model
+        early_stop = EarlyStopping(monitor='loss', patience=10, restore_best_weights=True) # Prevent overfitting
+
+        # Train the model
+        model.fit(X_train, Y_train, epochs=500, batch_size=64, callbacks=[checkpoint, early_stop])  # Increase epochs potentially
+
+        self.model = model       
+
+    def inference(self, symbol, timeframe, start_date, end_date):
+        data = self.get_stock_data(symbol, timeframe, start_date, end_date)
+
+        if data is None:
+            return "Error fetching data", None
+
+        data = self.calculate_indicators(data)
+
+        if len(data) < 20: 
+            return "Insufficient data for prediction & chart", None
+
+        predicted_price = self.predict_next_movement(data)
+        chart = self.plot_candlestick_chart(data, symbol, timeframe)
+        return f"Predicted price: ${predicted_price:.2f}", chart
 
 def main():
-    """Gradio Interface."""
-
-    symbol_input = gr.Textbox("AAPL", label="Symbol", interactive=True)  # Moved "AAPL" to the correct position
-    timeframe_input = gr.Dropdown(label="Timeframe", choices=["1m", "5m", "15m", "30m", "1h", "1d"], value="1d")
-
-    with gr.Blocks() as interface:
-        gr.Markdown("## Real-time Stock Market Analysis")
-        with gr.Row():
-            symbol_input = gr.Textbox("AAPL", label="Symbol", interactive=True)  
-            timeframe_input = gr.Dropdown(label="Timeframe", choices=["1m", "5m", "15m", "30m", "1h", "1d"], value="1d", interactive=True)
-
-        with gr.Row():
-            predict_button = gr.Button(value="Predict")
-            predicted_price = gr.Textbox(label="Predicted Price")
-
-        with gr.Row():
-            output_plot = gr.Plot(label="Candlestick Chart")
-
-        predict_button.click(fn=predict_next_day, inputs=[symbol_input, timeframe_input], outputs=predicted_price)
-
-        predicted_price.change(fn=plot_candlestick, inputs=[symbol_input, timeframe_input, predicted_price], outputs=output_plot)
-
-    interface.launch()
+    app = CandlestickApp()
+    iface = gr.Interface(
+        fn=app.inference,
+        inputs=[
+            gr.inputs.Textbox(lines=1, placeholder="Enter Stock Symbol (e.g., AAPL)", label="Stock Symbol"),
+            gr.inputs.Dropdown(["1m", "5m", "15m", "30m", "1h", "1d"], label="Timeframe"),
+            gr.inputs.DatePicker(label="Start Date"),  # New, for start date
+            gr.inputs.DatePicker(label="End Date"),    # New, for end date
+        ],
+        outputs=[
+            gr.outputs.Textbox(label="Prediction"),
+            gr.outputs.Image(label="Candlestick Chart"), 
+        ],
+        title="Stock Market Prediction & Analysis (Enhanced)",
+        description="Enter a stock symbol, timeframe, and date range to get a prediction and candlestick chart analysis.",
+    )
+
+    iface.launch()
 
 if __name__ == "__main__":
     main()