def test_convert_surge__on_rolled_dice_pool():
    rolled_dice_pool = dce.RolledDicePool.from_rolled_dice_list(
        [
            dse.RolledDouse(
                douse=att_dse.BlackAttackDouse(),
                symbol=sym.Surge(),
            ),
            dse.RolledDouse(
                douse=att_dse.RedAttackDouse(),
                symbol=sym.Hit(),
            ),
            dse.RolledDouse(
                douse=att_dse.WhiteAttackDouse(),
                symbol=sym.Crit(),
            ),
        ]
    )
    rolled_dice_pool_target = dce.RolledDicePool.from_rolled_dice_list(
        [
            dse.RolledDouse(
                douse=att_dse.RedAttackDouse(),
                symbol=sym.Hit(),
            ),
            dse.RolledDouse(
                douse=att_dse.BlackAttackDouse(),
                symbol=sym.Crit(),
            ),
            dse.RolledDouse(
                douse=att_dse.WhiteAttackDouse(),
                symbol=sym.Crit(),
            ),
        ]
    )
    assert actions.convert_all_surges_to_crit.on(rolled_dice_pool) == rolled_dice_pool_target
def test_convert_surge__on_rolled_douse__surge__wrong_conversion():
    rolled_douse = dse.RolledDouse(
        douse=att_dse.BlackAttackDouse(),
        symbol=sym.Surge(),
    )
    with pytest.raises(ValueError):
        actions.convert_all_surges_to_block.on(rolled_douse)
def run_examples_with_dice_conversion_before_calculation():
    rebel_troopers = dce.DicePool.from_dice_list(
        [
            att_dse.BlackAttackDouse(),
            att_dse.BlackAttackDouse(),
            att_dse.BlackAttackDouse(),
            att_dse.BlackAttackDouse(),
        ]
    )
    print('rebel troopers - before surges')
    pprint(count.CountSymbols().on(rebel_troopers).get_best(3))
    print('rebel troopers - after surge conversion')
    rebel_troopers = actions.convert_all_surges_to_hit.on(rebel_troopers)
    pprint(count.CountSymbols().on(rebel_troopers).get_best(3))
    print('rebel troopers - crits counted as hits')
    rebel_troopers = actions.convert_all_crits_to_hits_for_calculation.on(rebel_troopers)
    pprint(count.CountSymbols().on(rebel_troopers).get_best(3))
def run_examples_with_critical_and_aim__critical_applied_only_before_aim():
    rebel_troopers = dce.DicePool.from_dice_list(
        [
            att_dse.BlackAttackDouse(),
            att_dse.BlackAttackDouse(),
            att_dse.BlackAttackDouse(),
            att_dse.BlackAttackDouse(),
        ]
    )
    print('rebel troopers')
    pprint(count.CountSymbols().on(rebel_troopers).get_best(3))

    rebel_troopers = actions.critical_2.on(rebel_troopers)
    print('rebel troopers - after critical 2')
    pprint(count.CountSymbols().on(rebel_troopers).get_best(3))

    rebel_troopers = actions.aim.on(rebel_troopers)
    print('rebel troopers - after aim')
    pprint(count.CountSymbols().on(rebel_troopers).get_best(3))

    rebel_troopers = actions.convert_all_crits_to_hits_for_calculation.on(rebel_troopers)
    rebel_troopers = actions.remove_unconverted_surge_results_attack.on(rebel_troopers)
    print('rebel troopers - count')
    pprint(count.CountSymbols().on(rebel_troopers).get_best(10))
def test_roll__douse():
    douse = att_dse.BlackAttackDouse()
    rerolled_douse = rll.Roll(roll_policy=rll_pol.RollPolicyAttack()).on(douse)

    crit_prob = fractions.Fraction(1, 8) * fractions.Fraction(3, 2)
    assert rerolled_douse.as_dict[dse.RolledDouse(douse,
                                                  sym.Crit())] == crit_prob

    hit_prob = fractions.Fraction(3, 8) * fractions.Fraction(3, 2)
    assert rerolled_douse.as_dict[dse.RolledDouse(douse,
                                                  sym.Hit())] == hit_prob

    blank_prob = fractions.Fraction(3, 8) * fractions.Fraction(1, 2)
    assert rerolled_douse.as_dict[dse.RolledDouse(douse,
                                                  sym.Blank())] == blank_prob
def test_aim__black_attack_douse(output, prob):
    douse = att_dse.BlackAttackDouse()
    rerolled_douse_prob_dist: pd.ProbabilityDistribution = actions.aim.on(douse)
    print(rerolled_douse_prob_dist)
    assert rerolled_douse_prob_dist.as_dict[output] == prob
import pytest

from legion_dice_probs import actions
from legion_dice_probs.stochastic_objects import douse as dse
from legion_dice_probs.stochastic_objects import attack_douse as att_dse
from legion_dice_probs.stochastic_states import symbol as sym
from legion_dice_probs.stochastic_states import symbols as syms
from prob_dist_api import probability_distribution as pd


