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

class ParticleNadionReplicator(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2019, 12, 19)  # Set Start Date
        self.SetCash(100000)  # Set Strategy Cash
        
        self.SetBrokerageModel(BrokerageName.InteractiveBrokersBrokerage, AccountType.Margin)
        # self.AddEquity("SPY", Resolution.Minute)
        self.SetExecution(ImmediateExecutionModel())

        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel())

    
        self.SetRiskManagement(MaximumUnrealizedProfitPercentPerSecurity(0.03))
        tickers = ["SPY", "QQQQ", "TLT"]
        self.UniverseSettings.Resolution = Resolution.Minute
        symbols = [ Symbol.Create(ticker, SecurityType.Equity, Market.USA)  for ticker in tickers]
        self.SetUniverseSelection( ManualUniverseSelectionModel(symbols) )
        self.SetBenchmark("SPY")
        self.SetSecurityInitializer(lambda x: x.SetDataNormalizationMode(DataNormalizationMode.Raw))
        
        self.alpha = MyAlphaModel(2.5)#leverage
        self.AddAlpha(self.alpha)

class MyAlphaModel(AlphaModel):
    def __init__(self, leverage = 1):
        self.leveragePercentage = leverage
        self.equities = []
        
    def OnSecuritiesChanged(self, algorithm, changes):
        self.equities = [x for x in changes.AddedSecurities if str(type(x)) == "<class 'QuantConnect.Securities.Equity.Equity'>"]

        if algorithm.Portfolio.CashBook.TotalValueInAccountCurrency <= 25000 and self.leveragePercentage > 2:
            self.leveragePercentage = 2
        elif algorithm.Portfolio.CashBook.TotalValueInAccountCurrency > 25000 and self.leveragePercentage > 4:
            self.leveragePercentage = 4
        
        for symbol in self.equities:
            #Leverage
            algorithm.Securities[symbol.Symbol].SetLeverage(self.leveragePercentage)
            if self.leveragePercentage > 2:
                #PDT Margin Model for Short Trading... if any
                algorithm.Securities[symbol.Symbol].MarginModel = PatternDayTradingMarginModel()
                
    def Update(self, algorithm, data):
        insights = []
        
        if algorithm.IsWarmingUp or algorithm.Portfolio.Invested:
            return insights
        
        for symbol in self.equities:
            insights.append(Insight.Price(symbol.Symbol,timedelta(days=180), InsightDirection.Up))
        
        return insights