import unittest

import numpy as np
import pandas as pd

from primethai.config import AppConfig, StrategyConfig
from primethai.scanner import scan_candidates
from primethai.setups import setup_for


class SetupScannerTests(unittest.TestCase):
    def setUp(self):
        self.previous = pd.Series({
            "open": 10.0, "high": 11.0, "low": 9.5, "close": 10.0,
            "ema21_high": 10.0,
        })
        self.row = pd.Series({
            "open": 10.1, "high": 10.8, "low": 9.7, "close": 10.5,
            "ema21_high": 10.2, "ema21_close": 10.0, "ema21_low": 9.0,
            "trend_state": "NEUTRAL", "higher_low": False,
            "extension_atr": 0.5,
        })

    def test_scanner_setup_classifies_base_breakout_and_ema_reclaim(self):
        setup, entry, stop = setup_for(self.row, self.previous, prior_high=10.4)
        self.assertEqual(setup, "Base Breakout")
        self.assertEqual((entry, stop), (10.5, 9.0))

        # Keep the current close below yesterday's high so the EMA reclaim is distinct.
        prior = self.previous.copy()
        prior["high"] = 11.0
        prior["close"] = 10.0
        self.row["close"] = 10.5
        setup, _, _ = setup_for(self.row, prior, prior_high=11.0)
        self.assertEqual(setup, "21DMA High Reclaim")

    def test_extension_filter_blocks_scanner_setups(self):
        self.row["extension_atr"] = 1.2
        setup, entry, stop = setup_for(self.row, self.previous, prior_high=10.4)
        self.assertIsNone(setup)
        self.assertIsNone(entry)
        self.assertIsNone(stop)

    def test_watch_only_candidate_has_no_position_size(self):
        index = pd.bdate_range("2026-01-01", periods=25)
        frame = pd.DataFrame({
            "open": 10.0, "high": 10.6, "low": 9.6, "close": 10.0,
            "ema21_high": 9.5, "ema21_close": 9.4, "ema21_low": 9.0,
            "trend_state": "NEUTRAL", "higher_low": False,
            "extension_atr": 0.5, "avg_value_turnover20": 50_000_000.0,
            "rs_percentile": 80.0, "volume_contracting": False,
            "tightness_score": 50.0, "weekly_trend_state": "NEUTRAL",
            "rs_1m": 0.0, "rs_3m": 0.0, "rs_12m": 0.0,
            "rs_1m_vs_set": 0.0, "rs_3m_vs_set": 0.0,
            "rs_12m_vs_set": 0.0, "rs_1m_vs_set100ew": 0.0,
            "rs_3m_vs_set100ew": 0.0, "rs_12m_vs_set100ew": 0.0,
            "52w_proximity_pct": 90.0,
        }, index=index)
        market = pd.DataFrame({
            "regime": ["CONFIRMED_UPTREND"], "new_risk_allowed": [True],
        }, index=[index[-1]])
        config = AppConfig(strategy=StrategyConfig())

        result = scan_candidates({"TEST": frame}, market, config=config)

        self.assertEqual(result.candidates.iloc[0]["entry_type"], "WATCH")
        self.assertEqual(result.candidates.iloc[0]["shares"], 0)
        self.assertEqual(result.candidates.iloc[0]["risk_pct"], 0.0)

        frame.loc[index[-1], "close"] = 10.8
        frame.loc[index[-1], "high"] = 11.0
        market.loc[index[-1], "new_risk_allowed"] = False
        market.loc[index[-1], "regime"] = "CORRECTION"

        blocked = scan_candidates({"TEST": frame}, market, config=config)

        candidate = blocked.candidates.iloc[0]
        self.assertEqual(candidate["setup"], "Watch")
        self.assertEqual(candidate["entry_type"], "WATCH")
        self.assertEqual(candidate["shares"], 0)
        self.assertEqual(candidate["risk_pct"], 0.0)
        self.assertIn("market regime blocks new entries", candidate["reasons"])

        market.attrs["members_by_date"] = {pd.Timestamp(index[-1]): ["OTHER"]}
        nonmember = scan_candidates({"TEST": frame}, market, config=config)
        self.assertTrue(nonmember.candidates.empty)

        market.attrs["members_by_date"] = {pd.Timestamp(index[-1]): ["TEST"]}
        stale = scan_candidates({"TEST": frame.iloc[:-1]}, market, config=config)
        self.assertTrue(stale.candidates.empty)


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