@pytest.mark.parametrize(
    "output, prob",
    (
            (
                    dse.RolledDouse(douse=att_dse.BlackAttackDouse(), symbol=sym.Blank()),
                    fractions.Fraction(3, 8) * fractions.Fraction(1, 2),  # 3 / 16
            ),
            (
                    dse.RolledDouse(douse=att_dse.BlackAttackDouse(), symbol=sym.Surge()),
                    fractions.Fraction(1, 8) * fractions.Fraction(1, 2),  # 1 / 16
            ),
            (
                    dse.RolledDouse(douse=att_dse.BlackAttackDouse(), symbol=sym.Hit()),
                    fractions.Fraction(3, 8) * (1 + fractions.Fraction(1, 2)),  # 9 / 16
            ),
            (
                    dse.RolledDouse(douse=att_dse.BlackAttackDouse(), symbol=sym.Crit()),
                    fractions.Fraction(1, 8) * (1 + fractions.Fraction(1, 2)),  # 3 / 16
            ),
    )
def run_examples_with_raw_dice_probabilities():
    rebel_troopers = dce.DicePool.from_dice_list(
        [
            att_dse.BlackAttackDouse(),
            att_dse.BlackAttackDouse(),
            att_dse.BlackAttackDouse(),
            att_dse.BlackAttackDouse(),
        ]
    )
    print('rebel troopers')
    pprint(count.CountSymbols().on(rebel_troopers).get_best(3))

    rebel_troopers_with_z6_with_trooper = dce.DicePool.from_dice_list(
        [
            att_dse.WhiteAttackDouse(),
            att_dse.WhiteAttackDouse(),
            att_dse.WhiteAttackDouse(),
            att_dse.WhiteAttackDouse(),
            att_dse.WhiteAttackDouse(),
            att_dse.WhiteAttackDouse(),
            att_dse.BlackAttackDouse(),
            att_dse.BlackAttackDouse(),
            att_dse.BlackAttackDouse(),
            att_dse.BlackAttackDouse(),
            att_dse.BlackAttackDouse(),
        ]
    )
    print('rebel troopers with Z6 and trooper')
    pprint(count.CountSymbols().on(rebel_troopers_with_z6_with_trooper).get_best(3))

    droid_toopers_with_gun_with_trooper = dce.DicePool.from_dice_list(
        [
            att_dse.WhiteAttackDouse(),
            att_dse.WhiteAttackDouse(),
            att_dse.WhiteAttackDouse(),
            att_dse.WhiteAttackDouse(),
            att_dse.WhiteAttackDouse(),
            att_dse.WhiteAttackDouse(),
            att_dse.WhiteAttackDouse(),
            att_dse.WhiteAttackDouse(),
            att_dse.BlackAttackDouse(),
            att_dse.BlackAttackDouse(),
            att_dse.BlackAttackDouse(),
        ]
    )
    print('droid troopers')
    pprint(count.CountSymbols().on(droid_toopers_with_gun_with_trooper).get_best(3))
    rebel_troopers = actions.aim.on(rebel_troopers)
    print('rebel troopers - after aim')
    pprint(count.CountSymbols().on(rebel_troopers).get_best(3))

    rebel_troopers = actions.convert_all_crits_to_hits_for_calculation.on(rebel_troopers)
    rebel_troopers = actions.remove_unconverted_surge_results_attack.on(rebel_troopers)
    print('rebel troopers - count')
    pprint(count.CountSymbols().on(rebel_troopers).get_best(10))


if __name__ == '__main__':
    # run_examples_with_raw_dice_probabilities()
    # run_examples_with_dice_conversion_before_calculation()
    # run_examples_with_critical_and_aim__critical_applied_only_before_aim()

    rebel_troopers = dce_wtoks.DicePoolWithTokens.from_dice_list_and_tokens(
        [
            att_dse.BlackAttackDouse(),
            att_dse.BlackAttackDouse(),
            # att_dse.BlackAttackDouse(),
            # att_dse.BlackAttackDouse(),
        ],
        tokens=toks_spec.TokensAttack.from_attack_tokens(n_aim=1),
    )
    print('rebel troopers')
    pprint(count.CountSymbols().on(rebel_troopers).get_best(3))
    rebel_troopers = actions.aim.on(rebel_troopers)
    print('rebel troopers - after aim')
    pprint(rebel_troopers)
    pprint(count.CountSymbols().on(rebel_troopers).get_best(3))
])))
def test_convert_symbols__conversion_order__symbols(n_converts, symbols,
                                                    symbols_target):
    conversion = conv_syms.ConvertSymbols(
        conversion_policy=conv_pol.ConversionPolicyAttackAnyToCrit(),
        conversion_limit=n_converts,
    )
    assert conversion.on(syms.Symbols.from_symbols_list(
        symbols)) == syms.Symbols.from_symbols_list(symbols_target)


