"""Contracts of the sleeve portfolio (us_three_sleeve): causal targets, sleeves
kept to their own columns, shares that drift like separate accounts, and the
live trader's session timing."""

import unittest
from unittest import mock

import numpy as np
import pandas as pd

import trader
from strategies.permanent import ETF_UNIVERSE
from strategies.robust_momentum import RobustMomentumStrategy
from strategies.qqq_trend import QQQVolTrend
from strategies.sleeve_portfolio import (SleevePortfolioStrategy, decisions, drifting_shares,
                                         us_momentum_qqq_trend, us_three_sleeve)


def _bars(n_days: int = 520, n_stocks: int = 30, seed: int = 1) -> dict[str, pd.DataFrame]:
    rng = np.random.default_rng(seed)
    idx = pd.bdate_range("2024-01-01", periods=n_days)
    cols = [f"S{i:02d}" for i in range(n_stocks)] + sorted(set(ETF_UNIVERSE))
    r = rng.normal(0.0005, 0.015, (n_days, len(cols)))
    close = pd.DataFrame(100 * np.exp(np.cumsum(r, axis=0)), index=idx, columns=cols)
    opens = close.shift(1).bfill() * np.exp(rng.normal(0, 0.004, close.shape))
    return {"close": close, "open": opens}


class SleevePortfolioTests(unittest.TestCase):
    def test_targets_on_a_day_ignore_that_day_and_later_bars(self):
        bars = _bars()
        strat = us_three_sleeve()
        cut = bars["close"].index[450]
        base = strat.generate_signals(bars)
        shocked = {f: b.copy() for f, b in bars.items()}
        for f in shocked:
            shocked[f].loc[cut:] *= 1.7
        again = strat.generate_signals(shocked)
        pd.testing.assert_frame_equal(base.loc[:cut], again.loc[:cut])

    def test_stock_sleeves_hold_no_etfs_and_the_book_is_not_levered(self):
        bars = _bars()
        strat = us_three_sleeve()
        weights = strat.generate_signals(bars)
        etfs = set(ETF_UNIVERSE)
        self.assertLessEqual(weights.sum(axis=1).max(), 1.0 + 1e-9)
        # After warm-up the stock part is two sleeves' worth of whole names only.
        late = weights.iloc[-1]
        stocks = late[[c for c in late.index if c not in etfs]]
        self.assertEqual(int((stocks > 0).sum()) <= 30, True)
        self.assertGreater(stocks.sum(), 0.5)
        _, shares = strat.target_weights(bars)
        self.assertTrue(np.allclose(shares.sum(axis=1), 1.0))

    def test_decisions_are_the_strategy_weights_one_row_earlier(self):
        bars = _bars()
        strat = RobustMomentumStrategy("low_risk", top_n=5, rebal_freq=5, buffer=1.0)
        stock = {f: b[[c for c in b.columns if c.startswith("S")]] for f, b in bars.items()}
        d = decisions(strat, stock)
        w = strat.generate_signals(stock)
        pd.testing.assert_frame_equal(d.iloc[:-1], w.shift(-1).iloc[:-1].fillna(0.0))
        self.assertAlmostEqual(d.iloc[-1].sum(), 1.0)

    def test_shares_drift_with_sleeve_returns_and_reset_past_the_band(self):
        # Two sleeves fully in one asset each; A doubles on day 1, B is flat.
        returns = np.array([[0.0, 0.0], [1.0, 0.0], [0.0, 0.0], [0.0, 0.0]])
        wa = np.array([[1.0, 0.0]] * 4)
        wb = np.array([[0.0, 1.0]] * 4)
        base = np.array([0.5, 0.5])
        month_start = np.array([False, False, False, True])
        shares = drifting_shares([wa, wb], returns, base, month_start, band=0.25)
        np.testing.assert_allclose(shares[1], [2 / 3, 1 / 3])     # drifted like two accounts
        np.testing.assert_allclose(shares[2], [2 / 3, 1 / 3])     # no reset mid-month
        np.testing.assert_allclose(shares[3], [0.5, 0.5])         # 33% off base: reset
        kept = drifting_shares([wa, wb], returns, base, month_start, band=0.5)
        np.testing.assert_allclose(kept[3], [2 / 3, 1 / 3])       # inside a 50% band

    def test_trader_session_targets_use_bars_through_the_previous_day(self):
        bars = _bars()
        strat = us_three_sleeve()
        days = bars["close"].index[-2:]
        full = strat.generate_signals(bars).loc[days]
        got = trader.bar_session_weights(strat, bars, days)
        pd.testing.assert_frame_equal(got, full)

    def test_trader_truncates_all_bar_fields_before_padding_sleeves(self):
        bars = _bars()
        days = bars["close"].index[-4:-2]
        stocks = [c for c in bars["close"] if c not in ETF_UNIVERSE]
        member = pd.DataFrame(True, index=bars["close"].index, columns=stocks)
        expected = us_three_sleeve().generate_signals({f: b.loc[:days[-1]] for f, b in bars.items()}).loc[days]
        with mock.patch.object(trader, "index_members", lambda m, d: member.reindex(d)):
            actual = trader.session_signals("us", "us_three_sleeve", list(bars["close"]), days,
                                            bars["close"], bars["open"])
        pd.testing.assert_frame_equal(actual, expected.reindex(columns=actual.columns))

    def test_registered_as_a_us_mixed_strategy(self):
        self.assertIsInstance(trader.STRATEGY_REGISTRY["us_three_sleeve"](), SleevePortfolioStrategy)
        self.assertEqual(set(trader.strategy_extra_tickers("us_three_sleeve")), set(ETF_UNIVERSE))
        self.assertTrue(trader.strategy_supports_market("us_three_sleeve", "us"))
        self.assertFalse(trader.strategy_supports_market("us_three_sleeve", "cn"))
        self.assertEqual(trader.strategy_fields("us_three_sleeve"), ("close", "open"))


