from AlgorithmImports import *
from datetime import timedelta, datetime
from QuantConnect.Data.UniverseSelection import *
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel
from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel


class Third_Attempt(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 1, 1)  
        self.SetEndDate(2022, 3, 1)  
        self.SetCash(100000)  
        
        self.AddUniverseSelection(Highperformance())
        self.UniverseSettings.Resolution = Resolution.Daily

        self.AddAlpha(BuyPerformance())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(self.DateRules.Every(DayOfWeek.Monday)))
        self.SetExecution(ImmediateExecutionModel())
        self.SetWarmUp(90, Resolution.Daily)
        

class Highperformance (FundamentalUniverseSelectionModel):

    def __init__(self):
        super().__init__( True, None)
        self.lastMonth = -1
        

    def SelectCoarse(self, algorithm, coarse):
        if algorithm.Time.month == self.lastMonth:
            return Universe.Unchanged
        self.lastMonth = algorithm.Time.month

        sortedByVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)
        filteredByFundamentals = [x.Symbol for x in sortedByVolume if x.HasFundamentalData]

        return filteredByFundamentals
        
        
    def SelectFine(self, algorithm, fine):

        sorted_high = sorted([x for x in fine if x.MarketCap > 2e9
                            and 0.5 > x.OperationRatios.AVG5YrsROIC.FiveYears > 0.20
                            and 50 > x.ValuationRatios.PERatio > 20
                            and x.AssetClassification.MorningstarSectorCode != MorningstarSectorCode.FinancialServices
                            and x.AssetClassification.MorningstarSectorCode != MorningstarSectorCode.Healthcare], 
                            key = lambda x: x.ValuationRatios.PERatio, reverse=True)

        fundamental_universe = [x.Symbol for x in sorted_high[:5]]
        
        algorithm.Debug('Universe Selection:')
        algorithm.Debug(str(algorithm.Time))
        algorithm.Debug('/n  ')
        for security in fundamental_universe:
            algorithm.Debug(str(security.Value))

        return fundamental_universe
        

class BuyPerformance(AlphaModel):
    
    def __init__(self):
        self.lastMonth = -1
        self.newAdds = []
        self.newRemovals = []
        

    def Update(self, algorithm, data):
        
        insights = []
        for added in self.newAdds:
            if not algorithm.Securities[added].Invested and algorithm.Securities[added].HasData:
                insights.append(Insight(added, timedelta(days = 30), InsightType.Price, InsightDirection.Up))

        for removed in self.newRemovals:
            if removed not in data.Bars:
                continue
            insights.append(Insight(removed, timedelta(days = 30), InsightType.Price, InsightDirection.Flat))

        return insights
        

    def OnSecuritiesChanged(self, algorithm, changes):

        for security in changes.AddedSecurities:
            if security.Symbol not in self.newAdds and security.IsTradable:
                self.newAdds.append(security.Symbol)

        for security in changes.RemovedSecurities:
            if security.IsTradable:
                self.newRemovals.append(security.Symbol)