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=7),
                          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, 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

    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:
            if algorithm.Portfolio.Invested:
                algorithm.Debug(f'{algorithm.Time} Risk Model triggered with realized {drawdown*100:.2f}% drawdown. Liquidated.')
            self.isLiquidated = True
            return [ PortfolioTarget(target.Symbol, 0) for target in targets]
        return []