# region imports
from AlgorithmImports import *
# endregion

class DeterminedOrangeArmadillo(QCAlgorithm):

    _black_fridays = [
        datetime(1998, 11, 27),
        datetime(1999, 11, 26),
        datetime(2000, 11, 24),
        datetime(2001, 11, 23),
        datetime(2002, 11, 29),
        datetime(2003, 11, 28),
        datetime(2004, 11, 26),
        datetime(2005, 11, 25),
        datetime(2006, 11, 24),
        datetime(2007, 11, 23),
        datetime(2008, 11, 28),
        datetime(2009, 11, 27),
        datetime(2010, 11, 26),
        datetime(2011, 11, 25),
        datetime(2012, 11, 23),
        datetime(2013, 11, 29),
        datetime(2014, 11, 28),
        datetime(2015, 11, 27),
        datetime(2016, 11, 25),
        datetime(2017, 11, 24),
        datetime(2018, 11, 23),
        datetime(2019, 11, 29),
        datetime(2020, 11, 27),
        datetime(2021, 11, 26),
        datetime(2022, 11, 25),
        datetime(2023, 11, 24),
    ]

    _prime_days = [
        datetime(2015, 7, 15),
        datetime(2016, 7, 12),
        datetime(2017, 7, 11),
        datetime(2018, 7, 17),
        datetime(2019, 7, 15),
        datetime(2020, 10, 13),
        datetime(2021, 7, 21),
        datetime(2022, 7, 12),
        datetime(2023, 7, 11),
    ]
    _contract_symbol = None

    def initialize(self):
        self.set_start_date(2012, 11, 13)
        self.set_end_date(2024, 10, 1)
        self.set_cash(1_000_000)
        self.set_security_initializer(BrokerageModelSecurityInitializer(self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))
        self._amzn = self.add_equity('AMZN', data_normalization_mode=DataNormalizationMode.RAW)
        self._spy = self.add_equity('SPY')
        self._holding_period = self.get_parameter('holding_period', 14)
        for holidays in [self._black_fridays, self._prime_days]:
            for holiday in holidays:
                # Sell an AMZN put contract before the holiday.
                self.schedule.on(
                    self.date_rules.on(self._spy.exchange.hours.get_next_market_close(holiday - timedelta(self._holding_period), False)), 
                    self.time_rules.before_market_close(self._spy.symbol, 1), 
                    self._sell_put
                )
                # Liquidate the put contract after the holiday.
                self.schedule.on(
                    self.date_rules.on(self._spy.exchange.hours.get_next_market_close(holiday + timedelta(1), False)), 
                    self.time_rules.before_market_close(self._spy.symbol, 1), 
                    lambda: self.liquidate(self._contract_symbol) if self._contract_symbol else None
                )

    def on_data(self, data):
        # If the buyer exercises the Option, liquidate the underlying Equity position.
        if self._amzn.holdings.invested:
            self.liquidate(self._amzn.symbol)
            self._contract_symbol = None

    def _sell_put(self):
        chain = self.option_chain(self._amzn.symbol).data_frame
        if chain.empty:
            return
        expiry_threshold = self._amzn.exchange.hours.get_next_market_close(self.time + timedelta(self._holding_period), False)
        expiry = chain[chain.expiry > expiry_threshold].expiry.min()
        self._contract_symbol = chain[
            (chain.expiry == expiry) &
            (chain.right == OptionRight.PUT) &
            (chain.strike <= chain.underlyinglastprice)
        ].sort_values('openinterest').index[-1]
        self.add_option_contract(self._contract_symbol)
        self.set_holdings(self._contract_symbol, -0.2)