from AlgorithmImports import *
import numpy as np
import pandas as pd

# =============================================================================
# kinfo ENSEMBLE — 0379 (defense) + 0273 (offense), one daily algorithm
# Applies kinfo's #1 lesson (uncorrelated legs beat singles) to the two engines
# that VERIFIED on QC:
#   * 0379 symphony : DD-controlled multi-asset rotation (TLT/GLD/SHV + crash gate)
#                     -> OOS 32.6%/21.8%DD/Sh0.94, full 20.4/29.3/0.77
#   * 0273          : aggressive leveraged-Nasdaq regime rotator (QQQ/TQQQ/SQQQ,
#                     5-cond bear gate) -> kinfo full 33.6%/42.5%/Sh1.05
# Each sleeve is gross<=100%; blended at W_0379 / W_0273 (sum=1) so the combined
# book never borrows -> cash-faithful. QC-default margin (faithful proxy; the
# cash-mode MOO rotation artifact is avoided, see champ_lev_regime notes).
#
# IS = 2010-2020 | OOS = 2021-2025 | default full. Tune W_0379 / W_0273 by results.
# =============================================================================

class KinfoEnsemble(QCAlgorithm):

    # blend weights (sum to 1.0) — 70/30 chosen: best DD (25.2% full) + Sharpe 1.32/1.63, monotone in-sample
    W_0379 = 0.70
    W_0273 = 0.30

    USE_CASH = False   # False = margin proxy (gross<=100%, no borrowing). True = IBKR cash + 10bps + sells-before-buys
    REBAL_TOL = 0.08   # rebalance only when held names change or a weight moves >= this (cuts turnover/fees)

    T_0379 = ["QQQ","SPY","TLT","BSV","GLD","QLD","PSQ","SHV","IEF","QID","SMH","USD"]
    T_0273 = ["TQQQ","SQQQ"]              # QQQ shared with 0379 set; ONEQ + VIX added below
    HISTORY_BARS = 1500                   # 0273 needs ~6y for bear-streak; 0379 uses last 300

    # 0273 params
    STREAK_MIN = 40
    DIST_THRESHOLD = -0.08
    BOUNCE_ENTRY = 0.045
    MAX_SQQQ = 0.90

    def initialize(self):
        self.set_start_date(2021, 1, 1)
        self.set_end_date(2026, 6, 9)
        self.set_cash(100000)
        if self.USE_CASH:
            self.set_brokerage_model(BrokerageName.INTERACTIVE_BROKERS_BROKERAGE, AccountType.CASH)
            self.settings.minimum_order_margin_portfolio_percentage = 0.0

        self.syms = {}
        for t in self.T_0379 + self.T_0273:
            if t in self.syms:
                continue
            eq = self.add_equity(t, Resolution.DAILY)
            if self.USE_CASH:
                eq.set_fee_model(InteractiveBrokersFeeModel())
                eq.set_slippage_model(ConstantSlippageModel(0.001))
                eq.set_settlement_model(ImmediateSettlementModel())
            self.syms[t] = eq.symbol
        self.ixic = self.add_equity("ONEQ", Resolution.DAILY).symbol   # 0273 IXIC proxy
        try:
            self.vix = self.add_data(CBOE, "VIX", Resolution.DAILY).symbol
            self._has_vix = True
        except Exception:
            self.vix = None
            self._has_vix = False

        self.set_benchmark("QQQ")
        self.set_warm_up(self.HISTORY_BARS, Resolution.DAILY)
        self._last_key = None
        self._pending = None   # FILL-FIX: queued target weights, drained at next open
        self.schedule.on(self.date_rules.every_day("QQQ"),
                         self.time_rules.before_market_close("QQQ", 10), self.rebalance)

    # ---------- shared helpers ----------
    @staticmethod
    def _rsi_wilder(px, window):
        delta = px.diff()
        gain = delta.clip(lower=0.0)
        loss = (-delta).clip(lower=0.0)
        ag = gain.ewm(alpha=1.0/window, adjust=False).mean()
        al = loss.ewm(alpha=1.0/window, adjust=False).mean()
        rs = ag / al.replace(0.0, np.nan)
        return (100.0 - 100.0/(1.0+rs)).fillna(50.0)

    @staticmethod
    def _rsi_sma_seed(series, period):
        delta = series.diff()
        gain = delta.clip(lower=0.0).to_numpy()
        loss = (-delta.clip(upper=0.0)).to_numpy()
        n = len(series); ag = np.full(n, np.nan); al = np.full(n, np.nan)
        if n > period:
            ag[period] = np.nanmean(gain[1:period+1]); al[period] = np.nanmean(loss[1:period+1])
            ip = 1.0/period
            for i in range(period+1, n):
                ag[i] = ag[i-1]*(1-ip) + gain[i]*ip
                al[i] = al[i-1]*(1-ip) + loss[i]*ip
        rs = np.where(al > 0, ag/np.where(al == 0, np.nan, al), np.nan)
        return pd.Series(100.0 - 100.0/(1.0+rs), index=series.index).fillna(50.0)

    @staticmethod
    def _top(tickers, rsi10):
        best = tickers[0]; bv = rsi10[best]
        for t in tickers[1:]:
            if rsi10[t] > bv: best, bv = t, rsi10[t]
        return best

    # ---------- 0379 sleeve ----------
    def _decide_0379(self, close):
        rsi_assets = ["QQQ","SPY","TLT","BSV","GLD","PSQ","IEF","QLD","USD","SMH"]
        rsi10 = {t: float(self._rsi_sma_seed(close[t], 10).iloc[-1]) for t in rsi_assets}
        rsi60_spy = float(self._rsi_sma_seed(close["SPY"], 60).iloc[-1])
        spy = close["SPY"]; qqq = close["QQQ"]
        spy_ma200 = float(spy.rolling(200).mean().iloc[-1])
        spy_ma20  = float(spy.rolling(20).mean().iloc[-1])
        qqq_ma20  = float(qqq.rolling(20).mean().iloc[-1])
        spy_l = float(spy.iloc[-1]); qqq_l = float(qqq.iloc[-1])
        cr10  = float((qqq.iloc[-1]/qqq.iloc[-11]-1.0)*100.0)
        cr60  = float((qqq.iloc[-1]/qqq.iloc[-61]-1.0)*100.0)
        cr252 = float((qqq.iloc[-1]/qqq.iloc[-253]-1.0)*100.0)
        w = {}
        def add(t, x): w[t] = w.get(t, 0.0) + x
        if spy_l > spy_ma200:
            if rsi10["QQQ"] > 80.0: add("TLT", 1.0)
            elif rsi10["SPY"] > 80.0: add("TLT", 1.0)
            elif rsi60_spy > 60.0: add(self._top(["TLT","BSV","GLD"], rsi10), 1.0)
            else: add("QQQ", 1.0)
        else:
            if rsi10["QQQ"] < 30.0:
                add("QLD", 1.0)
            else:
                tlt_gt_psq = rsi10["TLT"] > rsi10["PSQ"]
                qqq_below = qqq_l < qqq_ma20
                if qqq_l > qqq_ma20:
                    if rsi10["PSQ"] < 30.0: add("PSQ", 0.50)
                    elif cr10 > 5.5: add("SHV", 0.50)
                    else: add("QQQ", 0.50)
                else:
                    add(self._top(["IEF","PSQ"], rsi10), 0.50)
                if cr252 < -20.0:
                    if qqq_below:
                        if cr60 <= -12.0:
                            sib1 = "SPY" if spy_l > spy_ma20 else ("QQQ" if tlt_gt_psq else "PSQ")
                            sib2 = "QQQ" if tlt_gt_psq else "PSQ"
                            add(sib1, 0.25); add(sib2, 0.25)
                        else:
                            add("QLD" if tlt_gt_psq else "QID", 0.50)
                    else:
                        if rsi10["PSQ"] < 31.0: add("PSQ", 0.50)
                        elif cr10 > 5.5: add("PSQ", 0.50)
                        else: add(self._top(["QQQ","SMH"], rsi10), 0.50)
                else:
                    if qqq_below:
                        add("QLD" if tlt_gt_psq else "QID", 0.50)
                    else:
                        if rsi10["PSQ"] < 31.0: add("QID", 0.50)
                        elif cr10 > 5.5: add("QID", 0.50)
                        else: add(self._top(["QLD","USD"], rsi10), 0.50)
        return w

    # ---------- 0273 sleeve ----------
    def _decide_0273(self, df, vix_series):
        df = df.copy()
        df["qqq_sma_50"] = df["QQQ"].rolling(50).mean()
        df["qqq_sma_200"] = df["QQQ"].rolling(200).mean()
        df["qqq_dist_sma_50"] = (df["QQQ"] - df["qqq_sma_50"]) / df["qqq_sma_50"]
        df["qqq_dist_sma_200"] = (df["QQQ"] - df["qqq_sma_200"]) / df["qqq_sma_200"]
        df["qqq_sma_200_slope_20"] = (df["qqq_sma_200"] - df["qqq_sma_200"].shift(20)) / df["qqq_sma_200"].shift(20)
        df["qqq_roc_21"] = df["QQQ"].pct_change(21)
        df["ixic_sma_50"] = df["IXIC"].rolling(50).mean()
        df["ixic_sma_100"] = df["IXIC"].rolling(100).mean()
        df["ixic_sma_175"] = df["IXIC"].rolling(175).mean()
        df["ixic_sma_200"] = df["IXIC"].rolling(200).mean()
        df["ixic_sma_200_dist"] = (df["IXIC"] - df["ixic_sma_200"]) / df["ixic_sma_200"]
        df["ixic_above_sma_50"] = (df["IXIC"] > df["ixic_sma_50"]).astype(float)
        df["ixic_above_sma_100"] = (df["IXIC"] > df["ixic_sma_100"]).astype(float)
        df["ixic_above_sma_175"] = (df["IXIC"] > df["ixic_sma_175"]).astype(float)
        df["qqq_rsi_14"] = self._rsi_wilder(df["QQQ"], 14)
        df["tqqq_vol20"] = df["TQQQ"].pct_change().rolling(20).std() * np.sqrt(252)
        if vix_series is not None and not vix_series.empty:
            v = vix_series.reindex(df.index).ffill()
            df["vix_ratio_20"] = (v / v.rolling(20).mean()).fillna(1.0)
        else:
            df["vix_ratio_20"] = 1.0

        feat = df.shift(1).iloc[-1]
        req = ["qqq_dist_sma_50","qqq_dist_sma_200","qqq_sma_200_slope_20","qqq_roc_21",
               "tqqq_vol20","qqq_rsi_14","ixic_sma_200_dist","ixic_above_sma_50",
               "ixic_above_sma_100","ixic_above_sma_175","vix_ratio_20"]
        if feat[req].isna().any():
            return {}

        above175 = bool(feat["ixic_above_sma_175"] > 0.5)
        above100 = bool(feat["ixic_above_sma_100"] > 0.5)
        above50 = bool(feat["ixic_above_sma_50"] > 0.5)
        rsi14 = float(feat["qqq_rsi_14"]); roc21 = float(feat["qqq_roc_21"])
        dist_200_ixic = float(feat["ixic_sma_200_dist"]); vix_ratio = float(feat["vix_ratio_20"])
        strong_bull = above175 and above50 and (rsi14 > 55.0) and (roc21 > 0.0)
        full_momentum = (rsi14 > 60.0) and (roc21 > 0.0)
        regime = 4
        if above100: regime = 3
        if above175 and (not strong_bull): regime = 2
        if above175 and strong_bull and (not full_momentum): regime = 5
        if above175 and strong_bull and full_momentum: regime = 1
        vix_scale = 1.0
        if vix_ratio >= 2.2: vix_scale = 0.38
        elif vix_ratio >= 1.5: vix_scale = 0.62
        elif vix_ratio >= 1.1: vix_scale = 0.88
        w_qqq, w_tqqq, w_sqqq = 0.0, 0.0, 0.0
        very_extended = (regime == 1) and (dist_200_ixic > 0.15)
        normal_bull = (regime == 1) and (not very_extended)
        if very_extended: w_tqqq = float(np.clip(0.84*vix_scale, 0.35, 0.84))
        elif normal_bull: w_tqqq = float(np.clip(0.90*vix_scale, 0.35, 0.90))
        elif regime == 5: w_tqqq = float(np.clip(0.60*vix_scale, 0.28, 0.60))
        elif regime == 2: w_tqqq = float(np.clip(0.70*vix_scale, 0.28, 0.70))
        elif regime == 4: w_qqq = 0.45
        elif regime == 3:
            w_tqqq = 0.62 if rsi14 > 55 else (0.48 if rsi14 > 45 else 0.34)
        t_vol = float(feat["tqqq_vol20"])
        if above175 and (t_vol < 0.60): w_tqqq = min(0.95, w_tqqq*1.05)
        if above175 and (t_vol >= 0.80): w_tqqq = min(0.95, w_tqqq*0.75)
        if very_extended: w_tqqq = min(0.95, w_tqqq*0.92)

        now = df.iloc[-1]
        dist50 = float(now["qqq_dist_sma_50"]); dist200 = float(now["qqq_dist_sma_200"])
        slope200 = float(now["qqq_sma_200_slope_20"])
        if any(np.isnan([dist50, dist200, slope200])):
            return {}
        below_200 = (df["qqq_dist_sma_200"] < 0).astype(int).values
        streak = 0
        for b in below_200[:-1][::-1]:
            if b == 1: streak += 1
            else: break
        bear_gate = (streak >= self.STREAK_MIN) and (dist200 < dist50) and (dist200 < self.DIST_THRESHOLD) and (slope200 < 0.0)
        roc3 = float(df["QQQ"].pct_change(3).shift(1).iloc[-1])
        if bear_gate and (roc3 >= self.BOUNCE_ENTRY):
            w_qqq, w_tqqq, w_sqqq = 0.0, 0.0, self.MAX_SQQQ
        gross = w_qqq + w_tqqq + w_sqqq
        if gross > 1.0:
            w_qqq /= gross; w_tqqq /= gross; w_sqqq /= gross
        return {"QQQ": w_qqq, "TQQQ": w_tqqq, "SQQQ": w_sqqq}

    # ---------- combined rebalance ----------
    def rebalance(self):
        if self.is_warming_up:
            return
        all_syms = [self.syms[t] for t in self.syms] + [self.ixic]
        hist = self.history(all_syms, self.HISTORY_BARS, Resolution.DAILY)
        if hist is None or hist.empty or "close" not in hist.columns:
            return
        try:
            wide = hist["close"].unstack(level=0)
        except Exception:
            return
        col = {}
        for t in self.syms:
            s = self.syms[t]
            if s in wide.columns: col[t] = wide[s]
        if self.ixic in wide.columns: col["IXIC"] = wide[self.ixic]
        if any(t not in col for t in self.T_0379) or "IXIC" not in col:
            return
        close = pd.DataFrame(col).dropna()
        if len(close) < 260:
            return

        # vix series
        vix_series = None
        if self._has_vix and self.vix is not None:
            try:
                vh = self.history(self.vix, self.HISTORY_BARS, Resolution.DAILY)
                if vh is not None and not vh.empty:
                    s = vh["value"] if "value" in vh.columns else (vh["close"] if "close" in vh.columns else None)
                    if s is not None:
                        if isinstance(s.index, pd.MultiIndex):
                            s.index = s.index.get_level_values(-1)
                        s.index = pd.to_datetime(s.index).tz_localize(None)
                        vix_series = s[~pd.Index(s.index).duplicated(keep="last")].sort_index().astype(float)
            except Exception:
                vix_series = None

        w379 = self._decide_0379(close[self.T_0379])
        df0273 = close[["QQQ","TQQQ","SQQQ","IXIC"]]
        if vix_series is not None:
            df0273.index = close.index
        w273 = self._decide_0273(df0273, vix_series)

        target = {}
        for t, x in w379.items():
            target[t] = target.get(t, 0.0) + self.W_0379 * x
        for t, x in w273.items():
            if x > 0:
                target[t] = target.get(t, 0.0) + self.W_0273 * x

        target = {t: round(w, 3) for t, w in target.items() if w > 0.005}
        # rebalance-on-change with tolerance band (cuts turnover/fees on cash):
        # only trade if the held names change OR some weight moves >= REBAL_TOL.
        if self._last_key is not None:
            same_names = set(target) == set(self._last_key)
            moves = [abs(target.get(t, 0.0) - self._last_key.get(t, 0.0))
                     for t in set(target) | set(self._last_key)]
            if same_names and (max(moves) if moves else 0.0) < self.REBAL_TOL:
                return
        self._last_key = target

        # FILL-FIX (ports the 476/504cg discipline): do NOT execute intraday on Daily
        # resolution. An intraday market order / set_holdings on Daily fills at the
        # PRIOR bar's close (yesterday) -> inflates the backtest vs live. Instead QUEUE
        # the target weights; on_data drains them via MarketOnOpenOrder so fills land at
        # the NEXT session open — a price that only exists AFTER the decision.
        self._pending = target

    def on_data(self, data):
        # FILL-FIX drain: execute the queued rebalance at the next open via MOO.
        if self.is_warming_up or self._pending is None:
            return
        target = self._pending
        self._pending = None
        keep = set(self.syms[t] for t in target)
        # close names no longer targeted
        for kvp in list(self.portfolio):
            sym, h = kvp.key, kvp.value
            if h.invested and sym not in keep and sym != self.ixic and h.quantity != 0:
                self.market_on_open_order(sym, -h.quantity)
        # size off current equity & last price; submit reductions before increases
        equity = self.portfolio.total_portfolio_value
        def cur_w(t):
            h = self.portfolio[self.syms[t]]
            return (float(h.holdings_value) / equity) if equity > 0 else 0.0
        for t, w in sorted(target.items(), key=lambda kv: kv[1] - cur_w(kv[0])):
            sym = self.syms[t]
            price = float(self.securities[sym].price)
            if price <= 0:
                continue
            target_qty = int(equity * w / price)
            cur = self.portfolio[sym].quantity if self.portfolio.contains_key(sym) else 0
            delta = target_qty - cur
            if delta != 0:
                self.market_on_open_order(sym, delta)