Save to CSV and plot resuts added

requirements.txt

@@ -1,3 +1,4 @@
 pandas
+matplotlib
 python-dotenv
 alpaca-py

stock_analysis.py

@@ -2,6 +2,8 @@ from datetime import datetime
 import os
 
 import pandas as pd
+import matplotlib.pyplot as plt
+import matplotlib.dates as mdates
 from dotenv import load_dotenv
 
 from alpaca.data.historical import StockHistoricalDataClient
@@ -32,27 +34,149 @@ def fetch_stock_data(symbol, start_date, end_date):
 
 
 def analyze_stock(df):
+    df["timestamp"] = pd.to_datetime(df["timestamp"])
+    df.set_index("timestamp", inplace = True)
     df["Return"] = df["close"].pct_change()
     df["MA_3"] = df["close"].rolling(window=3).mean()
 
-    print("\nFirst rows:")
+ 
-    print(df.head())
 
-    print("\nStock analysis:")
-    print(f"Highest price: {df['close'].max():.2f}")
-    print(f"Lowest price: {df['close'].min():.2f}")
-    print(f"Average price: {df['close'].mean():.2f}")
-    print(f"Average return: {df['Return'].mean():.4f}")
 
-    print("\nClose, Return, MA_3:")
-    print(df[["timestamp", "close", "Return", "MA_3"]])
 
+    stats = {
+    "highest_price": df["close"].max(),
+    "lowest_price": df["close"].min(),
+    "average_price": df["close"].mean(),
+    "average_return": df["Return"].mean(),
+    }
+
+    return df, stats
+
+def save_to_csv(df, symbol, start_date, end_date):
+
+    os.makedirs("./data", exist_ok=True)
+
+    df.to_csv(f"./data/{symbol}_{end_date}_{start_date}.csv", index = True)
+
+def plot_results(df):
+
+    os.makedirs("./data/figures", exist_ok=True)
+
+    plt.figure(figsize=(20,5))
+    plt.plot(df.index, df["close"], label = "Close", linestyle = "-", marker = "o", color = "blue")
+    plt.plot(df.index,df["MA_3"], label = "MA_3", linestyle = ":", color = "orange")
+    plt.grid(True)
+
+
+    
+    plt.xticks(df.index, rotation = 45)
+
+    plt.tight_layout()
+        
+    plt.legend()
+
+    plt.savefig(f"./data/figures/Close&MA_3_{symbol}_{end_date}_{start_date}.png", dpi=300, bbox_inches="tight")
+
+    plt.close()
+
+    plt.figure(figsize=(20,5))
+
+    plt.scatter(df.index, df["Return"])
+
+    plt.xticks(df.index, rotation=45)
+
+    plt.tight_layout()
+    
+    plt.savefig(f"./data/figures/Return_{symbol}_{end_date}_{start_date}.png", dpi=300, bbox_inches="tight")
+
+    plt.close()
+
+def plot_results(df, symbol, start_date, end_date):
+    os.makedirs("./data/figures", exist_ok=True)
+
+    # Close price + moving average
+    plt.figure(figsize=(20, 5))
+
+    plt.plot(df.index, df["close"], label="Close", linestyle="-", marker="o")
+    plt.plot(df.index, df["MA_3"], label="MA_3", linestyle=":")
+
+    plt.title(f"{symbol} Close Price and 3-Day Moving Average")
+    plt.xlabel("Date")
+    plt.ylabel("Price")
+
+    plt.grid(True)
+    plt.xticks(df.index, rotation=45)
+
+    # ax = plt.gca()
+    # ax.xaxis.set_major_locator(mdates.DayLocator(interval=1))
+    # ax.xaxis.set_major_formatter(mdates.DateFormatter("%d/%m/%Y"))
+
+    plt.legend()
+    plt.tight_layout()
+
+    plt.savefig(
+        f"./data/figures/Close_MA3_{symbol}_{start_date}_{end_date}.png",
+        dpi=300,
+        bbox_inches="tight",
+    )
+
+    plt.close()
+
+    # Daily returns
+    plt.figure(figsize=(20, 5))
+
+    plt.plot(df.index, df["Return"], marker="o", linestyle="-")
+    plt.axhline(0, linestyle="--")
+
+    plt.title(f"{symbol} Daily Returns")
+    plt.xlabel("Date")
+    plt.ylabel("Return")
+
+    plt.grid(True)
+    plt.xticks(df.index, rotation=45)
+    plt.tight_layout()
+
+    plt.savefig(
+        f"./data/figures/Return_{symbol}_{start_date}_{end_date}.png",
+        dpi=300,
+        bbox_inches="tight",
+    )
+
+    plt.close()
 
 if __name__ == "__main__":
+    
     symbol = input("Enter stock symbol, for example AAPL: ").upper()
 
     start_date = input("Start date YYYY-MM-DD: ")
     end_date = input("End date YYYY-MM-DD: ")
 
     df = fetch_stock_data(symbol, start_date, end_date)
-    analyze_stock(df)
+
+    analyzed_df, stats = analyze_stock(df)
+
+    print("\nFirst rows:")
+    print(df.head())
+
+    print("\nStock analysis:")
+    print(f"Highest price: {stats['highest_price']:.2f}")
+    print(f"Lowest price: {stats['lowest_price']:.2f}")
+    print(f"Average price: {stats['average_price']:.2f}")
+    print(f"Average return: {stats['average_return']:.4f}")
+
+    print("\nClose, Return, MA_3:")
+    print(df[["close", "Return", "MA_3"]])
+
+    save = input("Save to CSV? [Y/n]: ").lower()
+
+    if save == "n":
+        pass
+    else:
+        save_to_csv(analyzed_df, symbol, start_date, end_date)
+
+    plot = input("Plot figures to png [Y/n]: ").lower()
+
+    if plot == "n":
+        pass
+    else:
+        plot_results(analyzed_df, symbol, start_date, end_date)