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


class ImprovedCoveredCallAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2025, 6, 1)
        self.set_end_date(2025, 12, 31)
        self.set_cash(100_000)

        self.settings.seed_initial_prices = True
        seeder = FuncSecuritySeeder(self.get_last_known_prices)
        self.add_security_initializer(lambda s: seeder.seed_security(s))

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

        # Strategy parameters
        self._target_weight       = 1.0
        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   # buy back when 75% of premium captured
        self._stop_loss_mult      = 3.0    # buy back if call 3x'd in value (cap loss)

        # State
        self._short_call            = None
        self._short_call_entry_price = None
        self._short_call_strike     = None
        self._short_call_expiry     = None
        self._spy_fully_allocated   = False  # avoid redundant set_holdings calls

        # Trend filter: skip selling calls when SPY is in a downtrend
        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)

        self.schedule.on(
            self.date_rules.every(DayOfWeek.MONDAY),
            self.time_rules.after_market_open(self._spy, 5),
            self.sell_covered_call
        )

        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

        # Already have an open short call — don't stack
        if self._short_call is not None and self.portfolio[self._short_call].invested:
            return

        underlying_price = self.securities[self._spy].price
        if underlying_price <= 0:
            return

        # Trend filter: skip selling calls when SPY is BELOW its slow SMA
        # (downtrend — premium doesn't offset drawdown risk)
        if not self._sma_fast.is_ready or not self._sma_slow.is_ready:
            return

        if self._sma_fast.current.value < self._sma_slow.current.value * 0.99:
            self.debug("Skipping covered call: SPY in downtrend.")
            return

        # Buy SPY only if not already allocated (avoid repeated market orders)
        if not self._spy_fully_allocated:
            self.set_holdings(self._spy, self._target_weight)
            self._spy_fully_allocated = True

        # Pull option chain
        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
            exp = contract.expiry.date()
            if exp < min_expiry or exp > max_expiry:
                continue
            # OTM only — at least 1% above spot
            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

        def contract_score(c):
            # Prefer contracts closest to target delta; fall back to moneyness
            if c.greeks is not None and c.greeks.delta is not None:
                return abs(abs(c.greeks.delta) - self._target_delta)
            # Fallback: target ~3% OTM as a proxy for ~0.30 delta
            return abs((c.strike / underlying_price) - 1.03) * 10  # scaled to same range

        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

        self.market_order(option_symbol, -contracts_to_sell)

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

        self.debug(
            f"Sold {contracts_to_sell} covered call(s): "
            f"strike={selected.strike}, expiry={selected.expiry.date()}, mid={mid:.2f}"
        )

    # ------------------------------------------------------------------
    def manage_position(self):
        if self.is_warming_up or 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

        entry = self._short_call_entry_price

        # --- Profit take: option decayed to (1 - profit_take_pct) of entry ---
        if option_price <= entry * (1 - self._profit_take_pct):
            self.debug(f"Profit target hit: option at {option_price:.2f} vs entry {entry:.2f}")
            self.liquidate(self._short_call, tag="Profit target")
            self.reset_short_call_tracking()
            return

        # --- Stop loss: option has 3x'd — cap the pain ---
        if option_price >= entry * self._stop_loss_mult:
            self.debug(f"Stop loss hit: option at {option_price:.2f} vs entry {entry:.2f}")
            self.liquidate(self._short_call, tag="Stop loss")
            self.reset_short_call_tracking()
            return

        # --- Close ITM calls at expiration to avoid assignment ---
        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 call at expiration.")
            self.liquidate(self._short_call, tag="ITM at expiry")
            self.reset_short_call_tracking()
            return

        # --- Roll early: if DTE <= 2 and call is OTM, close and re-sell Monday ---
        if days_to_expiry <= 2 and underlying_price < self._short_call_strike:
            self.debug("Rolling OTM call early (low DTE).")
            self.liquidate(self._short_call, tag="Early roll OTM")
            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