# region imports
from AlgorithmImports import *
# endregion


TICKERS = [
    "BIL", "UVXY", "TQQQ", "TECL", "UPRO", "SQQQ", "SOXL",
    "SPY", "SPXL", "XLK", "DBC", "TLT", "NVDA",
]

BLOCKS = [
    (0.9, 0.1), (0.8, 0.2), (0.7, 0.3), (0.6, 0.4), (0.5, 0.5),
    (0.4, 0.6), (0.3, 0.7), (0.2, 0.8), (0.1, 0.9),
]

BIL_BIAS = 0.4  # Increased cash allocation
MAX_LEVERAGE = 0.55  # Reduced leverage for lower drawdown
VOL_LOOKBACK = 20  # Days to measure volatility
DD_STOP_THRESHOLD = 0.10  # Lower drawdown threshold for earlier exit


class TQQQFTLTQC(QCAlgorithm):
    """TQQQ FTLT Yolo - top 14 RSI Sort. Updates RSI/TI at 15:40 only."""

    def initialize(self):
        self.set_start_date(2020, 1, 1)
        self.set_end_date(2025, 12, 31)
        self.set_cash(100000)
        self.symbols = {}
        self._rsi = {}
        self._ma20 = {}
        self._ma200 = {}
        self._nvda_roc = RateOfChangePercent(1)
        self._nvda_stdev5 = StandardDeviation(5)
        self._spy_returns = RollingWindow[float](VOL_LOOKBACK)
        self._peak_equity = self.portfolio.total_portfolio_value
        for t in TICKERS:
            sym = self.add_equity(t, Resolution.MINUTE).symbol
            self.symbols[t] = sym
            self._rsi[t] = {
                2: RelativeStrengthIndex(2, MovingAverageType.WILDERS),
                5: RelativeStrengthIndex(5, MovingAverageType.WILDERS),
                10: RelativeStrengthIndex(10, MovingAverageType.WILDERS),
                14: RelativeStrengthIndex(14, MovingAverageType.WILDERS),
            }
            self._ma20[t] = SimpleMovingAverage(20)
            self._ma200[t] = SimpleMovingAverage(200)
        self.set_warm_up(200, Resolution.MINUTE)

    def _get_volatility_scalar(self) -> float:
        """Calculate position size scalar based on volatility and drawdown."""
        if self._spy_returns.count < VOL_LOOKBACK:
            return 0.4
        
        returns = [self._spy_returns[i] for i in range(self._spy_returns.count)]
        vol = np.std(returns) * np.sqrt(252)
        
        # Balanced vol scaling
        if vol > 0.30:
            vol_scalar = 0.3
        elif vol > 0.25:
            vol_scalar = 0.5
        elif vol > 0.20:
            vol_scalar = 0.7
        else:
            vol_scalar = 1.0
        
        # Balanced drawdown protection
        current_equity = self.portfolio.total_portfolio_value
        if current_equity > self._peak_equity:
            self._peak_equity = current_equity
        drawdown = (self._peak_equity - current_equity) / self._peak_equity
        
        # Hard stop: exit all risky positions
        if drawdown > DD_STOP_THRESHOLD:
            return 0.0
        elif drawdown > 0.10:
            dd_scalar = 0.4
        elif drawdown > 0.07:
            dd_scalar = 0.6
        else:
            dd_scalar = 1.0
        
        return min(vol_scalar, dd_scalar)
    
    def _update_indicators(self, slice):
        """Update RSI and TI at 15:40 using bar closes."""
        spy_sym = self.symbols["SPY"]
        if spy_sym in slice.bars and self._spy_returns.count > 0:
            prev_price = self.securities[spy_sym].price
            curr_price = slice.bars[spy_sym].close
            if prev_price > 0:
                daily_return = (curr_price - prev_price) / prev_price
                self._spy_returns.add(daily_return)
        for t in TICKERS:
            sym = self.symbols[t]
            if sym not in slice.bars:
                continue
            bar = slice.bars[sym]
            self._rsi[t][2].update(bar.end_time, bar.close)
            self._rsi[t][5].update(bar.end_time, bar.close)
            self._rsi[t][10].update(bar.end_time, bar.close)
            self._rsi[t][14].update(bar.end_time, bar.close)
            self._ma20[t].update(bar.end_time, bar.close)
            self._ma200[t].update(bar.end_time, bar.close)
            if t == "NVDA":
                self._nvda_roc.update(bar.end_time, bar.close)
                if self._nvda_roc.is_ready:
                    self._nvda_stdev5.update(bar.end_time, self._nvda_roc.current.value)

    def _get_ti(self, ticker: str) -> dict:
        """Get current TI values. Returns dict with rsi_2, rsi_5, rsi_10, rsi_14, ma_20, ma_200, close, stdev5."""
        out = {}
        sym = self.symbols[ticker]
        close = self.securities[sym].price
        out["close"] = close
        out["rsi_2"] = self._rsi[ticker][2].current.value
        out["rsi_5"] = self._rsi[ticker][5].current.value
        out["rsi_10"] = self._rsi[ticker][10].current.value
        out["rsi_14"] = self._rsi[ticker][14].current.value
        out["ma_20"] = self._ma20[ticker].current.value
        out["ma_200"] = self._ma200[ticker].current.value
        out["stdev5"] = self._nvda_stdev5.current.value * 100 if ticker == "NVDA" else 0.0
        return out

    def _bil_replace(self) -> str:
        ti = self._get_ti("NVDA")
        if not self._nvda_stdev5.is_ready:
            return "BIL"
        return "SOXL" if ti["stdev5"] > ti["rsi_2"] else "BIL"

    def _tqqq_ftlt_reddit(self) -> str:
        spy = self._get_ti("SPY")
        if not self._ma200["SPY"].is_ready:
            return "BIL"
        spy_above_ma200 = spy["close"] > spy["ma_200"]
        if spy_above_ma200:
            tqqq = self._get_ti("TQQQ")
            if tqqq["rsi_10"] > 79:
                return "UVXY"
            spxl = self._get_ti("SPXL")
            if spxl["rsi_10"] > 80:
                return "UVXY"
            xlk = self._get_ti("XLK")
            dbc = self._get_ti("DBC")
            if xlk["rsi_10"] > dbc["rsi_10"]:
                return "TQQQ"
            return "BIL"
        tqqq = self._get_ti("TQQQ")
        if tqqq["rsi_10"] < 31:
            return "TECL"
        if spy["rsi_10"] < 30:
            return "UPRO"
        if self._ma20["TQQQ"].is_ready and tqqq["close"] < tqqq["ma_20"]:
            sqqq = self._get_ti("SQQQ")
            tlt = self._get_ti("TLT")
            xlk = self._get_ti("XLK")
            dbc = self._get_ti("DBC")
            if sqqq["rsi_10"] > tlt["rsi_10"]:
                if xlk["rsi_10"] > dbc["rsi_10"]:
                    return "BIL"
                return "SQQQ"
            return "BIL"
        sqqq = self._get_ti("SQQQ")
        if sqqq["rsi_10"] < 31:
            return "SQQQ"
        xlk = self._get_ti("XLK")
        dbc = self._get_ti("DBC")
        if xlk["rsi_10"] > dbc["rsi_10"]:
            if dbc["close"] < dbc["ma_20"]:
                return "TQQQ"
            return "BIL"
        return "BIL"

    def _decide_weights(self) -> dict:
        merged = {}
        for w_bil, w_ftlt in BLOCKS:
            t_bil = self._bil_replace()
            t_ftlt = self._tqqq_ftlt_reddit()
            bw = 1.0 / len(BLOCKS)
            merged[t_bil] = merged.get(t_bil, 0.0) + bw * w_bil
            merged[t_ftlt] = merged.get(t_ftlt, 0.0) + bw * w_ftlt
        candidates = [(t, merged[t]) for t in merged if merged[t] > 0]
        if not candidates:
            return {"BIL": 1.0}
        top2 = sorted(candidates, key=lambda x: -self._get_ti(x[0])["rsi_14"])[:2]
        top2_weights = {t: 1.0 / len(top2) for t, _ in top2}
        out = {t: 0.0 for t in TICKERS}
        total = sum(top2_weights.values())
        if total > 0:
            for t, w in top2_weights.items():
                out[t] = w / total
        return out

    def on_data(self, slice):
        if self.time.hour != 15 or self.time.minute != 40:
            return
        self._update_indicators(slice)
        if not self._rsi["SPY"][14].is_ready or not self._ma200["SPY"].is_ready:
            return
        if self.is_warming_up:
            return
        weights = self._decide_weights()
        if "BIL" in weights:
            weights["BIL"] += BIL_BIAS
        
        # Apply volatility/drawdown scaling
        vol_scalar = self._get_volatility_scalar()
        
        # If in severe drawdown mode, go 100% cash
        if vol_scalar == 0.0:
            for t in TICKERS:
                sym = self.symbols[t]
                if t == "BIL":
                    self.set_holdings(sym, 1.0)
                else:
                    self.liquidate(sym)
            return
        
        # Scale down risky assets, increase cash
        risky_assets = ["TQQQ", "TECL", "UPRO", "SQQQ", "SOXL", "UVXY"]
        for t in risky_assets:
            if t in weights:
                weights[t] *= vol_scalar
        
        # Add remaining to BIL
        reduction = 1.0 - vol_scalar
        weights["BIL"] = weights.get("BIL", 0) + reduction
        
        # Apply max leverage constraint
        total = sum(weights.values())
        if total > MAX_LEVERAGE:
            scale = MAX_LEVERAGE / total
            for t in weights:
                weights[t] *= scale
            weights["BIL"] = weights.get("BIL", 0) + (1.0 - MAX_LEVERAGE)
        
        total = sum(weights.values())
        if total <= 0:
            return
        
        for t in TICKERS:
            sym = self.symbols[t]
            if sym not in slice.bars:
                continue
            w = weights.get(t, 0) / total if total > 0 else 0.0
            self.set_holdings(sym, float(w))