"""Tests for the capital-allocation simulation engine."""

from __future__ import annotations

import unittest

from app.simulation import Candidate, allocate_capital


class AllocationTest(unittest.TestCase):
    def make_candidates(self):
        # dividend + high score
        return [
            Candidate("A", 10.0, 1.5, True, 5.0),
            Candidate("B", 20.0, 1.2, True, 3.0),
            Candidate("C", 5.0, 1.0, False, 0.0),
            Candidate("D", 2.0, 0.5, False, 0.0),
            Candidate("E", 50.0, 0.9, True, 8.0),
        ]

    def test_allocates_three_buckets(self):
        res = allocate_capital(200_000, self.make_candidates())
        bucket_notional = res.bucket_notional
        self.assertIn(1, bucket_notional)
        self.assertIn(2, bucket_notional)
        self.assertIn(3, bucket_notional)
        # total invested + unallocated == capital
        self.assertAlmostEqual(res.invested + res.unallocated_cash, 200_000, places=2)
        # every order is a multiple of 100
        for o in res.orders:
            self.assertEqual(o.qty % 100, 0)
            self.assertGreaterEqual(o.qty, 100)

    def test_min_100_shares_respected(self):
        # tiny capital: bucket1 50% still must afford 100 shares
        res = allocate_capital(2_000, self.make_candidates())
        for o in res.orders:
            self.assertGreaterEqual(o.qty, 100)
            self.assertEqual(o.notional, o.qty * o.price)
        # never over-invest beyond capital
        self.assertLessEqual(res.invested, 2_000)

    def test_bucket3_excludes_already_bought(self):
        cands = [Candidate("A", 10.0, 1.5, True, 5.0),
                 Candidate("B", 10.0, 0.1, True, 8.0)]
        res = allocate_capital(300_000, cands)
        # A (bought in bucket1) must NOT also appear in bucket3
        bucket3_syms = [o.symbol for o in res.orders if o.bucket == 3]
        bucket1_syms = [o.symbol for o in res.orders if o.bucket == 1]
        overlap = set(bucket3_syms) & set(bucket1_syms)
        self.assertEqual(overlap, set())

    def test_bucket3_ranks_by_yield_ignoring_score(self):
        # The spec: bucket3 = highest dividend yield, IGNORING score.
        # A low-score but high-yield name must rank above a high-score low-yield name.
        cands = [
            Candidate("HIGH_SCORE", 10.0, 8.0, True, 2.0),   # score 8, yield 2%
            Candidate("HIGH_YIELD", 10.0, 0.1, True, 7.0),   # score 0.1, yield 7%
            Candidate("MID", 10.0, 5.0, True, 3.0),
        ]
        # Big capital so bucket1 consumes only the top score name, leaving
        # HIGH_YIELD (not HIGH_SCORE) to be the bucket3 top pick.
        res = allocate_capital(1_000_000, cands)
        b3 = [o.symbol for o in res.orders if o.bucket == 3]
        # HIGH_YIELD (7%) should be selected in bucket3 before HIGH_SCORE (2%)
        self.assertIn("HIGH_YIELD", b3)
        if "HIGH_SCORE" in b3:
            hi = b3.index("HIGH_SCORE")
            hy = b3.index("HIGH_YIELD")
            self.assertLess(hy, hi)

    def test_invalid_capital_raises(self):
        with self.assertRaises(Exception):
            allocate_capital(0, self.make_candidates())
        with self.assertRaises(Exception):
            allocate_capital(-100, self.make_candidates())

    def test_no_candidates_raises(self):
        with self.assertRaises(Exception):
            allocate_capital(100_000, [])

    def test_cash_fallback_when_nothing_fits(self):
        # only very high-priced names, tiny capital -> cannot buy 100 shares -> cash
        cands = [Candidate("X", 1000.0, 2.0, True, 3.0)]
        res = allocate_capital(50_000, cands)  # 50% = 25k, can't buy 100*1000
        self.assertEqual(res.orders, [])
        self.assertAlmostEqual(res.unallocated_cash, 50_000, places=2)


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