import unittest

import numpy as np
import pandas as pd

from strategies.permanent import TrendRiderV4


class TrendRiderV4Tests(unittest.TestCase):
    def test_v4_builds_capped_multi_asset_portfolio(self):
        dates = pd.bdate_range("2023-01-02", periods=320)
        trend = np.linspace(100.0, 180.0, len(dates))
        prices = pd.DataFrame(
            {
                "SPY": trend,
                "QQQ": trend * 1.10,
                "SSO": trend * 1.55,
                "QLD": trend * 1.65,
                "UPRO": trend * 2.00,
                "TQQQ": trend * 2.20,
                "SHY": np.linspace(100.0, 103.0, len(dates)),
                "IEF": np.linspace(100.0, 104.0, len(dates)),
                "TLT": np.linspace(100.0, 105.0, len(dates)),
                "GLD": np.linspace(100.0, 115.0, len(dates)),
                "DBC": np.linspace(90.0, 105.0, len(dates)),
            },
            index=dates,
        )

        strategy = TrendRiderV4(max_single_weight=0.35, max_leveraged_weight=0.50)
        weights = strategy.generate_signals(prices)
        active = weights[weights.sum(axis=1) > 0.99]

        self.assertFalse(active.empty)
        self.assertLessEqual(active.max(axis=1).max(), 0.350001)
        self.assertGreaterEqual((active > 0.001).sum(axis=1).min(), 4)
        leveraged = [c for c in ["SSO", "QLD", "UPRO", "TQQQ"] if c in active.columns]
        self.assertLessEqual(active[leveraged].sum(axis=1).max(), 0.500001)
        self.assertTrue(np.allclose(active.sum(axis=1), 1.0))


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