# https://quantpedia.com/Screener/Details/25
from QuantConnect.Data.UniverseSelection import *
import math
import numpy as np
import pandas as pd
import scipy as sp

class SmallCapInvestmentAlgorithm(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2016, 1, 1)   
        self.SetEndDate(2018, 2, 1)         
        self.SetCash(100000)            
        self.UniverseSettings.Resolution = Resolution.Daily
        self.filtered_fine = None
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        self.AddEquity("SPY", Resolution.Daily)
        self.Schedule.On(self.DateRules.MonthStart("SPY"), self.TimeRules.AfterMarketOpen("SPY"), self.rebalance)
        # Count the number of months that have passed since the algorithm starts
        self.months = -1
        self.yearly_rebalance = True
    
    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 (float(x.AdjustedPrice) > 5)]
            return self.filtered_coarse
        else: 
            return []      

    def FineSelectionFunction(self, fine):
        if self.yearly_rebalance:
            # Calculate the market cap and add the "MakretCap" property to fine universe object
            fine = [x for x in fine if (x.ValuationRatios.PERatio > 0) 
                                    and (x.EarningReports.BasicAverageShares.ThreeMonths > 0) 
                                    and (x.EarningReports.BasicEPS.TwelveMonths > 0)]
            for i in fine:
                i.MarketCap = float(i.EarningReports.BasicAverageShares.ThreeMonths * (i.EarningReports.BasicEPS.TwelveMonths*i.ValuationRatios.PERatio))
            sorted_market_cap = sorted(fine, key=lambda x: x.MarketCap)
            self.filtered_fine = [i.Symbol for i in sorted_market_cap[:10]]
            self.yearly_rebalance = False
            return self.filtered_fine
        else:
            return []

    def rebalance(self):
        # yearly rebalance
        self.months += 1
        if self.months%12 == 0:
            self.yearly_rebalance = True

    def OnData(self, data):
        if not self.yearly_rebalance: return 
        if self.filtered_fine:
            stocks_invested = [x.Key for x in self.Portfolio]
            # liquidate stocks not in the lowest market cap list 
            for i in stocks_invested:
                if i not in self.filtered_fine:
                    self.Liquidate(i) 
                elif i in self.filtered_fine:
                    self.SetHoldings(i, 1/len(self.filtered_fine))