from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel

class VentralModulatedThrustAssembly(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 7, 9)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        
        self.SetUniverseSelection(SelectionModel())
        self.SetAlpha(ConstantAlphaModel(InsightType.Price, InsightDirection.Up, timedelta(1), 0.025, None))
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())


class SelectionModel(FundamentalUniverseSelectionModel):

    def __init__(self,
        filterFineData = True,
        universeSettings = None,
        securityInitializer = None):

        super().__init__(filterFineData, universeSettings, securityInitializer)
        self.periodCheck = -1
        self.coarse = dict()

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

        self.coarse = {x.Symbol: x
            for x in coarse if x.HasFundamentalData and x.Volume > 0 and x.Price > 0}

        return [symbol for symbol,_ in self.coarse.items()]

    def SelectFine(self, algorithm, fine):

        DERatio = lambda x: x.OperationRatios.LongTermDebtEquityRatio.ThreeMonths
        
        fine = [x for x in fine if DERatio(x)
                               and x.CompanyReference.CountryId == "USA"
                               and (x.CompanyReference.PrimaryExchangeID in ["NYS", "NAS"])
                               and (algorithm.Time - x.SecurityReference.IPODate).days > 180
                               and x.EarningReports.BasicAverageShares.ThreeMonths * 
                                   x.EarningReports.BasicEPS.TwelveMonths *
                                   x.ValuationRatios.PERatio > 1.52e8]
        if len(fine) == 0:
            return Universe.Unchanged

        self.periodCheck = algorithm.Time.year

        sortedByDEratio = sorted(fine, key=DERatio, reverse=False)[:50]
        symbols = [x.Symbol for x in sortedByDEratio]

        averages = dict()
        history = algorithm.History(symbols, 200, Resolution.Daily).close.unstack(0)

        for symbol in symbols:
            # Remove NaN: symbol does not have 200 daily data points 
            df = history[symbol].dropna()
            if df.empty:
                continue

            mom = Momentum(126)
            for time, close in df.iteritems():
                mom.Update(time, close)

            # Adds Momentum to dict only if it is ready
            if mom.IsReady:
                averages[symbol] = mom

        # Update with current data
        for symbol, mom in averages.items():
            c = self.coarse.pop(symbol, None)
            mom.Update(c.EndTime, c.AdjustedPrice)

        sortedbyMomentum = sorted(averages.items(), key=lambda x: x[1], reverse=True)

        return [x[0] for x in sortedbyMomentum[:5]]