from AlgorithmImports import *

class ParameterizedOptionsAlgorithm(QCAlgorithm):
    def Initialize(self, 
                   start_date=(2020, 1, 1), 
                   end_date=(2025, 1, 20), 
                   initial_cash=100000, 
                   tickers=["QQQ"], # Tickers to trade during bullish months
                   macd_fast_period=24, 
                   macd_slow_period=120, 
                   macd_signal_period=12,
                   option_expiry_min_days=2, 
                   option_expiry_max_days=7,
                   option_strike_range=10,
                   contracts_to_trade=10,
                   hold_through_exp = True, # During bullish months for tech, whether to hold all the way to expiration or not
                   bullish_months=[6, 7, 11, 12],
                   neutral_months=[2, 8, 9],
                   momentum_months=[1, 3, 4, 5, 10]):
                   
        """
        QQQ Options Trading Algorithm
        
        Parameters:
        - start_date: Backtest start date (year, month, day)
        - end_date: Backtest end date (year, month, day)
        - initial_cash: Starting cash balance
        - tickers: List of tickers to trade
        - macd_fast_period: Fast period for MACD
        - macd_slow_period: Slow period for MACD
        - macd_signal_period: Signal period for MACD
        - option_expiry_min_days: Minimum days to expiry for options
        - option_expiry_max_days: Maximum days to expiry for options
        - option_strike_range: Number of strikes around current price
        - contracts_to_trade: Number of contracts to trade
        - bullish_months: Months with bullish strategy
        - neutral_months: Months with no trading
        - momentum_months: Months with momentum strategy
        """
        # Set the start and end date of the backtest
        self.SetStartDate(start_date[0], start_date[1], start_date[2])
        self.SetEndDate(end_date[0], end_date[1], end_date[2])
        
        # Set the initial cash balance
        self.SetCash(initial_cash)
        
        # Store configuration parameters
        self.tickers = tickers
        self.contracts_to_trade = contracts_to_trade
        self.bullish_months = bullish_months
        self.neutral_months = neutral_months
        self.momentum_months = momentum_months
        self.hold_through_exp = hold_through_exp

        # Option and equity tracking
        self.equities = {}
        self.options = {}

        # Initialize securities
        for ticker in tickers:
            self.equities[ticker] = self.AddEquity(ticker, Resolution.HOUR)
            self.options[ticker] = self.AddOption(ticker, Resolution.Minute)
            
            # Set option universe filter
            self.options[ticker].SetFilter(lambda universe: 
                universe.IncludeWeeklys()
                .Expiration(timedelta(option_expiry_min_days), timedelta(option_expiry_max_days))
                .Strikes(-option_strike_range, option_strike_range)
            )
        
        # Create and initialize MACD with configurable parameters
        self.macd = self.MACD(
            self.equities["QQQ"].Symbol, 
            macd_fast_period, 
            macd_slow_period, 
            macd_signal_period, 
            MovingAverageType.Exponential
        )
        
        # MACD tracking variables
        self.previous_macd = None
        self.previous_signal = None

        # Track the option holding
        self.current_option = None
        self.current_tech_options = {}

    def OnData(self, data):
        for ticker in self.options:
            if not data.OptionChains.ContainsKey(self.options[ticker].Symbol):
                return

        # Skip trading during neutral months
        if self.Time.month in self.neutral_months:
            return
        
        # Momentum months strategy
        if self.Time.month in self.momentum_months:
            self._handle_momentum_strategy(data)

        # Bullish tech months strategy
        if self.Time.month in self.bullish_months:
            #self._handle_bullish_strategy(data)
            self._handle_momentum_strategy(data)

    def _handle_momentum_strategy(self, data):
        # Check if MACD is initialized
        if self.previous_macd is None or self.previous_signal is None:
            self._update_macd_values()
            return

        # MACD crossover logic
        if self._is_macd_bullish_crossover():
            self.Debug(f"MACD crossed above on {self.Time}.")
            self._buy_call_option(data)
        elif self._is_macd_bearish_crossover():
            self.Debug(f"MACD crossed below on {self.Time}.")
            self._sell_call_option()

        # Update MACD values
        self._update_macd_values()

    def _handle_bullish_strategy(self, data):
        # Check if MACD is initialized
        if self.previous_macd is None or self.previous_signal is None:
            self._update_macd_values()
            return

        # MACD crossover logic for bullish months
        if self._is_macd_bullish_crossover():
            self.Debug(f"MACD crossed above on {self.Time}.")
            self._buy_tech_options(data)
        elif self._is_macd_bearish_crossover():
            self.Debug(f"MACD crossed above on {self.Time}.")
            self._sell_tech_options()

        # Update MACD values
        self._update_macd_values()

    def _is_macd_bullish_crossover(self):
        return (self.previous_macd < self.previous_signal and 
                self.macd.Current.Value > self.macd.Signal.Current.Value)

    def _is_macd_bearish_crossover(self):
        return (self.previous_macd > self.previous_signal and 
                self.macd.Current.Value < self.macd.Signal.Current.Value)

    def _update_macd_values(self):
        self.previous_macd = self.macd.Current.Value
        self.previous_signal = self.macd.Signal.Current.Value

    def _buy_call_option(self, data):
        option_chain = data.OptionChains[self.options['QQQ'].Symbol]
        contract = self._select_call_option(option_chain)
        
        if contract:
            
            if self.current_option is None:
                self.Debug(f"Buying {self.contracts_to_trade} Option Contracts: {contract.Symbol}, Strike: {contract.Strike}, Expiry: {contract.Expiry}")
                self.MarketOrder(contract.Symbol, self.contracts_to_trade)
                self.current_option = contract.Symbol

    def _sell_call_option(self):
        if self.current_option is not None:
            contract = self.current_option
            self.Liquidate(contract)
            self.Debug(f"Closed QQQ Option Contract: {contract.Value}")
            self.current_option = None

    def _buy_tech_options(self, data):
        if not self.current_tech_options:
            for ticker in self.options:
                option_chain = data.OptionChains[self.options[ticker].Symbol]
                contract = self._select_call_option(option_chain, ticker)
                contracts = self.contracts_to_trade // len(self.tickers)
                if contract:
                    self.Debug(f"Buying {contracts} Call Option Contracts: {contract.Symbol}, Strike: {contract.Strike}, Expiry: {contract.Expiry}")
                    self.MarketOrder(contract.Symbol, contracts)
                    self.current_tech_options[self.options[ticker].Symbol] = contract.Symbol

    def _sell_tech_options(self):
        if not self.hold_through_exp:
            for ticker, option_symbol in self.current_tech_options.items():
                self.Debug(f"Closing {ticker} Option Contract: {option_symbol}")
            
            self.liquidate()
            self.Debug("Closed All Option Contracts.")
            
            self.current_tech_options = {}  # Reset after liquidation
        else:
            self.Debug("Holding through expiration.")



    def _select_call_option(self, option_chain, ticker='QQQ'):
        # Filter options within expiry range
        one_week_from_now = self.Time.date() + timedelta(days=5)
        one_week_and_two_days_from_now = self.Time.date() + timedelta(days=7)

        contracts = [
            contract for contract in option_chain 
            if one_week_from_now <= contract.Expiry.date() <= one_week_and_two_days_from_now
        ]
        
        # Sort by proximity to current price
        price = self.Securities[ticker].Price
        contracts = sorted(contracts, key=lambda x: abs(x.Strike - price))

        # Return first call option
        call_contracts = [c for c in contracts if c.Right == OptionRight.Call]
        return call_contracts[0] if call_contracts else None