stock-ta

# TA-Py
# Perform Technical Analysis on a given set of Yahoo! Finance data.

import datetime
import matplotlib.finance as fin
import urllib2
import matplotlib.pyplot as plt
import talib

# Define "asking for a symbol" function.

def ask_symbol():
    symbol = raw_input('Look up which symbol?\n>')
    print '\nFetching data for ' + symbol + '...' + '\n'
    return symbol

# Define "asking for a timeframe" function.

def ask_timeframe():
    while True:
        timeframe = raw_input('What timeframe? \n a: 1 Year \n b: 3 Months \n c: Daily \n>')
        if timeframe == 'a':
            past_date = datetime.date.today() - datetime.timedelta(days=365) # One year ago.
            print '\nBeginning from ' + str(past_date) + '...\n'
            return past_date
        elif timeframe == 'b':
            past_date = datetime.date.today() - datetime.timedelta(days=31*3) # One month ago.
            print '\nBeginning from ' + str(past_date) + '...\n'
            return past_date
        elif timeframe == 'c':
            # past_date = current_date - datetime.timedelta(days=1) # One day ago.
            # return past_date
            # break
            print '\nYahoo doesn\'t have enough data for that. Try a different timeframe.\n'
        else:
            print '\nInvalid option!\n'


# Define "getting the stock data" function.

def get_data(symbol, past_date):
    while True:
        try:
            data = fin.quotes_historical_yahoo(symbol, (past_date.year, past_date.month, past_date.day), (datetime.date.today().year, datetime.date.today().month, datetime.date.today().day), asobject=None)
            return data
        except urllib2.HTTPError:
            print 'Symbol not found!\n'
            symbol = ask_symbol()


# Define "process the raw data into closing price" function.

def process_data(data):
    date = data[:,0]
    close = data[:,2]
    sma = talib.SMA(close)
    rsi = talib.RSI(close)
    macd = talib.MACD(close)
    return (date, close, sma, rsi, macd)

# TODO: Define "determine if RSI is above 70 or below 30" function.

# TODO: Define "determine if MACD histogram is zero" function.

# Define "plot the stock data over time" function.

def plot_data(date, close, sma, rsi, macd):
    plt.figure(1)

    plt.subplot(311)
    plt.plot_date(x = date, y = close, fmt = 'r-')
    plt.plot_date(x = date, y = sma, fmt = 'b-')
    plt.ylabel("Closing price ($)")
    plt.title("Stock price to date")

    plt.subplot(312)
    plt.plot_date(x = date, y = rsi, fmt = 'g-')
    plt.axhline(y=70, xmin=0, xmax=1)
    plt.axhline(y=30, xmin=0, xmax=1)
    plt.ylabel("Percent (%)")

    plt.subplot(313)
    plt.plot_date(x = date, y = macd[0], fmt = 'k-')
    plt.plot_date(x = date, y = macd[1], fmt = 'r-')
    plt.axhline(y=0, xmin = 0, xmax = 1)
    plt.fill_between(date, macd[0]-macd[1], 0, alpha=0.5, facecolor='darkslategrey', edgecolor = 'darkslategrey')

    plt.show()

def menu():
    choice = raw_input("What would you like to do?\n\na: Lookup a symbol\nq: Quit\n\n>")
    if choice == "a":
        lookup()
    elif choice == "q":
        print "Goodbye!"
        quit()
    else:
        print "Invalid option!\n"
        menu()


def lookup():
    # Take stock symbol input.
    symbol = ask_symbol()
    # Establish the beginning date and timeframe.
    past_date = ask_timeframe()
    # Download date and price data from Yahoo. Ask the user to reinput the symbol if it doesn't exist.
    data = get_data(symbol, past_date)
    # Process the raw data.
    date, close, sma, rsi, macd = process_data(data)
    # Plot data.
    plot_data(date, close, sma, rsi, macd)
    # Return to main menu.
    menu()


if __name__ == "__main__":
    print "Welcome to TA-Py!"
    menu()