from AlgorithmImports import *
from datetime import timedelta
from typing import Dict, List, Optional

class EarningsStrangleAlgorithm(QCAlgorithm):
    def initialize(self) -> None:
        self.set_start_date(2023, 12, 28)
        self.set_end_date(2025, 1, 1)
        self.set_cash(100000)

        # Strategy parameters
        self.target_dte = 45  # Target days to expiration
        self.min_dte = 30     # Minimum DTE before forced close
        self.take_profit_pct = 0.30   # Take profit at 30%
        self.stop_loss_pct = 1.50     # Stop loss at 150%
        self.position_size = 0.02     # 2% of portfolio per trade
        self.target_delta = 0.25      # Target delta for both call and put
        
        # Tracking variables
        self._symbols: List[Symbol] = []
        self._earnings_dates: Dict[Symbol, datetime] = {}
        self._active_positions: Dict[Symbol, dict] = {}
        self._last_earnings_check: Dict[Symbol, datetime] = {}
        self._option_symbols: Dict[Symbol, Symbol] = {}  # equity_symbol -> option_canonical_symbol
        
        # Add equities and earnings data
        #for ticker in ["NVDA", "META", "MSFT", "AVGO", "GOOGL", "AMZN", "AAPL"]:
        for ticker in ["MSFT"]:
            equity = self.add_equity(ticker, Resolution.MINUTE).symbol
            self._symbols.append(equity)
            
            # Add options and track canonical symbol
            option = self.add_option(ticker)
            option_symbol = option.symbol
            option.set_filter(self.option_filter)
            
            # Map equity to option canonical symbol
            self._option_symbols[equity] = option_symbol
            
        # Add earnings dataset
        self._dataset_symbol = self.add_data(EODHDUpcomingEarnings, "earnings").symbol

    def option_filter(self, universe):
        """Filter options to reasonable strikes and expirations"""
        return universe.strikes(-20, 20).expiration(timedelta(20), timedelta(60))

    def on_data(self, slice: Slice) -> None:
        # Check for earnings announcements
        self.check_earnings_announcements(slice)
        
        # Manage existing positions
        self.manage_existing_positions(slice)
        
        # Check for new strangle opportunities (day after earnings)
        self.check_strangle_opportunities(slice)

    def check_earnings_announcements(self, slice: Slice) -> None:
        """Track upcoming earnings dates"""
        earnings_data = slice.get(EODHDUpcomingEarnings)
        if not earnings_data:
            return
            
        for symbol in self._symbols:
            earnings_for_symbol = earnings_data.get(symbol)
            if earnings_for_symbol and earnings_for_symbol.report_date:
                # Store earnings date
                earnings_date = earnings_for_symbol.report_date
                self._earnings_dates[symbol] = earnings_date
                
                self.debug(f"{symbol} earnings scheduled for {earnings_date} "
                          f"at {earnings_for_symbol.report_time} "
                          f"(Est EPS: {earnings_for_symbol.estimate})")

    def check_strangle_opportunities(self, slice: Slice) -> None:
        """Check for strangle opportunities the day after earnings"""
        current_date = self.time.date()
        
        for symbol in self._symbols:
            # Skip if already have position for this symbol
            if symbol in self._active_positions:
                continue
                
            # Check if yesterday was earnings day
            if symbol in self._earnings_dates:
                earnings_date = self._earnings_dates[symbol].date()
                days_since_earnings = (current_date - earnings_date).days
                
                # Enter strangle the day after earnings (allowing 1-2 days buffer)
                if 1 <= days_since_earnings <= 2:
                    self.enter_short_strangle(symbol, slice)

    def enter_short_strangle(self, symbol: Symbol, slice: Slice) -> None:
        """Enter a short strangle position"""
        try:
            # Get current stock price
            if not slice.bars.contains_key(symbol):
                return
                
            stock_price = slice.bars[symbol].close
            
            # Get canonical option symbol for the equity
            option_symbol = self._option_symbols.get(symbol)
            if not option_symbol:
                self.debug(f"No option symbol was stored for {symbol}")
                return

            chain = slice.option_chains.get(option_symbol)
            if not chain:
                self.debug(f"No option chain available for {option_symbol}")
                return

            # Filter contracts by DTE
            target_expiry = None
            valid_contracts = []
            
            for contract_symbol, contract in chain.contracts.items():
                dte = (contract.expiry.date() - self.time.date()).days
                if abs(dte - self.target_dte) <= 12:  # Within 12 days of target
                    valid_contracts.append(contract)
                    if not target_expiry:
                        target_expiry = contract.expiry

            if not valid_contracts:
                self.debug(f"No suitable contracts found for {symbol} on {self.time.date()}")
                self.log(f"Available contracts for {symbol}:")
                
                contracts_sorted = sorted(chain.contracts.values(), key=lambda c: (c.expiry, c.strike))

                for contract in contracts_sorted:
                    dte = (contract.expiry.date() - self.time.date()).days
                    self.debug(f" - {contract.symbol.Value}: Expiry={contract.expiry.date()}, Strike={contract.strike}, Right={contract.right}, DTE={dte}")

                    return

            # Find 25 delta call and put
            call_contract = None
            put_contract = None
            min_call_delta_diff = float('inf')
            min_put_delta_diff = float('inf')

            for contract in valid_contracts:
                if contract.expiry != target_expiry:
                    continue
                
                # Check if Greeks are available - access delta to trigger calculation
                try:
                    delta = contract.greeks.delta
                    if delta is None:
                        continue
                except:
                    continue
                    
                delta_abs = abs(delta)
                delta_diff = abs(delta_abs - self.target_delta)
                
                if contract.right == OptionRight.CALL and delta_diff < min_call_delta_diff:
                    # For calls, we want positive delta around 0.25
                    if 0.15 <= delta <= 0.35:
                        min_call_delta_diff = delta_diff
                        call_contract = contract
                        
                elif contract.right == OptionRight.PUT and delta_diff < min_put_delta_diff:
                    # For puts, we want negative delta around -0.25
                    if -0.35 <= delta <= -0.15:
                        min_put_delta_diff = delta_diff
                        put_contract = contract

            if not call_contract or not put_contract:
                self.debug(f"Could not find suitable 25-delta call/put pair for {symbol}")
                return

            # Calculate position size
            portfolio_value = self.portfolio.total_portfolio_value
            position_value = portfolio_value * self.position_size
            
            # Get option prices - use ask price for selling
            call_price = call_contract.ask_price if call_contract.ask_price > 0 else call_contract.last_price
            put_price = put_contract.ask_price if put_contract.ask_price > 0 else put_contract.last_price
            
            if call_price <= 0 or put_price <= 0:
                self.debug(f"Invalid option prices for {symbol}: Call={call_price}, Put={put_price}")
                return

            # Calculate contracts to trade (round down for safety)
            premium_per_strangle = (call_price + put_price) * 100  # Per contract
            contracts_to_trade = max(1, int(position_value / premium_per_strangle))

            # Place orders
            call_ticket = self.market_order(call_contract.symbol, -contracts_to_trade)
            put_ticket = self.market_order(put_contract.symbol, -contracts_to_trade)

            # Store position info
            entry_premium = (call_price + put_price) * contracts_to_trade * 100
            
            self._active_positions[symbol] = {
                'call_symbol': call_contract.symbol,
                'put_symbol': put_contract.symbol,
                'call_ticket': call_ticket,
                'put_ticket': put_ticket,
                'contracts': contracts_to_trade,
                'entry_date': self.time.date(),
                'entry_premium': entry_premium,
                'call_entry_price': call_price,
                'put_entry_price': put_price,
                'call_entry_delta': call_contract.greeks.delta,
                'put_entry_delta': put_contract.greeks.delta,
                'stock_price_at_entry': stock_price,
                'expiry': target_expiry,
                'take_profit_target': entry_premium * (1 - self.take_profit_pct),
                'stop_loss_target': entry_premium * (1 + self.stop_loss_pct),
            }

            self.debug(f"ENTERED SHORT STRANGLE for {symbol}:")
            self.debug(f"  Stock Price: ${stock_price:.2f}")
            self.debug(f"  Call Strike: ${call_contract.strike:.2f} @ ${call_price:.2f} (Delta: {call_contract.greeks.delta:.3f})")
            self.debug(f"  Put Strike: ${put_contract.strike:.2f} @ ${put_price:.2f} (Delta: {put_contract.greeks.delta:.3f})")
            self.debug(f"  Contracts: {contracts_to_trade}")
            self.debug(f"  Entry Premium: ${entry_premium:.2f}")
            self.debug(f"  DTE: {(target_expiry.date() - self.time.date()).days}")
            self.debug(f"  Strangle Width: ${abs(call_contract.strike - put_contract.strike):.2f}")

        except Exception as e:
            self.debug(f"Error entering strangle for {symbol}: {str(e)}")

    def manage_existing_positions(self, slice: Slice) -> None:
        """Manage existing strangle positions"""
        positions_to_close = []
        
        for symbol, position in self._active_positions.items():
            try:
                # Check if contracts still exist in portfolio
                call_holding = self.portfolio[position['call_symbol']]
                put_holding = self.portfolio[position['put_symbol']]
                
                if call_holding.quantity == 0 and put_holding.quantity == 0:
                    positions_to_close.append(symbol)
                    continue

                # Get current option prices and Greeks
                call_price, call_delta = self.get_current_option_data(position['call_symbol'], slice)
                put_price, put_delta = self.get_current_option_data(position['put_symbol'], slice)
                
                if call_price is None or put_price is None:
                    continue

                # Calculate current position value
                current_value = (call_price + put_price) * position['contracts'] * 100
                pnl = position['entry_premium'] - current_value
                pnl_pct = pnl / position['entry_premium']

                # Get current stock price
                stock_price = slice.bars[symbol].close if slice.bars.contains_key(symbol) else None
                
                # Calculate days to expiration
                dte = (position['expiry'].date() - self.time.date()).days
                
     
                # Debug output with delta information
                # self.debug(f"[{self.time.date()}] POSITION UPDATE - {symbol}:")
                # self.debug(f"  Stock: ${stock_price:.2f} (Entry: ${position['stock_price_at_entry']:.2f})")
                # self.debug(f"  Call: ${call_price:.2f} (Entry: ${position['call_entry_price']:.2f}) Delta: {call_delta:.3f} (Entry: {position['call_entry_delta']:.3f})")
                # self.debug(f"  Put: ${put_price:.2f} (Entry: ${position['put_entry_price']:.2f}) Delta: {put_delta:.3f} (Entry: {position['put_entry_delta']:.3f})")
                # self.debug(f"  PnL: ${pnl:.2f} ({pnl_pct:.1%})")
                # self.debug(f"  DTE: {dte}")
                # self.debug(f"  Delta Change: Call {call_delta - position['call_entry_delta']:.3f}, Put {put_delta - position['put_entry_delta']:.3f}")

                # Check exit conditions
                should_close = False
                close_reason = ""

                # Take profit
                if pnl >= position['entry_premium'] * self.take_profit_pct:
                    should_close = True
                    close_reason = f"TAKE PROFIT ({pnl_pct:.1%})"
                
                # Stop loss
                elif pnl <= -position['entry_premium'] * self.stop_loss_pct:
                    should_close = True
                    close_reason = f"STOP LOSS ({pnl_pct:.1%})"
                
                # Close before expiration
                elif dte <= self.min_dte:
                    should_close = True
                    close_reason = f"MIN DTE REACHED ({dte} days)"

                if should_close:
                    self.close_strangle_position(symbol, position, close_reason)
                    positions_to_close.append(symbol)

            except Exception as e:
                self.debug(f"Error managing position for {symbol}: {str(e)}")

        # Remove closed positions
        for symbol in positions_to_close:
            if symbol in self._active_positions:
                del self._active_positions[symbol]

    def get_current_option_data(self, option_symbol: Symbol, slice: Slice) -> tuple[Optional[float], Optional[float]]:
        """Get current option price and delta from slice data"""
        try:
            # Try to get from option chains first
            for kvp in slice.option_chains:
                chain = kvp.value
                for contract in chain:
                    if contract.symbol == option_symbol:
                        price = contract.bid_price if contract.bid_price > 0 else contract.last_price
                        delta = contract.greeks.delta if contract.greeks and contract.greeks.delta else None
                        return (price if price > 0 else None, delta)
            
            # Fallback to quote bars (price only)
            if slice.quote_bars.contains_key(option_symbol):
                quote = slice.quote_bars[option_symbol]
                price = (quote.bid.close + quote.ask.close) / 2 if quote.bid.close > 0 and quote.ask.close > 0 else None
                return (price, None)
                
            return (None, None)
        except:
            return (None, None)

    def get_current_option_price(self, option_symbol: Symbol, slice: Slice) -> Optional[float]:
        """Get current option price from slice data (backwards compatibility)"""
        price, _ = self.get_current_option_data(option_symbol, slice)
        return price

    def close_strangle_position(self, symbol: Symbol, position: dict, reason: str) -> None:
        """Close a strangle position"""
        try:
            # Close call position
            call_holding = self.portfolio[position['call_symbol']]
            if call_holding.quantity != 0:
                self.market_order(position['call_symbol'], -call_holding.quantity)

            # Close put position  
            put_holding = self.portfolio[position['put_symbol']]
            if put_holding.quantity != 0:
                self.market_order(position['put_symbol'], -put_holding.quantity)

            self.debug(f"CLOSED STRANGLE for {symbol} - Reason: {reason}")
            
        except Exception as e:
            self.debug(f"Error closing strangle for {symbol}: {str(e)}")

    def on_order_event(self, order_event: OrderEvent) -> None:
        """Handle order events"""
        if order_event.status == OrderStatus.FILLED:
            order = self.transactions.get_order_by_id(order_event.order_id)
            self.debug(f"ORDER FILLED: {order.symbol} {order.quantity} @ ${order_event.fill_price:.2f}")