# region imports
from AlgorithmImports import *
from datetime import timedelta
import math
# endregion


class ImprovedCoveredCallAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2024, 9, 1)
        self.set_end_date(2024, 12, 31)
        self.set_cash(100000)

        self.settings.seed_initial_prices = True

        seeder = FuncSecuritySeeder(self.get_last_known_prices)
        self.add_security_initializer(lambda security: seeder.seed_security(security))

        self._spy = self.add_equity(
            "SPY",
            data_normalization_mode=DataNormalizationMode.RAW
        ).symbol

        # Strategy parameters
        self._target_weight = 0.50
        self._target_delta = 0.30
        self._min_dte = 4
        self._max_dte = 10
        self._min_open_interest = 100
        self._max_bid_ask_spread = 0.15
        self._profit_take_pct = 0.75

        # Track short call
        self._short_call = None
        self._short_call_entry_price = None
        self._short_call_strike = None
        self._short_call_expiry = None

        # Trend filter
        self._sma_fast = self.sma(self._spy, 20, Resolution.DAILY)
        self._sma_slow = self.sma(self._spy, 50, Resolution.DAILY)

        self.set_warm_up(60, Resolution.DAILY)

        # Sell covered call weekly
        self.schedule.on(
            self.date_rules.every(DayOfWeek.MONDAY),
            self.time_rules.after_market_open(self._spy, 5),
            self.sell_covered_call
        )

        # Manage open short call daily
        self.schedule.on(
            self.date_rules.every_day(self._spy),
            self.time_rules.after_market_open(self._spy, 30),
            self.manage_position
        )

    def sell_covered_call(self):
        if self.is_warming_up:
            return

        # Do not sell a new call if one is already open
        if self._short_call is not None and self.portfolio[self._short_call].invested:
            return

        underlying_price = self.securities[self._spy].price

        # Trend filter: avoid selling calls during strong uptrend
        if self._sma_fast.current.value > self._sma_slow.current.value * 1.02:
            self.debug("Skipping covered call: SPY is in strong uptrend.")
            return

        # Buy/adjust SPY position to target allocation
        self.set_holdings(self._spy, self._target_weight)

        chain = self.option_chain(self._spy)
        if chain is None:
            return

        min_expiry = self.time.date() + timedelta(days=self._min_dte)
        max_expiry = self.time.date() + timedelta(days=self._max_dte)

        contracts = []

        for contract in chain:
            if contract.right != OptionRight.CALL:
                continue

            if contract.expiry.date() < min_expiry or contract.expiry.date() > max_expiry:
                continue

            # OTM calls only
            if contract.strike <= underlying_price * 1.01:
                continue

            bid = contract.bid_price
            ask = contract.ask_price

            if bid <= 0 or ask <= 0:
                continue

            if ask - bid > self._max_bid_ask_spread:
                continue

            if contract.open_interest < self._min_open_interest:
                continue

            contracts.append(contract)

        if not contracts:
            self.debug("No suitable covered call found.")
            return

        # Prefer 30-delta call if Greeks exist; otherwise choose nearest 2-5% OTM
        def contract_score(contract):
            delta_score = 10

            if contract.greeks is not None and contract.greeks.delta is not None:
                delta_score = abs(abs(contract.greeks.delta) - self._target_delta)

            moneyness_score = abs((contract.strike / underlying_price) - 1.03)

            return delta_score + moneyness_score

        selected = sorted(contracts, key=contract_score)[0]

        option_symbol = selected.symbol
        option_security = self.add_option_contract(option_symbol, Resolution.MINUTE)

        multiplier = option_security.contract_multiplier

        shares_held = self.portfolio[self._spy].quantity
        contracts_to_sell = math.floor(shares_held / multiplier)

        if contracts_to_sell <= 0:
            self.debug("Not enough SPY shares to sell covered call.")
            return

        quantity = -contracts_to_sell

        self.market_order(option_symbol, quantity)

        self._short_call = option_symbol
        self._short_call_entry_price = (selected.bid_price + selected.ask_price) / 2
        self._short_call_strike = selected.strike
        self._short_call_expiry = selected.expiry

        self.debug(
            f"Sold {abs(quantity)} covered call(s): "
            f"{option_symbol}, strike={selected.strike}, expiry={selected.expiry.date()}"
        )

    def manage_position(self):
        if self.is_warming_up:
            return

        if self._short_call is None:
            return

        if not self.portfolio[self._short_call].invested:
            self.reset_short_call_tracking()
            return

        if not self.securities.contains_key(self._short_call):
            return

        option_price = self.securities[self._short_call].price
        underlying_price = self.securities[self._spy].price

        if option_price <= 0 or self._short_call_entry_price is None:
            return

        # Take profit if option has lost 75% of its entry value
        profit_target_price = self._short_call_entry_price * (1 - self._profit_take_pct)

        if option_price <= profit_target_price:
            self.debug("Taking profit on short call.")
            self.liquidate(self._short_call, tag="Covered call profit target hit")
            self.reset_short_call_tracking()
            return

        # Roll/close if call is ITM near expiration
        days_to_expiry = (self._short_call_expiry.date() - self.time.date()).days

        if days_to_expiry <= 1 and underlying_price > self._short_call_strike:
            self.debug("Closing ITM short call near expiration.")
            self.liquidate(self._short_call, tag="Closing ITM covered call")
            self.reset_short_call_tracking()
            return

    def reset_short_call_tracking(self):
        self._short_call = None
        self._short_call_entry_price = None
        self._short_call_strike = None
        self._short_call_expiry = None