import math
from functools import lru_cache
import matplotlib.pyplot as plt

def count_max_less_than(m, n, k):

    @lru_cache(maxsize=None)
    def dp(people_left, cats_left):
        if cats_left == 0:
            return 1 if people_left == 0 else 0
        
        total = 0
        # Put x people in the current category (0 <= x < k)
        max_in_cat = min(people_left, k - 1)
        for x in range(max_in_cat + 1):
            # Choose which x people go into this category
            total += math.comb(people_left, x) * dp(people_left - x, cats_left - 1)
        return total
    
    return dp(n, m)

N_PEOPLE = 10
WEEKDAYS = 7
MONTHS = 12

k_values = list(range(2, N_PEOPLE + 1))
week_probs = []
month_probs = []

print("k    Weekday prob    Month prob")
print("-" * 35)

for k in k_values:
    # Count assignments where max frequency < k
    week_bad = count_max_less_than(WEEKDAYS, N_PEOPLE, k)
    month_bad = count_max_less_than(MONTHS, N_PEOPLE, k)
    
    # Total possible assignments
    week_total = WEEKDAYS ** N_PEOPLE
    month_total = MONTHS ** N_PEOPLE
    
    # P(at least k) = 1 - P(all frequencies < k)
    week_prob = 1 - week_bad / week_total
    month_prob = 1 - month_bad / month_total
    
    week_probs.append(week_prob)
    month_probs.append(month_prob)
    
    print(f"{k:2d}   {week_prob:.6f}    {month_prob:.6f}")

# ---- Plot ----
plt.figure(figsize=(10, 6))
plt.plot(k_values, week_probs, 'o-', label='Weekday (7 days)', linewidth=2, markersize=8)
plt.plot(k_values, month_probs, 's-', label='Month (12 months)', linewidth=2, markersize=8)
plt.xlabel('k (at least k people share the same category)', fontsize=12)
plt.ylabel('Probability', fontsize=12)
plt.title(f'Probability of at least k people sharing the same weekday vs. month\n(among {N_PEOPLE} people)', fontsize=14)
plt.legend(fontsize=12)
plt.grid(True, linestyle='--', alpha=0.6)
plt.xticks(k_values)
plt.ylim(0, 1.05)  # leave a little headroom

# Optional: uncomment the next line for a log-scale y-axis to see the tail better
# plt.yscale('log')

plt.tight_layout()
plt.show()