# =============================================================================
# "Balanced-v3" — 2-sleeve deployable book, ONE daily account, fund-of-funds
# netting.  Sleeves + book weights (sum=1.0, gross<=100%, never borrows):
#   531 momentum-breadth 0.50 (285-authoritative momentum + breadth + BIL sweep)
#   A4 reversal          0.50 (1-day cross-sectional low-reversal, mcap-weighted)
#   net_w[sym] = 0.50*531 + 0.50*a4   (gross<=1)
#
# Derived from "Balanced-v2" (A4 / gen263 / 531 / kinfo) by (1) KEEPING the 531
# and A4 sleeves plus the ENTIRE netting + T+1 MarketOnOpenOrder execution
# harness VERBATIM, (2) REMOVING the gen263 momentum sleeve and the kinfo regime
# sleeve entirely — their subscriptions, schedules, decision methods, target
# dicts, and their lines in the net-weight formula — and (3) RE-WEIGHTING to
# net = 0.50*531 + 0.50*A4.
#
# EXECUTION CONTRACT (all sleeves): a sleeve NEVER calls Liquidate/SetHoldings/
# market_order. Each only WRITES its in-bucket target dict (sum<=1.0) + flags
# self._dirty. Risk-off / empty branches set that sleeve's dict = {} (531 sweeps
# its bucket to BIL, its own yield hedge) — never a book-wide liquidate. The
# shared _execute_net harness nets the two dicts, submits sells before buys, and
# fills everything via MarketOnOpenOrder at the T+1 open (decision on T close, no
# look-ahead). The harness scales each dict by its book weight.
#
# 531 uses its OWN breadth latch (allow_531/was_risk_off_531/max_stress_531) and
# its OWN band-ceiling history (band_hist_531). It shares the read-only
# per-symbol EMA/ADX/stretch/close-window indicators (band_len=189, hist_len=126,
# adx_period=14) that OnData maintains — no new hot-path work.
#
# Dates: get_parameter -> default 2010-01-01 .. 2026-06-26.
# =============================================================================
from AlgorithmImports import *
from collections import defaultdict
import numpy as np
import pandas as pd


# ---- shared universe (VERBATIM); serves 531 momentum-breadth + A4 ----
class SectorTopUniverse(FundamentalUniverseSelectionModel):
    def __init__(self, algo, blacklist=None):
        self.algo = algo
        self.blacklist = set(blacklist or [])
        super().__init__(self._select)

    def _select(self, fundamentals):
        buckets = defaultdict(list)
        for f in fundamentals:
            if not f.has_fundamental_data:
                continue
            if f.symbol.Value in self.blacklist:
                continue
            if f.company_reference.primary_exchange_id not in ("NYS", "NAS", "ASE"):
                continue
            if f.price is None or f.price <= 5:
                continue
            if f.market_cap is None or f.market_cap < 5_000_000_000:
                continue
            sector = f.asset_classification.morningstar_sector_code
            if sector is None:
                continue
            buckets[sector].append(f)
        symbols = []
        for _, stocks in buckets.items():
            stocks.sort(key=lambda x: x.market_cap, reverse=True)
            symbols.extend(s.symbol for s in stocks[:100])
        return symbols


