from QuantConnect.Data.UniverseSelection import *
import math
import numpy as np
import pandas as pd
import scipy as sp

class PriceEarningsAnamoly(QCAlgorithm):

    def Initialize(self):

        self.SetStartDate(2016, 1, 1)   
        self.SetEndDate(2018, 1, 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)
        self._NumCoarseStocks = 200
        self._NumStocksInPortfolio = 10
        # 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
            CoarseWithFundamental = [x for x in coarse if x.HasFundamentalData and x.Price > 5]
            sortedByDollarVolume = sorted(CoarseWithFundamental, key=lambda x: x.DollarVolume, reverse=False) 
            top = sortedByDollarVolume[:self._NumCoarseStocks]
            return [i.Symbol for i in top]
        else: 
            return []      

    def FineSelectionFunction(self, fine):
        if self.yearly_rebalance:
            fine = [x for x in fine if (x.ValuationRatios.PERatio > 0)]
            for i in fine:
                i.PERatio = float(i.ValuationRatios.PERatio)
            sortedPERatio = sorted(fine, key=lambda x: x.PERatio)
            
            self.filtered_fine = [i.Symbol for i in sortedPERatio[:self._NumStocksInPortfolio]]
            self.yearly_rebalance = False
            return self.filtered_fine
        else:
            return []

    def rebalance(self):
        # Rebalance at the start of each year
        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 that are not in the list of stocks with lowest PE ratio
            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))