# region imports
from AlgorithmImports import *
# endregion


TICKERS = [
    "BIL", "UVXY", "TQQQ", "TECL", "UPRO", "SQQQ", "NVDL",
    "SPY", "SPXL", "XLK", "KMLM", "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),
]




class TQQQFTLTQC(QCAlgorithm):
    """TQQQ FTLT Yolo - top 14 RSI Sort."""

    def initialize(self):
        self.set_start_date(2021, 1, 1)
        self.set_cash(100000)
        self.symbols = {}
        self._rsi = {}
        self._ma20 = {}
        self._ma200 = {}
        self._nvda_prices = RollingWindow[float](6)  # Need 6 prices for 5 daily returns
        
        for t in TICKERS:
            sym = self.add_equity(t, Resolution.DAILY).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.DAILY)



    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
        if ticker == "NVDA":
            if self._nvda_prices.is_ready:
                # Calculate 5 most recent daily returns (close-to-close)
                returns = []
                for i in range(5):
                    if i + 1 < self._nvda_prices.count:
                        if self._nvda_prices[i] > 0 and self._nvda_prices[i+1] > 0:
                            ret = (self._nvda_prices[i] - self._nvda_prices[i+1]) / self._nvda_prices[i+1]
                            returns.append(ret)
                # Standard deviation of daily returns as percentage
                out["stdev5"] = np.std(returns) * 100 if len(returns) == 5 else 0.0
            else:
                out["stdev5"] = 0.0
        else:
            out["stdev5"] = 0.0
        return out

    def _bil_replace(self) -> str:
        ti = self._get_ti("NVDA")
        if not self._nvda_prices.is_ready:
            return "BIL"
        return "NVDL" 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")
            kmlm = self._get_ti("KMLM")
            if xlk["rsi_10"] > kmlm["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")
            kmlm = self._get_ti("KMLM")
            if sqqq["rsi_10"] > tlt["rsi_10"]:
                if xlk["rsi_10"] > kmlm["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")
        kmlm = self._get_ti("KMLM")
        if xlk["rsi_10"] > kmlm["rsi_10"]:
            if kmlm["close"] < kmlm["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]
        out = {t: 0.0 for t in TICKERS}
        for t, _ in top2:
            out[t] = 0.5
        return out

    def on_data(self, slice):
        # Update indicators with daily bar close (during warmup and live)
        for t in TICKERS:
            sym = self.symbols[t]
            if sym not in slice.bars:
                continue
            bar = slice.bars[sym]
            close = bar.close
            self._rsi[t][2].update(bar.end_time, close)
            self._rsi[t][5].update(bar.end_time, close)
            self._rsi[t][10].update(bar.end_time, close)
            self._rsi[t][14].update(bar.end_time, close)
            self._ma20[t].update(bar.end_time, close)
            self._ma200[t].update(bar.end_time, close)
            if t == "NVDA":
                self._nvda_prices.add(close)
        
        if self.is_warming_up:
            return
        
        if not self._rsi["SPY"][14].is_ready or not self._ma200["SPY"].is_ready:
            return

        weights = self._decide_weights()
        
        total = sum(weights.values())
        if total <= 0:
            return
        
        for t in TICKERS:
            sym = self.symbols[t]
            w = weights.get(t, 0) / total if total > 0 else 0.0
            self.set_holdings(sym, float(w))