import numpy as np
import pandas as pd


class MyAlgo(QCAlgorithm):
    
    def __init__(self):
    
        # Period of sma 
        self.sma_period = 60   # Base period for SMA
        self.comp_period = 1   # Number of the previous periods to compare
                               # Here, we compare current sma vs 1-period previous sma
        
    def Initialize(self):

        self.SetCash(1000000)
        # Start and end dates for the backtest.
        self.SetStartDate(2017,1,1)
        self.SetEndDate(2017,3,31)
        
        self.spy = self.AddEquity("SPY", Resolution.Minute).Symbol

        
        self.security_list = ["XLK", "XLF", "IVV", "XLE", "XLV"]
        
        for i in range(len(self.security_list)):
            self.AddEquity(self.security_list[i],Resolution.Minute)
        
        
        # Schedule the rebalance function
        self.Schedule.On(self.DateRules.EveryDay("SPY"), 
        self.TimeRules.AfterMarketOpen("SPY", 30), 
        Action(self.rebalance))
        
        
    
    def OnData(self, data):
        pass


    def rebalance(self):

        # ---------------------------------------------------------------
        # Stocks History               ----------------------------------
        # ---------------------------------------------------------------
        
        # Get history of stocks ('close') for the past N-days
        history = self.History(self.security_list, self.sma_period+self.comp_period, Resolution.Daily)
        price_history = history['close'].unstack(level=0)
        # Get the current close of each stock
        current = self.History(self.security_list, 1, Resolution.Minute)
        curr_history = current['close'].unstack(level=0)
        # Combine all in one
        price_history = price_history.append(curr_history.iloc[0])
        
        # Re-Create working dataframe to play with  --------------------
        # DF 
        df1 = pd.DataFrame(price_history.values, columns = price_history.columns)
        df1 = df1.dropna()
        #self.Log('Working_DF : ' +'\n' + str(df1) + '\n')
        
        ##############################################################
        # SMA   ------------------------------------------
        ##############################################################
        
        # Get current sma
        curr_sma = df1[-self.sma_period:].mean(0)
        # Get the previous sma 
        prev_sma = df1[-(self.sma_period+self.comp_period):-self.comp_period].mean(0)
        
        self.Log('\n'+'Curr_sma:'+'\n'+str(curr_sma)+'\n'+'Prev_sma:'+'\n'+str(prev_sma)+'\n')