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

class VerticalTachyonRegulators(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2011, 1, 1)  # Set Start Date
        self.SetEndDate(2019, 1, 1)
        #self.SetStartDate(2019, 1, 1)
        #self.SetEndDate(2019, 8, 16)
        self.SetCash(1000000)  # 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

        ## Coarse/Fine universe selection model
        self.UniverseSettings.Resolution = Resolution.Daily
        self.AddUniverse(self.CoarseSelectionFunction, self.FineSelectionFunction)

        ## Custom Alpha model where we can perform the trading logic and the filtering done based on Fundamental Data
        self.AddAlpha(FundamentalFactorAlphaModel(self.num_fine))


    def CoarseSelectionFunction(self, coarse):

        ## If not time to rebalance, keep the same universe
        if self.Time.month == self.lastMonth: 
            return Universe.Unchanged

        ## Else reassign the month variable and filter
        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)

        return [i.Symbol for i in selected[:self.num_coarse]]

    def FineSelectionFunction(self, fine):

        ## Filter the fine data for equities with non-zero/non-null Value,
        ## 1-month Price Change, and Book Value per Share
        filtered_fine = [x.Symbol for x in fine if x.OperationRatios.OperationMargin.Value > 0
                                        and x.ValuationRatios.PriceChange1M > 0
                                        and x.ValuationRatios.BookValuePerShare > 0]
        
        return filtered_fine


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

    def Update(self, algorithm, data):

        ## Return no insights if it's not the month to rebalance
        if algorithm.Time.month == self.lastMonth: 
            return []
        
        self.lastMonth = algorithm.Time.month

        ## Create empty list of insights
        insights = []

        ## We will liquidate any securities we're still invested in that we don't want to hold a position
        ## for the next month
        for kvp in algorithm.Portfolio:
            holding = kvp.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, None, None))
        
        ## Emit Up (buy) Insights for our desired long positions
        for symbol in self.longs:
            insights.append(Insight(symbol, timedelta(30), InsightType.Price, InsightDirection.Up, 0.01, None))
            
        ## Emit Down (sell) Insights for our desired short positions
        for symbol in self.shorts:
            insights.append(Insight(symbol, timedelta(30), InsightType.Price, InsightDirection.Down, 0.01, None))

        return insights


    def OnSecuritiesChanged(self, algorithm, changes):
        
        ## Get the securities
        added = [x for x in changes.AddedSecurities]
        
        ## Perform filtering/sorting based on Value, Quality, and Momentum
        sortedByfactor1 = sorted(added, key=lambda x: x.Fundamentals.OperationRatios.OperationMargin.Value, reverse=True)
        sortedByfactor2 = sorted(added, key=lambda x: x.Fundamentals.ValuationRatios.PriceChange1M, reverse=True)
        sortedByfactor3 = sorted(added, key=lambda x: x.Fundamentals.ValuationRatios.BookValuePerShare, reverse=True)
        
        ## Create dictionary to store scores
        scores = {}
        
        ## Assign a score to each stock.
        for i,ele in enumerate(sortedByfactor1):
            rank1 = i
            rank2 = sortedByfactor2.index(ele)
            rank3 = sortedByfactor3.index(ele)
            scores[ele] = rank1*0.2 + rank2*0.4 + rank3*0.4
        
        ## Sort the stocks by their scores
        sorted_stock = sorted(scores.items(), key=lambda d:d[1],reverse=False)
        sorted_symbol = [x[0] for x in 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:]]
        algorithm.Log('Long: ' + ', '.join(sorted([x.Value for x in self.longs])))
        algorithm.Log('Short: ' + ', '.join(sorted([x.Value for x in self.shorts])))
        

from pytz import utc
from itertools import groupby
from datetime import datetime, timedelta
UTCMIN = datetime.min.replace(tzinfo=utc)

