# region imports
from AlgorithmImports import *
from dataclasses import dataclass
from datetime import date, timedelta
from typing import Optional, List, Dict, Any
# endregion

PHASE_PUT_SELLING = "PUT_SELLING"
PHASE_COVERED_CALL = "COVERED_CALL"

MAX_IV_FOR_ENTRY = 0.45       # skip new put entries when ATM IV is extremely elevated
PROFIT_TAKE_FRACTION = 0.50   # buy back a short option at 50% of the premium received
ROLL_DTE = 21                 # roll OTM shorts at 21 DTE; ITM shorts ride to assignment
PUT_MIN_DTE = 25
PUT_MAX_DTE = 35
CALL_MIN_DTE = 20
CALL_MAX_DTE = 35
TARGET_PUT_DELTA = -0.30
TARGET_CALL_DELTA = 0.30
CASH_USAGE = 0.90             # fraction of cash reserved for cash-secured puts


@dataclass
class Campaign:
    start_date: Optional[date] = None
    premiums: float = 0.0
    share_cost_basis: float = 0.0
    share_count: int = 0


@dataclass
class ShortLeg:
    symbol: Symbol
    entry_premium: float   # per-contract premium at sale (reference price)
    contracts: int
    strike: float
    expiry: datetime


class WheelSpyAlgorithm(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2016, 1, 1)
        self.set_cash(100000)

        self._spy = self.add_equity("SPY", Resolution.DAILY).symbol
        option = self.add_option("SPY", Resolution.DAILY)
        option.set_filter(lambda u: u.strikes(-20, 20).expiration(timedelta(0), timedelta(45)))
        self._option_symbol = option.symbol

        self._phase = PHASE_PUT_SELLING
        self._campaign = Campaign()
        self._history: List[Dict[str, Any]] = []
        self._short_leg: Optional[ShortLeg] = None
        self._total_premiums = 0.0

    # ------------------------------------------------------------------ data

    def on_data(self, data: Slice):
        if self._option_symbol not in data.option_chains:
            return
        chain = data.option_chains[self._option_symbol]

        self._settle_leg()

        if self._short_leg is not None:
            self._manage_leg(chain)
        elif self._phase == PHASE_PUT_SELLING:
            self._try_sell_put(chain)
        elif self._phase == PHASE_COVERED_CALL:
            self._try_sell_covered_call(chain)

        self.plot("Wheel", "Portfolio Value", self.portfolio.total_portfolio_value)
        self.plot("Wheel", "Total Premiums", self._total_premiums)

    def on_order_event(self, order_event: OrderEvent):
        if order_event.status != OrderStatus.FILLED:
            return
        if not order_event.is_assignment:
            return
        leg = self._short_leg
        if leg is None or order_event.symbol != leg.symbol:
            return
        self._process_assignment(leg)

    def on_end_of_algorithm(self):
        self.log(f"Completed campaigns: {len(self._history)}")
        for i, record in enumerate(self._history, start=1):
            self.log(f"Campaign {i}: {record}")
        self.log(f"Total premiums collected (all campaigns): {self._total_premiums:.2f}")

    # -------------------------------------------------------------- position

    def _settle_leg(self):
        """Handle a tracked short leg whose position reached zero (expiry)."""
        leg = self._short_leg
        if leg is None:
            return
        if self.portfolio[leg.symbol].quantity != 0:
            return
        dte = (leg.expiry.date() - self.time.date()).days
        if dte > 0:
            # Flat before expiry (e.g. a rejected closing order): drop tracking.
            self.log(f"Leg {leg.symbol} flat before expiry; clearing tracking")
            self._short_leg = None
            return
        spot = self.securities[self._spy].price
        right = leg.symbol.id.option_right
        itm = spot < leg.strike if right == OptionRight.PUT else spot > leg.strike
        if itm:
            self._process_assignment(leg)
        else:
            self.log(
                f"Short {right} {leg.strike} expired worthless; "
                f"premium kept {leg.entry_premium * 100 * leg.contracts:.2f}"
            )
            self._realize_premium(leg, leg.entry_premium)
            self._short_leg = None

    def _manage_leg(self, chain):
        leg = self._short_leg
        if self.portfolio[leg.symbol].quantity <= 0:
            return
        contract = next((c for c in chain if c.symbol == leg.symbol), None)
        if contract is None:
            return
        dte = (leg.expiry.date() - self.time.date()).days
        if dte <= 0:
            # Expiry day: let assignment / worthless-expiry handling take over.
            return
        price = contract.close
        if price <= 0:
            price = (contract.bid_price + contract.ask_price) / 2
        if price <= 0:
            return
        right = leg.symbol.id.option_right
        if price <= PROFIT_TAKE_FRACTION * leg.entry_premium:
            self._buy_back(leg, price, "profit take at 50% of premium")
            return
        if dte <= ROLL_DTE:
            spot = self.securities[self._spy].price
            otm = spot > leg.strike if right == OptionRight.PUT else spot < leg.strike
            if otm:
                self._buy_back(leg, price, f"roll at {dte} DTE")
                # A new leg is sold on the next daily decision to avoid
                # simultaneous pending orders competing for buying power.

    def _buy_back(self, leg: ShortLeg, price: float, reason: str):
        realized = (leg.entry_premium - price) * 100 * leg.contracts
        self.market_order(leg.symbol, leg.contracts, tag=f"Buy to close: {reason}")
        self._realize_premium(leg, leg.entry_premium - price)
        self.log(f"Bought back {leg.symbol} ({reason}) at {price:.2f}; realized {realized:.2f}")
        self._short_leg = None

    # ------------------------------------------------------------ new entries

    def _try_sell_put(self, chain):
        spot = self.securities[self._spy].price
        if spot <= 0:
            return
        if not self._iv_ok(chain, spot):
            self.log(f"Put entry skipped: ATM IV above {MAX_IV_FOR_ENTRY}")
            return
        candidates = [
            c for c in chain
            if c.symbol.id.option_right == OptionRight.PUT
            and PUT_MIN_DTE <= (c.expiry.date() - self.time.date()).days <= PUT_MAX_DTE
        ]
        if not candidates:
            return
        expiry = min(c.expiry for c in candidates)
        same_expiry = [c for c in candidates if c.expiry == expiry]
        target = self._pick_by_delta(same_expiry, TARGET_PUT_DELTA, spot * 0.95)
        strike = target.symbol.id.strike_price
        contracts = int(self.portfolio.cash * CASH_USAGE / (strike * 100))
        if contracts < 1:
            return
        self.market_order(target.symbol, -contracts, tag="Sell cash-secured put")
        if self._campaign.start_date is None:
            self._campaign.start_date = self.time.date()
        self._short_leg = ShortLeg(target.symbol, target.close, contracts, strike, target.expiry)
        self.log(
            f"Sold {contracts} SPY {strike} put exp {target.expiry.date()} "
            f"at {target.close:.2f} (spot {spot:.2f})"
        )

    def _try_sell_covered_call(self, chain):
        shares = self.portfolio[self._spy].quantity
        contracts = int(shares // 100)
        if contracts < 1:
            return
        spot = self.securities[self._spy].price
        candidates = [
            c for c in chain
            if c.symbol.id.option_right == OptionRight.CALL
            and CALL_MIN_DTE <= (c.expiry.date() - self.time.date()).days <= CALL_MAX_DTE
        ]
        if not candidates:
            return
        expiry = min(c.expiry for c in candidates)
        same_expiry = [c for c in candidates if c.expiry == expiry]
        target = self._pick_by_delta(same_expiry, TARGET_CALL_DELTA, spot * 1.05)
        self.market_order(target.symbol, -contracts, tag="Sell covered call")
        self._short_leg = ShortLeg(
            target.symbol, target.close, contracts, target.symbol.id.strike_price, target.expiry
        )
        self.log(
            f"Sold {contracts} SPY {target.symbol.id.strike_price} call exp {target.expiry.date()} "
            f"at {target.close:.2f} (spot {spot:.2f})"
        )

    # ------------------------------------------------------------- assignment

    def _process_assignment(self, leg: ShortLeg):
        right = leg.symbol.id.option_right
        if right == OptionRight.PUT and self._phase == PHASE_PUT_SELLING:
            self._realize_premium(leg, leg.entry_premium)
            self._campaign.share_cost_basis = leg.strike * 100 * leg.contracts
            self._campaign.share_count = leg.contracts * 100
            self._phase = PHASE_COVERED_CALL
            self.log(
                f"PUT assigned at {leg.strike}: holding {self._campaign.share_count} shares, "
                f"cost basis {self._campaign.share_cost_basis:.2f}"
            )
        elif right == OptionRight.CALL and self._phase == PHASE_COVERED_CALL:
            self._realize_premium(leg, leg.entry_premium)
            self._complete_campaign("called away", leg.strike)
        self._short_leg = None

    def _complete_campaign(self, outcome: str, call_strike: Optional[float] = None):
        share_pnl = 0.0
        if call_strike is not None:
            share_pnl = call_strike * 100 * (self._campaign.share_count // 100) \
                - self._campaign.share_cost_basis
        record = {
            "start": self._campaign.start_date,
            "end": self.time.date(),
            "premiums": round(self._campaign.premiums, 2),
            "share_cost_basis": round(self._campaign.share_cost_basis, 2),
            "share_pnl": round(share_pnl, 2),
            "total_pnl": round(self._campaign.premiums + share_pnl, 2),
            "outcome": outcome,
        }
        self._history.append(record)
        self.log(f"Campaign complete ({outcome}): {record}")
        self._campaign = Campaign()
        self._phase = PHASE_PUT_SELLING

    # ---------------------------------------------------------------- helpers

    def _realize_premium(self, leg: ShortLeg, per_contract_premium: float):
        amount = per_contract_premium * 100 * leg.contracts
        self._campaign.premiums += amount
        self._total_premiums += amount

    def _iv_ok(self, chain, spot: float) -> bool:
        ivs = [
            c.implied_volatility for c in chain
            if abs(c.symbol.id.strike_price - spot) <= 10 and c.implied_volatility > 0
        ]
        if not ivs:
            return True
        return (sum(ivs) / len(ivs)) <= MAX_IV_FOR_ENTRY

    def _pick_by_delta(self, contracts, target_delta: float, fallback_strike: float):
        def delta_of(c):
            greeks = c.greeks
            return greeks.delta if greeks is not None else 0.0

        with_delta = [c for c in contracts if delta_of(c) != 0.0]
        if with_delta:
            return min(with_delta, key=lambda c: abs(delta_of(c) - target_delta))
        return min(contracts, key=lambda c: abs(c.symbol.id.strike_price - fallback_strike))