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
from alpaca.data.requests import StockBarsRequest
from alpaca.data.timeframe import TimeFrame


load_dotenv()

API_KEY = os.getenv("ALPACA_API_KEY")
SECRET_KEY = os.getenv("ALPACA_SECRET_KEY")


def fetch_stock_data(symbol, start_date, end_date):
    client = StockHistoricalDataClient(API_KEY, SECRET_KEY)

    request = StockBarsRequest(
        symbol_or_symbols=symbol,
        timeframe=TimeFrame.Day,
        start=datetime.fromisoformat(start_date),
        end=datetime.fromisoformat(end_date),
    )

    bars = client.get_stock_bars(request)
    df = bars.df.reset_index()

    return df


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()

 



    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)

    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)