from AlgorithmImports import *
import math


class FactorModelingWOptions(QCAlgorithm):
    def Initialize(self) -> None:
        # Backtest parameters
        self.SetStartDate(2013, 1, 1)
        self.SetEndDate(2014, 1, 1)
        self.SetCash(100000)

        # Universe & valuation signal (smaller universe reduces fine + options load)
        self.UNIVERSE_SIZE = 300
        self.spy = self.AddEquity("SPY", Resolution.DAILY).Symbol
        # Options require RAW normalization; keep equities ADJUSTED
        self.SetSecurityInitializer(
            lambda sec: sec.SetDataNormalizationMode(
                DataNormalizationMode.RAW if sec.Type == SecurityType.Option else DataNormalizationMode.ADJUSTED
            )
        )

        # Option selection rules
        self.OPT_MAX_SIDE_DEVIATION = 0.08  # 8% max deviation from mid on either side
        # Tighter filter: around ATM strikes and 25-60 DTE approximates second monthly
        self.OPT_EXP_MIN = 25
        self.OPT_EXP_MAX = 60
        self.OPT_STRIKES_RANGE = 3  # +/- strikes around ATM
        # Cap number of underlyings per leg to avoid too many option chains
        self.MAX_UNDERLYINGS_PER_LEG = 25

        # State
        self.option_by_underlying = {}
        self.fine_last_sorted = []
        self.rebalance_flag = False

        # Universe selection: coarse/fine
        self.AddUniverse(self.CoarseSelection, self.FineSelection)

        # Universe settings to reduce data churn
        self.UniverseSettings.Resolution = Resolution.DAILY
        self.UniverseSettings.ExtendedMarketHours = False
        self.UniverseSettings.FillForward = True
        self.UniverseSettings.MinimumTimeInUniverse = 5  # at least a week

        # Weekly on Friday after market open
        self.Schedule.On(self.DateRules.Every(DayOfWeek.Friday), self.TimeRules.AfterMarketOpen(self.spy, 30), self.FlagRebalance)

        # Warm-up
        self.SetWarmUp(5, Resolution.DAILY)

        # Caches
        self._second_monthly_cache = {}
        self._second_monthly_cache_key = None

    def CoarseSelection(self, coarse):
        filtered = [c for c in coarse if c.HasFundamentalData and c.Price is not None and c.Price > 5]
        filtered.sort(key=lambda c: c.DollarVolume, reverse=True)
        return [c.Symbol for c in filtered[: self.UNIVERSE_SIZE]]

    def FineSelection(self, fine):
        rows = []
        for f in fine:
            try:
                # Exclude Financials sector
                is_financial = False
                ac = getattr(f, 'AssetClassification', None)
                sector_code = getattr(ac, 'MorningstarSectorCode', None) if ac is not None else None
                if sector_code is not None:
                    try:
                        # Prefer enum comparison when available
                        from QuantConnect.Data.Fundamental import MorningstarSectorCode  # type: ignore
                        if sector_code == MorningstarSectorCode.FinancialServices:
                            is_financial = True
                    except Exception:
                        # Fallback: common code for Financial Services is 103
                        try:
                            if int(sector_code) == 103:
                                is_financial = True
                        except Exception:
                            pass
                if is_financial:
                    continue

                # Use EV/EBIT as valuation metric (lower is cheaper)
                vr = getattr(f, 'ValuationRatios', None)
                ev_ebit = None
                if vr is not None:
                    ev_ebit = getattr(vr, 'EVToEBIT', None)
                    if ev_ebit is None:
                        ev_ebit = getattr(vr, 'EVToEbit', None)
                if ev_ebit is None or not math.isfinite(ev_ebit) or ev_ebit <= 0 or ev_ebit > 100:
                    continue
                rows.append((f.Symbol, ev_ebit))
            except Exception:
                continue
        if not rows:
            return []
        # Sort ascending: lowest EV/EBIT (cheapest) first
        rows.sort(key=lambda t: t[1])
        self.fine_last_sorted = [sym for sym, _ in rows]
        # Keep a manageable tradable set: top and bottom quartiles union
        n = len(rows)
        q = max(1, n // 4)
        top_syms = [sym for sym, _ in rows[:q]]      # cheapest quartile
        bot_syms = [sym for sym, _ in rows[-q:]]     # most expensive quartile
        return list(dict.fromkeys(top_syms + bot_syms))

    def FlagRebalance(self):
        self.rebalance_flag = True

    def OnData(self, data: Slice) -> None:
        if self.IsWarmingUp or not self.rebalance_flag:
            return
        self.rebalance_flag = False

        ranks = list(getattr(self, 'fine_last_sorted', []))
        if len(ranks) < 8:
            return

        # Quartiles
        n = len(ranks)
        q = n // 4
        top_quartile = ranks[:q]      # cheap by valuation (high earnings yield)
        bottom_quartile = ranks[-q:]  # expensive by valuation
        # Cap number of underlyings traded per leg to keep it fast
        top_quartile = top_quartile[:self.MAX_UNDERLYINGS_PER_LEG]
        bottom_quartile = bottom_quartile[:self.MAX_UNDERLYINGS_PER_LEG]

        # Build desired straddles
        desired = set()
        pv = self.Portfolio.TotalPortfolioValue
        # Equal budget split: 50% for short straddles, 50% for long straddles
        budget_short = pv * 0.5
        budget_long = pv * 0.5
        per_short = budget_short / max(1, len(top_quartile))
        per_long = budget_long / max(1, len(bottom_quartile))

        # Ensure subscriptions for all candidates
        for u in set(top_quartile + bottom_quartile):
            self._EnsureOptionSubscription(u)

        # Helper to select ATM straddle on second monthly expiry
        def select_straddle(chain, underlying):
            if chain is None:
                return None
            # Collect unique expiries
            expiries = sorted({c.Expiry.date() for c in chain})
            # Filter to standard monthly expiries (3rd Friday)
            def is_third_friday(d):
                return d.weekday() == 4 and 15 <= d.day <= 21
            # Cache key per-week to avoid recomputing
            week_key = (self.Time.year, self.Time.isocalendar()[1])
            if self._second_monthly_cache_key != week_key:
                self._second_monthly_cache = {}
                self._second_monthly_cache_key = week_key
            target_expiry = self._second_monthly_cache.get(underlying)
            if target_expiry is None:
                monthly = [d for d in expiries if is_third_friday(d)]
                if len(monthly) >= 2:
                    target_expiry = monthly[1]
                elif len(expiries) >= 2:
                    target_expiry = expiries[1]
                else:
                    return None
                self._second_monthly_cache[underlying] = target_expiry
            # Spot and ATM strike
            spot = self.Securities[underlying].Price
            # Find strikes on target expiry
            contracts_exp = [c for c in chain if c.Expiry.date() == target_expiry]
            if not contracts_exp:
                return None
            # Prefer contracts with valid quotes to reduce filtering later
            contracts_exp = [c for c in contracts_exp if (c.BidPrice is not None and c.AskPrice is not None)]
            strikes = sorted({c.Strike for c in contracts_exp})
            if not strikes:
                return None
            atm_strike = min(strikes, key=lambda k: abs(k - spot))
            call = next((c for c in contracts_exp if c.Right == OptionRight.CALL and c.Strike == atm_strike), None)
            put  = next((c for c in contracts_exp if c.Right == OptionRight.PUT  and c.Strike == atm_strike), None)
            if call is None or put is None:
                return None
            # Bid/ask filter
            okc, midc = self._BidAskOk(call.BidPrice, call.AskPrice)
            okp, midp = self._BidAskOk(put.BidPrice, put.AskPrice)
            if not (okc and okp):
                return None
            return (call.Symbol, midc, put.Symbol, midp)

        # Short straddles on top quartile
        for u in top_quartile:
            opt_sym = self.option_by_underlying.get(u)
            if opt_sym is None:
                continue
            chain = data.OptionChains.get(opt_sym)
            sel = select_straddle(chain, u)
            if sel is None:
                continue
            call_sym, call_mid, put_sym, put_mid = sel
            prem = (call_mid + put_mid) * 100
            if prem <= 0:
                continue
            qty = int(max(0, per_short // prem))
            if qty <= 0:
                continue
            desired.update([call_sym, put_sym])
            # Place/adjust orders (sell straddle)
            self._EnsurePosition(call_sym, -qty, tag="Short straddle call")
            self._EnsurePosition(put_sym, -qty, tag="Short straddle put")

        # Long straddles on bottom quartile
        for u in bottom_quartile:
            opt_sym = self.option_by_underlying.get(u)
            if opt_sym is None:
                continue
            chain = data.OptionChains.get(opt_sym)
            sel = select_straddle(chain, u)
            if sel is None:
                continue
            call_sym, call_mid, put_sym, put_mid = sel
            prem = (call_mid + put_mid) * 100
            if prem <= 0:
                continue
            qty = int(max(0, per_long // prem))
            if qty <= 0:
                continue
            desired.update([call_sym, put_sym])
            # Place/adjust orders (buy straddle)
            self._EnsurePosition(call_sym, qty, tag="Long straddle call")
            self._EnsurePosition(put_sym, qty, tag="Long straddle put")

        # Clean up undesired options
        for kv in self.Portfolio:
            sym = kv.Key
            if sym.SecurityType == SecurityType.Option and self.Portfolio[sym].Invested and sym not in desired:
                self.Liquidate(sym, tag="Weekly roll")

    # Utilities
    def _EnsureOptionSubscription(self, underlying: Symbol):
        if underlying in self.option_by_underlying:
            return self.option_by_underlying[underlying]
        opt = self.AddOption(underlying, Resolution.DAILY)
        # Keep chains small: near-ATM (+/- OPT_STRIKES_RANGE) and target 2nd-month window (EXP_MIN..EXP_MAX)
        opt.SetFilter(lambda u: u.Strikes(-self.OPT_STRIKES_RANGE, self.OPT_STRIKES_RANGE)
                                 .Expiration(self.OPT_EXP_MIN, self.OPT_EXP_MAX))
        self.option_by_underlying[underlying] = opt.Symbol
        return opt.Symbol

    def _BidAskOk(self, bid: float, ask: float):
        if bid is None or ask is None or bid <= 0 or ask <= 0:
            return False, 0.0
        mid = 0.5 * (bid + ask)
        if mid <= 0:
            return False, 0.0
        dev = max(mid - bid, ask - mid) / mid
        return dev <= self.OPT_MAX_SIDE_DEVIATION, mid

    def _EnsurePosition(self, symbol: Symbol, target_qty: int, tag: str):
        held = self.Portfolio[symbol].Quantity if symbol in self.Portfolio else 0
        delta = target_qty - held
        if delta != 0:
            self.MarketOrder(symbol, delta, tag=tag)