class EqualWeightingPortfolioConstructionModel(PortfolioConstructionModel):
    '''Provides an implementation of IPortfolioConstructionModel that gives equal weighting to all securities.
    The target percent holdings of each security is 1/N where N is the number of securities.
    For insights of direction InsightDirection.Up, long targets are returned and
    for insights of direction InsightDirection.Down, short targets are returned.'''

    def __init__(self, resolution = Resolution.Daily):
        '''Initialize a new instance of EqualWeightingPortfolioConstructionModel
        Args:
            resolution: Rebalancing frequency'''
        self.insightCollection = InsightCollection()
        self.removedSymbols = []
        self.nextExpiryTime = UTCMIN
        self.rebalancingTime = UTCMIN
        self.rebalancingPeriod = TimeSpan.FromDays(30)

    def ShouldCreateTargetForInsight(self, insight):
        '''Method that will determine if the portfolio construction model should create a
        target for this insight
        Args:
            insight: The insight to create a target for'''
        return True

    def DetermineTargetPercent(self, activeInsights):
        '''Will determine the target percent for each insight
        Args:
            activeInsights: The active insights to generate a target for'''
        result = {}

        # give equal weighting to each security
        count = sum(x.Direction != InsightDirection.Flat for x in activeInsights)
        percent = 0 if count == 0 else 1.0 / count
        for insight in activeInsights:
            result[insight] = insight.Direction * percent
        return result

    def CreateTargets(self, algorithm, insights):
        '''Create portfolio targets from the specified insights
        Args:
            algorithm: The algorithm instance
            insights: The insights to create portfolio targets from
        Returns:
            An enumerable of portfolio targets to be sent to the execution model'''

        targets = []

        if (algorithm.UtcTime <= self.nextExpiryTime and
            algorithm.UtcTime <= self.rebalancingTime and
            len(insights) == 0 and
            self.removedSymbols is None):
            return targets

        for insight in insights:
            if self.ShouldCreateTargetForInsight(insight):
                self.insightCollection.Add(insight)

        # Create flatten target for each security that was removed from the universe
        if self.removedSymbols is not None:
            universeDeselectionTargets = [ PortfolioTarget(symbol, 0) for symbol in self.removedSymbols ]
            targets.extend(universeDeselectionTargets)
            self.removedSymbols = None

        # Get insight that haven't expired of each symbol that is still in the universe
        activeInsights = self.insightCollection.GetActiveInsights(algorithm.UtcTime)

        # Get the last generated active insight for each symbol
        lastActiveInsights = []
        for symbol, g in groupby(activeInsights, lambda x: x.Symbol):
            lastActiveInsights.append(sorted(g, key = lambda x: x.GeneratedTimeUtc)[-1])

        # Determine target percent for the given insights
        percents = self.DetermineTargetPercent(lastActiveInsights)

        errorSymbols = {}
        for insight in lastActiveInsights:
            target = PortfolioTarget.Percent(algorithm, insight.Symbol, percents[insight])
            if not target is None:
                targets.append(target)
            else:
                errorSymbols[insight.Symbol] = insight.Symbol

        # Get expired insights and create flatten targets for each symbol
        expiredInsights = self.insightCollection.RemoveExpiredInsights(algorithm.UtcTime)

        expiredTargets = []
        for symbol, f in groupby(expiredInsights, lambda x: x.Symbol):
            if not self.insightCollection.HasActiveInsights(symbol, algorithm.UtcTime) and not symbol in errorSymbols:
                expiredTargets.append(PortfolioTarget(symbol, 0))
                continue

        targets.extend(expiredTargets)

        self.nextExpiryTime = self.insightCollection.GetNextExpiryTime()
        if self.nextExpiryTime is None:
            self.nextExpiryTime = UTCMIN

        self.rebalancingTime = algorithm.UtcTime + self.rebalancingPeriod

        return targets

    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 removed symbol and invalidate them in the insight collection
        self.removedSymbols = [x.Symbol for x in changes.RemovedSecurities]
        self.insightCollection.Clear(self.removedSymbols)