from __future__ import annotations

import unittest

import numpy as np

from cbpred.experiments import (
    _CROSS_COLUMNS,
    _GLOBAL_COLUMNS,
    _cross_sectional_context,
    experiment_variants,
    feature_names,
)
from cbpred.pipeline import aggregate_8h, compute_indicators

from tests.test_pipeline import make_hour_rows


class ExperimentFeatureTests(unittest.TestCase):
    def test_variant_feature_counts(self):
        counts = {variant.name: len(feature_names(variant)) for variant in experiment_variants()}
        self.assertEqual(counts["baseline"], 182)
        self.assertEqual(counts["compact7"], 164)
        self.assertEqual(counts["compact6"], 146)
        self.assertEqual(counts["rsi_sma"], 236)
        self.assertEqual(counts["global6"], 188)
        self.assertEqual(counts["cross4"], 186)
        self.assertEqual(counts["global_cross"], 192)

    def test_cross_sectional_context_is_ranked_by_timestamp(self):
        frames = {
            f"COIN{index}": aggregate_8h(make_hour_rows(0, 8 * 90, index * 0.1))
            for index in range(4)
        }
        indicators = {symbol: compute_indicators(frame) for symbol, frame in frames.items()}
        profiles = {
            symbol: {
                "asset_atr_price_top20_median": 0.05,
                "asset_quote_volume_top20_median_log": 12.0 + index,
            }
            for index, symbol in enumerate(frames)
        }
        global_context, cross_context = _cross_sectional_context(indicators, profiles)
        self.assertEqual(list(global_context.columns), list(_GLOBAL_COLUMNS))
        self.assertEqual(set(next(iter(cross_context.values())).columns), set(_CROSS_COLUMNS))
        self.assertEqual(len(cross_context), 4)
        for frame in cross_context.values():
            values = frame.iloc[20].to_numpy(dtype=float)
            self.assertTrue(np.isfinite(values).all())
            self.assertTrue(np.all((0.0 < values) & (values <= 1.0)))


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