from AlgorithmImports import *
from datetime import timedelta


class ChainedUniverseAlgorithm(QCAlgorithm):
    def initialize(self):
        # Setup
        self.set_start_date(2021, 2, 1)
        self.set_end_date(2021, 4, 1)
        self.set_cash(100_000)
        self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        self.set_benchmark("SPY")

        # Parameters
        self._min_price = 10

        # Universe
        self.universe_settings.asynchronous = True
        self.universe_settings.data_normalization_mode = DataNormalizationMode.RAW
        self.universe_settings.resolution = Resolution.DAILY
        # self.universe_settings.schedule.on = self.date_rules.every(DayOfWeek.Tuesday)
        self.set_security_initializer(CustomSecurityInitializer(self))
        universe = self.add_universe(self._fundamental_function)
        self.add_universe_options(universe, self._option_filter_function)
        
        # Initialize variables
        self._option_symbols = set()
        self._open_straddles = {}
        self.day = 0

        # Warm up so Greeks data are populated
        self.set_warm_up(5, Resolution.DAILY)

    def _fundamental_function(self, fundamental: List[Fundamental]) -> List[Symbol]:
        if self.time.weekday() != 1 and self.time.weekday() != 2:
            return Universe.UNCHANGED

        filtered = (f for f in fundamental if f.price > self._min_price)
        sorted_by_dollar_volume = sorted(filtered, key=lambda f: f.dollar_volume, reverse=True)
        return [f.symbol for f in sorted_by_dollar_volume[:30]]

    def _option_filter_function(self, option_filter_universe: OptionFilterUniverse) -> OptionFilterUniverse:
        if self.time.weekday() != 1 and self.time.weekday() != 2:
            return
        # return option_filter_universe.straddle(60)
        return option_filter_universe.expiration(timedelta(20), timedelta(200)).strikes(-50, +50).delta(-0.65, 0.65).implied_volatility(0.3, 2)

    def on_data(self, data: Slice) -> None:
        if self.is_warming_up or self.day == self.time.day:
            return

        # Check if it is Tuesday
        if self.time.weekday() != 1 :
            return

        # Get chain and return if it is None
        chains = data.option_chains
        if not chains:
            return

        # Gather slopes for each symbol
        symbol_slopes = []  # list to store tuples (symbol, slope, call_contract, put_contract)
        
        # Trade
        iv_30 = None
        iv_180 = None
        for symbol, chain in chains.items():
            underlying_price = chain.underlying.price

            # Separate calls/puts for convenience
            calls = [x for x in chain if x.right == OptionRight.CALL and x.greeks is not None and x.greeks.delta > 0.35]
            puts  = [x for x in chain if x.right == OptionRight.PUT  and x.greeks is not None and x.greeks.delta < -0.35]
            
            if not calls or not puts:
                continue

            # Pick the call option closest to 30 days
            sorted_calls_30 = sorted(calls, key=lambda c: abs((c.expiry - self.time).days - 30))
            if not sorted_calls_30:
                continue
            expiry30_call = sorted_calls_30[-1].expiry
            calls_30day   = [c for c in calls if c.expiry == expiry30_call]
            if not calls_30day:
                continue
            contract_call = sorted(calls_30day, key=lambda c: abs(c.strike - underlying_price))[0]

            # Pick the put option closest to 30 days
            sorted_puts_30 = sorted(puts, key=lambda p: abs((p.expiry - self.time).days - 30))
            if not sorted_puts_30:
                continue
            expiry30_put = sorted_puts_30[-1].expiry
            puts_30day   = [p for p in puts if p.expiry == expiry30_put]
            if not puts_30day:
                continue
            contract_put = sorted(puts_30day, key=lambda p: abs(p.strike - underlying_price))[0]

            # For the 180-day calls
            sorted_calls_180 = sorted(calls, key=lambda c: abs((c.expiry - self.time).days - 180))
            if not sorted_calls_180:
                continue
            expiry180_call = sorted_calls_180[-1].expiry
            calls_180day   = [c for c in calls if c.expiry == expiry180_call]
            if not calls_180day:
                continue
            contract_call_180 = sorted(calls_180day, key=lambda c: abs(c.strike - underlying_price))[0]

            # For the 180-day puts
            sorted_puts_180 = sorted(puts, key=lambda p: abs((p.expiry - self.time).days - 180))
            if not sorted_puts_180:
                continue
            expiry180_put = sorted_puts_180[-1].expiry
            puts_180day   = [p for p in puts if p.expiry == expiry180_put]
            if not puts_180day:
                continue
            contract_put_180 = sorted(puts_180day, key=lambda p: abs(p.strike - underlying_price))[0]
            
            # Ensure we have IV data
            if (contract_call.implied_volatility is None or 
                contract_put.implied_volatility is None or
                contract_call_180.implied_volatility is None or
                contract_put_180.implied_volatility is None):
                continue
            
            # Calculate short-term IV (30 day) and long-term IV (180 day)
            iv_short = (contract_call.implied_volatility + contract_put.implied_volatility) / 2
            iv_long  = (contract_call_180.implied_volatility + contract_put_180.implied_volatility) / 2

            # Edge case: avoid division by zero if iv_long is 0 or None
            if iv_long <= 0:
                continue


        ############# HERE I REMOVED SOME CODE #############
            # Buy straddle
            self.market_order(contract_call.symbol, 1)
            self.market_order(contract_put.symbol, 1)

        # Update day so we don't trade multiple times
        self.day = self.time.day


class CustomSecurityInitializer(BrokerageModelSecurityInitializer):
    def __init__(self, algorithm: QCAlgorithm) -> None:
        super().__init__(algorithm.brokerage_model, FuncSecuritySeeder(algorithm.get_last_known_prices))
        self.algorithm = algorithm

    def initialize(self, security: Security) -> None:
        # First, call the superclass definition
        # This method sets the reality models of each security using the default reality models of the brokerage model
        super().initialize(security)

        # Overwrite the price model        
        if security.type == SecurityType.OPTION: # Option type
            security.price_model = OptionPriceModels.crank_nicolson_fd()