from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import InsightWeightingPortfolioConstructionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel
from Risk.MaximumDrawdownPercentPerSecurity import MaximumDrawdownPercentPerSecurity
from Risk.TrailingStopRiskManagementModel import TrailingStopRiskManagementModel
from AlphaModel import FundamentalFactorAlphaModel

class VerticalTachyonRegulators(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 7, 1)
        #self.SetEndDate(2021, 8, 1)
        self.startCash = 100000
        self.SetCash(self.startCash)
        self.averages = {}
        self.slowPeriod = int(self.GetParameter("slowPeriod"))
        if self.slowPeriod is None:
            self.slowPeriod = 160
        self.fastPeriod = int(self.GetParameter("fastPeriod"))
        if self.fastPeriod is None:
            self.fastPeriod = 15

        # Benchmark
        #self.bench = self.AddEquity("SPY", Resolution.Daily).Symbol
        #self.benchmark = self.AddEquity("SPY", Resolution.Daily)

        #self.benchStartPrice = self.benchmark.Price
        #self.SetBenchmark(self.bench)

        # Execution model
        self.SetExecution(ImmediateExecutionModel())
        
        # Portfolio construction model
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel())
        #self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())

        # Risk model
        stopRisk = self.GetParameter("stopRisk")
        if stopRisk is None:
            stopRisk = 0.25
        self.SetRiskManagement(TrailingStopRiskManagementModel(float(stopRisk)))
        #self.SetRiskManagement(MaximumDrawdownPercentPerSecurity(float(stopRisk)))
        
        # Universe selection
        self.num_coarse = 4000
        self.num_fine = 30
        self.UniverseSettings.Leverage = 1
        self.UniverseSettings.Resolution = Resolution.Daily
        #self.UniverseSettings.MinimumTimeInUniverse = timedelta(days=30)
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        
        # rebalancing
        self.lastMonth = -1

        # Set factor weighting        
        quality_weight = self.GetParameter("quality_weight")
        if quality_weight is None:
            quality_weight = 15

        size_weight = self.GetParameter("size_weight")
        if size_weight is None:
            size_weight = 1
            
        value_weight = self.GetParameter("value_weight")
        if value_weight is None:
            value_weight = 1

        # Alpha Model
        self.AddAlpha(FundamentalFactorAlphaModel(self.num_fine, quality_weight, value_weight, size_weight))
        
        # schedule weekly plotting
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday), 
                        self.TimeRules.At(10, 30),
                        self.Plotting)


    def Plotting(self):
        self.Plot("Positions", "Num", len([x.Symbol for x in self.Portfolio.Values if self.Portfolio[x.Symbol].Invested]))
        
        #benchPerformance = self.benchmark.Price / self.benchStartPrice
        #self.Plot("Performance", "Performance", benchPerformance)


    def CoarseSelectionFunction(self, coarse):
        selector = []
        
        # If not time to rebalance, keep the same universe
        if self.Time.month == self.lastMonth: 
            return Universe.Unchanged
        
        # Else reassign the month variable
        self.lastMonth = self.Time.month
        
        # Select only those with fundamental data and a sufficiently large price
        # Sort by top dollar volume: most liquid to least liquid
        selected = sorted([x for x in coarse if x.HasFundamentalData and x.Price > 5],
            key = lambda x: x.DollarVolume, reverse=True)

        # Only add uptrending securities (EMA based)
        for finer in selected:
            symbol = finer.Symbol
            #self.Debug(symbol)
            
            if symbol not in self.averages:
                # Call history to get an array of X days of history data
                history = self.History(symbol, self.slowPeriod, Resolution.Daily)
                
                #Pass in the history result to SelectionData
                self.averages[symbol] = SelectionData(history, self.slowPeriod, self.fastPeriod)

            self.averages[symbol].update(self.Time, finer.AdjustedPrice)
            
            if  self.averages[symbol].is_ready() and self.averages[symbol].fast > self.averages[symbol].slow:
                selector.append(finer)

        return [x.Symbol for x in selector[:self.num_coarse]]


    def FineSelectionFunction(self, fine):
        # Filter the fine data for equities with non-zero/non-null values
        filtered_fine = [x.Symbol for x in fine if x.OperationRatios.ROE.SixMonths > 0.05
                                                and x.ValuationRatios.PriceChange1M > 0
                                                and x.ValuationRatios.FCFYield > 0
                                                and x.MarketCap > 0]

        return filtered_fine

class SelectionData():
    def __init__(self, history, slowPeriod, fastPeriod):
        #Save the fast and slow ExponentialMovingAverage
        self.slow = ExponentialMovingAverage(slowPeriod)
        self.fast = ExponentialMovingAverage(fastPeriod)
        
        #Loop over the history data and update the indicators
        for bar in history.itertuples():
            self.fast.Update(bar.Index[1], bar.close)
            self.slow.Update(bar.Index[1], bar.close)
    
    #Check if our indicators are ready
    def is_ready(self):
        return self.slow.IsReady and self.fast.IsReady
    
    #Use the "indicator.Update" method to update the time and price of both indicators
    def update(self, time, price):
        self.fast.Update(time, price)
        self.slow.Update(time, price)

