from QuantConnect import Resolution
from QuantConnect.Algorithm import QCAlgorithm


class NCAVsimple(QCAlgorithm):
    def Initialize(self):
        #rebalancing should occur in July
        self.SetStartDate(2000, 1, 1)  
        self.SetCash(100000) 
        self.UniverseSettings.Resolution = Resolution.Daily

        self.filtered_fine = None
        self.filtered_coarse = None
        self.symbol = self.AddEquity('SPY', Resolution.Daily).Symbol

        self.coarse_count = 3000
        # self.SetSecurityInitializer(lambda x: x.SetMarketPrice(self.GetLastKnownPrice(x)))
        
        self.yearly_rebalance = False
        self.AddUniverse(self.CoarseSelectionFunction,self.FineSelectionFunction)
        self.Schedule.On(self.DateRules.MonthEnd(self.symbol), self.TimeRules.AfterMarketOpen(self.symbol), self.rebalance)
        
    def CoarseSelectionFunction(self, coarse):
        if self.yearly_rebalance:
            # drop stocks which have no fundamental data or have low price
            self.filtered_coarse = [x.Symbol for x in coarse if (x.HasFundamentalData) and x.Market == 'usa']
            return self.filtered_coarse
        else: 
            return Universe.Unchanged      
    
    def FineSelectionFunction(self, fine):
        if self.yearly_rebalance:
            # #calculate the NCAV/MV and add the property to fine universe object
            # #filters out the companies in the financial sector as suggested
            # fine = [x for x in fine if (float(x.FinancialStatements.BalanceSheet.CurrentAssets.Value) > 0) 
            #                         and (float(x.EarningReports.BasicAverageShares.Value) > 0)
            #                         and (float(x.FinancialStatements.BalanceSheet.CurrentLiabilities.Value) > 0)
            #                         and (float(x.MarketCap) > 0)
            #                         #  This indicator will denote which one of the six industry data collection templates applies to the company. 
            #                         # Each industry data collection template includes data elements that are commonly reported by companies in that industry. 
            #                         # N=Normal (Manufacturing), M=Mining, U=Utility, T=Transportation, B=Bank, I=Insurance 
            #                         and (x.CompanyReference.IndustryTemplateCode!="B")
            #                         and (x.CompanyReference.IndustryTemplateCode!="I")]
            # fine = sorted(fine, key = lambda x:x.MarketCap, reverse=True)
            # self.Debug("fine len: " + str(len(fine)))
            
            # for i in fine:
            #     #calculates the net current asset value per share
            #     self.Debug(i.Symbol)
            #     total_liabilities = i.FinancialStatements.BalanceSheet.TotalLiabilitiesAsReported.TwelveMonths
            #     current_assets = i.FinancialStatements.BalanceSheet.CurrentAssets.TwelveMonths
            #     i.ncav = (current_assets - total_liabilities)/float(i.MarketCap)
            #     if i.ncav > 1.5:
            #         self.Debug("-total_liabilities: " + str(total_liabilities))
            #         self.Debug("-current_assets: " + str(current_assets))
            #         self.Debug("-market_cap: " + str(i.MarketCap))
            #         self.Debug("-ncav: " + str(i.ncav))
                            
            # # keeps all symbols that have a NCAV/MV higher than 1.5
            # # sorted_fine = sorted(fine, lambda x: x.ncav)
            # self.filtered_fine = [i.Symbol for i in fine if (i.ncav > 1.5)]
            fine = [x for x in fine if x.EarningReports.BasicAverageShares.ThreeMonths > 0 and x.MarketCap != 0 and x.ValuationRatios.WorkingCapitalPerShare != 0]
            sorted_by_market_cap = sorted(fine, key = lambda x:x.MarketCap, reverse=True)
            top_by_market_cap = [x for x in sorted_by_market_cap[:self.coarse_count]]
            
            # NCAV/MV calc.
            # self.filtered_fine = [x.Symbol for x in top_by_market_cap if ((x.ValuationRatios.WorkingCapitalPerShare * x.EarningReports.BasicAverageShares.ThreeMonths) / x.MarketCap) > 1.5]
        
            self.filtered_fine = [x.Symbol for x in top_by_market_cap if ((x.FinancialStatements.BalanceSheet.CurrentAssets.TwelveMonths - x.FinancialStatements.BalanceSheet.TotalLiabilitiesAsReported.TwelveMonths) / x.MarketCap) > 1.5]

            return self.filtered_fine
        else:
            return []
    
    def rebalance(self):
        #yearly rebalance
        if self.Time.month == 6:
            self.Debug("Rebalance" + str(self.Time))
            self.yearly_rebalance = True
        
    def OnData(self, data):
        if not self.yearly_rebalance: return 

        stocks_invested = [x.Key for x in self.Portfolio if x.Value.Invested]
        for symbol in stocks_invested:
            if symbol not in self.filtered_fine:
                self.Liquidate(symbol)

        for symbol in self.filtered_fine:
            if self.Securities[symbol].Price != 0 and self.Securities[symbol].IsTradable:  # Prevent error message.
                self.SetHoldings(symbol, 1 / len(self.filtered_fine))

        self.yearly_rebalance = False