%matplotlib inline

import pandas as pd
import zipline
from zipline import TradingAlgorithm
from zipline.data.loader import load_bars_from_yahoo

# Uncomment and run these lines to create a cache file of benchmarks 
# and historical treasury rates in your ~/.zipline directory.
# You should only have to do this once unless you delete the cache.

# zipline.data.loader.dump_treasury_curves()
# zipline.data.loader.dump_benchmarks('SPY')

start = pd.Timestamp('2008-01-01', tz='UTC')
end = pd.Timestamp('2013-01-01', tz='UTC')

input_data = load_bars_from_yahoo(
    stocks=['AAPL', 'MSFT'],
    start=start,
    end=end,
)
input_data

input_data.loc[:,:,'price'].plot()

volumes = input_data.loc[:,:,'volume']
volumes.plot()

# Quarterly volumes.  Resampling is awesome!
volumes.resample('1Q', how='sum').plot(kind='bar', stacked=True)

# A very simple example algo, using the TradingAlgorithm subclass interface.
class BuyAndHoldAlgorithm(TradingAlgorithm):
    
    def initialize(self):
        self.has_ordered = False
    
    def handle_data(self, data):
        """
        Buy 100 shares of every stock in our universe at the start of 
        the simulation.
        """
        if not self.has_ordered:
            for stock in data:
                self.order(stock, 100)
            self.has_ordered = True

my_algo = BuyAndHoldAlgorithm()
results = my_algo.run(input_data)

# Results has very fine-grained info on what your algorithm did.
# These are the raw values that we used to create our displays
# on Quantopian.
list(results.columns)

# My algo's positions, on days 0 and 1.
list(results.positions[[0,1]])

results.portfolio_value.plot()

%%zipline --symbols=AAPL --start=2009-01-01 --end=2013-01-01 -o outvar
# This is an IPython cell magic.  It's essentially a way to pass the contents
# of a cell into another program.  The %%zipline cell magic runs a a simulation using
# the initialize and handle_data functions defined in the cell, binding its output
# to the name passed to the -o flag.

# Unlike on Quantopian, you need to import magic functions into your namespace.
from zipline.api import (
    add_history,
    history,
    order_target,
    record,
    symbol,
)

def initialize(context):
    # Register 2 histories that track daily prices,
    # one with a 100 window and one with a 300 day window
    add_history(20, '1d', 'price')
    add_history(80, '1d', 'price')

    context.i = 0


def handle_data(context, data):
    # Skip first 300 days to get full windows
    context.i += 1
    if context.i < 80:
        return

    # Compute averages
    # history() has to be called with the same params
    # from above and returns a DataFrame with a DatetimeIndex
    # and columns given by the securities in the backtest.
    short_mavg = history(20, '1d', 'price').mean()
    long_mavg = history(80, '1d', 'price').mean()

    # Trading logic
    if short_mavg['AAPL'] > long_mavg['AAPL']:
        # order_target orders as many shares as needed to
        # achieve the desired number of shares.
        order_target('AAPL', 100)
        
    elif short_mavg['AAPL'] < long_mavg['AAPL']:
        order_target('AAPL', 0)

    # Save values for later inspection
    record(AAPL=data['AAPL'].price,
           short_mavg=short_mavg['AAPL'],
           long_mavg=long_mavg['AAPL'])

output = outvar.dropna(how='any')

import matplotlib.pyplot as plt
fig = plt.figure()
aapl_subplot = fig.add_subplot('211', xlabel='Date', ylabel='Price')
position_value_subplot = fig.add_subplot('212', xlabel='Date', ylabel='Value')

output['AAPL'].plot(ax=aapl_subplot)
output['portfolio_value'].plot(ax=position_value_subplot)

plt.gcf().set_size_inches(14, 10)