import math
import numpy as np
import pandas as pd
import statistics

from datetime import datetime, timedelta


class BasicTemplateAlgorithm(QCAlgorithm):

    def Initialize(self):

        self.SetCash(100000)
        self.SetStartDate(2015, 8, 1)
        self.SetEndDate(2015, 9, 30)
            
        # Add securities and get the data
        self.symbols = ["SPY","IWM"]

        for s in self.symbols:
            self.AddEquity(s, Resolution.Minute)

        # Schedule trades at 10am    
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
                 self.TimeRules.AfterMarketOpen("SPY", 5),       
                 Action(self.Rebalance))
        
        # Days to warm up the indicators
        self.SetWarmup(timedelta(20))
        
    def OnData(self, slice):
        pass
    
    def Rebalance(self):
        # Get 21 previous days closes plus the last minute close
        history = self.History(self.symbols, 21, Resolution.Daily)
        last_minute_data = self.History(self.symbols, 1, Resolution.Minute)

        # Append the last minute data
        history = history.append(last_minute_data)

        # Extract just the close prices and 
        # Use the 'unstack' method to make a column for each equity
        close_prices = history.close.unstack(level=0)

        # Log what we have to make sure we added the last price to the end
        self.Log("{}".format(close_prices))
        
        # Use the built in 'pct_change' and 'std' to get the volatility 
        # The calculation will include the latest price
        annl_stdev_series = (close_prices.
                            pct_change(axis=0).
                            std(axis=0, ddof=0) * (252.0 ** 0.5))
        
        # Iterate through the self.symbols list to order
        for stock in self.symbols:
            # do whatever calculation to find weight
            # here as an example it's just the ratio of std dev
            weight = annl_stdev_series[stock] / annl_stdev_series.sum()
            self.SetHoldings(stock, weight)
            # Log the weights to see what we ordered
            self.Log("{}  {}".format(stock, weight))