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

class OptimizedDeltaHedgingStrategy(QCAlgorithm):
    def Initialize(self):
        # Basic setup
        self.SetStartDate(2021, 12, 1)
        self.SetEndDate(2021, 12, 17)
        self.SetCash(2000000)
        
        # Add securities
        self.qqq = self.AddEquity("QQQ", Resolution.Hour)
        
        # Initialize strategy parameters first
        self.SetupStrategyParameters()
        
        # Add option with appropriate filter
        self.qqq_option = self.AddOption("QQQ", Resolution.Hour)
        self.qqq_option.SetFilter(self.OptionFilterFunc)
        
        # Initialize variables and set up schedules
        self.InitializeTrackingVariables()
        self.CalculateInitialHistoricalVolatility()
        
        # Set up schedules
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(9, 30), self.UpdateHistoricalVolatility)
        self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(16, 0), self.LogDailySummary)
        
        # Strategy rotation if testing multiple strategies
        if len(self.strategy_combinations) > 1:
            self.Debug("Testing multiple strategy combinations")
            self.Schedule.On(self.DateRules.EveryDay(), self.TimeRules.At(9, 31), self.RotateStrategy)
    
    def SetupStrategyParameters(self):
        """Set up strategy testing parameters"""
        # Create strategy testing combinations
        self.strategy_combinations = [
            # Base strategy - ATM calls
            {
                "strike_range": (-10, 10),
                "expiry": datetime(2021, 12, 17),
                "option_type": "calls",
                "strike_selection": "atm",
                "iv_comparison": True
            },
            # ITM calls
            {
                "strike_range": (-10, 10),
                "expiry": datetime(2021, 12, 17),
                "option_type": "calls",
                "strike_selection": "itm",
                "iv_comparison": True  
            },
            # OTM calls
            {
                "strike_range": (-10, 10),
                "expiry": datetime(2021, 12, 17),
                "option_type": "calls",
                "strike_selection": "otm",
                "iv_comparison": True
            },
            # ATM puts
            {
                "strike_range": (-10, 10),
                "expiry": datetime(2021, 12, 17),
                "option_type": "puts",
                "strike_selection": "atm",
                "iv_comparison": True
            },
            # ITM puts
            {
                "strike_range": (-10, 10),
                "expiry": datetime(2021, 12, 17),
                "option_type": "puts",
                "strike_selection": "itm",
                "iv_comparison": True
            },
            # OTM puts
            {
                "strike_range": (-10, 10),
                "expiry": datetime(2021, 12, 17),
                "option_type": "puts",
                "strike_selection": "otm",
                "iv_comparison": True
            },
            # Earlier expiry test - ATM calls
            {
                "strike_range": (-10, 10),
                "expiry": datetime(2021, 12, 10),
                "option_type": "calls",
                "strike_selection": "atm",
                "iv_comparison": True
            }
        ]
        
        # Set initial strategy
        self.current_strategy_index = 0
        self.current_strategy = self.strategy_combinations[0]
        self.strategy_performance = {}
        self.strategy_start_value = 0
    
    def InitializeTrackingVariables(self):
        """Initialize variables for tracking positions and performance"""
        # Option position tracking
        self.option_symbol = None
        self.option_strike = 0
        self.is_position_open = False
        self.position_type = None
        self.option_position_quantity = 0
        
        # Hedging parameters
        self.last_delta_hedge_time = datetime.min
        self.hedge_frequency = timedelta(hours=1)
        self.option_expiry = self.current_strategy["expiry"]
        self.start_trading_time = datetime(2021, 12, 1, 10, 0, 0)
        
        # Risk management parameters
        self.max_capital_usage = 0.05
        self.max_margin_usage = 0.70
        self.margin_buffer = 0.30
        
        # Volatility and PnL tracking
        self.historical_volatility = 0
        self.daily_pnl_start = 0
        self.today_date = None
    
    def OptionFilterFunc(self, universe):
        """Filter options based on strategy parameters"""
        strike_range = self.current_strategy["strike_range"]
        option_type = self.current_strategy["option_type"]
        expiry_range = self.option_expiry - self.Time
        
        # Verify expiry is valid
        if expiry_range.total_seconds() <= 0:
            return universe.Strikes(-5, 5).Expiration(timedelta(0), timedelta(20))
        
        # Apply filter based on option type
        if option_type == "calls":
            return universe.Strikes(strike_range[0], strike_range[1]).Expiration(timedelta(0), expiry_range).CallsOnly()
        elif option_type == "puts":
            return universe.Strikes(strike_range[0], strike_range[1]).Expiration(timedelta(0), expiry_range).PutsOnly()
        else:  # both
            return universe.Strikes(strike_range[0], strike_range[1]).Expiration(timedelta(0), expiry_range)
    
    def CalculateInitialHistoricalVolatility(self):
        """Calculate historical volatility at start"""
        hist_start = self.Time - timedelta(days=40)
        hist_end = self.Time - timedelta(days=1)
        history = self.History(self.qqq.Symbol, hist_start, hist_end, Resolution.Daily)
        
        if not history.empty and len(history) >= 25:
            closes = history['close'].values[-25:]
            log_returns = np.diff(np.log(closes))
            self.historical_volatility = np.std(log_returns) * np.sqrt(252)
    
    def UpdateHistoricalVolatility(self):
        """Daily update of historical volatility"""
        end_date = self.Time.date() - timedelta(days=1)
        start_date = end_date - timedelta(days=40)
        
        history = self.History(self.qqq.Symbol, start_date, end_date, Resolution.Daily)
        
        if len(history) >= 25:
            closes = history['close'].values[-25:]
            log_returns = np.diff(np.log(closes))
            self.historical_volatility = np.std(log_returns) * np.sqrt(252)
            
            self.daily_pnl_start = self.Portfolio.TotalPortfolioValue
            self.today_date = self.Time.date()
    
    def RotateStrategy(self):
        """Switch to the next strategy to test"""
        # Record performance of current strategy
        current_performance = self.Portfolio.TotalPortfolioValue - self.strategy_start_value
        strategy_key = self.FormatStrategyKey(self.current_strategy)
        self.strategy_performance[strategy_key] = current_performance
        
        self.Debug(f"Strategy {strategy_key} performance: ${current_performance:,.2f}")
        
        # Rotate to next strategy
        self.current_strategy_index = (self.current_strategy_index + 1) % len(self.strategy_combinations)
        self.current_strategy = self.strategy_combinations[self.current_strategy_index]
        
        # Reset for new strategy
        self.CloseAllPositions()
        self.InitializeTrackingVariables()
        self.strategy_start_value = self.Portfolio.TotalPortfolioValue
        
        self.Debug(f"Switching to strategy: {self.FormatStrategyKey(self.current_strategy)}")
    
    def FormatStrategyKey(self, strategy):
        """Create readable key for strategy"""
        return (f"{strategy['option_type']}_{strategy['strike_selection']}_"
                f"exp{strategy['expiry'].strftime('%m-%d')}")
    
    def SelectOption(self, chain):
        """Select appropriate option based on strategy"""
        underlying_price = chain.Underlying.Price
        strike_selection = self.current_strategy["strike_selection"]
        option_type = self.current_strategy["option_type"]
        
        # Filter by option type
        if option_type == "calls":
            contracts = [c for c in chain if c.Right == OptionRight.Call]
        elif option_type == "puts":
            contracts = [c for c in chain if c.Right == OptionRight.Put]
        else:
            # Select based on largest IV vs HV spread
            calls = [c for c in chain if c.Right == OptionRight.Call]
            puts = [c for c in chain if c.Right == OptionRight.Put]
            
            if not calls or not puts:
                return None
                
            # Find closest ATM options
            calls.sort(key=lambda x: abs(x.Strike - underlying_price))
            puts.sort(key=lambda x: abs(x.Strike - underlying_price))
            
            call_candidate = calls[0]
            put_candidate = puts[0]
            
            # Select best option type based on IV spread
            if abs(call_candidate.ImpliedVolatility - self.historical_volatility) > abs(put_candidate.ImpliedVolatility - self.historical_volatility):
                contracts = calls
            else:
                contracts = puts
        
        if not contracts:
            return None
            
        # Sort by strike price
        contracts.sort(key=lambda x: x.Strike)
        
        # Select appropriate strike based on strategy
        selected_contract = None
        
        if option_type == "calls" or (option_type == "both" and contracts[0].Right == OptionRight.Call):
            # Call selection logic
            if strike_selection == "atm":
                atm = [c for c in contracts if c.Strike >= underlying_price]
                selected_contract = atm[0] if atm else None
            elif strike_selection == "itm":
                itm = [c for c in contracts if c.Strike < underlying_price]
                selected_contract = itm[-1] if itm else None
            elif strike_selection == "otm":
                otm = [c for c in contracts if c.Strike > underlying_price]
                selected_contract = otm[1] if len(otm) > 1 else otm[0] if otm else None
        else:
            # Put selection logic
            if strike_selection == "atm":
                atm = [c for c in contracts if c.Strike <= underlying_price]
                selected_contract = atm[-1] if atm else None
            elif strike_selection == "itm":
                itm = [c for c in contracts if c.Strike > underlying_price]
                selected_contract = itm[0] if itm else None
            elif strike_selection == "otm":
                otm = [c for c in contracts if c.Strike < underlying_price]
                selected_contract = otm[-2] if len(otm) > 1 else otm[-1] if otm else None
        
        # Default to closest ATM if no match found
        if not selected_contract and contracts:
            selected_contract = min(contracts, key=lambda x: abs(x.Strike - underlying_price))
            
        return selected_contract
    
    def OpenOptionPosition(self, option_contract):
        """Open new option position based on strategy"""
        option_price = (option_contract.BidPrice + option_contract.AskPrice) / 2
        
        # Calculate safe position size
        underlying_price = self.Securities[self.qqq.Symbol].Price
        estimated_margin = (underlying_price * 0.2 + option_price) * 100
        
        available_capital = self.Portfolio.Cash * self.max_capital_usage
        max_by_capital = int(available_capital / option_price / 100) 
        max_by_margin = int(self.Portfolio.Cash * self.max_margin_usage / estimated_margin)
        
        self.option_position_quantity = max(1, min(max_by_capital, max_by_margin) // 2)
        self.option_symbol = option_contract.Symbol
        self.option_strike = option_contract.Strike
        
        # Determine position direction based on IV vs HV
        implied_vol = option_contract.ImpliedVolatility
        
        if implied_vol > self.historical_volatility:
            self.position_type = "short"
            self.Sell(self.option_symbol, self.option_position_quantity)
            self.Debug(f"SHORT {option_contract.Symbol}, Strike: {option_contract.Strike}, IV: {implied_vol:.2f}, HV: {self.historical_volatility:.2f}")
        else:
            self.position_type = "long"
            self.Buy(self.option_symbol, self.option_position_quantity)
            self.Debug(f"LONG {option_contract.Symbol}, Strike: {option_contract.Strike}, IV: {implied_vol:.2f}, HV: {self.historical_volatility:.2f}")
            
        self.is_position_open = True
    
    def PerformDeltaHedge(self, option_contract):
        """Execute delta hedging"""
        try:
            delta = option_contract.Greeks.Delta
            
            # Calculate required hedge quantity
            target_shares = -delta * self.option_position_quantity * 100
            current_shares = self.Portfolio[self.qqq.Symbol].Quantity
            shares_to_trade = int(target_shares - current_shares)
            
            # Ensure we don't exceed margin limits
            available_cash = self.Portfolio.Cash * self.margin_buffer
            safe_quantity = int(min(abs(shares_to_trade), available_cash / self.Securities[self.qqq.Symbol].Price))
            safe_shares = safe_quantity if shares_to_trade > 0 else -safe_quantity
            
            # Execute trade if needed
            if abs(safe_shares) > 0:
                if safe_shares > 0:
                    self.Buy(self.qqq.Symbol, abs(safe_shares))
                else:
                    self.Sell(self.qqq.Symbol, abs(safe_shares))
            
            self.last_delta_hedge_time = self.Time
            
        except Exception as e:
            self.Debug(f"Delta hedging error: {str(e)}")

    def CheckMarginStatus(self):
        """Monitor and manage margin usage"""
        margin_remaining = self.Portfolio.MarginRemaining
        margin_total = self.Portfolio.TotalMarginUsed + margin_remaining

        if margin_remaining < self.margin_buffer * margin_total:
            # Reduce position if margin gets low
            reduction_factor = 3

            if self.position_type == "long":
                qty_to_reduce = self.option_position_quantity // reduction_factor
                if qty_to_reduce > 0:
                    self.Sell(self.option_symbol, qty_to_reduce)
                    self.option_position_quantity -= qty_to_reduce
                    self.Debug(f"Margin protection: Reduced long position by {qty_to_reduce} contracts")
            else:  # short position
                qty_to_reduce = self.option_position_quantity // reduction_factor
                if qty_to_reduce > 0:
                    self.Buy(self.option_symbol, qty_to_reduce)
                    self.option_position_quantity -= qty_to_reduce
                    self.Debug(f"Margin protection: Reduced short position by {qty_to_reduce} contracts")

    def LogDailySummary(self):
        """Log daily performance metrics"""
        if self.today_date is None:
            return

        daily_pnl = self.Portfolio.TotalPortfolioValue - self.daily_pnl_start
        total_pnl = self.Portfolio.TotalProfit

        self.Debug(f"Date: {self.today_date} | Strategy: {self.FormatStrategyKey(self.current_strategy)} | " +
                  f"Daily PnL: ${daily_pnl:,.2f} | Total PnL: ${total_pnl:,.2f}")

        self.daily_pnl_start = self.Portfolio.TotalPortfolioValue

    def CloseAllPositions(self):
        """Close all open positions"""
        if self.is_position_open:
            if self.position_type == "long":
                self.Sell(self.option_symbol, self.option_position_quantity)
            else:
                self.Buy(self.option_symbol, self.option_position_quantity)

            # Close underlying hedge position
            underlying_position = self.Portfolio[self.qqq.Symbol].Quantity
            if underlying_position > 0:
                self.Sell(self.qqq.Symbol, underlying_position)
            elif underlying_position < 0:
                self.Buy(self.qqq.Symbol, abs(underlying_position))

            self.is_position_open = False
            self.Debug(f"Closed all positions for strategy: {self.FormatStrategyKey(self.current_strategy)}")

    def OnData(self, slice):
        """Main algorithm event handler"""
        if self.Time < self.start_trading_time:
            return

        if slice.OptionChains.Count == 0:
            return

        chain = list(slice.OptionChains.Values)[0]
        if not chain:
            return

        # Open new position if none exists
        if not self.is_position_open:
            selected_option = self.SelectOption(chain)
            if selected_option:
                self.OpenOptionPosition(selected_option)
            return

        # Perform delta hedging at scheduled intervals
        if (self.Time - self.last_delta_hedge_time) >= self.hedge_frequency:
            option_contract = next((contract for contract in chain 
                              if contract.Symbol == self.option_symbol), None)

            if option_contract is None:
                return

            self.CheckMarginStatus()
            self.PerformDeltaHedge(option_contract)

    def OnEndOfAlgorithm(self):
        """Final analysis at end of algorithm run"""
        # Record final strategy performance
        current_performance = self.Portfolio.TotalPortfolioValue - self.strategy_start_value
        strategy_key = self.FormatStrategyKey(self.current_strategy)
        self.strategy_performance[strategy_key] = current_performance

        # Find best performing strategy
        best_strategy = max(self.strategy_performance.items(), key=lambda x: x[1]) if self.strategy_performance else (None, 0)

        # Print summary of all strategies
        self.Debug("\n=== STRATEGY PERFORMANCE SUMMARY ===")
        for strat, pnl in sorted(self.strategy_performance.items(), key=lambda x: x[1], reverse=True):
            self.Debug(f"{strat}: ${pnl:,.2f}")

        # Print final results and analysis
        self.Debug("\n=== FINAL RESULTS ===")
        if best_strategy[0]:
            self.Debug(f"Best Strategy: {best_strategy[0]} with PnL: ${best_strategy[1]:,.2f}")

            # Parse which parameters were used in best strategy
            parts = best_strategy[0].split('_')
            option_type = parts[0]
            strike_selection = parts[1]
            expiry = parts[2].replace('exp', '')

            self.Debug(f"\nOptimal Parameters:")
            self.Debug(f"- Option Type: {option_type}")
            self.Debug(f"- Strike Selection: {strike_selection}")
            self.Debug(f"- Expiration: {expiry}")

            # Explain why this strategy performed best
            self.Debug("\nExplanation:")
            if strike_selection == "itm":
                self.Debug("ITM options have higher delta exposure, capturing more directional movement")
            elif strike_selection == "otm":
                self.Debug("OTM options have higher vega exposure, benefiting from volatility changes")
            else:
                self.Debug("ATM options have balanced delta/gamma exposure and good liquidity")

            if option_type == "calls":
                self.Debug("Calls were effective given the underlying price movement direction")
            else:
                self.Debug("Puts were effective due to downward price movement or volatility spikes")

            self.Debug("IV-HV comparison successfully identified options with mispriced volatility")