Update app.py

app.py

@@ -85,55 +85,27 @@ def plot_to_image(plt, title, market_cap):
     buf.seek(0)
     return Image.open(buf)
 
-def plot_macd_strategy(data, company_name, ticker, market_cap):
-    """Plot MACD strategy for a single company."""
+def plot_indicator(data, company_name, ticker, indicator, market_cap):
+    """Plot selected technical indicator for a single company."""
     plt.figure(figsize=(16, 10))
-    exp1 = data['Close'].ewm(span=12, adjust=False).mean()
-    exp2 = data['Close'].ewm(span=26, adjust=False).mean()
-    macd = exp1 - exp2
-    signal = macd.ewm(span=9, adjust=False).mean()
-    
-    buy_price, sell_price, macd_signal = implement_macd_strategy(data['Close'], macd, signal)
-    
-    plt.plot(data.index, data['Close'], color='gray', linewidth=2, label="Close Price")
-    plt.plot(data.index, macd, label='MACD', color='skyblue')
-    plt.scatter(data.index, buy_price, marker='^', color='green', s=100, label='BUY SIGNAL')
-    plt.scatter(data.index, sell_price, marker='v', color='red', s=100, label='SELL SIGNAL')
-    plt.ylabel('MACD Strategy', fontsize=FONT_SIZE)
-    
-    return plot_to_image(plt, f'{company_name} ({ticker}) MACD Strategy', market_cap)
-
-def implement_macd_strategy(prices, macd, signal):
-    buy_price = []
-    sell_price = []
-    macd_signal = []
-    signal_state = 0
-    for i in range(len(macd)):
-        if macd[i] > signal[i]:
-            if signal_state != 1:
-                buy_price.append(prices[i])
-                sell_price.append(np.nan)
-                signal_state = 1
-                macd_signal.append(signal_state)
-            else:
-                buy_price.append(np.nan)
-                sell_price.append(np.nan)
-                macd_signal.append(0)
-        elif macd[i] < signal[i]:
-            if signal_state != -1:
-                buy_price.append(np.nan)
-                sell_price.append(prices[i])
-                signal_state = -1
-                macd_signal.append(signal_state)
-            else:
-                buy_price.append(np.nan)
-                sell_price.append(np.nan)
-                macd_signal.append(0)
-        else:
-            buy_price.append(np.nan)
-            sell_price.append(np.nan)
-            macd_signal.append(0)
-    return buy_price, sell_price, macd_signal
+    if indicator == "SMA":
+        sma_55 = data['Close'].rolling(window=55).mean()
+        sma_200 = data['Close'].rolling(window=200).mean()
+        plt.plot(data.index, data['Close'], label='Close')
+        plt.plot(data.index, sma_55, label='55-day SMA')
+        plt.plot(data.index, sma_200, label='200-day SMA')
+        plt.ylabel('Price', fontsize=FONT_SIZE)
+    elif indicator == "MACD":
+        exp1 = data['Close'].ewm(span=12, adjust=False).mean()
+        exp2 = data['Close'].ewm(span=26, adjust=False).mean()
+        macd = exp1 - exp2
+        signal = macd.ewm(span=9, adjust=False).mean()
+        plt.plot(data.index, macd, label='MACD')
+        plt.plot(data.index, signal, label='Signal Line')
+        plt.bar(data.index, macd - signal, label='MACD Histogram')
+        plt.ylabel('MACD', fontsize=FONT_SIZE)
+
+    return plot_to_image(plt, f'{company_name} ({ticker}) {indicator}', market_cap)
 
 def plot_indicators(company_names, indicator_types):
     """Plot the selected indicators for the selected companies."""
@@ -156,20 +128,19 @@ def plot_indicators(company_names, indicator_types):
             data, market_cap = future.result()
             if data is None:
                 continue
-            if indicator == "MACD Strategy":
-                images.append(plot_macd_strategy(data, company, ticker, market_cap))
+            images.append(plot_indicator(data, company, ticker, indicator, market_cap))
 
     return images, ""
 
 def select_all_indicators(select_all):
     """Select or deselect all indicators based on the select_all flag."""
-    indicators = ["MACD Strategy"]
+    indicators = ["SMA", "MACD"]
     return indicators if select_all else []
 
 def launch_gradio_app():
     """Launch the Gradio app for interactive plotting."""
     company_choices = list(COMPANY_TICKERS.keys())
-    indicators = ["MACD Strategy"]
+    indicators = ["SMA", "MACD"]
 
     def fetch_and_plot(company_names, indicator_types):
         images, error_message = plot_indicators(company_names, indicator_types)