"""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

import numpy as np
import pandas as pd

import trader
from strategies.permanent import ETF_UNIVERSE
from strategies.robust_momentum import RobustMomentumStrategy
from strategies.sleeve_portfolio import (SleevePortfolioStrategy, decisions, drifting_shares,
                                         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_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"))


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