import unittest

import numpy as np
import pandas as pd

from research.lab.futures_stock_signals import configurations, overnight_exposure, risk_exposure


class FuturesTimingTests(unittest.TestCase):
    def setUp(self):
        dates = pd.bdate_range("2020-01-01", periods=180)
        t = np.arange(len(dates))
        equity = 100 * np.exp(.001 * t + .002 * np.sin(t))
        rates = 100 * np.exp(.0001 * t + .0002 * np.sin(t))
        self.prices = pd.DataFrame({"ES=F": equity, "NQ=F": equity, "ZN=F": rates,
                                    "SPY": equity, "QQQ": equity, "IEF": rates}, index=dates)

    def test_future_prices_cannot_change_earlier_decisions(self):
        changed = self.prices.copy()
        changed.iloc[110:, :3] *= .7
        for config in configurations():
            original, known = risk_exposure(self.prices, config)
            modified, modified_known = risk_exposure(changed, config)
            last = 110 + config.get("lag", 1)
            with self.subTest(config=config["name"]):
                pd.testing.assert_series_equal(original.iloc[:last], modified.iloc[:last])
                pd.testing.assert_series_equal(known.iloc[:last], modified_known.iloc[:last])

    def test_same_day_futures_crash_is_not_usable_for_entry_or_preceding_night(self):
        config = next(c for c in configurations() if c["name"] == "futures_equity_shock_default")
        changed = self.prices.copy()
        changed.iloc[110:, :2] *= .7
        entry, _ = risk_exposure(changed, config)
        self.assertEqual(entry.iloc[110], 1.0)
        self.assertEqual(entry.iloc[111], .5)
        night = overnight_exposure(entry)
        self.assertEqual(night.iloc[111], 1.0)
        self.assertEqual(night.iloc[112], .5)

    def test_identical_etf_and_futures_inputs_have_identical_rules(self):
        configs = {c["name"]: c for c in configurations()}
        for name, config in configs.items():
            if not name.startswith("futures_"):
                continue
            with self.subTest(config=name):
                a, av = risk_exposure(self.prices, config)
                b, bv = risk_exposure(self.prices, configs[name.replace("futures_", "etf_", 1)])
                pd.testing.assert_series_equal(a, b, check_names=False)
                pd.testing.assert_series_equal(av, bv, check_names=False)

    def test_unavailable_prices_are_not_backward_filled_from_future(self):
        config = next(c for c in configurations() if c["name"] == "futures_equity_shock_default")
        missing = self.prices.copy()
        missing.iloc[110, :2] = np.nan
        exposure, known = risk_exposure(missing, config)
        self.assertFalse(known.iloc[111])
        self.assertEqual(exposure.iloc[111], 1.0)
        prefix, prefix_known = risk_exposure(missing.iloc[:112], config)
        pd.testing.assert_series_equal(exposure.iloc[:112], prefix)
        pd.testing.assert_series_equal(known.iloc[:112], prefix_known)


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