'''
Dynamic Version of the Arbitrage stock model
Includes contract swapping when onhand AR + cost to exit + optimization < market AR
Strategy Design by Thomas Deng
Programmed by Thomas Deng & Eli Webster
'''

from AlgorithmImports import *
from collections import defaultdict
from typing import Dict, Any


class DynamicAR(QCAlgorithm):

    def initialize(self) -> None:
        self.set_start_date(2023, 1, 1)
        self.set_end_date(2026, 6, 1)
        self.set_cash(10_000_000)
        self.settings.seed_initial_prices = True
        # Add a "risk-free" asset.
        self._shv = self.add_equity('SHV')
        # Add the Equities and Options to trade.
        tickers = ["AAPL", "MSFT", "NVDA", "AMZN", "META", "GOOGL", "TSLA"]  # Mag 7 stocks only
        self._equities = []
        for ticker in tickers:
            equity = self.add_equity(ticker)
            option = self.add_option(ticker)
            option.set_filter(lambda u: u.include_weeklys().expiration(7, 36).strikes(-30, 0))
            equity.option_symbol = option.symbol
            self._equities.append(equity)
        # Add a member to track our current trade.
        self._trade = None
        # Define some parameters.
        self._min_ar = self.get_parameter('min_ar', 5)  # Minimum AR threshold to buy contract
        self._swap_optimization = 3 # Minimum AR improvement required to justify a swap
        # Add a Scheduled Event to scan for new trades.
        # Scan every 5 minutes starting from market open, for 30 mins.
        for i in range(self.get_parameter('entry_scans', 7)):
            self.schedule.on(
                self.date_rules.every_day('SPY'),
                self.time_rules.after_market_open('SPY', max(5*i, 1)), 
                self._check_for_entry
            )
        # Add a Scheduled Event to scan for exits.
        # Liquidate the portfolio on the day the contracts expiry.
        self.schedule.on(
            self.date_rules.every_day('SPY'),
            self.time_rules.before_market_close('SPY', 2),
            self._check_expiry
        )

    def _check_for_entry(self) -> None:
        # Find the best available opportunity.
        best_trade = self._find_best_trade()
        # If holding a position, evaluate whether to swap or hold.
        if self._trade:
            # If expired, liquidate.
            if (self.time.date() >= self._trade.expiry or
                # Detect early assignment: stock position gone but expiry not reached
                not self._trade.underlying.invested and self.time.date() < self._trade.expiry):
                self.liquidate()
                self._trade = None
            # If there is a better opportunity and its worth swapping to it, go for it.
            elif best_trade and best_trade["ar"] > self._get_hand_ar() + self._get_exit_cost() + self._swap_optimization:
                self.liquidate()
                self._trade = None
                self._enter_position(best_trade)
        # If there are no opportunities that passed our filters, do nothing.
        elif not best_trade:
            return
        # If the best opportunity doesn't pass our threshold, hold the risk-free asset.
        elif best_trade["ar"] < self._min_ar:
            if not self._shv.invested:
                self.set_holdings(self._shv, 1)
        # Otherwise, sell the risk-free asset and buy the conversion.
        elif self._shv.invested:
            self.liquidate(self._shv)
            self._enter_position(best_trade)

    # If the contracts expire today, liquidate them.
    def _check_expiry(self) -> None:
        if self._trade and self.time.date() == self._trade.expiry:
            self.liquidate()
            self._trade = None

    # Scan all tickers and return the best opportunity for annualized returns.
    def _find_best_trade(self) -> Dict[str, Any]:
        opportunities = []
        # Scan all contract chains in selected tickers, finding best contract
        for underlying in self._equities:
            # Organize the current Option chain into a dictionary, grouping mirror contracts together.
            chain = self.current_slice.option_chains.get(underlying.option_symbol)
            if not chain:
                continue
            pairs = defaultdict(dict)
            for contract in chain:
                key = (contract.strike, contract.expiry.date())
                pairs[key][contract.right] = contract
            # Collect all unique strikes to calculate strike distance later.
            all_strikes = sorted(set(contract.strike for contract in chain))
            # A = Ask Price of Stock (use Price if Ask not available)
            underlying_price = underlying.ask_price or underlying.price
            # Loop through each pair on contracts.
            today = self.time.date()
            for (strike, expiry), legs in pairs.items():
                # Skip contracts expiring today.
                if expiry == today:
                    continue
                # Get the call and put.
                call = legs.get(OptionRight.CALL)
                put = legs.get(OptionRight.PUT)
                if not call or not put:
                    continue
                # S = Strike Price
                # C = Bid Price of Call
                # P = Ask Price of Put
                s = strike
                c = call.bid_price
                p = put.ask_price
                net_premium = c - p
                # Apply filters:
                #  1) Net premium must be positive.
                if (net_premium < 0 or 
                    # 2) P > 0 and C > 0.
                    c <= 0 or p <= 0 or 
                    # Strike must be within 10% of the underlying price.
                    s < 0.90 * underlying_price or s > 1.10 * underlying_price):
                    continue
                # Save the arbitrage opportunity.
                opportunities.append({
                    "underlying": underlying,
                    "call": call,
                    "put": put,
                    "ar": self._annualized_return(underlying, call, put)
                })
        # Select the arbitrage opportunity with the greatest annualized return.
        return max(opportunities, key=lambda t: t['ar']) if opportunities else {}

    def _annualized_return(self, underlying, call, put):
        # C = Bid Price of Call (what we can sell our held call for)
        c = self.securities[call.symbol].bid_price 
        # P = Ask Price of Put (what it costs to close our short put)
        p = self.securities[put.symbol].ask_price
        # A = Ask Price of Stock
        a = underlying.ask_price
        s = call.strike
        t = self._years_to_expiry(call.expiry.date())
        return ((c - p) + (s - a)) * 100 / (a * t)

    def _years_to_expiry(self, expiry):
        days_to_expiry = (expiry - self.time.date()).days
        return max(days_to_expiry / 365.25, 1 / 365)

    # Execute entry into a new position.
    def _enter_position(self, trade: Dict[str, Any]) -> None:
        underlying = trade["underlying"]
        call = trade["call"]
        put = trade["put"]
        # Calculate the cost per combo: 
        #   Buy stock at ask, buy put at ask, and sell call at bid
        combo_price = 100 * (underlying.ask_price + put.ask_price - call.bid_price)
        combos = int(self.portfolio.margin_remaining * .95 / combo_price)
        if combos < 1:
            return
        # Build the trade legs.
        legs = [
            Leg.create(underlying, 100),
            Leg.create(put.symbol, 1),
            Leg.create(call.symbol, -1)
        ]
        # Plat the combo order.
        self.combo_market_order(legs, combos)
        # Keep a reference to the trade.
        self._trade = ComboTrade(underlying, call, put, call.expiry.date())

    # AR of the position currently held (using current market prices)
    def _get_hand_ar(self) -> float:
        return self._annualized_return(self._trade.underlying, self._trade.call, self._trade.put)

    # Annualized cost of exiting the current position (bid-ask spread)
    # Formula: (put_spread + call_spread) * 100 / (T * A_bid)
    def _get_exit_cost(self) -> float:
        # Get the securities in the current trade.
        underlying = self._trade.underlying
        call = self.securities[self._trade.call.symbol]
        put = self.securities[self._trade.put.symbol]
        # Get the entry/exit prices.
        call_ask = call.ask_price
        call_bid = call.bid_price
        put_ask = put.ask_price
        put_bid = put.bid_price
        stock_bid = underlying.bid_price
        # Ensure the calculation will be valid.
        if not all([call_ask, call_bid, put_ask, put_bid, stock_bid]):
            return np.inf
        # Calculate annualized return of exiting the position.
        put_spread = put_ask - put_bid
        call_spread = call_ask - call_bid
        t = self._years_to_expiry(self._trade.expiry)
        return (put_spread + call_spread) * 100 / (t * stock_bid)


class ComboTrade:
    
    def __init__(self, underlying, call, put, expiry):
        self.underlying = underlying
        self.call = call
        self.put = put
        self.expiry = expiry