#from datetime import timedelta
from QuantConnect.Data.UniverseSelection import *
from Selection.FundamentalUniverseSelectionModel import FundamentalUniverseSelectionModel
from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel

class FirstAttempt(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 1, 1)  # Set Start Date
        self.SetEndDate(2022,1,1) # Setting end date
        self.SetCash(100000)  # Set Strategy Cash
        
        #self.SetBenchmark('SPY')
        self.lastMonth = -1
        
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverseSelection(HighRoICHighPEUniverse())
        #self.AddUniverseSelection(HighRoICHighPEUniverse())
        
        self.AddAlpha(BuyValue())
        
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel(lambda time: None))
        self.SetExecution(ImmediateExecutionModel())
        

class HighRoICHighPEUniverse(FundamentalUniverseSelectionModel):
            
    def __init__(self):
        super().__init__(True, None) #took it from the bootcamp
        self.lastMonth = -1

    
    def SelectCoarse(self, algorithm, coarse):
        
        if algorithm.Time.month == self.lastMonth:
            return Universe.Unchanged
            
        self.lastMonth = algorithm.Time.month
       
        filtered = [ x for x in coarse if x.HasFundamentalData ]
        
        return [x.Symbol for x in filtered]
        
    def SelectFine(self,algorithm, fine):
        
        sorted_high = sorted([x for x in fine if x.MarketCap < 2e9 
                        and x.OperationRatios.ROIC.OneYear > 0.2
                        and x.ValuationRatios.PERatio > 20
                        and x.AssetClassification.MorningstarSectorCode != MorningstarSectorCode.FinancialServices
                        and x.AssetClassification.MorningstarSectorCode != MorningstarSectorCode.Healthcare],
                        key=lambda x: x.ValuationRatios.PERatio, reverse=True)
        
        
        algorithm.Log(f"High Universe consists of {len(sorted_high)} securities")
        algorithm.Log(f"High Universe #1 {sorted_high[0]} securities")
        algorithm.Log(f"High Universe #2 {sorted_high[1]} securities")
        algorithm.Log(f"High Universe #3 {sorted_high[2]} securities")
        algorithm.Log(f"High Universe #4 {sorted_high[3]} securities")
        algorithm.Log(f"High Universe #5 {sorted_high[4]} securities")
        
       
        sorted_low = sorted([x for x in fine if x.MarketCap < 2e9 
                        and x.OperationRatios.ROIC.OneYear < 0.05
                        and x.ValuationRatios.PERatio < 10
                        and x.AssetClassification.MorningstarSectorCode != MorningstarSectorCode.FinancialServices
                        and x.AssetClassification.MorningstarSectorCode != MorningstarSectorCode.Healthcare],
        key=lambda x: x.ValuationRatios.PERatio, reverse=True)
        
        #and x.ValuationRatios.PERatio < 10
        #key=lambda x: x.OperationRatios.OperationMargin.OneYear, reverse=True
        
        algorithm.Log(f"Low Universe consists of {len(sorted_low)} securities")
        algorithm.Log(f"Low Universe #1 {sorted_low[0]} securities")
        algorithm.Log(f"Low Universe #2 {sorted_low[1]} securities")
        algorithm.Log(f"Low Universe #3 {sorted_low[2]} securities")
        algorithm.Log(f"Low Universe #4 {sorted_low[3]} securities")
        algorithm.Log(f"Low Universe #5 {sorted_low[4]} securities")
        
        Universe = sorted_high[:5] + sorted_low[-5:]
        
        return [z.Symbol for z in Universe]
     
        
class BuyValue (AlphaModel):
    
    
    def __init__(self):
        
        self.lastMonth = -1
        self.longs = []
        self.shorts = []
    
    def Update (self, algorithm, data):
        
        if algorithm.Time.month == self.lastMonth:
            return []
            
        self.lastMonth = algorithm.Time.month
        
        insights = []
        
        #Check if existing algorithm portfolio contains any securities in long or short directories
        #if we are invested on securities that are not in long or shorts, then we go 'flat'
        for x in algorithm.Portfolio:
            holding = x.Value
            symbol = holding.Symbol
            if holding.Invested and symbol not in self.longs and symbol not in self.shorts:
                insights.append(Insight(symbol, timedelta(30), InsightType.Price, InsightDirection.Flat))
                
        #generate insight-Up for long positions
        for symbol in self.longs:
            insights.append(Insight(symbol, timedelta(30), InsightType.Price, InsightDirection.Up))
        
        #generate insights for short positions
        for symbol in self.shorts:
            insights.append(Insight(symbol, timedelta(30), InsightType.Price, InsightDirection.Down))
            
        return insights
        
        
    def OnSecuritiesChanged(self, algorithm, changes):
        
        added = [x for x in changes.AddedSecurities]
        
        sortedByPE = sorted(added, key=lambda x: x.Fundamentals.ValuationRatios.PERatio, reverse=True)
        
        self.longs = [x.Symbol for x in sortedByPE[:5]]
        algorithm.Log('LONGS: ' + ', '.join(sorted([x.Value for x in self.longs])))
        
        self.shorts = [x.Symbol for x in sortedByPE[-5:]]
        algorithm.Log('SHORTS: ' + ', '.join(sorted([x.Value for x in self.shorts])))