# region imports
from AlgorithmImports import *
# endregion


class VolumeSpikesInOilFuturesAlgorithm(QCAlgorithm):

    # Define some parameters
    SPIKE_MULTIPLE_ETH = 3  # x weekly-average minute volume, outside regular hours
    SPIKE_MULTIPLE_RTH = 6  # x weekly-average minute volume, during regular hours
    REFERENCE_MULTIPLE = 2  # x the same minute's volume on the busiest reference day
    TAKE_PROFIT = 0.1       # Close positions at +10% profit.
    SIGNAL_COOLDOWN = timedelta(hours=2)  # Don't enter new positions more frequently than this.

    def initialize(self):
        self.set_start_date(2026, 3, 20)
        self.set_end_date(2026, 7, 8)
        self.set_cash(100_000)
        # Add the universe of Oil Future contracts.
        self._contracts = []
        future = self.add_future(Futures.Energy.CRUDE_OIL_WTI, fill_forward=False, extended_market_hours=True)
        future.set_filter(0, 60)
        # Add the SPY ETF.
        self._spy = self.add_equity("SPY")
        # Add some members to assist with the trading logic.
        self._invested_contract = None
        self._trading_armed = False
        self._last_signal_time = datetime.min
        self._max_average_volume = None
        self._bars_of_most_liquid_contract = None
        # Add Scheduled Events to reset the daily volume tally, to refresh the weekly snapshot, and to arm trading.
        self.schedule.on(self.date_rules.every_day(future), self.time_rules.at(17, 0), self._reset_total_volume)
        self.schedule.on(self.date_rules.week_end(future), self.time_rules.at(17, 0), self._reset_weekly_snapshot)
        self.schedule.on(self.date_rules.week_start(future), self.time_rules.at(19, 0), self._arm_trading)
        # Add a warm-up period to prime the factors and historical bars.
        self.set_warm_up(timedelta(14))

    def _reset_total_volume(self):
        for contract in self._contracts:
            contract.total_volume_today = 0

    def _arm_trading(self):
        self._trading_armed = True

    def _reset_weekly_snapshot(self):
        # Get the max of each contract's average minute volume during RTH.
        self._max_average_volume = max(sum(c.volume_in_rth) / len(c.volume_in_rth) for c in self._contracts)
        # Get the minute bars of the contract that traded the most this week.
        self._bars_of_most_liquid_contract = max(self._contracts, key=lambda c: sum(bar.volume for bar in c.bars)).bars
        for contract in self._contracts:
            contract.volume_in_rth = []
            contract.bars = []
        self._trading_armed = False
            
    def on_securities_changed(self, changes):
        # As contracts enter the universe, use duck-typing to attach some custom members.
        for security in changes.added_securities:
            if security.type == SecurityType.EQUITY or security.symbol.is_canonical():
                continue
            security.total_volume_today = 0
            security.volume_in_rth = []
            security.bars = []
            self._contracts.append(security)
        for security in changes.removed_securities:
            self._contracts.remove(security)

    def on_data(self, data: Slice):
        # Apply the take-profit logic if we're holding a contract.
        if self._invested_contract:
            holding = self._invested_contract.holdings
            price = self._invested_contract.close
            if (holding.is_long and price > (1+self.TAKE_PROFIT) * holding.average_price or 
                holding.is_short and price < (1-self.TAKE_PROFIT) * holding.average_price):
                self.liquidate(tag="TAKE PROFIT")
                self._invested_contract = None
        # Record the volume and minute-bar of each contract.
        is_rth = self.is_market_open(self._spy)
        tradable_contracts = []
        for contract in self._contracts:
            bar = data.bars.get(contract)
            if not bar:
                continue
            contract.total_volume_today += bar.volume
            if is_rth:
                contract.volume_in_rth.append(bar.volume)
            contract.bars.append(bar)
            tradable_contracts.append(contract)
        # Check if we should scan for a new trade.
        if self.is_warming_up or not self._trading_armed or not tradable_contracts:
            return
        # Get the contract with the most volume today.
        contract = max(tradable_contracts, key=lambda contract: contract.total_volume_today)
        # Check if the current volume is abnormally high by comparing against last week's trading activity:
        #  - the max average minute volume in regular trading hours
        #  - the max volume of the most liquid contract during the same time (hour:minute combos)
        bar = data.bars.get(contract)
        if (bar.volume <= self._max_average_volume * (self.SPIKE_MULTIPLE_RTH if is_rth else self.SPIKE_MULTIPLE_ETH) or
            bar.volume <= self.REFERENCE_MULTIPLE * max([b.volume for b in self._bars_of_most_liquid_contract if b.end_time.time() == self.time.time()])):
            return
        # Check if the signal cool-down period has ended.
        if self.time <= self._last_signal_time + self.SIGNAL_COOLDOWN:
            return
        self._last_signal_time = self.time
        # Open a one-lot position (direction +1 long / -1 short), flipping any opposite one.
        # Green bar -> long, red bar -> short.
        direction = 1 if bar.close > bar.open else -1
        if not self._invested_contract:
            self.market_order(contract, direction)
            self._invested_contract = contract
            return
        holding = self._invested_contract.holdings
        if holding.is_short if direction == 1 else holding.is_long:
            self.liquidate()
            self.market_order(contract, direction)
            self._invested_contract = contract