from AlgorithmImports import *
from datetime import timedelta
from typing import List, Tuple
import datetime as dt

class UpcomingEarningsExampleAlgorithm(QCAlgorithm):
    options_by_symbol: dict = {}        # { underlying: (callSym, putSym) }
    pending_trade: set = set()          # underlyings awaiting liquidity check & entry
    active_positions: set = set()       # underlyings with open positions (protected from removal)
    max_spread_pct: float = 0.10        # 10%
    target_delta: float = 0.20          # 20 delta for strangle legs
    min_market_cap: float = 600e6       # Minimum $600M market cap
    min_dollar_volume: float = 50e6     # Minimum $50M daily dollar volume
    min_stock_price: float = 10.0       # Minimum $10 stock price
    min_credit: float = 1.00            # Minimum $1.00 credit per strangle

    def initialize(self) -> None:
        self.set_start_date(2021, 1, 1)
        self.set_end_date(2025, 10, 31)
        self.set_cash(100000)
        self.set_time_zone("America/New_York")

        self.set_security_initializer(BrokerageModelSecurityInitializer(
            self.brokerage_model, FuncSecuritySeeder(self.get_last_known_prices)))

        self.universe_settings.resolution = Resolution.MINUTE
        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
        self.add_universe(EODHDUpcomingEarnings, self.selection)
        
        # Schedule liquidation at 9:45 AM ET
        self.schedule.on(
            self.date_rules.every_day(),
            self.time_rules.at(9, 45),
            self.liquidate_all_positions
        )

    def selection(self, earnings: List[EODHDUpcomingEarnings]) -> List[Symbol]:
        """
        Select stocks that report:
        - AfterMarket TODAY (we trade today before close, hold through tonight's earnings)
        - BeforeMarket TOMORROW (we trade today before close, hold through tomorrow's pre-market earnings)

        reportDate is a date object.
        reportTime is a string: "BeforeMarket" or "AfterMarket".
        Returns list(Symbol) as expected by the rest of the algorithm.
        """
        today = self.Time.date()
        tomorrow = today + timedelta(days=1)
        symbols = set()

        for e in earnings:
            # Get date and time
            report_date = getattr(e, "ReportDate", None) or getattr(e, "report_date", None)
            report_time = getattr(e, "ReportTime", None) or getattr(e, "report_time", None)

            if not report_date or not report_time:
                continue

            # Normalize date (ensure it's a date object)
            report_date = report_date.date() if hasattr(report_date, "date") else report_date

            # Selection logic - FIXED
            # Case 1: AfterMarket today (earnings tonight after close)
            if report_time == "AfterMarket" and report_date == today:
                symbols.add(e.Symbol)
                # self.Debug(f"[SELECTED] {e.Symbol.Value}: AfterMarket TODAY ({report_date})")
            
            # Case 2: BeforeMarket tomorrow (earnings tomorrow morning before open)
            elif report_time == "BeforeMarket" and report_date == tomorrow:
                symbols.add(e.Symbol)
                # self.Debug(f"[SELECTED] {e.Symbol.Value}: BeforeMarket TOMORROW ({report_date})")

        # Plot and return
        if symbols:
            self.Plot("Universe", "Count", len(symbols))

        return list(symbols)


    def on_securities_changed(self, changes: SecurityChanges) -> None:
        # Handle equity additions: pick 20-delta call/put, subscribe to contracts, but DO NOT trade yet
        for added in [s for s in changes.added_securities if s.type == SecurityType.EQUITY]:
            # Skip if underlying is not tradable (delisted, halted, etc.)
            if not added.is_tradable:
                # self.debug(f"[SKIP] {added.symbol.Value} is not tradable, skipping")
                continue
            
            # Filter by stock price
            if added.price < self.min_stock_price:
                # self.debug(f"[SKIP] {added.symbol.Value} price ${added.price:.2f} < ${self.min_stock_price:.0f}")
                continue
            
            # Filter by market cap
            if hasattr(added.fundamentals, 'market_cap') and added.fundamentals.market_cap:
                if added.fundamentals.market_cap < self.min_market_cap:
                    # self.debug(f"[SKIP] {added.symbol.Value} market cap ${added.fundamentals.market_cap/1e9:.2f}B < ${self.min_market_cap/1e9:.0f}B")
                    continue
            
            # Filter by dollar volume (use 30-day average)
            if hasattr(added, 'volume_model') and added.volume_model:
                dollar_volume = added.volume_model.value * added.price
                if dollar_volume < self.min_dollar_volume:
                    # self.debug(f"[SKIP] {added.symbol.Value} dollar volume ${dollar_volume/1e6:.1f}M < ${self.min_dollar_volume/1e6:.0f}M")
                    continue
            
            call, put = self.select_option_contracts(added.symbol)
            if not call or not put:
                continue

            # Cache selection and subscribe to get live quotes for liquidity check
            # Wrap in try-except to handle any issues with adding options
            try:
                self.options_by_symbol[added.symbol] = (call, put)
                call_sub = self.add_option_contract(call).symbol
                put_sub = self.add_option_contract(put).symbol
                self.pending_trade.add(added.symbol)
            except Exception as e:
                # self.debug(f"[ERROR] Failed to add options for {added.symbol.Value}: {str(e)}")
                # Clean up if partially added
                self.options_by_symbol.pop(added.symbol, None)
                continue

        # Handle equity removals: ONLY remove if NOT in active positions
        for removed in [s for s in changes.removed_securities if s.type == SecurityType.EQUITY]:
            # KEY FIX: Don't remove securities with open positions
            if removed.symbol in self.active_positions:
                # self.debug(f"[UNIVERSE] {removed.symbol.Value} removed from universe but has active positions - keeping until liquidation")
                continue
            
            # Safe to remove - no active positions
            self.liquidate(removed.symbol)
            contracts = self.options_by_symbol.pop(removed.symbol, None)
            self.pending_trade.discard(removed.symbol)
            if contracts:
                for c in contracts:
                    self.remove_option_contract(c)
            # self.debug(f"[UNIVERSE] Removed {removed.symbol.Value}, flattened & unsubscribed.")

    def on_data(self, slice: Slice) -> None:
        # Only enter trades between 3:50 PM and 4:00 PM ET (10 min before close)
        current_time = self.time.time()
        entry_window_start = dt.time(15, 50)  # 3:50 PM ET
        entry_window_end = dt.time(16, 0)     # 4:00 PM ET
        
        # Check if we're in the entry window
        if not (entry_window_start <= current_time <= entry_window_end):
            return
        
        # Try to enter for each pending underlying once quotes are available & liquid
        to_remove = []
        for underlying in list(self.pending_trade):
            call_sym, put_sym = self.options_by_symbol.get(underlying, (None, None))
            if not call_sym or not put_sym:
                to_remove.append(underlying)
                continue

            call_sec = self.securities.get(call_sym, None)
            put_sec  = self.securities.get(put_sym,  None)

            # Need valid quotes
            call_spread = self._spread_pct(call_sec)
            put_spread  = self._spread_pct(put_sec)
            if call_spread is None or put_spread is None:
                # quotes not ready yet
                continue

            # If either leg is illiquid (> max_spread_pct), skip trading this underlying
            if call_spread > self.max_spread_pct or put_spread > self.max_spread_pct:
                # Unsubscribe the contracts to free resources
                for c in (call_sym, put_sym):
                    self.remove_option_contract(c)
                to_remove.append(underlying)
                self.options_by_symbol.pop(underlying, None)
                continue

            # Passed liquidity check — calculate position size as 0.5% of portfolio
            portfolio_value = self.portfolio.total_portfolio_value
            target_position_value = portfolio_value * 0.005  # 0.5% of portfolio
            
            # Estimate the credit received per strangle (mid prices)
            call_mid = (call_sec.bid_price + call_sec.ask_price) / 2.0
            put_mid = (put_sec.bid_price + put_sec.ask_price) / 2.0
            credit_per_contract = call_mid + put_mid  # Credit per share
            
            # Check minimum credit requirement
            if credit_per_contract < self.min_credit:
                # self.debug(f"[SKIP] {underlying.Value} credit ${credit_per_contract:.2f} < ${self.min_credit:.2f} minimum")
                # Unsubscribe the contracts to free resources
                for c in (call_sym, put_sym):
                    self.remove_option_contract(c)
                to_remove.append(underlying)
                self.options_by_symbol.pop(underlying, None)
                continue
            
            credit_per_strangle = credit_per_contract * 100  # *100 for contract multiplier
            
            # Calculate number of strangles to sell (minimum 1)
            if credit_per_strangle > 0:
                quantity = max(1, int(target_position_value / credit_per_strangle))
            else:
                quantity = 1
            
            # Place the short strangle (as combo)
            self.combo_market_order([Leg.create(call_sym, -1), Leg.create(put_sym, -1)], quantity)
            self.debug(f"[ENTRY] Short {quantity}x strangle on {underlying.Value} (call spr={call_spread:.1%}, put spr={put_spread:.1%}, credit=${credit_per_contract:.2f}/share)")
            
            # Mark as active position to protect from universe removal
            self.active_positions.add(underlying)
            to_remove.append(underlying)

        # Clean up pending set for processed symbols
        for u in to_remove:
            self.pending_trade.discard(u)

    def select_option_contracts(self, underlying: Symbol) -> Tuple[Symbol, Symbol]:
        """
        Choose a same-expiry 20-delta strangle (shortest expiration):
        - Put: ~20 delta (OTM put)
        - Call: ~20 delta (OTM call)
        
        Also filters for option liquidity by checking spreads.
        """
        # Get all tradable option contracts for filtering
        option_contract_list = self.option_chain(underlying)
        if not option_contract_list:
            return None, None

        # Nearest expiry at least 1 week out
        min_expiry = self.time + timedelta(days=7)
        valid = [x for x in option_contract_list if x.id.date >= min_expiry]
        if len(valid) < 2:
            return None, None
        
        # Get the shortest expiry
        expiry = min(x.id.date for x in valid)
        same_expiry = [x for x in valid if x.id.date == expiry]

        # Get spot price
        spot = self.securities[underlying].price
        
        # Split by right - for strangle we want OTM options
        puts = [x for x in same_expiry if x.id.option_right == OptionRight.PUT and x.id.strike_price < spot]
        calls = [x for x in same_expiry if x.id.option_right == OptionRight.CALL and x.id.strike_price > spot]

        if not puts or not calls:
            return None, None

        # Find the option contracts closest to target delta
        put_contract = self._find_delta_contract(underlying, puts, -self.target_delta)
        call_contract = self._find_delta_contract(underlying, calls, self.target_delta)

        if not put_contract or not call_contract:
            return None, None

        # Check option liquidity by spread (if securities already exist)
        # This provides early filtering before we subscribe in OnSecuritiesChanged
        if put_contract in self.securities and call_contract in self.securities:
            put_spread = self._spread_pct(self.securities[put_contract])
            call_spread = self._spread_pct(self.securities[call_contract])
            
            if put_spread is not None and call_spread is not None:
                if put_spread > self.max_spread_pct or call_spread > self.max_spread_pct:
                    # self.debug(f"[SKIP] {underlying.Value} options too illiquid at selection (call={call_spread:.1%}, put={put_spread:.1%})")
                    return None, None

        return call_contract, put_contract

    def _find_delta_contract(self, underlying: Symbol, contracts: List[Symbol], target_delta: float) -> Symbol:
        """
        Find the contract closest to the target delta.
        If greeks are not available, estimate using a simple strike distance heuristic.
        """
        if not contracts:
            return None
        
        spot = self.securities[underlying].price
        
        # Try to use actual greeks if available
        best_contract = None
        best_delta_diff = float('inf')
        
        for contract in contracts:
            # Try to get the option contract security
            if contract in self.securities:
                option_sec = self.securities[contract]
                # Check if greeks are available
                if hasattr(option_sec, 'greeks') and option_sec.greeks is not None:
                    delta = option_sec.greeks.delta
                    delta_diff = abs(delta - target_delta)
                    if delta_diff < best_delta_diff:
                        best_delta_diff = delta_diff
                        best_contract = contract
        
        # If we found a contract using greeks, return it
        if best_contract is not None:
            return best_contract
        
        # Fallback: use strike distance as proxy for delta
        target_distance_pct = 0.10  # 10% OTM as approximation for 20 delta
        
        if target_delta > 0:  # Call
            target_strike = spot * (1 + target_distance_pct)
        else:  # Put
            target_strike = spot * (1 - target_distance_pct)
        
        # Find contract with strike closest to target
        best_contract = min(contracts, key=lambda x: abs(x.id.strike_price - target_strike))
        
        return best_contract

    def liquidate_all_positions(self) -> None:
        """
        Liquidate all option positions at 9:45 AM ET
        """
        if self.portfolio.invested:
            self.liquidate()
            # self.debug(f"[EXIT] Liquidated all positions at {self.time}")
            
            # Clear tracking dictionaries and active positions
            for underlying in list(self.options_by_symbol.keys()):
                contracts = self.options_by_symbol.pop(underlying, None)
                if contracts:
                    for c in contracts:
                        if c in self.securities:
                            self.remove_option_contract(c)
            
            self.pending_trade.clear()
            self.active_positions.clear()  # Clear active positions after liquidation

    def _spread_pct(self, sec: Security) -> float:
        """
        Returns bid/ask spread as a fraction of mid.
        None if quotes are not available or invalid.
        """
        if sec is None:
            return None
        bid = float(sec.bid_price) if sec.bid_price is not None else 0.0
        ask = float(sec.ask_price) if sec.ask_price is not None else 0.0
        if bid <= 0 or ask <= 0:
            return None
        mid = 0.5 * (bid + ask)
        if mid <= 0:
            return None
        return (ask - bid) / mid