# region imports
from AlgorithmImports import *
# endregion

class FocusedYellowGorilla(QCAlgorithm):
    def Initialize(self):
        self.SetStartDate(2024, 1, 8)
        self.SetEndDate(2024, 1, 8)
        self.SetCash(10000000)

        # 添加股票数据，以5分钟为间隔合并数据
        self.aapl = self.AddEquity("AAPL", Resolution.Minute, extendedMarketHours = True).Symbol
        fiveMinuteConsolidator = TradeBarConsolidator(timedelta(minutes=5))
        fiveMinuteConsolidator.DataConsolidated += self.FiveMinuteBarHandler
        self.SubscriptionManager.AddConsolidator(self.aapl, fiveMinuteConsolidator)


        # Calculate the EMA with 5 5-minute bars
        self.fast_ema = ExponentialMovingAverage(5)
        
        # Calculate the EMA with 20 5-minut bars
        self.slow_ema = ExponentialMovingAverage(20)
        
        self.SetWarmup(100)

        chart = Chart("Trade")
        self.AddChart(chart)
        chart.AddSeries(Series("EMA-5", SeriesType.Line,"", Color.Purple))
        chart.AddSeries(Series("EMA-20", SeriesType.Line, "", Color.Black))
        chart.AddSeries(Series('Candle', SeriesType.Bar))    
        chart = Chart("Candlestick")
        self.AddChart(chart)
        chart.AddSeries(CandlestickSeries("AAPL", "$"))

    def FiveMinuteBarHandler(self, sender, consolidated):
    #     '''This is our event handler for our 5 minute trade bar defined above in Initialize(). So each time the
    #     consolidator produces a new 5 minute bar, this function will be called automatically. The 'sender' parameter
    #      will be the instance of the IDataConsolidator that invoked the event, but you'll almost never need that!'''

        self.fast_ema.Update(consolidated.Time, consolidated.Close)
        self.slow_ema.Update(consolidated.Time, consolidated.Close)

        # 获取当前快慢均线值
        fast_ema_value = self.fast_ema.Current.Value
        slow_ema_value = self.slow_ema.Current.Value

        self.Plot("Trade", 'EMA-5', fast_ema_value)
        self.Plot("Trade", 'EMA-20', slow_ema_value)
        self.Plot("Trade", "AAPL", consolidated)

        # 判断快慢均线交叉
        if fast_ema_value > slow_ema_value and self.fast_ema.Previous.Value <= self.slow_ema.Previous.Value:
            # 发出买入信号
            self.SetHoldings("AAPL", 1)
            self.Debug("买入 AAPL")

        elif fast_ema_value < slow_ema_value and self.fast_ema.Previous.Value >= self.slow_ema.Previous.Value:
            # 发出卖出信号
            self.Liquidate("AAPL")
            self.Debug("卖出 AAPL")



    def OnData(self, slice):
        pass