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 PortfolioTarget
from QuantConnect.Algorithm.Framework.Risk import RiskManagementModel

class ScheduledRiskTest(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2015, 1, 1)
        self.SetEndDate(2015, 2, 1)
        self.SetCash(1000)  
        self.SetBrokerageModel(BrokerageName.OandaBrokerage, AccountType.Margin)
        self.AddForex("EURUSD", Resolution.Minute, Market.Oanda)
        
        # Emit an insight every Monday
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Monday),
                         self.TimeRules.AfterMarketOpen('EURUSD',15),
                         self.TradeStuff)
                         
        self.SetPortfolioConstruction(CustomPortfolioModel())
        self.SetExecution(ImmediateExecutionModel())
        self.AddRiskManagement(PortfolioDrawdownRiskModel(1e-7))    # Some very small drawdown threshold
        
    def OnData(self, data):
        pass
    
    def TradeStuff(self):
        # Long EURUSD
        insight = Insight('EURUSD',  
                          timedelta(days=2),
                          InsightType.Price,
                          InsightDirection.Up,
                          None,
                          None,
                          None,
                          1
                        )
        self.EmitInsights(insight)

class CustomPortfolioModel(PortfolioConstructionModel):
    def __init__(self):
        pass
        
    def CreateTargets(self, algorithm, insights):
        
        targets = []
        for insight in insights:
            targ = PortfolioTarget(insight.Symbol, 1) #insight.Direction*insight.Weight)
            targets.append(targ)
        return targets        

class PortfolioDrawdownRiskModel(RiskManagementModel):
    """
    Tracks and limits portfolio drawdowns from highest equity value. Resets trailing stops on new targets.
    """
    def __init__(self, maxDrawdown=0.1):
        self.maxDrawdown = maxDrawdown
        self.currentTargets = set()                                                                    # Tracks states of new targets
        self.portfolioHigh = 0                                                                         # Tracks portfolio highs
        self.isLiquidated = False          # Flag for whether risk model liquidated
        self.logged = False

    def ManageRisk(self, algorithm, targets):

        # Reset trackers on new targets
        if set(targets) != set(self.currentTargets):
            self.currentTargets = set(targets)
            self.portfolioHigh = 0
            self.isLiquidated = False
        
        # Update trailing highs of portfolio value
        portfolioValue = algorithm.Portfolio.TotalPortfolioValue
        if portfolioValue > self.portfolioHigh:
            self.portfolioHigh = portfolioValue
            return []
        
        if self.portfolioHigh == 0:
            return []
        
        # Liquidate portfolio if exceed drawdown
        drawdown = (portfolioValue / self.portfolioHigh) - 1.0
        if (drawdown < -self.maxDrawdown) or self.isLiquidated:
            self.isLiquidated = True
            
            portfolio_targets = []
            for kvp in algorithm.Securities:
                security = kvp.Value
                if security.Invested:
                    portfolio_targets.append(PortfolioTarget(security.Symbol, 0))
                    algorithm.Debug(f'{algorithm.Time} Risk Model triggered with realized {drawdown*100:.2f}% drawdown. Liquidated.')
            return portfolio_targets
        return []