from pandas import DataFrame, read_csv, set_option, load, read_clipboard, pivot
set_option('display.max_columns', 150)
set_option('display.max_rows', 500)
set_option('display.date_dayfirst', True)
set_option('notebook_repr_html',True)
set_option('display.width',500)

from zipline.algorithm import TradingAlgorithm
from zipline.transforms import MovingAverage
from zipline.utils.factory import load_from_yahoo
from datetime import datetime
import pytz

class DualMovingAverage(TradingAlgorithm):
    """Dual Moving Average Crossover algorithm.
    This algorithm buys stock once its short moving average crosses
    its long moving average (indicating upwards momentum) and sells
    it once the averages cross again (indicating downwards
    momentum).

    """
    def initialize(self, short_window=3, long_window=10, starting_cash = 1000):
       
        self.add_transform(MovingAverage, 'short_mavg', ['price'],
                           window_length=short_window)

        self.add_transform(MovingAverage, 'long_mavg', ['price'],
                           window_length=long_window)

        self.invested = False

    def handle_data(self, data):
        self.short_mavg = data[stock].short_mavg['price']
        self.long_mavg = data[stock].long_mavg['price']
        self.buy = False
        self.sell = False

        if self.short_mavg > self.long_mavg and not self.invested:
            self.order(stock, tr_volume)
            self.invested = True
            self.buy = True
        elif self.short_mavg < self.long_mavg and self.invested:
            self.order(stock, -1*tr_volume)
            self.invested = False
            self.sell = True

        self.record(short_mavg=self.short_mavg,
                    long_mavg=self.long_mavg,
                    buy=self.buy,
                    sell=self.sell)
        

if __name__ == '__main__':
    start = datetime(2013, 1, 1, 0, 0, 0, 0, pytz.utc)
    end = datetime(2013, 6, 8, 0, 0, 0, 0, pytz.utc)
    stock = "WOW.AX"
    tr_volume = 10000

    data = load_from_yahoo(stocks=[stock], indexes={}, start=start, end=end)
    dma = DualMovingAverage()
    results = dma.run(data)

fig = plt.figure(figsize=(20, 12), dpi=320, facecolor="#342F2F", edgecolor='w')

ax1 = fig.add_subplot(211, axisbg='k')

results['short_mavg'].plot(ax=ax1, color='#00FF26', ls=":", lw=2)
results['long_mavg'].plot(ax=ax1, color='#FFA500')

data[stock].plot(ax=ax1, lw=3, color='#226dc9')

ax1.plot(results.ix[results.buy].index, results.short_mavg[results.buy],
         '^', markersize=10, color='w')
ax1.plot(results.ix[results.sell].index, results.short_mavg[results.sell],
         'v', markersize=10, color='w')
ax1.grid(color='#D9D9D9', alpha=1, lw=0.5)


ax2 = fig.add_subplot(212, axisbg='k')
results.portfolio_value.plot(ax=ax2, lw=3, color='#179349')
ax2.grid(color='#D9D9D9', alpha=1, lw=0.5)
x = np.arange(0,len(results.portfolio_value))
ax2.fill_between(results.portfolio_value.index,results.portfolio_value, color="#179349", alpha=0.5)

ax1.set_title('Price and Moving Avarage(s), $\n', color="w")
ax2.set_title('Portfolio, $\n', color="w")

ax1.tick_params(axis="x", colors="#BFBFBF")
ax1.tick_params(axis="y", colors="#BFBFBF")
ax2.tick_params(axis="x", colors="#BFBFBF")
ax2.tick_params(axis="y", colors="#BFBFBF")

ax1.set_xlabel("")


subplots_adjust(hspace = 0.4)

leg = ax1.legend(loc=0, fancybox=True, shadow=True)
leg.get_frame().set_alpha(0.01)
for text in leg.get_texts(): text.set_color("#BFBFBF")

results.portfolio_value.tail()