# region imports
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
# endregion

class Third_Attempt(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 1, 1)  # Set Start Date
        self.SetEndDate(2022, 1, 1)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        
        self.AddUniverseSelection(Highperformance())
        self.UniverseSettings.Resolution = Resolution.Daily

        self.AddAlpha(BuyPerformance())
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(lambda time: None))
        self.SetExecution(ImmediateExecutionModel())

class Highperformance (FundamentalUniverseSelectionModel):

    def __init__(self):
        super().__init__( True, None)
        self.lastMonth = -1
        self.spy = Symbol.Create('SPY', SecurityType.Equity, Market.USA)

    def SelectCoarse(self, algorithm, coarse):
        #run the algorithm once a month, return Universe.Unchanged in case we are looking at exactly the same month
        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]] + [self.spy]

        algorithm.Debug('Universe Selection:')
        algorithm.Debug(str(algorithm.Time))
        algorithm.Debug('/n  ')
        for security in fundamental_universe:
            algorithm.Debug(str(security))

        #return [x.Symbol for x in sorted_high[:5]] + [self.spy]
        return [x.Symbol for x in sorted_high[:5]]

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

        #removals to be removed
        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(30), InsightType.Price, InsightDirection.Up))

        # to be removed at a later stage
        for removed in self.newRemovals:
            if removed not in data.Bars:
                continue
            insights.append(Insight(removed, timedelta(30), InsightType.Price, InsightDirection.Flat))

        return insights

    def OnSecuritiesChanged(self, algorithm, changes):
        #When assets are added to the universe, they will trigger OnSecuritiesChanged() event.
        #From there, you can initialize any state or history required for the Alpha Model
        for security in changes.AddedSecurities:
            #if security.Symbol != self.spy and security.Symbol not in self.newAdds:
            if security.Symbol not in self.newAdds and security.IsTradable:
                self.newAdds.append(security.Symbol)

        for security in changes.RemovedSecurities:
            #if security.Symbol != self.spy:
            if security.IsTradable:
                self.newRemovals.append(security.Symbol)