import math
import numpy as np
import pandas as pd
import talib

from datetime import datetime, timedelta

class BasicTemplateAlgorithm(QCAlgorithm):

    def Initialize(self):

        self.SetCash(100000)
        self.SetStartDate(2017, 1, 1)
        self.SetEndDate(2017, 1, 31)
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
            
        # Add securities and get the data
        self.equities = ["SPY","QQQ"]
        self.sma5 = []
        self.rsi30 = []
        for s in self.equities:
            self.AddEquity(s, Resolution.Minute)
            self.sma5.append(self.SMA(s, 5, Resolution.Daily))

        # Schedule trades at 10am    
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.AfterMarketOpen("SPY", 30),       
                 Action(self.Rebalance))
  
        # 5 days to warm up the indicators
        self.SetWarmup(timedelta(5))
        
    def OnData(self, slice):
        pass
    
    def Rebalance(self):
        
        # global variables
        allocation = 1.0
        vol = 100+1
        long_mult = 1.0
        short_mult = -1.0
        target_weight = 0.9
        security_ratio = 1.0
        
        for i in range(len(self.equities)):
            # position sizing variables
            #prices = self.History(self.equities[i], vol)['close']
            #price_stdev = (prices.pct_change()+1).apply(np.log).std()*(252**0.5)  
            #target_weight = (allocation/price_stdev/100)*security_ratio
            
            # trading signal variables
            price = self.Securities[self.equities[i]].Price
            sma5 = self.sma5[i].Current.Value
            #self.sma5[i].Update(IndicatorDataPoint(Time, price))
            shares = self.Portfolio[self.equities[i]].Quantity
                        
            # trading logic
            if shares <= 0 and price >= sma5:
                self.SetHoldings(self.equities[i], target_weight*long_mult)
            if shares > 0 and price < sma5:
                self.SetHoldings(self.equities[i], target_weight*short_mult)
                #self.Liquidate()
            self.Debug("{0}:  price: {1}  sma5: {2}  ".format(self.equities[i], round(price,2), round(sma5,2)))
            self.Plot(self.equities[i], price)
            self.Plot(self.equities[i], sma5)