from scipy.stats import binom, hypergeom
import pandas as pd

shop_odds = {
    3: [0.75, 0.25, 0, 0, 0],
    4: [0.55, 0.3, 0.15, 0, 0],
    5: [0.45, 0.33, 0.20, 0.02, 0],
    6: [0.3, 0.4, 0.25, 0.05, 0],
    7: [0.19, 0.3, 0.4, 0.1, 0.01],
    8: [0.18, 0.25, 0.32, 0.22, 0.03],
    9: [0.15, 0.2, 0.25, 0.3, 0.1],
    10: [0.05, 0.1, 0.2, 0.4, 0.25],
}

pool_sizes = {
    1: 30,
    2: 25,
    3: 18,
    4: 10,
    5: 9,
    6: 9,
}

num_champs = {
    1: 14,
    2: 13,
    3: 13,
    4: 12,
    5: 8,
    6: 3,
}


def prob_x(x, attempts, cost_odd, total_pool, target_pool):
    """
    Calculates P(X in all attempts)
    """
    prob = 0

    for k in range(attempts + 1):  # k = number of costs in all attempts
        P_D_k = binom.pmf(k, attempts, cost_odd)  # Binomial probability of drawing k times
        P_x_given_k = hypergeom.pmf(
            x, total_pool, target_pool, k
        )  # Hypergeometric probability of exactly x hits in k draws

        prob += P_D_k * P_x_given_k  # Weighted sum over all possible k values

    return prob


def prob_x_first_shop_y_after(x, y, attempts, cost_odd, total_pool, target_pool, shop_slots=5):
    """
    Calculates P(X in first shop & Y in remaining shops)
    """
    prob = 0
    other_shops = attempts - shop_slots
    for k_first_shop in range(shop_slots + 1):  # k_first_shop = number of costs in the first shop
        # Binomial probability of drawing k_first_shop times in first shop
        P_D_k_first_shop = binom.pmf(k_first_shop, shop_slots, cost_odd)
        # Probability of exactly x hits in k_first_shop draws
        P_x_given_k_first_shop = hypergeom.pmf(x, total_pool, target_pool, k_first_shop)

        for k_after in range(other_shops + 1):  # k_next_35 = number of costs afterwards
            # Binomial probability of drawing k_after times after
            P_D_k_after = binom.pmf(k_after, other_shops, cost_odd)
            # Probability of exactly y hits in k_after draws
            P_y_given_k_after = hypergeom.pmf(y, total_pool - k_first_shop, target_pool - x, k_after)

            prob += P_D_k_first_shop * P_x_given_k_first_shop * P_D_k_after * P_y_given_k_after

    return prob

def prob_y_after_x_first_shop(x: int, y: int, attempts:int , cost_odd: int, total_pool: int, target_pool: int, shop_slots: int=5):
    """
    Calculates P(Y in remaining shops | X in first shop)
    """
    a_n_b = prob_x_first_shop_y_after(x, y, attempts, cost_odd, total_pool, target_pool, shop_slots)
    b = prob_x(x, shop_slots, cost_odd, total_pool, target_pool)
    return a_n_b / b


shop_slots = 5
num_players = 8
attempts = num_players * shop_slots
cost = 4
level = 9
shop_odd = shop_odds[level]
cost_odd = shop_odd[cost - 1]
target_pool = pool_sizes[cost]
num_champ = num_champs[cost]
total_pool = target_pool * num_champ

result_table = pd.DataFrame(index=range(5), columns=range(5))

for x in range(5):
    for y in range(5):
        prob = prob_y_after_x_first_shop(x, y, attempts, cost_odd, total_pool, target_pool)
        result_table.at[x, y] = prob

# Add a row at the top for prob_x(y, attempts-shop_slots)
top_row = [prob_x(y, attempts - shop_slots, cost_odd, total_pool, target_pool) for y in range(5)]
result_table.loc[-1] = top_row
result_table.index = result_table.index + 1
result_table = result_table.sort_index()

# Label the rows
result_table.index = ["P(other Y)"] + [f"P(other Y | me {x})" for x in range(5)]

# # # Cast to string for printing
str_table = result_table.astype(float).map(lambda x: f"{x:.1%}")
print(str_table)


def expected_cost(y, cost_odd, total_pool, target_pool):
    p = 0
    for x in range(1):
        p += prob_x(x+1, 5, cost_odd, total_pool - y, target_pool - y)
    expected_rolls = 1 / p
    return expected_rolls * 2


expected_cost_per_Y = {col: expected_cost(col, cost_odd, total_pool, target_pool) for col in result_table.columns}
result_table["Expected Cost"] = result_table.mul(pd.Series(expected_cost_per_Y), axis=1).sum(axis=1)
print(result_table)