def _qqq_bars(n_days: int = 520, drift: float = 0.0005, seed: int = 2) -> dict[str, pd.DataFrame]:
    """Stocks plus the combo's ETFs only (its live universe: S&P 500 + QQQ, TQQQ, SHY)."""
    other = [c for c in ETF_UNIVERSE if c not in ("TQQQ", "SHY")]
    bars = {f: b.drop(columns=other) for f, b in _bars(n_days, seed=seed).items()}
    rng = np.random.default_rng(seed + 1)
    idx = bars["close"].index
    q = pd.Series(100 * np.exp(np.cumsum(rng.normal(drift, 0.012, n_days))), index=idx)
    for f in bars:
        bars[f]["QQQ"] = q if f == "close" else q.shift(1).bfill()
    return bars


class QQQTrendTests(unittest.TestCase):
    def test_weights_on_a_day_ignore_that_day_and_later_closes(self):
        close = _qqq_bars()["close"]
        base = QQQVolTrend().generate_signals(close)
        shocked = close.copy()
        cut = close.index[400]
        shocked.loc[cut:, "QQQ"] *= 0.5
        pd.testing.assert_frame_equal(base.loc[:cut], QQQVolTrend().generate_signals(shocked).loc[:cut])

    def test_out_below_the_average_and_capped_above_it(self):
        up = _qqq_bars(drift=0.002)["close"]
        down = _qqq_bars(drift=-0.002)["close"]
        w_up = QQQVolTrend(cap=1.5).generate_signals(up).iloc[260:]
        w_down = QQQVolTrend().generate_signals(down).iloc[260:]
        self.assertLessEqual(w_up["TQQQ"].max(), 0.5 + 1e-12)
        self.assertGreater(w_up["TQQQ"].mean(), 0.3)
        np.testing.assert_allclose(w_up[["TQQQ", "SHY"]].sum(axis=1), 1.0)
        self.assertLess(w_down["TQQQ"].mean(), 0.05)

    def test_momentum_qqq_trend_is_a_us_mixed_strategy(self):
        self.assertIsInstance(trader.STRATEGY_REGISTRY["us_momentum_qqq_trend"](), SleevePortfolioStrategy)
        self.assertEqual(set(trader.strategy_extra_tickers("us_momentum_qqq_trend")), {"QQQ", "SHY", "TQQQ"})
        self.assertFalse(trader.strategy_supports_market("us_momentum_qqq_trend", "cn"))
        bars = _qqq_bars()
        strat = us_momentum_qqq_trend()
        w = strat.generate_signals(bars)
        self.assertLessEqual(w.sum(axis=1).max(), 1.0 + 1e-9)
        stocks = [c for c in w.columns if c[1:].isdigit()]
        self.assertLessEqual(int((w[stocks].iloc[-1] > 0).sum()), 10)
        self.assertGreater(w[["TQQQ", "SHY"]].iloc[-1].sum(), 0.3)
        days = bars["close"].index[-2:]
        pd.testing.assert_frame_equal(trader.bar_session_weights(strat, bars, days), w.loc[days])


if __name__ == "__main__":
    unittest.main()