import pandas as pd
import numpy as np
from clr import AddReference
AddReference("QuantConnect.Research")
from QuantConnect.Research import QuantBook


class TestingAlgorithm(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021,4,26)  
        self.SetEndDate(2021,4,28)
        self.SetCash(10000) 
        self.UniverseSettings.Resolution = Resolution.Minute
        self.AddUniverse(self.CoarseFilterFunction, self.FineFilterFunction)
 
        self.spy = self.AddEquity("SPY")
        
        self.Schedule.On(self.DateRules.EveryDay("SPY"), self.TimeRules.AfterMarketOpen("SPY", 60), self.MainFunction)
        
        # Set display option for logging purposes
        pd.set_option('display.float_format', '{:.7E}'.format)



    # Coarse filter
    def CoarseFilterFunction(self, coarse):
        self.Log("Running coarse filter function")
        self.coarsefilter = [x.Symbol for x in coarse if x.DollarVolume > 100000000 and x.HasFundamentalData]
        return self.coarsefilter


    #  Fine filter
    def FineFilterFunction(self, fine):
        self.Log("Running fine filter function")
        self.stocks = sorted(fine, key=lambda x: x.MarketCap, reverse=True)
        finefilter = [x.Symbol for x in self.stocks]
        self.tickers = [x.Symbol.Value for x in self.stocks]
        return finefilter

    def MainFunction(self):
        self.Log("Running main function")
        qb = QuantBook()
            
        shares = [x.FinancialStatements.BalanceSheet.OrdinarySharesNumber.ThreeMonths for x in self.stocks]
        mktcap = [x.MarketCap for x in self.stocks]
        revenue = []
        for x in self.stocks:
            revenue.append(qb.GetFundamental(x.Symbol, "FinancialStatements.IncomeStatement.TotalRevenue.ThreeMonths", self.Time - timedelta(1), self.Time).iloc[0,0])
        
        self.df = pd.DataFrame({
            'shares': shares,
            'mktcap': mktcap,
            'revenue': revenue
        },
            index = self.tickers)
        self.df = self.df.dropna()


        self.LoggingFunction()
        return


    def LoggingFunction(self):
        for index, row in self.df.head(10).iterrows():
            self.Log(str(row))
        return