@pytest.mark.parametrize(
    "conversion_policy, n_converts, rolled_dice, rolled_dice_target", (
        (conv_pol.ConversionPolicyAttackAnyToCrit(), 2, [
            dse.RolledDouse(att_dse.RedAttackDouse(), sym.Blank()),
            dse.RolledDouse(att_dse.BlackAttackDouse(), sym.Surge()),
            dse.RolledDouse(att_dse.WhiteAttackDouse(), sym.Hit()),
        ], [
            dse.RolledDouse(att_dse.RedAttackDouse(), sym.Crit()),
            dse.RolledDouse(att_dse.BlackAttackDouse(), sym.Crit()),
            dse.RolledDouse(att_dse.WhiteAttackDouse(), sym.Hit()),
        ]),
        (conv_pol.ConversionPolicyAttackAnyToCrit(), 1, [
            dse.RolledDouse(att_dse.RedAttackDouse(), sym.Blank()),
            dse.RolledDouse(att_dse.BlackAttackDouse(), sym.Surge()),
            dse.RolledDouse(att_dse.WhiteAttackDouse(), sym.Hit()),
        ], [
            dse.RolledDouse(att_dse.RedAttackDouse(), sym.Crit()),
            dse.RolledDouse(att_dse.BlackAttackDouse(), sym.Surge()),
            dse.RolledDouse(att_dse.WhiteAttackDouse(), sym.Hit()),
        ]),
def test_convert_surge__on_dice_pool():
    dice_pool = dce.DicePool.from_dice_list(
        [
            att_dse.WhiteAttackDouse(),
            att_dse.RedAttackDouse(),
            att_dse.BlackAttackDouse(),
        ]
    )
    result_all_surge = dce.RolledDicePool.from_rolled_dice_list(
        [
            dse.RolledDouse(
                douse=att_dse.WhiteAttackDouse(),
                symbol=sym.Surge()
            ),
            dse.RolledDouse(
                douse=att_dse.BlackAttackDouse(),
                symbol=sym.Surge()
            ),
            dse.RolledDouse(
                douse=att_dse.RedAttackDouse(),
                symbol=sym.Surge()
            ),
        ]
    )
    assert dice_pool.get_probability_distribution().as_dict[result_all_surge] == fractions.Fraction(
        numerator=1,
        denominator=8 ** 3,
    ), "No conversion sanity check."
    result_all_hit = dce.RolledDicePool.from_rolled_dice_list(
        [
            dse.RolledDouse(
                douse=att_dse.WhiteAttackDouse(),
                symbol=sym.Hit()
            ),
            dse.RolledDouse(
                douse=att_dse.BlackAttackDouse(),
                symbol=sym.Hit()
            ),
            dse.RolledDouse(
                douse=att_dse.RedAttackDouse(),
                symbol=sym.Hit()
            ),
        ]
    )
    assert dice_pool.get_probability_distribution().as_dict[result_all_hit] == fractions.Fraction(
        numerator=1 * 3 * 5,
        denominator=8 ** 3,
    ), "No conversion sanity check."

    dice_pool_converted_to_hit_prob_dist = actions.convert_all_surges_to_hit.on(dice_pool)
    assert dice_pool_converted_to_hit_prob_dist.as_dict[result_all_surge] == 0
    assert dice_pool_converted_to_hit_prob_dist.as_dict[result_all_hit] == fractions.Fraction(
        numerator=2 * 4 * 6,
        denominator=8 ** 3,
    )

    dice_pool_converted_to_hit_prob_dist = actions.convert_all_surges_to_crit.on(dice_pool)
    assert dice_pool_converted_to_hit_prob_dist.as_dict[result_all_surge] == 0
    assert dice_pool_converted_to_hit_prob_dist.as_dict[result_all_hit] == fractions.Fraction(
        numerator=1 * 3 * 5,
        denominator=8 ** 3,
    )
from legion_dice_probs.stochastic_objects import douse as dse
from legion_dice_probs.stochastic_states import symbol as sym


@pytest.mark.parametrize("policy, douse_cls, symbol, target_prob_dist", (
    (
        rll_pol.RollPolicyAttack(),
        att_dse.WhiteAttackDouse,
        sym.Blank(),
        att_dse.WhiteAttackDouse().get_probability_distribution(),
    ),
    (
        rll_pol.RollPolicyAttack(),
        att_dse.BlackAttackDouse,
        sym.Surge(),
        att_dse.BlackAttackDouse().get_probability_distribution(),
    ),
    (
        rll_pol.RollPolicyAttackBlanksOnly(),
        att_dse.RedAttackDouse,
        sym.Surge(),
        dse.RolledDouse(att_dse.RedAttackDouse(),
                        sym.Surge()).get_probability_distribution(),
    ),
    (
        rll_pol.RollPolicyAttackCritFish(),
        att_dse.RedAttackDouse,
        sym.Hit(),
        att_dse.RedAttackDouse().get_probability_distribution(),
    ),
))