# region imports
from AlgorithmImports import *
# endregion


class OneDTEIndexOptionUniverseAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2022, 1, 1)
        self.set_end_date(2022, 9, 1)
        self.set_cash(1_000_000)
        self.settings.automatic_indicator_warm_up = True
        self.portfolio.margin_call_model = MarginCallModel.NULL
        self.portfolio.set_positions(SecurityPositionGroupModel.NULL)
        self.set_security_initializer(MySecurityInitializer(self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))
        # Add 1DTE SPY contracts.
        self._index = self.add_index('SPX')
        self._option = self.add_index_option(self._index.symbol, 'SPXW')
        self._option.set_filter(lambda u: u.include_weeklys().expiration(1, 1).strikes(-20, 20))
        # Create a member to track when the algorithm should trade.
        self._can_trade = False
        date_rule = self.date_rules.every_day(self._option.symbol)
        self.schedule.on(date_rule, self.time_rules.at(15, 45), lambda: setattr(self, '_can_trade', True))
        self.schedule.on(date_rule, self.time_rules.at(16, 15), lambda: setattr(self, '_can_trade', False))
        self.schedule.on(date_rule, self.time_rules.after_market_open(self._index.symbol, 180), self.liquidate)
        # Load some parameters.
        self._bb_period = self.get_parameter('bb_period', 15)
        self._bb_std = self.get_parameter('bb_std', 3)
        self._position_size = self.get_parameter('position_size', 0.01) # 0.01 => 1%
        self._minimum_dollar_volume = self.get_parameter('minimum_dollar_volume', 1_000_000)

    def on_data(self, data):
        # Only process data during last 30 minutes of the day.
        if not self._can_trade:
            return
        # Get the Option chain.
        chain = data.option_chains.get(self._option.symbol)
        if not chain:
            return
        # Select contracts that have a volume above the upper Bollinger Band and a dollar volume above the threshold.
        contracts = [self.securities[c.symbol] for c in chain]
        contracts = [c for c in contracts if not c.invested and c.volume > c.volume_bb.upper_band.current.value and c.volume * c.price * 100 > self._minimum_dollar_volume]
        if not contracts:
            return
        # Select a contract (for example, the contract with volume furthest above its upper volume Bollinger Band).
        contract = sorted(contracts, key=lambda c: c.volume - c.volume_bb.current.value)[-1]
        # Place the trade.
        self.market_order(
            contract.symbol, 
            -max(1, int(contract.volume*self._position_size)), 
            tag=f'{contract.volume_bb.upper_band.current.value}; {contract.volume}; {self._index.price}'
        )
    
    def on_securities_changed(self, changes):
        # Add a Bollinger Band indicator to each 1DTE contract.
        for security in changes.added_securities:
            security.volume_bb = self.bb(security.symbol, self._bb_period, self._bb_std, resolution=Resolution.MINUTE, selector=Field.VOLUME)


class MySecurityInitializer(BrokerageModelSecurityInitializer):

    def __init__(self, brokerage_model: IBrokerageModel, security_seeder: ISecuritySeeder) -> None:
        super().__init__(brokerage_model, security_seeder)    

    def initialize(self, security: Security) -> None:
        super().initialize(security)
        security.set_buying_power_model(BuyingPowerModel.NULL)