from QuantConnect.Algorithm.Framework.Portfolio import PortfolioTarget
class SMATrading(QCAlgorithm):

    def Initialize(self):
        self.SetStartDate(2021, 4, 25)  #设置回测起始时间
        self.SetEndDate(2021, 4, 27) #设置回测结束时间
        self.SetCash(100000000)  #设置金额
     
        symbols = [Symbol.Create('SPY', SecurityType.Equity, Market.USA)] #获取SPY数据
        #1 使用手动资产集选择模型
        self.AddUniverseSelection(ManualUniverseSelectionModel(symbols)) 
        self.UniverseSettings.Resolution = Resolution.Minute  # 设置资产使用min分辨率
        self.UniverseSettings.ExtendedMarketHours = True  #提取盘前盘后交易数据
        #self.SetSecurityInitializer(lambda x: x.SetDataNormalizationMode(DataNormalizationMode.Raw)) #请求资产集证券的原始价格
        # self.SetWarmUp(30) #预热30笔数据
        self.AddAlpha(SmaCrossAlphaModel(5 , 30 , Resolution.Minute)) # 使用smaAlpha模型初始化框架
        self.SetPortfolioConstruction(EqualWeightingPortfolioConstructionModel()) # 使用EqualWeightingPortfolio模型初始化框架
        self.SetExecution(MyExecutionModel()) # 使用Execution模型初始化框架
        self.AddRiskManagement( MaximumDrawdownPercentPerSecurity(0.05) ) # 使用MaximumDrawdown模型初始化框架
        
        #self.SetWarmUp(30)
 # 2创建sma Alpha 模型
class SmaCrossAlphaModel(AlphaModel):
    '''Alpha model that uses an SMA cross to create insights'''

    def __init__(self,fastPeriod = 5 ,slowPeriod = 30,resolution = Resolution.Minute):
        self.fastPeriod = fastPeriod
        self.slowPeriod = slowPeriod
        self.resolution = resolution
        self.predictionInterval = Time.Multiply(Extensions.ToTimeSpan(resolution), fastPeriod)
        self.symbolDataBySymbol = {}
        
        resolutionString = Extensions.GetEnumString(resolution, Resolution)
        
    def OnSecuritiesChanged(self, algorithm, changes):
        for added in changes.AddedSecurities:
            symbolData = self.symbolDataBySymbol.get(added.Symbol)
            if symbolData is None:
                symbolData = SymbolData(algorithm , added , self.fastPeriod , self.slowPeriod)
                self.symbolDataBySymbol[added.Symbol] = symbolData
            else:
                # a security that was already initialized was re-added, reset the indicators
                symbolData.Fast.Reset()
                symbolData.Slow.Reset()
                
    def Update(self, algorithm, data):
        insights = []
        for symbol, symbolData in self.symbolDataBySymbol.items():
            if symbolData.consolidated:
                if symbolData.Fast.IsReady and symbolData.Slow.IsReady and symbolData.five_minutes_close.IsReady: #algorithm.Time.minute%5 == 0
                    if symbolData.SlowIsOverFast and symbolData.FastIsOverClose :
                        insights.append(Insight.Price(symbolData.Symbol, self.predictionInterval, InsightDirection.Down))
                    elif symbolData.FastIsOverSlow and symbolData.CloseIsOverFast:
                        insights.append(Insight.Price(symbolData.Symbol, self.predictionInterval, InsightDirection.Up))
            symbolData.consolidated = False
        return insights

class SymbolData:
    '''Contains data specific to a symbol required by this model'''
        
    def __init__(self, algorithm, security ,fastPeriod , slowPeriod ):
        self.algo = algorithm
        self.Security = security
        self.Symbol = security.Symbol
        self.consolidator = TradeBarConsolidator(timedelta(minutes=5))
        self.consolidator.DataConsolidated += self.OnDataConsolidated
        algorithm.SubscriptionManager.AddConsolidator(self.Symbol, self.consolidator)
        self.five_minutes_close = RollingWindow[float](2)
        # create fast/slow SMAs
        self.Fast = SimpleMovingAverage(fastPeriod)
        self.Slow = SimpleMovingAverage(slowPeriod)
        algorithm.RegisterIndicator(self.Symbol, self.Fast, self.consolidator)
        algorithm.RegisterIndicator(self.Symbol, self.Slow, self.consolidator)
        
        hist = algorithm.History(self.Symbol, 30, Resolution.Minute).loc[self.Symbol]
        if not hist.empty:
            for idx, row in hist.iterrows():
                tradebar = TradeBar(idx, self.Symbol, row.open, row.high, row.low, row.close, row.volume)
                self.consolidator.Update(tradebar)
                #self.Fast.Update(idx, row.close)
                #self.Slow.Update(idx, row.close)
                #self.five_minutes_close.Add(row.Close)
        self.consolidated = False
        
    # def Update(self, time, value):
    #     self.Fast.Update(time, value)
    #     self.Slow.Update(time, value)
    def OnDataConsolidated(self, sender, bar):
        #self.algo.Log(f"OnDataConsolidated called at {self.algo.Time}")
        self.consolidated = True
        self.five_minutes_close.Add(bar.Close)

    @property
    def FastIsOverSlow(self): return self.Fast.Current.Value > self.Slow.Current.Value
    @property
    def SlowIsOverFast(self): return self.Fast.Current.Value < self.Slow.Current.Value
    @property
    def CloseIsOverFast(self): return self.Fast.Current.Value < self.five_minutes_close[0]
    @property
    def FastIsOverClose(self): return self.Fast.Current.Value > self.five_minutes_close[0]
 
 # 3创建投资组合模型       
class EqualWeightingPortfolioConstructionModel(PortfolioConstructionModel):
    
    def CreateTargets(self, algorithm, insights):
        targets = []
        for insight in insights:
            targets.append(PortfolioTarget(insight.Symbol, insight.Direction * 100))
        return targets

# 5 创建执行模型
class MyExecutionModel(ExecutionModel):

    def Execute(self, algorithm, targets):
        for target in targets:
            open_quantity = sum([x.Quantity for x in algorithm.Transactions.GetOpenOrders(target.Symbol)])
            existing = algorithm.Securities[target.Symbol].Holdings.Quantity + open_quantity
            quantity = target.Quantity - existing
            # bidprice = algorithm.Securities[target.Symbol].BidPrice
            # askprice = algorithm.Securities[target.Symbol].AskPrice
            if quantity > 0: algorithm.MarketOrder(target.Symbol, quantity)
            if quantity < 0: algorithm.MarketOrder(target.Symbol, quantity)
        return []

# 4 创建MaximumDrawdown风险管理模型
class MaximumDrawdownPercentPerSecurity(RiskManagementModel):
    def __init__(self, maximumDrawdownPercent = 0.05):
        self.maximumDrawdownPercent = -abs(maximumDrawdownPercent)
    def ManageRisk(self, algorithm, targets):
        targets = []
        for kvp in algorithm.Securities:
            security = kvp.Value
            if not security.Invested: continue
            pnl = security.Holdings.UnrealizedProfitPercent
            if pnl < self.maximumDrawdownPercent:
                targets.append(PortfolioTarget(security.Symbol, 0)) # liquidate
        return targets