import math
import unittest
from datetime import date

from quant60.features import (
    TrainOnlyFeaturePreprocessor,
    build_executable_label,
    compute_price_features,
)
from quant60.synthetic import generate_synthetic_bars


class FeatureTests(unittest.TestCase):
    def setUp(self):
        self.bars = generate_synthetic_bars(
            ["600000.XSHG"],
            start_date="2024-01-02",
            trading_days_count=150,
            seed=7,
            include_market_events=False,
        )

    def test_feature_as_of_does_not_change_when_future_bars_are_appended(self):
        decision = self.bars[99].trading_date
        early = compute_price_features(self.bars[:100], as_of=decision)
        with_future = compute_price_features(self.bars, as_of=decision)
        self.assertEqual(early.values, with_future.values)
        self.assertEqual(early.history_end, decision)

    def test_long_window_features_are_explicitly_missing(self):
        snapshot = compute_price_features(self.bars[:30])
        self.assertIsNone(snapshot.values["momentum_60"])
        self.assertIsNotNone(snapshot.values["momentum_20"])

    def test_preprocessor_uses_training_statistics_and_missing_indicator(self):
        processor = TrainOnlyFeaturePreprocessor(winsor_fraction=0).fit(
            [{"a": 1.0}, {"a": 2.0}, {"a": 3.0}]
        )
        transformed = processor.transform_one({"a": None})
        self.assertEqual(transformed["a__missing"], 1.0)
        self.assertAlmostEqual(transformed["a"], 0.0)
        extreme = processor.transform_one({"a": 1_000_000.0})
        expected = processor.transform_one({"a": 3.0})
        self.assertEqual(extreme["a"], expected["a"])

    def test_executable_label_uses_sessions_after_decision(self):
        decision_index = 80
        label = build_executable_label(
            self.bars,
            decision_date=self.bars[decision_index].trading_date,
            horizon_sessions=5,
        )
        self.assertEqual(label.status, "OK")
        self.assertEqual(
            label.entry_date,
            self.bars[decision_index + 1].trading_date,
        )
        self.assertEqual(
            label.exit_date,
            self.bars[decision_index + 6].trading_date,
        )
        self.assertTrue(math.isfinite(label.value))

    def test_suspended_entry_is_censored_not_dropped(self):
        bars = generate_synthetic_bars(
            ["600000.XSHG"],
            start_date=date(2024, 1, 2),
            trading_days_count=40,
            seed=7,
            include_market_events=True,
        )
        decision = bars[10].trading_date
        label = build_executable_label(
            bars,
            decision_date=decision,
            horizon_sessions=5,
        )
        self.assertEqual(label.status, "NON_EXECUTABLE_CENSORED")
        self.assertIsNone(label.value)


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