from datetime import timedelta
from QuantConnect.Data.UniverseSelection import * 
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class SectorBalancedPortfolioConstruction(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2017, 9, 19) 
        self.SetEndDate(2017, 10, 15) 
        self.SetCash(10000000) 
        self.UniverseSettings.Resolution = Resolution.Hour
        self.SetUniverseSelection(MyUniverseSelectionModel())
        self.SetAlpha(ConstantAlphaModel(InsightType.Price, InsightDirection.Up, timedelta(1), 0.025, None))
        self.SetPortfolioConstruction(MySectorWeightingPortfolioConstructionModel())
        self.SetExecution(ImmediateExecutionModel())

class MyUniverseSelectionModel(FundamentalUniverseSelectionModel):

    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 > 10]
        sortedByDollarVolume = sorted(filtered, key=lambda x: x.DollarVolume, reverse=True)
        return [x.Symbol for x in sortedByDollarVolume][:1000]

    def SelectFine(self, algorithm, fine):
        for f in fine:
            try:
                f.AssetClassification
            except:
                algorithm.Quit(f"Error with {f.Symbol}")
                return []
        
        filtered = [f for f in fine if f.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.Technology]
        self.technology = sorted(filtered, key=lambda f: f.MarketCap, reverse=True)[:50]
        filtered = [f for f in fine if f.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.FinancialServices]
        self.financialServices = sorted(filtered, key=lambda f: f.MarketCap, reverse=True)[:50]
        filtered = [f for f in fine if f.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.ConsumerDefensive]
        self.consumerDefensive = sorted(filtered, key=lambda f: f.MarketCap, reverse=True)[:50]
        return [x.Symbol for x in self.technology + self.financialServices + self.consumerDefensive]
        #filtered = [f for f in fine if f.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.Industrials]
        #self.industrial = sorted(filtered, key=lambda f: f.MarketCap, reverse=True)[:5]
        #filtered = [f for f in fine if f.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.Energy]
        #self.energy = sorted(filtered, key=lambda f: f.MarketCap, reverse=True)[:5]
        #filtered = [f for f in fine if f.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.BasicMaterials]
        #self.basicmaterial = sorted(filtered, key=lambda f: f.MarketCap, reverse=True)[:5]
        #filtered = [f for f in fine if f.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.CommunicationServices]
        #self.communication = sorted(filtered, key=lambda f: f.MarketCap, reverse=True)[:5]
        #filtered = [f for f in fine if f.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.Utilities]
        #self.utility = sorted(filtered, key=lambda f: f.MarketCap, reverse=True)[:5]
        #filtered = [f for f in fine if f.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.Healthcare]
        #self.health = sorted(filtered, key=lambda f: f.MarketCap, reverse=True)[:5]
        #filtered = [f for f in fine if f.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.ConsumerCyclical]
        #self.consumerCyclical = sorted(filtered, key=lambda f: f.MarketCap, reverse=True)[:5]
        #filtered = [f for f in fine if f.AssetClassification.MorningstarSectorCode == MorningstarSectorCode.RealEstate]
        #self.realestate = sorted(filtered, key=lambda f: f.MarketCap, reverse=True)[:5]

      
        
        
class MySectorWeightingPortfolioConstructionModel(EqualWeightingPortfolioConstructionModel):

    def __init__(self, rebalance = Resolution.Daily):
        super().__init__()
        self.symbolBySectorCode = dict()
        self.result = dict()
        self.sector_code_by_symbol = dict()

    def DetermineTargetPercent(self, activeInsights):
        #1. Set the self.sectorBuyingPower before by dividing one by the length of self.symbolBySectorCode
        self.sectorBuyingPower = 1/len(self.symbolBySectorCode)
            
        for sector, symbols in self.symbolBySectorCode.items():
            #2. Search for the active insights in this sector. Save the variable self.insightsInSector
            self.insightsInSector = [insight for insight in activeInsights if insight.Symbol in symbols]
        
            #3. Divide the self.sectorBuyingPower by the length of self.insightsInSector to calculate the variable percent
            # The percent is the weight we'll assign the direction of the insight
            self.percent = self.sectorBuyingPower / len(self.insightsInSector)
        
            #4. For each insight in self.insightsInSector, assign each insight an allocation. 
            # The allocation is calculated by multiplying the insight direction by the self.percent 
            for insight in self.insightsInSector:
                self.result[insight] = insight.Direction * self.percent
        
        return self.result


    def OnSecuritiesChanged(self, algorithm, changes):
        for security in changes.AddedSecurities:
            sectorCode = security.Fundamentals.AssetClassification.MorningstarSectorCode
            if sectorCode not in self.symbolBySectorCode:
                self.symbolBySectorCode[sectorCode] = list()
            self.symbolBySectorCode[sectorCode].append(security.Symbol) 

        for security in changes.RemovedSecurities:
            algorithm.Debug(f"Removed {security.Symbol}; {type(security.Fundamentals)}")
            #sectorCode = security.Fundamentals.AssetClassification.MorningstarSectorCode
            #if sectorCode in self.symbolBySectorCode:
            #    symbol = security.Symbol
            #    if symbol in self.symbolBySectorCode[sectorCode]:
            #        self.symbolBySectorCode[sectorCode].remove(symbol)

        super().OnSecuritiesChanged(algorithm, changes)