#region imports
from AlgorithmImports import *
#endregion
from datetime import timedelta
import math

class FundamentalFactorAlphaModel(AlphaModel):
    
    
    def __init__(self, num_fine, quality_weight, value_weight, size_weight):
        
        # Initialize the various variables/helpers we'll need
        self.lastYear = -1
        self.longs = []
        self.shorts = []
        self.num_fine = num_fine
        self.period = timedelta(253)
        
        # normalize quality, value, size weights
        weights = [quality_weight, value_weight, 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 len([x.Symbol for x in algorithm.Portfolio.Values if algorithm.Portfolio[x.Symbol].Invested]) > 1:
            return []

        # 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 or shorts
        for security in algorithm.Portfolio.Values:
            if security.Invested and security.Symbol not in self.longs and security.Symbol not in self.shorts:
                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 ))
                
        shorts_length = len(self.shorts)
        
        for i in range(shorts_length):
            insights.append(Insight(self.shorts[i], self.period, InsightType.Price, 
                                    InsightDirection.Down, None, (shorts_length - i)**2, None, (shorts_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
        added = [x for x in changes.AddedSecurities]
        
        # Assign quality, value, size score to each stock
        quality_scores = self.Scores(added, [ 
                                            (lambda x: x.Fundamentals.OperationRatios.DebttoAssets.Value, False, 1)])
        
        value_scores = self.Scores(added, [
                                            (lambda x: (x.Fundamentals.FinancialStatements.BalanceSheet.CurrentAssets.TwelveMonths - x.Fundamentals.FinancialStatements.BalanceSheet.CurrentLiabilities.TwelveMonths - x.Fundamentals.FinancialStatements.BalanceSheet.LongTermCapitalLeaseObligation.TwelveMonths - x.Fundamentals.FinancialStatements.BalanceSheet.LongTermDebt.TwelveMonths - x.Fundamentals.FinancialStatements.BalanceSheet.NonCurrentDeferredLiabilities.TwelveMonths) / x.Fundamentals.MarketCap, True, 1),
                                            ])
        
        size_scores = self.Scores(added, [(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)
        long_symbols = [tup[0] for tup in sorted_stock][:self.num_fine]
        # short_symbols = [tup[0] for tup in sorted_stock][-math.floor(self.num_fine/2):]

        # Sort the top stocks into the long_list
        self.longs = [security.Symbol for security in long_symbols]
        # self.shorts = [security.Symbol for security in short_symbols]
        
        # Log symbols and their score
        algorithm.Log(", ".join([str(x.Symbol.Value) + ": " + str(scores[x]) for x in long_symbols]))
        # algorithm.Log(", ".join([str(x.Symbol.Value) + ": " + str(scores[x]) for x in short_symbols]))


    def Scores(self, added, 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(added, 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

#region imports
from AlgorithmImports import *
#endregion
from Execution.ImmediateExecutionModel import ImmediateExecutionModel
from Portfolio.EqualWeightingPortfolioConstructionModel import InsightWeightingPortfolioConstructionModel
from Risk.MaximumDrawdownPercentPerSecurity import MaximumDrawdownPercentPerSecurity
from AlphaModel import FundamentalFactorAlphaModel


class VerticalTachyonRegulators(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(1999, 1, 1)
        self.SetEndDate(2023, 4, 30)
        self.SetCash(1000000)
        
        # Execution model
        self.SetExecution(ImmediateExecutionModel())
        
        # Portfolio construction model
        self.Settings.RebalancePortfolioOnSecurityChanges = False
        self.SetPortfolioConstruction(InsightWeightingPortfolioConstructionModel(self.RebalanceFunction))

        # Universe selection
        self.num_coarse = 50_000
        self.num_fine = 35
        self.lastYear = -1
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)
        
        # Set factor weighting        
        quality_weight = 0
        size_weight = 0
        value_weight = 1
        
        # Alpha Model
        self.AddAlpha(FundamentalFactorAlphaModel(self.num_fine, quality_weight, value_weight, size_weight))
        
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday), 
                        self.TimeRules.At(10, 30),
                        self.Plotting)
        self.lastRebalanceTime = self.StartDate

    def RebalanceFunction(self, time):
        if len([x.Symbol for x in self.Portfolio.Values if self.Portfolio[x.Symbol].Invested]) > 1:
            return None
        return time

    def Plotting(self):
        self.Plot("Positions", "Num", len([x.Symbol for x in self.Portfolio.Values if self.Portfolio[x.Symbol].Invested]))


    def CoarseSelectionFunction(self, coarse):
        if len([x.Symbol for x in self.Portfolio.Values if self.Portfolio[x.Symbol].Invested]) > 1:
            return Universe.Unchanged
            
        # Sort by dollar volume: most liquid to least liquid
        selected = sorted([x for x in coarse if x.HasFundamentalData], key = lambda x: x.DollarVolume, reverse=True)
        
        return [x.Symbol for x in selected[:self.num_coarse]]


    def FineSelectionFunction(self, fine):
        filtered_fine = [x for x in fine if x.MarketCap > 50_000_000
                                        and x.FinancialStatements.BalanceSheet.CurrentAssets.TwelveMonths - x.FinancialStatements.BalanceSheet.CurrentLiabilities.TwelveMonths - x.FinancialStatements.BalanceSheet.LongTermCapitalLeaseObligation.TwelveMonths - x.FinancialStatements.BalanceSheet.LongTermDebt.TwelveMonths - x.FinancialStatements.BalanceSheet.NonCurrentDeferredLiabilities.TwelveMonths > 0
        ]
        fine_symbols = [x.Symbol for x in filtered_fine]
        return fine_symbols

    def OnEndOfAlgorithm(self):
    	self.Liquidate()