class BalancedV3(QCAlgorithm):

    # ---- book weights (sum=1) ----
    W_531 = 0.50          # 531 momentum-breadth sleeve bucket weight
    W_A4 = 0.50           # A4 reversal sleeve bucket weight

    DEBUG_FILLS = True

    WARMUP_BARS = 260     # primes shared EMA(189)/ADX(14)/close-window(189) for 531

    # ---- A4 params (VERBATIM) ----
    A4_TOP_N = 100
    A4_TOP_K = 10
    A4_MAX_W = 0.30
    A4_WIN = 9

    # ---- shared per-symbol indicator params (maintained in OnData; used by 531) ----
    band_len = 189
    hist_len = 126
    adx_limit = 35
    adx_period = 14

    # ---- 531 momentum-breadth params (from S531_285_authoritative.py) ----
    lookbacks531 = [21, 63, 126, 189, 252]     # equal-weight avg (1/3/6/9/12-mo)
    stock_count_531 = 10
    max_weight531 = 0.20

    # =====================================================================
    def Initialize(self):
        sy = int(self.get_parameter("start_year", "2010"))
        sm = int(self.get_parameter("start_month", "1"))
        sd = int(self.get_parameter("start_day", "1"))
        ey = int(self.get_parameter("end_year", "2026"))
        em = int(self.get_parameter("end_month", "6"))
        ed = int(self.get_parameter("end_day", "26"))
        self.SetStartDate(sy, sm, sd)
        self.SetEndDate(ey, em, ed)
        self.SetCash(100_000)
        self.SetBrokerageModel(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.MARGIN)
        self.Settings.MinimumOrderMarginPortfolioPercentage = 0.0

        self._n_fills = 0; self._max_moo = 0.0; self._n_moo = 0
        self._max_liq = 0.0; self._n_liq = 0; self._n_inv = 0; self._max_gross = 0.0

        # 531's yield hedge (fixed ETF, outside the fundamental universe)
        self.bil = self.add_equity("BIL", Resolution.DAILY).symbol
        self.fixed_syms = {self.bil}       # fixed ETFs -> NOT in momentum/A4 universe

        # ---- momentum/A4 shared universe ----
        self.UniverseSettings.Resolution = Resolution.Daily
        self.UniverseSettings.DataNormalizationMode = DataNormalizationMode.TOTAL_RETURN
        self.UniverseSettings.Leverage = 4
        self.SetUniverseSelection(SectorTopUniverse(self, blacklist={"GME", "AMC"}))

        # shared per-symbol indicator state (read-only for both sleeves)
        self.symbols = set()
        self.ma = {}; self.adx = {}; self.close_win = {}
        self.stretch_ema = {}; self.stretch_max = {}
        self.current_band_idx = {}
        self.BOTTOM_LEVELS = {0, 1, 2, 3, 4}

        # ---- sleeve target dicts + dirty flag ----
        self.tgt_a4 = {}; self.tgt_531 = {}
        self._dirty = False

        # ---- 531 sleeve's OWN breadth-latch + band-ceiling state ----
        self.allow_531 = True
        self.was_risk_off_531 = False
        self.max_stress_531 = 0.0
        self.band_hist_531 = {}

        self.SetBenchmark("SPY")
        self.SetWarmUp(self.WARMUP_BARS, Resolution.Daily)

        # A4 + 531 monthly decision at month end (write target dicts; T-close -> T+1 open)
        self.Schedule.On(self.DateRules.MonthEnd("SPY"),
                         self.TimeRules.BeforeMarketClose("SPY", 5), self.RebalanceA4)
        self.Schedule.On(self.DateRules.MonthEnd("SPY"),
                         self.TimeRules.BeforeMarketClose("SPY", 5), self.Rebalance531)

    # =====================================================================
    def OnSecuritiesChanged(self, changes):
        for sec in changes.AddedSecurities:
            s = sec.Symbol
            if s in self.fixed_syms:
                continue
            sec.SetFeeModel(InteractiveBrokersFeeModel())
            sec.SetSlippageModel(ConstantSlippageModel(0.001))
            self.symbols.add(s)
            self.stretch_max[s] = 0.0
            self.ma[s] = self.EMA(s, self.band_len, Resolution.Daily)
            self.adx[s] = self.ADX(s, self.adx_period, Resolution.Daily)
            self.stretch_ema[s] = self.EMA(s, self.band_len, Resolution.Daily)
            self.close_win[s] = RollingWindow[float](self.band_len)
            self.band_hist_531[s] = RollingWindow[int](self.hist_len)   # 531's own ceiling history
        for sec in changes.RemovedSecurities:
            s = sec.Symbol
            if s in self.fixed_syms:
                continue
            self.symbols.discard(s)
            self.ma.pop(s, None); self.adx.pop(s, None); self.stretch_ema.pop(s, None)
            self.close_win.pop(s, None)
            self.band_hist_531.pop(s, None)
            self.current_band_idx.pop(s, None); self.stretch_max.pop(s, None)

    def _band_index(self, price, bands):
        for i in range(len(bands) - 1):
            if bands[i] <= price < bands[i + 1]:
                return i
        return len(bands) - 2

    # =====================================================================
    def OnData(self, data):
        for s in list(self.symbols):
            if not data.ContainsKey(s):
                continue
            bar = data[s]
            if bar is None:
                continue
            close = bar.Close
            self.close_win[s].Add(close)
            if not self.close_win[s].IsReady or not self.ma[s].IsReady:
                continue
            dev = np.std(list(self.close_win[s]))
            if dev <= 0:
                continue
            mid = self.ma[s].Current.Value
            stretch = abs(close - mid) / dev
            self.stretch_ema[s].Update(self.Time, stretch)
            if stretch > self.stretch_max[s]:
                self.stretch_max[s] = stretch
            bands = [mid - dev * 1.618, mid - dev * 1.382, mid - dev, mid - dev * 0.809,
                     mid - dev * 0.5, mid - dev * 0.382, mid, mid + dev * 0.382, mid + dev * 0.5,
                     mid + dev * 0.809, mid + dev, mid + dev * 1.382, mid + dev * 1.618]
            self.current_band_idx[s] = self._band_index(close, bands)

        if self.IsWarmingUp:
            return

        if self._dirty:
            self._execute_net()
            self._dirty = False

    # =====================================================================
    def RebalanceA4(self):
        # A4: 1-day cross-sectional low-reversal on the top-100-by-mcap subset of the universe.
        if self.IsWarmingUp or len(self.symbols) < 20:
            return
        # top-100 by market cap among universe members (== A4's standalone universe)
        cand = []
        for s in self.symbols:
            f = self.Securities[s].Fundamentals
            if f is None or f.MarketCap is None:
                continue
            cand.append((s, float(f.MarketCap)))
        if len(cand) < 20:
            return
        cand.sort(key=lambda x: x[1], reverse=True)
        a4_syms = [s for s, _ in cand[:self.A4_TOP_N]]
        hist = self.History(a4_syms, self.A4_WIN + 2, Resolution.Daily)
        if hist.empty or "low" not in hist.columns:
            return
        lows = hist["low"].unstack(level=0)
        cs_rank = lows.rank(axis=1, pct=True)
        ts_rank = cs_rank.rolling(self.A4_WIN).apply(lambda x: x.iloc[-1] / len(x) if len(x) > 0 else 0).iloc[-1]
        score = (-ts_rank).dropna()
        if score.empty:
            return
        top = score.nlargest(self.A4_TOP_K).index.tolist()
        if not top:
            return
        mcaps = {}
        for s in top:
            f = self.Securities[s].Fundamentals
            if f is None or f.MarketCap is None:
                continue
            mcaps[s] = float(f.MarketCap)
        if not mcaps:
            return
        total = sum(mcaps.values())
        target = {s: min(self.A4_MAX_W, mcaps[s] / total) for s in mcaps}
        sw = sum(target.values())
        self.tgt_a4 = {s: v / sw for s, v in target.items()} if sw > 0 else {}
        self._dirty = True

    # =====================================================================
    def Rebalance531(self):
        # 531 momentum-breadth sleeve (285-authoritative). Ported faithfully from
        # S531_285_authoritative.py but it only WRITES self.tgt_531 (in-bucket,
        # sum<=1.0) + flags _dirty; it NEVER calls Liquidate/SetHoldings. Risk-off
        # / no-signal branches sweep the WHOLE 531 bucket to BIL (531's own yield
        # hedge), not a book-wide liquidate. Uses its OWN breadth latch + band
        # ceiling history. Reuses OnData's shared read-only indicators.
        if self.IsWarmingUp:
            return
        idxs = list(self.current_band_idx.values())
        if len(idxs) < 50:
            return
        bottom_frac = sum(i in self.BOTTOM_LEVELS for i in idxs) / len(idxs)
        self.max_stress_531 = max(self.max_stress_531, bottom_frac)
        if bottom_frac >= 0.45:
            self.allow_531 = False
            self.was_risk_off_531 = True
        elif self.was_risk_off_531:
            denominator = max(self.max_stress_531, 0.10)
            improvement = (self.max_stress_531 - bottom_frac) / denominator
            if improvement >= 0.60 or bottom_frac < 0.15:
                for s in self.symbols:
                    if s in self.band_hist_531:
                        self.band_hist_531[s] = RollingWindow[int](self.hist_len)
                self.allow_531 = True
                self.was_risk_off_531 = False
                self.max_stress_531 = 0.0
        else:
            self.allow_531 = True
        # Risk-off: sweep 531's bucket 100% to the Treasury hedge (BIL)
        if not self.allow_531:
            self.tgt_531 = {self.bil: 1.0}; self._dirty = True
            return
        hist = self.History(list(self.symbols), max(self.lookbacks531) + 1, Resolution.Daily)
        if hist.empty:
            return
        closes = hist["close"].unstack(0)
        momentum = {}
        for s in self.symbols:
            if s not in closes:
                continue
            px = closes[s]
            if len(px) < max(self.lookbacks531) + 1:
                continue
            if not self.adx[s].IsReady or self.adx[s].Current.Value > self.adx_limit:
                continue
            mom = np.mean([px.iloc[-1] / px.iloc[-lb - 1] - 1 for lb in self.lookbacks531])
            if not self.ma[s].IsReady:
                continue
            price = self.Securities[s].Price
            ema = self.ma[s].Current.Value
            if price <= ema:
                continue
            if mom > 0:
                momentum[s] = mom
        if not momentum:
            self.tgt_531 = {self.bil: 1.0}; self._dirty = True
            return
        top = sorted(momentum, key=momentum.get, reverse=True)[:self.stock_count_531]
        scaled = {}
        for s in top:
            if not self.ma[s].IsReady or not self.stretch_ema[s].IsReady:
                continue
            dev = np.std(list(self.close_win[s]))
            if dev <= 0:
                continue
            mid = self.ma[s].Current.Value
            lm = self.stretch_ema[s].Current.Value
            lm2 = lm / 2.0
            lm3 = lm2 * 0.38196601
            lm4 = lm * 1.38196601
            lm5 = lm * 1.61803399
            lm6 = (lm + lm2) / 2.0
            bands = [
                mid - dev * lm5, mid - dev * lm4, mid - dev * lm,
                mid - dev * lm6, mid - dev * lm2, mid - dev * lm3, mid,
                mid + dev * lm3, mid + dev * lm2, mid + dev * lm6,
                mid + dev * lm, mid + dev * lm4, mid + dev * lm5
            ]
            price = self.Securities[s].Price
            idx = self._band_index(price, bands)
            self.band_hist_531[s].Add(idx)
            hist_idx = list(self.band_hist_531[s])
            historical_high = max(hist_idx) if hist_idx else idx
            if historical_high <= 0:
                scale = 1.0
            elif idx >= historical_high:
                scale = 0.0
            else:
                scale = max(0.2, 1.0 - idx / historical_high)
            current_stretch = self.stretch_ema[s].Current.Value
            peak_stretch = self.stretch_max.get(s, 0.0)
            if idx >= 10 and peak_stretch > 0 and current_stretch < (peak_stretch * 0.80):
                scale = 0.2
            scaled[s] = (momentum[s] * self.adx[s].Current.Value) * scale
        if not scaled:
            self.tgt_531 = {self.bil: 1.0}; self._dirty = True
            return
        min_stress = 0.15
        max_stress = 0.45
        target_exposure = float(round(np.interp(bottom_frac, [min_stress, max_stress], [1.0, 0.0]), 2))
        total_scaled = sum(scaled.values())
        raw_weights = {s: v / total_scaled for s, v in scaled.items()}
        capped_weights = {s: min(self.max_weight531, w) for s, w in raw_weights.items()}
        current_sum = sum(capped_weights.values())
        final_weights = {}
        if current_sum > 0:
            for s, w in capped_weights.items():
                final_weights[s] = (w / current_sum) * target_exposure
        # Sweep all unallocated bucket capital into BIL to eliminate cash drag
        hedge_allocation = round(1.0 - target_exposure, 2)
        if hedge_allocation > 0:
            final_weights[self.bil] = hedge_allocation
        self.tgt_531 = {s: w for s, w in final_weights.items() if w > 0}
        self._dirty = True

    # =====================================================================
    def _execute_net(self):
        net = {}
        for sym, w in self.tgt_a4.items():
            net[sym] = net.get(sym, 0.0) + self.W_A4 * w
        for sym, w in self.tgt_531.items():
            net[sym] = net.get(sym, 0.0) + self.W_531 * w
        net = {s: w for s, w in net.items() if w > 0.0005}

        equity = self.Portfolio.TotalPortfolioValue
        keep = set(net)
        for kvp in list(self.Portfolio):
            sym, h = kvp.Key, kvp.Value
            if h.Invested and sym not in keep and h.Quantity != 0:
                self.MarketOnOpenOrder(sym, -h.Quantity, tag=f"EXIT now={self.Time.date()}")
        def cur_w(sym):
            return (float(self.Portfolio[sym].HoldingsValue) / equity) if equity > 0 else 0.0
        for sym, w in sorted(net.items(), key=lambda kv: kv[1] - cur_w(kv[0])):
            if not self.Securities.ContainsKey(sym):
                continue
            price = float(self.Securities[sym].Price)
            if price <= 0:
                continue
            tq = int(equity * w / price)
            cur = self.Portfolio[sym].Quantity if self.Portfolio.ContainsKey(sym) else 0
            d = tq - cur
            if d != 0:
                self.MarketOnOpenOrder(sym, d, tag=f"ENTRY w={w:.3f} px={price:.2f} now={self.Time.date()}")

    # =====================================================================
    def OnOrderEvent(self, oe):
        if not self.DEBUG_FILLS:
            return
        if oe.Status == OrderStatus.INVALID:
            self._n_inv += 1
            self.SetRuntimeStatistic("Invalid", str(self._n_inv))
            msg = str(oe.Message)[:50] if oe.Message else "?"
            self.SetRuntimeStatistic("InvMsg", f"{oe.Symbol.Value}:{msg}")
            return
        if oe.Status != OrderStatus.FILLED or oe.FillQuantity == 0:
            return
        sec = self.Securities[oe.Symbol]; fp = float(oe.FillPrice); o = float(sec.Open)
        dev = abs(fp - o) / o if o > 0 else 0.0
        try:
            od = self.Transactions.GetOrderById(oe.OrderId); ot = int(od.Type) if od is not None else -1
        except Exception:
            ot = -1
        if ot == 4:
            self._max_moo = max(self._max_moo, dev); self._n_moo += 1
        else:
            self._max_liq = max(self._max_liq, dev); self._n_liq += 1
        self._n_fills += 1
        gross = sum(abs(p.HoldingsValue) for p in self.Portfolio.Values) / self.Portfolio.TotalPortfolioValue
        self._max_gross = max(self._max_gross, gross)
        self.SetRuntimeStatistic("Fills", str(self._n_fills))
        self.SetRuntimeStatistic("MaxVsOpen_MOO_pct", f"{self._max_moo*100:.4f}")
        self.SetRuntimeStatistic("nMOO", str(self._n_moo))
        self.SetRuntimeStatistic("nLiq", str(self._n_liq))
        self.SetRuntimeStatistic("MaxGross_pct", f"{self._max_gross*100:.1f}")