"""
This cell is going to create the basic framework of the algorithm
"""

import zipline
import pytz
from datetime import datetime
import matplotlib.pyplot as pyplot
from collections import defaultdict

from zipline import TradingAlgorithm
#from zipline.api import order_target, record, symbol, history, add_history, order_target_percent, order
from zipline.api import order_target, record, symbol, history, order_target_percent, order
import numpy as np
import pandas as pd

#: NOTICE HOW THIS IS OUTSIDE INITIALIZE, BECAUSE IT IS, WE CAN REDFINE IT EVERYTIME WE REDINE INITIALIZE
min_rsi = 30
max_rsi = 70
    schedule_function(my_rebalance, date_rules.week_start(), time_rules.market_open(hours=1))
     
    # Record tracking variables at the end of each day.
    schedule_function(my_record_vars, date_rules.every_day(), time_rules.market_close())
     
    # Create our dynamic stock selector.
    attach_pipeline(make_pipeline(context), 'my_pipeline')
         
def make_pipeline(context):
    """
    A function to create our dynamic stock selector (pipeline). Documentation on
    pipeline can be found here: https://www.quantopian.com/help#pipeline-title
    """
    
    context.min_rsi = 30
    context.max_rsi = 70
    context.stock_size = 3
    
    # Create a dollar volume factor.
    dollar_volume = AverageDollarVolume(window_length=1)
 
    # Pick the top 1% of stocks ranked by dollar volume.
    high_dollar_volume = dollar_volume.percentile_between(99, 100)

    rsi = RSI()
    
    pipe = Pipeline(
        screen = high_dollar_volume,
        columns = {
            'dollar_volume': dollar_volume,
            'rsi': rsi,
        }
    )
    return pipe
 
def before_trading_start(context, data):
    """
    Called every day before market open.
    """
    context.output = pipeline_output('my_pipeline').dropna()
#    log.info("Original DF:\n%s" %context.output)
  
    context.long_secs = context.output.sort_values('rsi', ascending = True).head(context.stock_size)
    context.short_secs = context.output.sort_values('rsi', ascending = False).head(context.stock_size)
    
    context.security_list = context.long_secs.index.union(context.short_secs.index).tolist()

    context.security_set = set(context.security_list)
    
def my_assign_weights(context, data):
    """
    Assign weights to securities that we want to order.
    """
    pass
 
def my_rebalance(context,data):
    """
    Execute orders according to our schedule_function() timing. 
    """
    for stock in context.portfolio.positions:
        if stock not in context.security_set:
            order_target_percent(stock, 0)
            
    for stock in context.long_secs.index:
        if data.can_trade(stock):
            order_target_percent(stock, 1.0/context.stock_size)
            
    for stock in context.short_secs.index:
        if data.can_trade(stock):
            order_target_percent(stock, -1.0/context.stock_size)
    
    
def my_record_vars(context, data):
    """
    Plot variables at the end of each day.
    """
    pass
 
def handle_data(context,data):
    """
    Called every minute.
    """
    pass

"""
This cell is going to load in the data, run the algorithm, and print out the Sharpe Ratio
"""

data = get_pricing(
    ['AAPL', 'SPY'],
    start_date='2015-01-01',
    end_date = '2015-12-31',
    frequency='daily'
)

algo_obj = TradingAlgorithm(
    initialize=initialize, 
#    handle_data=handle_data
)

#: See this perf_manual object? I'm going to use that to get ending portfolio value
perf_manual = algo_obj.run(data.transpose(2,1,0))

#: Get the sharpe ratio
sharpe = (perf_manual.returns.mean()*252)/(perf_manual.returns.std() * np.sqrt(252))
print "The Sharpe ratio is %0.6f" % sharpe

# Step Two: Running a parameter search

min_dollar_vols = [weight for weight in np.arange(20, 40, 5)]
max_dollar_vols = [weight for weight in np.arange(60, 80, 5)]

#: Create a dictionary to hold all the results of our algorithm run
all_sharpes = defaultdict(dict)

for min_dollar_vol in min_dollar_vols:
    for max_dollar_vol in max_dollar_vols:
        
        #: Redfine initialize with new weights
        def initialize(context):
            context.min_rsi = 30
            context.max_rsi = 70
            context.stock_size = 3

            
        algo_obj = TradingAlgorithm(
            initialize=initialize, 
#            handle_data=handle_data
        )
        perf_manual = algo_obj.run(data.transpose(2,1,0))
        sharpe = (perf_manual.returns.mean()*252)/(perf_manual.returns.std() * np.sqrt(252))
        
        #: Add the result to our dict
        all_sharpes[aapl_weight][spy_weight] = sharpe

all_sharpes = pd.DataFrame(all_sharpes)
all_sharpes.index.name = "SPY Weight"
all_sharpes.columns.name = "AAPL Weight"

# Step Three: Plot our results

import matplotlib.pyplot as pyplot

def heat_map(df):
    """
    This creates our heatmap using our sharpe ratio dataframe
    """
    fig = pyplot.figure()
    ax = fig.add_subplot(111)
    axim = ax.imshow(df.values,cmap = pyplot.get_cmap('RdYlGn'), interpolation = 'nearest')
    ax.set_xlabel(df.columns.name)
    ax.set_xticks(np.arange(len(df.columns)))
    ax.set_xticklabels(list(df.columns))
    ax.set_ylabel(df.index.name)
    ax.set_yticks(np.arange(len(df.index)))
    ax.set_yticklabels(list(df.index))
    ax.set_title("Sharpe Ratios")
    pyplot.colorbar(axim)
    
#: Plot our heatmap
heat_map(all_sharpes)

print all_sharpes

# Step Four: Use your own algorithm¶