#from System.Collections.Generic import List
import decimal as d
import numpy as np
import pandas as pd
import time

class BatchProblem(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2017, 5, 24)  #Set Start Date
        self.SetEndDate(  2017, 6, 24)    #Set End Date
        self.SetCash(100000)           #Set Strategy Cash

        self.UniverseSettings.Resolution = Resolution.Minute
        #self.UniverseSettings.Leverage   = 4

        self.max_per_side_assets     = 20 # so 2*N is total assets at L=2
        self.coarse_count            = 300 #3000

        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        self.AddUniverse(self.CoarseSelectionFunction)
        
        self.AddEquity("SPY", Resolution.Minute)
        
            
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
            #self.TimeRules.At(8,45),
            self.TimeRules.AfterMarketOpen("SPY", 0),
            Action(self.Downloader))
        
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
            #self.TimeRules.At(8,55),
            self.TimeRules.AfterMarketOpen("SPY", 1),
            Action(self.Downloader))
        
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
            #self.TimeRules.At(9,5),
            self.TimeRules.AfterMarketOpen("SPY", 2),
            Action(self.Downloader))
        
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
            #self.TimeRules.At(9,15),
            self.TimeRules.AfterMarketOpen("SPY", 3),
            Action(self.Downloader))
        
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
            #self.TimeRules.At(9,25),
            self.TimeRules.AfterMarketOpen("SPY", 4),
            Action(self.Strategy))
        
        self.Schedule.On(self.DateRules.EveryDay("SPY"),
            #self.TimeRules.AfterMarketOpen("SPY", 0),
            self.TimeRules.AfterMarketOpen("SPY", 5),
            Action(self.Rebalance))

        self.Schedule.On(self.DateRules.EveryDay("SPY"),
            self.TimeRules.BeforeMarketClose("SPY", 3),
            Action(self.FireSale))
            
        self.universe              = []
        self.Reset_Baskets()

        self.n_hist_items = 3 #100
        
        self.splotName = 'Strategy Info'
        sPlot = Chart(self.splotName)
        sPlot.AddSeries(Series('N_Universe', SeriesType.Line, 0)) # Label only
        self.AddChart(sPlot)
        
    def Reset_Baskets(self):
        
        self.current_symbols_long  = []
        self.current_symbols_short = []

        self.split_universe        = []
        self.current_subset        = 0
        self.asset_hist            = {}
        


    def Downloader(self):
        
        self.Log("Downloader %d "%self.current_subset + str(self.Time))
        
        if len(self.split_universe) == 0:
            return
        
        # Download subset of universe to prevent timeout via History server
        for symbol in self.split_universe[self.current_subset]:
            self.asset_hist[symbol] = self.History([symbol,], self.n_hist_items, Resolution.Daily).astype(np.float32)
        self.current_subset += 1
        
        
    def ProcessAsset(self, symbol):
        
        asset_hist   = self.asset_hist[symbol]

        #compute something with history to show difference
        dp = asset_hist.open.pct_change().values[-1]
        
        return symbol, dp
        
    def Strategy(self):
        
        self.Log("Strategy " + str(self.Time))
        if len(self.universe) == 0:
            return
        
        # Collect the info we need per assets to filter/mask
        res = [self.ProcessAsset(symbol) for symbol in self.universe]# if str(symbol) in self.asset_hist]

        symbols, std = zip(*res)
        symbols      = np.array(symbols)      
        std          = np.array(std)
        
        self.current_symbols_long  = symbols[np.argsort(std)][-self.max_per_side_assets:]
        self.current_symbols_short = symbols[np.argsort(std)][:self.max_per_side_assets]
        
    
    def Rebalance(self):
        
        self.Log("Rebalance " + str(self.Time))
        
        # Enter short positions
        if len(self.current_symbols_short) > 0:
            for sym in self.current_symbols_short:
                self.SetHoldings(sym, -0.1)
        
        # Enter long positions
        if len(self.current_symbols_long) > 0:
            for sym in self.current_symbols_long:
                self.SetHoldings(sym, 0.1)
    
        self.Plot(self.splotName,'N_Universe', len(self.universe))
        
    def FireSale(self):
        self.Liquidate()
        
    def CoarseSelectionFunction(self, coarse):
        
        
        self.Log("CoarseSelectionFunction " + str(self.Time))
        
        # Filter the values of the dict: by price
        values = filter(lambda x: (x.Price >= 5.0) and (x.Price <= 250.0), coarse)

        # sort descending by daily dollar volume
        sortedByDollarVolume = sorted(values, key=lambda x: x.DollarVolume, reverse=True) 

        # Assets selected
        # we need to return only the symbol objects
        result = [ x.Symbol for x in sortedByDollarVolume[:self.coarse_count] ]
        
        # Clear the baskets and prepare for a new universe
        self.Reset_Baskets()
        
        # we need to call Strategy n times to hit our uni target
        self.split_universe = np.array_split(result, 4) # universe is ~2k so use 4 updated taking ~5-8 min each?
        
        return result

    # this event fires whenever we have changes to our universe
    
    # this event fires whenever we have changes to our universe
    def OnSecuritiesChanged(self, changes):
        
        # liquidate removed securities
        for security in changes.RemovedSecurities:
            if security.Symbol in self.universe:
                self.universe.remove(security.Symbol)
            if security.Invested:
                self.Log(str("OnSecuritiesChanged Liquidate of %s"%str(security.Symbol)))
                self.Liquidate(security.Symbol)

        # we want equal allocation in each security in our universe
        for security in changes.AddedSecurities:
            security.MarginModel   =  PatternDayTradingMarginModel()
            if security.Symbol not in self.universe:
                if str(security.Symbol) != "SPY":
                    self.universe.append(security.Symbol)