from datetime import timedelta

class FundamentalFactorAlphaModel(AlphaModel):
    
    def __init__(self, num_fine, quality_weight, value_weight, size_weight):
        
        # Initialize the various variables/helpers we'll need
        self.lastMonth = -1
        self.longs = []
        self.num_fine = num_fine
        self.period = timedelta(31)

        # normalize quality, value, size weights
        weights = [float(quality_weight), float(value_weight), float(size_weight)]
        weights = [float(i)/sum(weights) for i in weights]
        
        self.quality_weight = weights[0]
        self.value_weight = weights[1]
        self.size_weight = weights[2]


    def Update(self, algorithm, data):
        '''Updates this alpha model with the latest data from the algorithm.
        This is called each time the algorithm receives data for subscribed securities
        Args:
            algorithm: The algorithm instance
            data: The newa available
        Returns:
            New insights'''
        
        # Return no insights if it's not time to rebalance
        if algorithm.Time.month == self.lastMonth: 
            return []
        self.lastMonth = algorithm.Time.month
        
        # List of insights
        # Insights of the form: Insight(symbol, timedelta, type, direction, magnitude, confidence, sourceModel, weight)
        insights = []
        
        # Close old positions if they aren't in longs
        for security in algorithm.Portfolio.Values:
            if security.Invested and security.Symbol not in self.longs:
                insights.append(Insight(security.Symbol, self.period, InsightType.Price, 
                                            InsightDirection.Flat, None, None, None, None))
        
        length = len(self.longs)
        
        for i in range(length):
            insights.append(Insight(self.longs[i], self.period, InsightType.Price, 
                                    InsightDirection.Up, None, (length - i)**2, None, (length - i)**2 ))
        
        return insights


    def OnSecuritiesChanged(self, algorithm, changes):
        '''Event fired each time the we add/remove securities from the data feed
        Args:
            algorithm: The algorithm instance that experienced the change in securities
            changes: The security additions and removals from the algorithm'''

        # Get the added securities
        addedSecurities = [x for x in changes.AddedSecurities]
        algorithm.Debug(len(addedSecurities))
        
        if len(addedSecurities) == 0:
            return
        
        # Assign quality, value, size score to each stock
        quality_scores = self.Scores(addedSecurities, [(lambda x: x.Fundamentals.OperationRatios.ROE.SixMonths, True, 1)]) 
                                            #(lambda x: x.Fundamentals.OperationRatios.QuickRatio.Value, True, 1)]) 
                                            #(lambda x: x.Fundamentals.OperationRatios.DebttoAssets.Value, False, 1)])
        
        value_scores = self.Scores(addedSecurities, [(lambda x: x.Fundamentals.ValuationRatios.FCFYield, True, 1)])
        
        # larger caps = True, smaller caps = False
        size_scores = self.Scores(addedSecurities, [(lambda x: x.Fundamentals.MarketCap, False, 1)])
        
        scores = {}
        # Assign a combined score to each stock 
        for symbol,value in quality_scores.items():
            quality_rank = value
            value_rank = value_scores[symbol]
            size_rank = size_scores[symbol]
            scores[symbol] = quality_rank*self.quality_weight + value_rank*self.value_weight + size_rank*self.size_weight
        
        # Sort the securities by their scores
        sorted_stock = sorted(scores.items(), key=lambda tup : tup[1], reverse=False)
        sorted_symbol = [tup[0] for tup in sorted_stock][:self.num_fine]
        
        # Sort the top stocks into the long_list
        self.longs = [security.Symbol for security in sorted_symbol]
        
        # Log longs symbols and their score
        #algorithm.Log(", ".join([str(x.Symbol.Value) + ": " + str(scores[x]) for x in sorted_symbol]))


    def Scores(self, addedSecurities, fundamentals):
        '''Assigns scores to each stock in added
        Args: 
            added: list of sceurities 
            fundamentals: list of 3-tuples (lambda function, bool, float)
        Returns:
            Dictionary with score for each security'''
        
        length = len(fundamentals)
        
        if length == 0:
            return {}
        
        # Initialize helper variables
        scores = {}
        sortedBy = []
        rank = [0 for _ in fundamentals]
        
        # Normalize weights
        weights = [tup[2] for tup in fundamentals]
        weights = [float(i)/sum(weights) for i in weights]
        
        # Create sorted list for each fundamental factor passed
        for tup in fundamentals:
            sortedBy.append(sorted(addedSecurities, key=tup[0], reverse=tup[1]))
        
        # Create and save score for each symbol
        for index,symbol in enumerate(sortedBy[0]):
            
            # Save symbol's rank for each fundamental factor
            rank[0] = index
            for j in range(1, length):
                rank[j] = sortedBy[j].index(symbol)
            
            # Save symbol's total score
            score = 0
            for i in range(length):
                score += rank[i] * weights[i]
            scores[symbol] = score
            
        return scores