from __future__ import annotations

import sys
import unittest
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
if str(ROOT) not in sys.path:
    sys.path.insert(0, str(ROOT))

from tools.build_distribution_risk_score_v2 import calculate_distribution_risk


class TestDistributionRiskParity(unittest.TestCase):

    def test_distribution_risk_parity_scenarios(self):
        # Scenario 1: Smart Money Outflow only
        row_1 = {
            "close": 10000,
            "ma20": 10000,
            "frg_5d": -100,
            "inst_5d": -200,
        }
        res_1 = calculate_distribution_risk(row_1, kospi_ret_5d=0.0)
        self.assertEqual(res_1["distribution_risk_score"], 30)
        self.assertIn("smart_money_outflow", res_1["reason_codes"])
        self.assertEqual(res_1["anti_distribution_state"], "PASS")

        # Scenario 2: High upper wick and low flow credit under priceAboveMa20
        row_2 = {
            "close": 12000,
            "ma20": 10000, # priceAboveMa20 = True
            "high": 15000,
            "low": 10000,
            # upperWickRatio = (15000-12000)/5000 = 3000/5000 = 0.60 >= 0.45
            "flow_credit": 0.35, # flow_credit < 0.40
        }
        res_2 = calculate_distribution_risk(row_2, kospi_ret_5d=0.0)
        self.assertIn("upper_wick_distribution", res_2["reason_codes"])
        self.assertIn("flow_credit_low", res_2["reason_codes"])
        # score = 15 (upper wick) + 20 (flow credit low) = 35
        self.assertEqual(res_2["distribution_risk_score"], 35)

        # Scenario 3: Trim Review threshold (score >= 55)
        row_3 = {
            "close": 10000,
            "ma20": 10000,
            "frg_5d": -100,
            "inst_5d": -200, # +30
            "flow_credit": 0.30, # +20
            "volume": 70,
            "avg_volume_5d": 100, # volume < 80% of avg_vol_5d -> +20
        }
        res_3 = calculate_distribution_risk(row_3, kospi_ret_5d=0.0)
        # score = 30 + 20 + 20 = 70 (BLOCK_BUY)
        self.assertEqual(res_3["distribution_risk_score"], 70)
        self.assertEqual(res_3["anti_distribution_state"], "BLOCK_BUY")

    def test_distribution_risk_early_warning_signals(self):
        # Early warning signal 1: New high volume contraction
        row_4 = {
            "close": 9800,
            "high_52w": 10000, # close >= 97% of 52w high -> nearNewHigh = True
            "volume": 70,
            "avg_volume_5d": 100, # volume < 80% -> +12
        }
        res_4 = calculate_distribution_risk(row_4, kospi_ret_5d=0.0)
        self.assertIn("new_high_volume_contraction", res_4["reason_codes"])
        self.assertEqual(res_4["pre_distribution_warning"], "EARLY_WARNING")

        # Early warning signal 2: Surge weak flow
        row_5 = {
            "close": 10000,
            "ret_5d": 6.0, # ret5d >= 5
            "flow_credit": 0.40, # flow_credit < 0.45 -> +10
        }
        res_5 = calculate_distribution_risk(row_5, kospi_ret_5d=0.0)
        self.assertIn("surge_weak_flow", res_5["reason_codes"])
        self.assertEqual(res_5["pre_distribution_warning"], "EARLY_WARNING")


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