from AlgorithmImports import *

class ESFuturesTrendFollowingAlgorithm(QCAlgorithm):

    threshold = 0.01 # 1%
    
    def Initialize(self) -> None:
        self.SetCash(1000000)
        self.SetStartDate(2015, 1, 1)  # Start date

        # Define the futures contract for ES
        self.es_continuous_contract = self.AddFuture(Futures.Metals.Gold,
                                          dataNormalizationMode=DataNormalizationMode.BackwardsRatio,
                                          dataMappingMode=DataMappingMode.OpenInterest,
                                          contractDepthOffset=0)

        self.symbol = self.es_continuous_contract.Symbol

        # Historical data
        history = self.History(self.symbol, 100, Resolution.Daily)
        self.Debug(f"We got {len(history)} items from our history request")

        self.ema = self.EMA(self.symbol, 100, Resolution.Daily)
        if not history.empty:
            for time, row in history.droplevel(0).loc[self.symbol].iterrows():
                self.ema.Update(IndicatorDataPoint(time, row.close))

        self.Debug("Initialization complete. EMA100 is set up.")


    def OnData(self, slice: Slice) -> None:
        # Accessing data and handling roll overs
        for symbol, changed_event in slice.SymbolChangedEvents.items():
            self.HandleRollover(changed_event)

        mapped_symbol = self.es_continuous_contract.Mapped

        if not (slice.Bars.ContainsKey(self.symbol) and self.ema.IsReady and mapped_symbol):
            return

        # If you have subscribed to a single futures contract:
        if self.symbol in slice.Bars:
            bar = slice.Bars[self.symbol]
            # Now you can access open, high, low, close, and volume of the bar
            open_price = bar.Open
            high_price = bar.High
            low_price = bar.Low
            close_price = bar.Close
            price = bar.Price
            volume = bar.Volume

        # Trend-following logic based on 100 EMA
        if price > self.ema.Current.Value * (1+self.threshold) and not self.Portfolio[mapped_symbol].IsLong:
            self.MarketOrder(mapped_symbol, 1)  # Adjust position size as needed
            self.Log(f"Going Long on {mapped_symbol}. Price: {price}, EMA100: {self.ema.Current.Value}, at Time: {self.Time}.")
        elif price < self.ema.Current.Value * (1+self.threshold) and not self.Portfolio[mapped_symbol].IsShort:
            self.MarketOrder(mapped_symbol, -1) # Adjust position size as needed
            self.Log(f"Going Short on {mapped_symbol}. Price: {price}, EMA100: {self.ema.Current.Value}, at Time: {self.Time}.")


    def HandleRollover(self, changed_event):
        old_symbol = changed_event.OldSymbol
        new_symbol = changed_event.NewSymbol
        tag = f"Rollover at {self.Time}: {old_symbol} -> {new_symbol}"
        quantity = self.Portfolio[old_symbol].Quantity

        # Rolling over: to liquidate any position of the old mapped contract and switch to the newly mapped contract
        self.Liquidate(old_symbol, tag=tag)
        if quantity != 0: self.MarketOrder(new_symbol, quantity, tag = tag)
        self.Log(tag)