from AlgorithmImports import *
from Portfolio.RiskParityPortfolioOptimizer import RiskParityPortfolioOptimizer


class RiskParityBugDemoAlgorithm(QCAlgorithm):
    """
    Demonstrates two bugs in RiskParityPortfolioOptimizer:

    Bug 1 — Bounds are silently ignored during optimization:
        Newton-CG doesn't support bounds, so maximum_weight is never
        enforced during the solve. The constraint only applies as a
        post-hoc np.clip(), which is too late to affect the solution.

    Bug 2 — Clipping after normalization breaks sum-to-1:
        weights = np.clip(solver["x"] / sum, min_w, max_w)
        If any weight is clipped, the portfolio no longer sums to 1.0
        and the risk parity property is silently violated.
    """

    MAX_WEIGHT = 0.20  # Tight cap

    def initialize(self):
        self.set_start_date(self.end_date - timedelta(5 * 365))
        self.set_cash(100_000)
        self.settings.rebalance_portfolio_on_security_changes = False

        tickers = ["IEF", "TLT", "SPY", "EFA", "EEM", "JPXN", "XLK"]
        self._securities = [self.add_equity(t, Resolution.DAILY) for t in tickers]
        # Build a ticker -> Symbol map for set_holdings calls later
        self._symbol_by_ticker = {s.symbol.value: s.symbol for s in self._securities}

        # Positional args only — keyword names are not exposed in the stubs
        self._constrained_optimizer   = RiskParityPortfolioOptimizer(0.0, self.MAX_WEIGHT)
        self._unconstrained_optimizer = RiskParityPortfolioOptimizer(1e-5, 1.0)

        self.set_warm_up(timedelta(252))

    def on_warmup_finished(self):
        time_rule = self.time_rules.at(8, 0)
        self.schedule.on(self.date_rules.month_start("SPY"), time_rule, self._rebalance)
        if self.live_mode:
            self._rebalance()
        else:
            self.schedule.on(self.date_rules.today, time_rule, self._rebalance)

    def _rebalance(self):
        history = self.history(self._securities, 253, Resolution.DAILY)
        if history.empty:
            return

        returns = history["close"].unstack(level=0).pct_change().dropna()
        if returns.empty or returns.shape[1] < 2:
            return

        # columns are ticker strings after unstack; NOT Symbol objects
        tickers = list(returns.columns)

        unconstrained = self._unconstrained_optimizer.optimize(returns)
        constrained   = self._constrained_optimizer.optimize(returns)

        weight_sum    = float(np.sum(constrained))
        any_violation = any(float(w) > self.MAX_WEIGHT + 1e-6 for w in constrained)

        #  Bug 1: did any weight exceed the declared bound? 
        self.plot("Bug1 - Bound Violations", "bound_violated",  int(any_violation))
        self.plot("Bug1 - Bound Violations", "max_weight_cap",  self.MAX_WEIGHT)

        #  Bug 2: constrained weights must sum to 1.0 
        self.plot("Bug2 - Weight Sum", "constrained_sum",   weight_sum)
        self.plot("Bug2 - Weight Sum", "expected_sum_1_0",  1.0)

        for i, ticker in enumerate(tickers):
            c_w = float(constrained[i])
            u_w = float(unconstrained[i])
            # ticker is already a plain string — no .value needed
            self.plot("Weights - Constrained",   ticker, c_w)
            self.plot("Weights - Unconstrained",  ticker, u_w)

            if c_w > self.MAX_WEIGHT + 1e-6:
                self.log(
                    f"[BUG 1] {self.time:%Y-%m-%d}: {ticker} weight={c_w:.4f} "
                    f"exceeds max_weight={self.MAX_WEIGHT} — Newton-CG ignored the bound."
                )

        if abs(weight_sum - 1.0) > 1e-4:
            self.log(
                f"[BUG 2] {self.time:%Y-%m-%d}: constrained weights sum to "
                f"{weight_sum:.6f}, not 1.0 — post-clip normalization is broken."
            )

        # Trade on unconstrained weights (known-good baseline)
        for i, ticker in enumerate(tickers):
            symbol = self._symbol_by_ticker.get(ticker)
            if symbol:
                self.set_holdings(symbol, float(unconstrained[i]))