'''
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 *
import math
from collections import defaultdict

class DynamicAR(QCAlgorithm):

    def Initialize(self):
        # Simulation date, currently 2023 - 2026 End of Q1
        self.set_start_date(2023, 1, 1)
        self.set_end_date(2026, 3, 31)
        self.set_cash(10000000)

        # Disable fees and margin call (Auto Liquidate)
        self.portfolio.margin_call_model = MarginCallModel.NULL
        #self.set_security_initializer(lambda s: s.SetFeeModel(ConstantFeeModel(0)))

        # Money Market Variables
        self.shv = self.add_equity("SHV").Symbol
        self.in_money_market = False

        # Mag 7 stocks only
        self.tickers = ["AAPL", "MSFT", "NVDA", "AMZN", "META", "GOOGL", "TSLA"]
        self.underlyings = {}
        self.option_symbols = {}

        self.current_position = None

        # Minimum AR threshold to buy contract
        self.min_ar = 5

        # Optimization hurdle: minimum AR improvement required to justify a swap
        self.swap_optimization = 3.0

        self.trade_history = []

        # Go through selected tickers and get contracts according to filters
        for ticker in self.tickers:
            equity = self.add_equity(ticker)
            self.underlyings[ticker] = equity.Symbol

            option = self.add_option(ticker)
            option.SetFilter(lambda u: u.include_weeklys().Expiration(7, 36))
            self.option_symbols[ticker] = option.Symbol

        # Morning entry window
        times = [(9, 30), (9, 35), (9, 40), (9, 45), (9, 50), (9, 55), (10, 0)]

        for h, m in times:
            self.schedule.on(
                self.date_rules.every_day(),
                self.time_rules.at(h, m),
                self.CheckForEntry
            )
        
        # Expiry-day liquidation at 3:58pm (2 min before close)
        # Then market orders have time to fill before 4pm
        self.schedule.on(
            self.date_rules.every_day(),
            self.time_rules.before_market_close("AAPL", 2),
            self.CheckExpiry
        )
    

    # Check if contract expiring
    def CheckExpiry(self):
        if self.current_position is None:
            return

        today = self.time.date()
        expiry = self.current_position["expiry"]

        if today == expiry:
            #self.debug(f"Expiry day: liquidating at 3:58pm before close")
            self.liquidate()
            self.current_position = None


    def CheckForEntry(self):
        # Find the best available market opportunity
        best_trade = self.FindBestTrade()

        # If holding a position, evaluate whether to swap or hold
        if self.current_position is not None:
            # Auto-liquidate if expired
            if self.time.date() >= self.current_position["expiry"]:
                self.liquidate()
                #self.debug("Position expired: liquidated all holdings")
                self.current_position = None
                return

            ticker = self.current_position["trade_record"]["ticker"]
            underlying = self.underlyings[ticker]

            # Detect early assignment: stock position gone but expiry not reached
            stock_qty = self.portfolio[underlying].quantity
            if stock_qty == 0 and self.time.date() < self.current_position["expiry"]:
                #self.debug(f"Early assignment detected on {ticker}: liquidating remaining legs")
                self.liquidate()
                self.current_position = None
                return

            # Compute hand AR (AR of the position we currently hold)
            hand_ar = self.GetHandAR()

            # Compute exit cost (bid-ask spread cost of unwinding, annualized)
            exit_cost = self.GetExitCost()

            market_ar = 0.0

            if best_trade:
                market_ar = best_trade["ar"]

            total_current = hand_ar + exit_cost + self.swap_optimization

            # self.debug(
            #     f"Hand AR: {hand_ar:.2f}% | Exit Cost: {exit_cost:.2f}% | "
            #     f"Optimization Hurdle: {self.swap_optimization:.2f}% | "
            #     f"Total Required: {total_current:.2f}% | Market AR: {market_ar:.2f}%"
            # )

            # Check if best contract beats optimized total
            if best_trade is not None and market_ar > total_current:
                improvement = market_ar - total_current
                #self.debug(f"SWAP - Improvement: {improvement:.2f}%")
                self.SwapPosition(best_trade)
            #else:
                #self.debug("HOLD - current position is still optimal")

            return

        # No contracts
        if best_trade is None:
            return

        # Buy money market if cant reach threshold
        if best_trade["ar"] < self.min_ar:
            if not self.in_money_market:
                cash = self.portfolio.margin_remaining
                shv_price = self.securities[self.shv].AskPrice
                if shv_price <= 0:
                    shv_price = self.securities[self.shv].Price
                shares = int(cash / shv_price)
                if shares > 0:
                    self.market_order(self.shv, shares)
                    self.in_money_market = True
                    #self.debug(f"No qualifying AR: parked {shares} shares in SHV @ ${shv_price:.2f}")
            return
        else:
            # Beating threshold, sell money market
            if self.in_money_market:
                self.liquidate(self.shv)
                self.in_money_market = False
                #self.debug("AR threshold met: exited SHV")

        self.EnterPosition(best_trade)


    # Scan all tickers and return the best (call, put, ar) combo dict
    def FindBestTrade(self):
        best_trade = None

        # Scan all contract chains in selected tickers, finding best contract
        for ticker in self.tickers:
            chain = self.current_slice.option_chains.get(self.option_symbols[ticker])
            if not chain:
                continue

            result = self.GetBestContract(chain, ticker)
            if not result:
                continue

            call, put, ar = result

            # Save best contract
            if best_trade is None or ar > best_trade["ar"]:
                best_trade = {
                    "ticker": ticker,
                    "call": call,
                    "put": put,
                    "ar": ar
                }

        return best_trade


    # Execute entry into a new position
    def EnterPosition(self, trade):
        ticker  = trade["ticker"]
        call = trade["call"]
        put = trade["put"]
        underlying = self.underlyings[ticker]

        # A = Ask Price of Stock
        underlying_price = self.securities[underlying].AskPrice
        if underlying_price <= 0:
            underlying_price = self.securities[underlying].Price

        # C = Bid Price of Call
        call_bid = call.BidPrice

        # P = Ask Price of Put
        put_ask = put.AskPrice

        cash = self.portfolio.margin_remaining

        # Cost per combo: Buy stock at Ask + Buy put at Ask - Sell call at Bid
        combo_price = (underlying_price * 100) + (put_ask * 100) - (call_bid * 100)

        combos = int(cash * .95 / combo_price)

        if combos < 1:
            self.debug("Insufficient cash to enter position")
            return

        legs = [
            Leg.create(underlying, 100),
            Leg.create(put.Symbol, 1),
            Leg.create(call.Symbol, -1)
        ]

        self.combo_market_order(legs, combos)

        self.debug(
            f"POST-FILL | Stock: {self.portfolio[underlying].quantity} shares | "
            f"Call qty: {self.portfolio[call.Symbol].quantity} | "
            f"Put qty: {self.portfolio[put.Symbol].quantity}"
        )

        trade_record = {
            "entry_time": self.time,
            "ticker": ticker,
            "call_symbol": call.Symbol,
            "put_symbol": put.Symbol,
            "strike": call.Strike,
            "expiry": call.Expiry.date(),
            "ar": trade["ar"],
            "combos": combos,
            "entry_stock_price": underlying_price,
            "entry_call_bid": call_bid,
            "entry_put_ask": put_ask,
            "combo_cost": combo_price
        }
        self.trade_history.append(trade_record)

        self.current_position = {
            "expiry": call.Expiry.date(),
            "trade_record": trade_record
        }

        self.debug(
            f"Trade: {self.time} | {ticker} | "
            f"Strike: {call.Strike} | Expiry: {call.Expiry.date()} | "
            f"AR: {trade['ar']:.2f}% | Combos: {combos}"
        )
        self.debug(
            f"  -> Call: {call.Symbol} @ ${call_bid:.2f} | "
            f"Put: {put.Symbol} @ ${put_ask:.2f} | "
            f"Stock Ask: ${underlying_price:.2f}"
        )


    # Liquidate current position and enter a new one
    def SwapPosition(self, new_trade):
        old_ticker = self.current_position["trade_record"]["ticker"]
        new_ticker = new_trade["ticker"]
        #self.debug(f"Swapping from {old_ticker} to {new_ticker}")

        self.liquidate()
        self.current_position = None

        self.EnterPosition(new_trade)

        # Mark latest trade history entry as a swap
        if self.trade_history:
            self.trade_history[-1]["action"] = "SWAP"


    # AR of the position currently held (using current market prices)
    # Formula mirrors GetBestContract: ((C - P) + (S - A)) * 100 / (A * T)
    def GetHandAR(self):
        if self.current_position is None:
            return 0.0

        record  = self.current_position["trade_record"]
        ticker = record["ticker"]
        call = record["call_symbol"]
        put = record["put_symbol"]
        underlying = self.underlyings[ticker]

        if (call not in self.securities or
                put not in self.securities or
                underlying not in self.securities):
            return 0.0

        # C = Bid Price of Call (what we can sell our held call for)
        C = self.securities[call].BidPrice
        # P = Ask Price of Put (what it costs to close our short put)
        P = self.securities[put].AskPrice
        # A = Ask Price of Stock
        A = self.securities[underlying].AskPrice
        if A <= 0:
            A = self.securities[underlying].Price

        S = record["strike"]

        if C <= 0 or P <= 0 or A <= 0:
            return 0.0

        expiry = self.current_position["expiry"]
        days_to_expiry  = (expiry - self.time.date()).days
        T = max(days_to_expiry / 365.25, 1 / 365)

        hand_ar = ((C - P) + (S - A)) * 100 / (A * T)
        return hand_ar


    # Annualized cost of exiting the current position (bid-ask spread)
    # Formula: (put_spread + call_spread) * 100 / (T * A_bid)
    def GetExitCost(self):
        if self.current_position is None:
            return 0.0

        record = self.current_position["trade_record"]
        ticker = record["ticker"]
        call_sym = record["call_symbol"]
        put_sym = record["put_symbol"]
        underlying = self.underlyings[ticker]

        if (call_sym not in self.securities or
                put_sym not in self.securities or
                underlying not in self.securities):
            return 0.0

        call = self.securities[call_sym]
        put = self.securities[put_sym]
        stock = self.securities[underlying]

        call_ask = call.AskPrice
        call_bid = call.BidPrice
        put_ask = put.AskPrice
        put_bid = put.BidPrice
        stock_bid = stock.BidPrice

        if call_ask <= 0 or call_bid <= 0 or put_ask <= 0 or put_bid <= 0 or stock_bid <= 0:
            return 0.0

        expiry = self.current_position["expiry"]
        days_to_expiry = (expiry - self.time.date()).days
        T = max(days_to_expiry / 365.25, 1 / 365)

        put_spread  = put_ask - put_bid
        call_spread = call_ask - call_bid

        exit_cost = (put_spread + call_spread) * 100 / (T * stock_bid)
        return exit_cost

    # For a given option chain, return (call, put, ar) with highest AR
    def GetBestContract(self, chain, ticker):

        pairs = defaultdict(dict)
        today = self.time.date()

        underlying = self.underlyings[ticker]

        # A = Ask Price of Stock (use Price if Ask not available)
        underlying_price = self.securities[underlying].AskPrice
        if underlying_price <= 0:
            underlying_price = self.securities[underlying].Price

        for contract in chain:
            key = (contract.Strike, contract.Expiry.date())
            pairs[key][contract.Right] = contract

        # Initalize starting variables
        best = None
        best_ar = -math.inf

        for (strike, expiry), legs in pairs.items():
            # Skip contracts expiring today
            if expiry == today:
                continue

            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.BidPrice
            P = put.AskPrice

            # Net premium must be positive
            net_premium = C - P

            if net_premium < 0:
                continue

            # Exclude P = 0.00 and C = 0.00
            if C <= 0 or P <= 0:
                continue

            # Filter strikes within +-10% of underlying
            if S < 0.90 * underlying_price or S > 1.10 * underlying_price:
                continue

            # T = Time Till Expiry (Years)
            days_to_expiry = (expiry - today).days
            T = max(days_to_expiry / 365.25, 1 / 365)

            # AR = Annualized Return (%)
            A  = underlying_price
            ar = ((C - P) + (S - A)) * 100 / (A * T)

            if ar > best_ar:
                best_ar = ar
                best = (call, put, ar)

        return best


    # Log trades, summary, and add to txt file
    def OnEndOfAlgorithm(self):
        lines = []
        lines.append("=" * 80)
        lines.append("TRADE HISTORY")
        lines.append("=" * 80)

        for i, trade in enumerate(self.trade_history, 1):
            action = trade.get("action", "ENTER")
            lines.append(
                f"Trade {i} [{action}]: {trade['entry_time']} | {trade['ticker']} | "
                f"Strike: {trade['strike']} | Expiry: {trade['expiry']} | "
                f"AR: {trade['ar']:.2f}% | Combos: {trade['combos']}"
            )
            lines.append(
                f"  -> Call: {trade['call_symbol']} @ ${trade['entry_call_bid']:.2f} | "
                f"Put: {trade['put_symbol']} @ ${trade['entry_put_ask']:.2f} | "
                f"Stock Ask: ${trade['entry_stock_price']:.2f}"
            )

        lines.append("=" * 80)
        lines.append(f"Total Trades: {len(self.trade_history)}")

        if self.trade_history:
            avg_ar = sum(t["ar"] for t in self.trade_history) / len(self.trade_history)
            max_ar = max(t["ar"] for t in self.trade_history)
            min_ar = min(t["ar"] for t in self.trade_history)
            lines.append(f"Average AR: {avg_ar:.2f}%")
            lines.append(f"Max AR: {max_ar:.2f}%")
            lines.append(f"Min AR: {min_ar:.2f}%")

        lines.append(f"Final Portfolio Value: ${self.portfolio.total_portfolio_value:,.2f}")
        lines.append(f"Total Return: {((self.portfolio.total_portfolio_value - 10000000) / 10000000 * 100):.2f}%")
        lines.append(f"Total Fees: ${self.portfolio.total_fees:.2f}")
        lines.append("=" * 80)

        content = "\n".join(lines)
        self.object_store.save("trade_history.txt", content)
        self.debug("Trade history saved to object store: trade_history.txt")

        for line in lines:
            self.debug(line)