from QuantConnect import *
from QuantConnect.Algorithm import *
from QuantConnect.Algorithm.Framework import *
from QuantConnect.Algorithm.Framework.Portfolio import PortfolioTarget
from QuantConnect.Algorithm.Framework.Risk import RiskManagementModel
from Portfolio.EqualWeightingPortfolioConstructionModel import EqualWeightingPortfolioConstructionModel

class RiskTest(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 1, 1)      
        self.SetEndDate(2019,2,1)
        self.SetCash(1000)              
        self.AddEquity("SPY", Resolution.Minute)
        self.Securities['SPY'].FeeModel = ConstantFeeModel(0)
        
        # Emit new insights every Friday
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Friday),
                         self.TimeRules.AfterMarketOpen('SPY', 0),
                         self.PlaceOrder
                        )
        
        self.SetPortfolioConstruction(MyPortfolioModel())
        self.AddRiskManagement(MyRiskModel(1e-5))
        self.SetExecution(ImmediateExecutionModel())
        
    def OnData(self, data):
        pass
    
    def PlaceOrder(self):
        """
        Emit insight to short SPY every Friday
        """
        
        insight = Insight.Price('SPY', timedelta(days=30), InsightDirection.Down, None, None, None, 1)
        self.EmitInsights([insight])
        
class MyPortfolioModel(EqualWeightingPortfolioConstructionModel):
    def __init__(self):
        pass
        
    def CreateTargets(self, algorithm, insights):
        
        # Simple insight weighting PCM
        targets = []
        for insight in insights:
            targ = PortfolioTarget(insight.Symbol, insight.Direction*insight.Weight)
            targets.append(targ)
        return targets

class MyRiskModel(RiskManagementModel):

    def __init__(self, maxDrawdown=0.05):

        self.maxDrawdown = maxDrawdown
        self.liquidatedSymbols = set()                                  # Tracks symbols that have been liquidated
        self.currentTargets = []                                        # Tracks state of current targets

    def ManageRisk(self, algorithm, targets):
        
        # Reset trackers on new targets
        if (set(targets) != self.currentTargets) and len(targets)>0:
            algorithm.Log(f'New Targets. Quantity: {targets[0].Quantity}')
            self.liquidatedSymbols = set()
            self.currentTargets = set(targets)

        riskAdjustedTargets = []
        for _ in algorithm.Securities:
            symbol = _.Key                                              # Symbol object
            security = _.Value                                          # Security object
            ticker = symbol.Value                                       # String ticker

            symbolPnL = security.Holdings.UnrealizedProfitPercent       # Current PnL

            # Liquidate if exceed  drawdown
            if (symbolPnL < -self.maxDrawdown) or (ticker in self.liquidatedSymbols):
                riskAdjustedTargets.append(PortfolioTarget(symbol, 0))
                
                if algorithm.Securities[symbol].Invested:
                    self.liquidatedSymbols.add(ticker)
                    algorithm.Log(f'Trailing stop loss triggered for {ticker}.')

        return riskAdjustedTargets