#from QuantConnect.Data.UniverseSelection import * 
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class SectorBalancedPortfolioConstruction(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2016, 1, 1) 
        self.SetEndDate(2017, 1, 5) 
        self.SetCash(100000) 
        self.UniverseSettings.Resolution = Resolution.Daily
        self.SetUniverseSelection(MyUniverseSelectionModel())
        self.AddEquity('SPY', Resolution.Daily)
        self.sma = self.SMA('SPY', 200, Resolution.Daily)
        self.smaWindow = RollingWindow[float](2)
        
        # register sma update callback to our sma window
        self.sma.Updated += self.UpdateWindow
        
        self.SetWarmup(200)
        
    def UpdateWindow(self, sender, updated):
        if self.sma.IsReady:
            self.smaWindow.Add(self.sma.Current.Value)
            if self.smaWindow.Count == 2:
                slope = self.smaWindow[0] - self.smaWindow[1]
                self.Plot('Custom', 'SMA slope', slope)
        
class MyUniverseSelectionModel(FundamentalUniverseSelectionModel):
    # add parameter "algo" so our algorithm can be passed in
    def __init__(self):
        super().__init__(True, None, None)

    def SelectCoarse(self, algorithm, coarse):
        filtered = [x for x in coarse if x.HasFundamentalData and x.Price > 100]
        return [x.Symbol for x in filtered]

    def SelectFine(self, algorithm, fine):
        filtered = [f for f in fine if f.MarketCap > 2e10][:3]
        for x in filtered:
            algorithm.Debug(x.Symbol)
        return [x.Symbol for x in filtered]