from AlgorithmImports import *
import numpy as np
from datetime import datetime, timedelta

class DailyStrangleVIXStrategy(QCAlgorithm):
    
    def Initialize(self):
        # Set start and end dates
        self.SetStartDate(2025, 1, 1)
        self.SetEndDate(2025, 6, 1)
        
        # Set cash
        self.SetCash(25000)
        
        # Add equity data for SPY (our underlying)
        self.spy = self.add_equity("SPY", Resolution.Minute)
        self.spy.SetDataNormalizationMode(DataNormalizationMode.Raw)
        
        # Add VIX data
        self.vix = self.add_data(CBOE, "VIX", Resolution.Daily)
        
        # Add option data for SPY
        option = self.add_option("SPY", Resolution.Minute)
        self.spy_option = option.Symbol
        
        # Set option filter
        option.SetFilter(self.OptionFilter)
        
        # Initialize variables
        self.current_strangle = None
        self.entry_time = None
        self.vix_threshold = 20  # VIX threshold for favorable conditions
        self.spy_volatility_window = RollingWindow[float](20)  # 20-day rolling window for SPY volatility
        self.spy_price_window = RollingWindow[float](21)  # 21 days for volatility calculation
        self.vix_ma_window = RollingWindow[float](10)  # 10-day VIX moving average
        
        # Risk management for long strangle
        self.max_position_size = 0.05  # Max 5% of portfolio per trade (buying premium)
        self.profit_target = 1.0  # 100% profit target
        self.stop_loss = -0.5  # 50% stop loss (half premium paid)
        
        # Track daily P&L
        self.daily_pnl = 0
        self.trade_count = 0
        
        # Schedule functions
        self.schedule.on(self.date_rules.every_day("SPY"), 
                        self.time_rules.at(9, 35), 
                        self.CheckAndOpenStrangle)

        self.schedule.on(
            self.date_rules.every_day("SPY"), 
            self.time_rules.before_market_close("SPY", 10), 
            self.CloseStrangle
        )
        
        self.schedule.on(
            self.date_rules.every_day("SPY"), 
            self.time_rules.before_market_close("SPY", 5), 
            self.close_of_day
        )

        # debug and logging
        self.SetWarmUp(25)  # Warm up period for indicators
        
    def OptionFilter(self, universe):
        return universe.IncludeWeeklys().Strikes(-10, 10).Expiration(0, 7)
    
    def OnData(self, data):
        # Update price windows
        if self.spy.Symbol in data and data[self.spy.Symbol] is not None:
            self.spy_price_window.Add(data[self.spy.Symbol].Close)
        
        # Update VIX data
        if self.vix.Symbol in data and data[self.vix.Symbol] is not None:
            vix_value = data[self.vix.Symbol].Value
            self.vix_ma_window.Add(vix_value)
            
        # Calculate SPY realized volatility if we have enough data
        if self.spy_price_window.IsReady:
            returns = []
            for i in range(1, self.spy_price_window.Count):
                if self.spy_price_window[i] != 0:
                    daily_return = np.log(self.spy_price_window[i-1] / self.spy_price_window[i])
                    returns.append(daily_return)
            
            if len(returns) >= 19:  # Need at least 19 returns for 20-day volatility
                realized_vol = np.std(returns) * np.sqrt(252) * 100  # Annualized volatility in %
                self.spy_volatility_window.Add(realized_vol)
    
    def CheckAndOpenStrangle(self):
        """Check conditions and open strangle if favorable"""
        if self.is_warming_up:
            return
            
        # Don't open new position if we already have one
        if self.current_strangle is not None:
            return
            
        # Check if we have the required data
        if not self.vix_ma_window.IsReady or not self.spy_volatility_window.IsReady:
            return
            
        # Get current VIX and its moving average
        current_vix = self.vix_ma_window[0]
        vix_ma = sum(self.vix_ma_window) / self.vix_ma_window.Count
        
        # Get current realized volatility
        current_realized_vol = self.spy_volatility_window[0]
        avg_realized_vol = sum(self.spy_volatility_window) / self.spy_volatility_window.Count
        
        # Define favorable conditions for strangle
        favorable_conditions = self.CheckFavorableConditions(current_vix, vix_ma, 
                                                            current_realized_vol, avg_realized_vol)
        
        self.OpenStrangle()
    
    def CheckFavorableConditions(self, current_vix, vix_ma, current_realized_vol, avg_realized_vol):
        """
        Check if conditions are favorable for opening a strangle
        Returns True if conditions suggest high probability of movement
        """
        conditions_met = 0
        total_conditions = 0
        
        # Condition 1: VIX is elevated (above threshold)
        total_conditions += 1
        if current_vix > self.vix_threshold:
            conditions_met += 1
            self.debug(f"✓ VIX elevated: {current_vix:.2f} > {self.vix_threshold}")
        
        # Condition 2: VIX is above its recent average (trending up)
        total_conditions += 1
        if current_vix > vix_ma * 1.05:  # 5% above average
            conditions_met += 1
            self.debug(f"✓ VIX above MA: {current_vix:.2f} > {vix_ma:.2f}")
        
        # Condition 3: Realized volatility is picking up
        total_conditions += 1
        if current_realized_vol > avg_realized_vol * 1.1:  # 10% above average
            conditions_met += 1
            self.debug(f"✓ Realized vol elevated: {current_realized_vol:.2f} > {avg_realized_vol:.2f}")
        
        # Condition 4: VIX term structure (if available) - simplified check
        total_conditions += 1
        if current_vix > 15:  # Basic threshold indicating some market stress
            conditions_met += 1
            self.debug(f"✓ VIX shows market stress: {current_vix:.2f}")
        
        # Require at least 2 out of 4 conditions
        is_favorable = conditions_met >= 3
        
        self.debug(f"Conditions check: {conditions_met}/{total_conditions} met. Favorable: {is_favorable}")
        self.debug(f"VIX: {current_vix:.2f}, VIX MA: {vix_ma:.2f}, RV: {current_realized_vol:.2f}, RV MA: {avg_realized_vol:.2f}")
        
        return is_favorable
    
    def OpenStrangle(self):
        """Open a strangle position"""
        try:
            # Get option chain
            option_chain = self.CurrentSlice.OptionChains.get(self.spy_option)
            if option_chain is None or len(option_chain) == 0:
                return
            
            # Get current SPY price
            spy_price = self.Securities[self.spy.Symbol].Price
            
            # Find the best strikes for strangle
            target_call_strike = spy_price  
            target_put_strike = spy_price
            
            # Find closest options to target strikes
            calls = [x for x in option_chain if x.Right == OptionRight.Call]
            puts = [x for x in option_chain if x.Right == OptionRight.Put]
            
            if len(calls) == 0 or len(puts) == 0:
                return
                
            # Select the call closest to target strike
            best_call = min(calls, key=lambda x: abs(x.Strike - target_call_strike))
            
            # Select the put closest to target strike
            best_put = min(puts, key=lambda x: abs(x.Strike - target_put_strike))
            
            # Calculate position size based on premium and risk management
            call_premium = best_call.AskPrice
            put_premium = best_put.AskPrice
            total_premium = call_premium + put_premium
            
            if total_premium <= 0:
                return
                
            # Calculate maximum contracts based on premium cost (we're buying)
            max_risk = self.Portfolio.TotalPortfolioValue * self.max_position_size
            max_contracts = int(max_risk / (total_premium * 100))  # 100 shares per contract
            
            if max_contracts <= 0:
                return
                
            # Limit to reasonable size
            contracts = min(max_contracts, 10)
            
            # Execute the trades - BUY the strangle to profit from movement
            call_order = self.market_order(best_call.Symbol, -contracts)   # Buy call
            put_order = self.market_order(best_put.Symbol, -contracts)     # Buy put
            
            # Store strangle information
            self.current_strangle = {
                'call_symbol': best_call.Symbol,
                'put_symbol': best_put.Symbol,
                'call_strike': best_call.Strike,
                'put_strike': best_put.Strike,
                'contracts': contracts,
                'entry_premium': total_premium,
                'entry_time': self.Time,
                'call_order': call_order,
                'put_order': put_order
            }
            
            self.entry_time = self.Time
            self.trade_count += 1
            
            self.debug(f"Opened strangle #{self.trade_count}: Call {best_call.Strike}, Put {best_put.Strike}, " +
                      f"Contracts: {contracts}, Premium: ${total_premium:.2f}")
                      
        except Exception as e:
            self.debug(f"Error opening strangle: {str(e)}")
    
    def CloseStrangle(self):
        """Close the current strangle position"""
        if self.current_strangle is None:
            return
            
        try:
            call_symbol = self.current_strangle['call_symbol']
            put_symbol = self.current_strangle['put_symbol']
            contracts = self.current_strangle['contracts']
            
            # Close positions - Sell back the long strangle
            self.market_order(call_symbol, contracts)  # Sell call
            self.market_order(put_symbol, contracts)   # Sell put
            
            # Calculate P&L
            call_position = self.Portfolio[call_symbol]
            put_position = self.Portfolio[put_symbol]
            
            total_pnl = call_position.UnrealizedProfit + put_position.UnrealizedProfit
            
            self.debug(f"Closed strangle #{self.trade_count}: P&L: ${total_pnl:.2f}")
            
            # Reset strangle
            self.current_strangle = None
            self.entry_time = None
            
        except Exception as e:
            self.debug(f"Error closing strangle: {str(e)}")
    
    def OnOrderEvent(self, orderEvent):
        """Handle order events"""
        if orderEvent.Status == OrderStatus.Filled:
            self.debug(f"Order filled: {orderEvent.Symbol} {orderEvent.FillQuantity} @ ${orderEvent.FillPrice}")
    
    def close_of_day(self):
        """End of day processing"""
        # Force close any open positions
        if self.current_strangle is not None:
            self.CloseStrangle()
        
        # Log daily performance
        self.debug(f"End of day - Portfolio value: ${self.Portfolio.TotalPortfolioValue:.2f}")