# region imports
from AlgorithmImports import *
import numpy as np
import decimal
# endregion


class CustomVix(PythonData):
    def get_source(self, config, date, is_live):
        url = f"https://cdn.cboe.com/api/global/us_indices/daily_prices/{config.symbol.value}_History.csv"
        return SubscriptionDataSource(url, SubscriptionTransportMedium.REMOTE_FILE)

    def reader(self, config, line, date, is_live):
        try:
            if not line or not line[0].isdigit():
                return None
            data = line.split(',')
            v = CustomVix()
            v.symbol = config.symbol
            v.time = datetime.strptime(data[0], "%m/%d/%Y")
            v.value = decimal.Decimal(data[4])
            return v
        except:
            return None


class QuantumPortfolioConservativeV14(QCAlgorithm):

    def initialize(self):
        self.set_start_date(2018, 1, 1)
        self.set_end_date(2023, 12, 31)
        self.set_cash(1_000_000)
        self.set_warm_up(504)

        self.set_brokerage_model(
            BrokerageName.CHARLES_SCHWAB,
            AccountType.MARGIN
        )

        # Assets
        self.spxl = self.add_equity("SPXL", Resolution.DAILY).symbol
        self.tqqq = self.add_equity("TQQQ", Resolution.DAILY).symbol

        self.vix = self.add_data(CustomVix, "VIX", Resolution.DAILY).symbol
        self.vix3m = self.add_data(CustomVix, "VIX3M", Resolution.DAILY).symbol
        self.vvix = self.add_data(CustomVix, "VVIX", Resolution.DAILY).symbol

        # Indicators
        self.ema200_spxl = self.ema(self.spxl, 200)
        self.ema200_tqqq = self.ema(self.tqqq, 200)
        self.ema20_spxl = self.ema(self.spxl, 20)
        self.ema20_tqqq = self.ema(self.tqqq, 20)

        # Risk state
        self.peak_equity = self.portfolio.total_portfolio_value
        self.profit_peak = self.portfolio.total_portfolio_value
        self.equity_cap = self.portfolio.total_portfolio_value

        self.cooldown_days = 0
        self.ema_break_count = 0

        self.schedule.on(
            self.date_rules.every_day(self.spxl),
            self.time_rules.after_market_open(self.spxl, 30),
            self.rebalance
        )

    def rebalance(self):
        if self.is_warming_up:
            return

        equity = self.portfolio.total_portfolio_value

        # ==================================================
        # ① 组合级硬回撤兜底
        # ==================================================
        self.peak_equity = max(self.peak_equity, equity)
        drawdown = (self.peak_equity - equity) / self.peak_equity

        if drawdown > 0.20:
            self.liquidate()
            self.cooldown_days = 20
            self.peak_equity = equity
            self.profit_peak = equity
            self.equity_cap = equity
            self.log(f"FORCE EXIT | DD={drawdown:.2%}")
            return

        if self.cooldown_days > 0:
            self.cooldown_days -= 1
            return

        # ==================================================
        # ② EMA200 + VVIX → 趋势死亡 → 全清仓
        # ==================================================
        spxl_price = self.securities[self.spxl].price
        tqqq_price = self.securities[self.tqqq].price
        vvix = self.securities[self.vvix].price

        below_ema = (
            spxl_price < self.ema200_spxl.current.value and
            tqqq_price < self.ema200_tqqq.current.value
        )

        if below_ema and vvix > 100:
            self.ema_break_count += 1
        else:
            self.ema_break_count = 0

        if self.ema_break_count >= 3:
            self.liquidate()
            self.cooldown_days = 20
            self.ema_break_count = 0
            self.log("EMA200 + VVIX CONFIRMED → FULL EXIT")
            return

        # ==================================================
        # ③ VVIX → 杠杆软上限
        # ==================================================
        if vvix > 120:
            max_exposure = 0.5
        elif vvix > 100:
            max_exposure = 1.0
        else:
            max_exposure = 2.0

        # ==================================================
        # ④ VIX Panic → 系统性风险
        # ==================================================
        vix = self.securities[self.vix].price
        vix3m = self.securities[self.vix3m].price
        if vix <= 0 or vix3m <= 0:
            return

        if vix3m / vix < 1.05 or vix > 25:
            self.liquidate()
            return

        # ==================================================
        # ⑤ Equity Cap（⭐压 45% 回撤的关键）
        # ==================================================
        self.equity_cap = max(self.equity_cap, equity)
        equity_ratio = equity / self.equity_cap

        if equity_ratio < 0.90:
            risk_multiplier = 0.5
        elif equity_ratio < 0.95:
            risk_multiplier = 0.7
        else:
            risk_multiplier = 1.0

        # ==================================================
        # ⑥ 基础目标仓位
        # ==================================================
        target_spxl = max_exposure * 0.5 * risk_multiplier
        target_tqqq = max_exposure * 0.5 * risk_multiplier

        # ==================================================
        # ⑦ 杠杆 ETF 专用保护
        # ==================================================
        self.profit_peak = max(self.profit_peak, equity)
        profit_dd = (self.profit_peak - equity) / self.profit_peak

        if profit_dd > 0.08:
            target_spxl *= 0.5
            target_tqqq *= 0.5

        def scale_by_extension(price, ema):
            dist = (price - ema) / ema
            if dist > 0.45:
                return 0.5
            elif dist > 0.30:
                return 0.7
            return 1.0

        target_spxl *= scale_by_extension(spxl_price, self.ema200_spxl.current.value)
        target_tqqq *= scale_by_extension(tqqq_price, self.ema200_tqqq.current.value)

        vol_spxl = self.history(self.spxl, 20, Resolution.DAILY)["close"].pct_change().std() * np.sqrt(252)
        vol_tqqq = self.history(self.tqqq, 20, Resolution.DAILY)["close"].pct_change().std() * np.sqrt(252)

        if vol_spxl > 0.60 and spxl_price < self.ema20_spxl.current.value:
            target_spxl = 0

        if vol_tqqq > 0.60 and tqqq_price < self.ema20_tqqq.current.value:
            target_tqqq = 0

        # ==================================================
        # ⑧ 执行
        # ==================================================
        self.set_holdings(self.spxl, target_spxl)
        self.set_holdings(self.tqqq, target_tqqq)

        self.log(
            f"VVIX={vvix:.1f} | DD={drawdown:.2%} | "
            f"EquityRatio={equity_ratio:.2f} | "
            f"SPXL={target_spxl:.2f} | TQQQ={target_tqqq:.2f}"
        )