from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel

class VerticalTachyonRegulators(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2018, 1, 1)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        
        ## Set execution model to mimic market orders
        self.SetExecution(ImmediateExecutionModel())
        
        ## set equal weighting protfolio construction to mimic intital algorithm weighting scheme
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())

        ## Helper variables for universe selection and checking for conditions to update only at the
        ## beginning of each month
        self.num_coarse = 250
        self.num_fine = 10
        self.lastMonth = -1
        self.symbols = []
        
        ## Coarse/Fine universe selection model
        self.UniverseSettings.Resolution = Resolution.Daily
        self.SetUniverseSelection(FineFundamentalUniverseSelectionModel(self.CoarseSelectionFunction, self.FineSelectionFunction, None, None))

        ## Custom Alpha model where we can perform the trading logic and much of the filtering previously done in the Fine function and OnData methods
        self.AddAlpha(CustomAlphaModel(self.lastMonth, self.num_fine, self.num_coarse))


    def CoarseSelectionFunction(self, coarse):

        ## If not time to rebalance, return an empty list
        if self.Time.month == self.lastMonth: 
            return []

        selected = [x for x in coarse if (x.HasFundamentalData) 
                    and (float(x.Price) > 5)]
        
        sortedByDollarVolume = sorted(selected, key=lambda x: x.DollarVolume, reverse=True) 
        top = sortedByDollarVolume[:self.num_coarse]
        self.symbols = [i.Symbol for i in top]

        return self.symbols
    
    def FineSelectionFunction(self, fine):
        
        ## If not time to rebalance, return an empty list
        if self.Time.month == self.lastMonth: 
            return []
            
        ## Else reassign the month variable and filter
        self.lastMonth = self.Time.month       
        
        filtered_fine = [x.Symbol for x in fine if x.OperationRatios.OperationMargin.Value
                                        and x.ValuationRatios.PriceChange1M 
                                        and x.ValuationRatios.BookValuePerShare]
        
        ## Sanity check to examine PB Ratios
        for x in fine:
            self.Log(x.Symbol.Value + ' PB Ratio: ' + str(x.ValuationRatios.PBRatio))

        self.symbols = filtered_fine
        return self.symbols


class CustomAlphaModel:
    
    def __init__(self, lastMonth, num_fine, num_coarse):
        
        ## Initialize the various variables/helpers we'll need
        self.lastMonth = lastMonth
        self.longs = []
        self.shorts = []
        self.num_fine = num_fine
        self.num_coarse = num_coarse
        
    def Update(self, algorithm, data):
        
        ## Create empty list of insights
        insights = []
        
        ## Return no insights if it's not the month to rebalance
        if algorithm.Time.month == self.lastMonth: 
            return []
        
        ## We will liquidate any securities we're still invested in that we don't want to hold a position
        ## for the next month
        for i in algorithm.ActiveSecurities:
            symbol = i.Value.Symbol
            if algorithm.Securities[symbol].Invested and (symbol not in self.longs) and (symbol not in self.shorts):
                insights.append(Insight(symbol, TimeSpan.FromDays(30), InsightType.Price, InsightDirection.Flat, None, None))
        
        ## Emit Up (buy) Insights for our desired long positions
        for symbol in self.longs:
            insights.append(Insight(symbol, TimeSpan.FromDays(30), InsightType.Price, InsightDirection.Up, 0.01, None))
            
        ## Emit Down (sell) Insights for our desired short positions
        for symbol in self.shorts:
            algorithm.Log(symbol)
            insights.append(Insight(symbol, TimeSpan.FromDays(30), InsightType.Price, InsightDirection.Down, 0.01, None))
            
        
        return insights

    def OnSecuritiesChanged(self, algorithm, changes):
        
        ## Get symbols from the universe
        symbols = [ x.Symbol for x in changes.AddedSecurities ]
        
        ## Get the security objects
        filtered_fine = [algorithm.Securities[symbol] for symbol in symbols]
        
        ## Perform filtering/sorting that was previously done in the Coarse/Fine functions -- these lists could not
        ## be passed into the Alpha model as before, so we access the Fundamental data here instead
        sortedByfactor1 = sorted(filtered_fine, key=lambda x: x.Fundamentals.OperationRatios.OperationMargin.Value, reverse=True)
        sortedByfactor2 = sorted(filtered_fine, key=lambda x: x.Fundamentals.ValuationRatios.PriceChange1M, reverse=True)
        sortedByfactor3 = sorted(filtered_fine, key=lambda x: x.Fundamentals.ValuationRatios.BookValuePerShare, reverse=True)
        
        stock_dict = {}
        
        ## Assign a score to each stock, you can also change the rule of scoring here.
        for i,ele in enumerate(sortedByfactor1):
            rank1 = i
            rank2 = sortedByfactor2.index(ele)
            rank3 = sortedByfactor3.index(ele)
            score = sum([rank1*0.2,rank2*0.4,rank3*0.4])
            stock_dict[ele] = score
        
        ## Sort the stocks by their scores
        self.sorted_stock = sorted(stock_dict.items(), key=lambda d:d[1],reverse=False)
        sorted_symbol = [x[0] for x in self.sorted_stock]

        ## Sort the top stocks into the long_list and the bottom ones into the short_list
        self.longs = [x.Symbol for x in sorted_symbol[:self.num_fine]]
        self.shorts = [x.Symbol for x in sorted_symbol[-self.num_fine:]]