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.SetBrokerageModel(BrokerageName.AlphaStreams)
        self.SetStartDate(2011, 1, 1)
        #self.SetEndDate(2021, 8, 1)
        self.startCash = 100000
        self.averages = {}
        self.slowPeriod = self.GetParameter("slowPeriod")
        if self.slowPeriod is None:
            self.slowPeriod = 160
        self.fastPeriod = self.GetParameter("fastPeriod")
        if self.fastPeriod is None:
            self.fastPeriod = 15
        self.netMargin = self.GetParameter("netMargin")
        if self.netMargin is None:
            self.netMargin = 0.2
        self.netMargin = float(self.netMargin)

        # Benchmark
        self.SetBenchmark("SPY")

        # Execution model
        self.SetExecution(ImmediateExecutionModel())
        
        # Portfolio construction model
        # insightWeighting seems to outperform
        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 = 8000
        self.num_fine = 20
        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 is currently ROE
        quality_weight = self.GetParameter("quality_weight")
        if quality_weight is None:
            quality_weight = 8

        size_weight = self.GetParameter("size_weight")
        if size_weight is None:
            size_weight = 4
        
        # value is currently Net Profit Margin    
        value_weight = self.GetParameter("value_weight")
        if value_weight is None:
            value_weight = 16

        # 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]))
        self.Plot("Margin remaining", "USD", self.Portfolio.MarginRemaining)


    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)
            
        selected = selected[:self.num_coarse]

        # Only add up or down trending securities (EMA based)
        for finer in selected:
            symbol = finer.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)
            
            # EMA change direction here
            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]


    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.CompanyReference.IndustryTemplateCode == 'B'
                                                and x.OperationRatios.NetMargin.ThreeMonths > self.netMargin
                                                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

# For small/large cap preference, change line ca. 80
# For portfolio construction weights, line 56

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.long = {}
        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 = []
        sorted_symbol = [tup[0] for tup in self.long]
        longs = [security.Symbol for security in sorted_symbol]

        # Close old positions if they aren't in longs
        for security in algorithm.Portfolio.Values:
            if security.Invested and security.Symbol not in longs:

                insights.append(Insight(security.Symbol, self.period,
                                InsightType.Price, 
                               InsightDirection.Flat, None, None, None, None))

        length = len(self.long)
        for i, elem in enumerate(self.long):
            
            algorithm.Log(str(elem[0].Symbol)+ " investing")
            
            # the weight is the position in the self.long tuple divided with 
            # the length of the self.long tuple. The 0 entry is the highest
            j = length - i
            alt_weight = j**2
            weight = j/length
            
            # Insight(symbol, timedelta, type, direction, magnitude=None, confidence=None, sourceModel=None, weight=None)
            # so in our case the weights will be from the lenghts which again is max self.num_fine
            insights.append(Insight(elem[0].Symbol, self.period, InsightType.Price, 
                                    InsightDirection.Up, None, alt_weight,\
                                    None, alt_weight))
        
        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]

        # 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.OperationRatios.NetMargin.SixMonths, True, 1)])
        
        # Larger caps = True, smaller caps = False
        size_scores = self.Scores(addedSecurities, \
            [(lambda x: x.Fundamentals.MarketCap, True, 1)])
        
        scores = {}
        self.long = {}
        # 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
        # shouldnt this be sorted reverse = True??
        sorted_scores = sorted(scores.items(), key=lambda tup : tup[1], reverse=True)
        #algorithm.Log(sorted_scores)
        self.long = sorted_scores[:self.num_fine]


    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