from alphaMOM import MOMAlphaModel
# alpha model for insights only.  To be integrated
from riskTrailStop import TrailingStopRiskManagementModel

class DynamicUniverse(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2011, 1, 1)
        self.SetEndDate(2021, 5, 1)
        self.SetCash(1000000)
        self.rebalanceTime = datetime.min
        self.activeStocks = set()

        self.AddUniverse(self.CoarseFilter, self.FineFilter)
        self.UniverseSettings.Resolution = Resolution.Hour
        
        self.SetAlpha(MOMAlphaModel())
        
        # 10% trailing stop risk management
        self.SetRiskManagement(TrailingStopRiskManagementModel())
    
        self.portfolioTargets = []

    def CoarseFilter(self, coarse):
        # Rebalancing monthly
        if self.Time <= self.rebalanceTime:
            return self.Universe.Unchanged
        self.rebalanceTime = self.Time + timedelta(20)
        
        sortedByDollarVolume = sorted(coarse, key=lambda x: x.DollarVolume, reverse=True)
        return [x.Symbol for x in sortedByDollarVolume if x.Price > 10
                                                and x.DollarVolume > 1000000000][:100]

    def FineFilter(self, fine):
        sortedByPERatio = sorted(fine, key=lambda x: x.ValuationRatios.PERatio)
        return [x.Symbol for x in sortedByPERatio if x.ValuationRatios.PERatio > 0][10:]
        # takes top 10 by PERatio

    def OnSecuritiesChanged(self, changes):
        # close positions in removed securities
        for x in changes.RemovedSecurities:
            self.Liquidate(x.Symbol)
            self.activeStocks.remove(x.Symbol)
        
        # add new symbol to activeStocks
        for x in changes.AddedSecurities:
            self.activeStocks.add(x.Symbol)   

        # adjust targets if universe has changed
        self.portfolioTargets = [PortfolioTarget(symbol, 1/len(self.activeStocks)) 
                            for symbol in self.activeStocks]
                            
        # log the changes in the function 
        # self.Log(f"OnSecuritiesChanged({self.UtcTime}):: {changes}")

    def OnData(self, data):

        if self.portfolioTargets == []:
            return
        
        for symbol in self.activeStocks:
            if symbol not in data:
                return
        
        # apply alphaMOM criteria to constrain SetHoldings method to "Up" portfolioTargets only
        for symbol in self.portfolioTargets:
            if InsightDirection.Up:
                self.SetHoldings(self.portfolioTargets)
        
        self.portfolioTargets = []
        
    def OnOrderEvent(self, orderEvent):
        # log order data for text/csv output file
        order = self.Transactions.GetOrderById(orderEvent.OrderId)
        if orderEvent.Status == OrderStatus.Filled: 
            self.Log("{0}, {1}, {2}".format(self.Time, order.Type, orderEvent))

from datetime import timedelta

class MOMAlphaModel(AlphaModel):
    
    def __init__(self):
        self.mom = []
      
    def OnSecuritiesChanged(self, algorithm, changes):
        
        # Initialize a 30-day momentum indicator for each symbol
        for security in changes.AddedSecurities:
            symbol = security.Symbol
            self.mom.append({"symbol":symbol, "indicator":algorithm.MOM(symbol, 30, Resolution.Daily)})
        
    def Update(self, algorithm, data):

        # Sort the list of dictionaries by indicator in descending order
        ordered = sorted(self.mom, key=lambda kv: kv["indicator"].Current.Value, reverse=True)

        # Return a group of insights, emitting InsightDirection.Up for the first item of ordered, and InsightDirection.Flat for the second
        return Insight.Group(
            [
                Insight.Price(ordered[0]["symbol"], timedelta(1), InsightDirection.Up),
                Insight.Price(ordered[1]["symbol"], timedelta(1), InsightDirection.Flat)
            ])



# Your New Python File

from clr import AddReference
AddReference("System")
AddReference("QuantConnect.Common")
AddReference("QuantConnect.Algorithm")
AddReference("QuantConnect.Algorithm.Framework")

from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Algorithm.Framework import *
from QuantConnect.Algorithm.Framework.Portfolio import *
from QuantConnect.Algorithm.Framework.Risk import *

class TrailingStopRiskManagementModel(RiskManagementModel):
    '''Provides an implementation of IRiskManagementModel that limits the maximum possible loss
    measured from the highest unrealized profit'''
    def __init__(self, maximumDrawdownPercent = 0.10):
        '''Initializes a new instance of the TrailingStopRiskManagementModel class
        Args:
            maximumDrawdownPercent: The maximum percentage drawdown allowed for algorithm portfolio compared with the highest unrealized profit, defaults to 10% drawdown'''
        self.maximumDrawdownPercent = -abs(maximumDrawdownPercent)
        self.trailingHighs = dict()

    def ManageRisk(self, algorithm, targets):
        '''Manages the algorithm's risk at each time step
        Args:
            algorithm: The algorithm instance
            targets: The current portfolio targets to be assessed for risk'''
        riskAdjustedTargets = list()

        for kvp in algorithm.Securities:
            symbol = kvp.Key
            security = kvp.Value

            # Remove if not invested
            if not security.Invested:
                self.trailingHighs.pop(symbol, None)
                continue

            # Add newly invested securities
            if symbol not in self.trailingHighs:
                self.trailingHighs[symbol] = security.Holdings.AveragePrice   # Set to average holding cost
                continue

            # Check for new highs and update - set to tradebar high
            if self.trailingHighs[symbol] < security.High:
                self.trailingHighs[symbol] = security.High
                continue

            # Check for securities past the drawdown limit
            securityHigh = self.trailingHighs[symbol]
            drawdown = (security.Low / securityHigh) - 1

            if drawdown < self.maximumDrawdownPercent:
                # liquidate
                riskAdjustedTargets.append(PortfolioTarget(symbol, 0))

        return riskAdjustedTargets