probability.py

# src/tft_analyzer/probability.py

import json
from pathlib import Path

# --- Static Data ---

# Fallback constants used when the config file is missing or unreadable.
_FALLBACK_POOL_SIZES = {1: 29, 2: 22, 3: 18, 4: 12, 5: 10}
_FALLBACK_SHOP_ODDS = {
    1:  {1: 100, 2: 0,  3: 0,  4: 0,  5: 0},
    2:  {1: 100, 2: 0,  3: 0,  4: 0,  5: 0},
    3:  {1: 75,  2: 25, 3: 0,  4: 0,  5: 0},
    4:  {1: 55,  2: 30, 3: 15, 4: 0,  5: 0},
    5:  {1: 45,  2: 33, 3: 20, 4: 2,  5: 0},
    6:  {1: 25,  2: 40, 3: 30, 4: 5,  5: 0},
    7:  {1: 19,  2: 30, 3: 35, 4: 15, 5: 1},
    8:  {1: 18,  2: 25, 3: 32, 4: 22, 5: 3},
    9:  {1: 10,  2: 20, 3: 25, 4: 35, 5: 10},
    10: {1: 5,   2: 10, 3: 20, 4: 40, 5: 25},
}

# Initialize active constants with fallback values.
POOL_SIZES = dict(_FALLBACK_POOL_SIZES)
SHOP_ODDS = {k: dict(v) for k, v in _FALLBACK_SHOP_ODDS.items()}

def load_configurations(config_path: Path = Path(__file__).parent / "data" / "game_constants.json"):
    """Loads game constants from a JSON file if it exists. Falls back to hardcoded defaults."""
    global POOL_SIZES, SHOP_ODDS
    if config_path.exists():
        try:
            with open(config_path, 'r') as f:
                data = json.load(f)

                # Convert string keys from JSON back to integers for Python logic
                if "pool_sizes" in data:
                    POOL_SIZES = {int(k): v for k, v in data["pool_sizes"].items()}

                if "shop_odds" in data:
                    SHOP_ODDS = {
                        int(lvl): {int(t): prob for t, prob in odds.items()}
                        for lvl, odds in data["shop_odds"].items()
                    }
        except Exception as e:
            print(f"Warning: Could not load game constants from {config_path}. Using fallback defaults. Error: {e}")
            POOL_SIZES = dict(_FALLBACK_POOL_SIZES)
            SHOP_ODDS = {k: dict(v) for k, v in _FALLBACK_SHOP_ODDS.items()}

# Attempt to load configurations on module import
load_configurations()

def calculate_shop_odds(level: int, tier: int, units_gone_for_champion: int, total_tier_units_gone: int, num_champions_in_tier: int) -> float:
    """
    Calculates the probability of finding a specific unit in the next shop.

    Args:
        level: The player's current level (1-10).
        tier: The cost of the unit (1-5).
        units_gone_for_champion: The number of copies of that specific unit already taken from the pool.
        total_tier_units_gone: The total number of cards of that tier removed from the pool (including the specific unit).
        num_champions_in_tier: The number of unique champions in that tier.

    Returns:
        The probability (0.0 to 1.0) of the unit appearing in the next 5 shop slots.
    """
    if not (1 <= level <= 10):
        raise ValueError("Level must be between 1 and 10.")
    if not (1 <= tier <= 5):
        raise ValueError("Tier must be between 1 and 5.")

    # 1. Determine Tier Odds
    tier_odds = SHOP_ODDS.get(level, {}).get(tier, 0) / 100.0
    if tier_odds <= 0:
        return 0.0

    # 2. Determine Pool State
    total_copies_per_champ = POOL_SIZES.get(tier, 0)
    if total_copies_per_champ == 0:
        return 0.0

    # 3. Calculate Remaining Copies of Desired Unit
    remaining_copies = total_copies_per_champ - units_gone_for_champion
    if remaining_copies <= 0:
        return 0.0

    # 4. Calculate Total Remaining Cards in Tier Pool
    total_pool_size_initial = total_copies_per_champ * num_champions_in_tier
    current_pool_size = total_pool_size_initial - total_tier_units_gone

    if current_pool_size <= 0:
        return 0.0

    # 5. Calculate Probability per Slot
    # P(slot = unit) = P(slot = tier) * P(unit | tier)
    prob_unit_given_tier = remaining_copies / current_pool_size
    prob_per_slot = tier_odds * prob_unit_given_tier

    # 6. Calculate Probability for Shop (5 slots)
    # P(at least one) = 1 - (1 - P_slot)^5
    prob_shop = 1 - (1 - prob_per_slot) ** 5

    return prob_shop

def calculate_shop_odds_simple(level: int, tier: int, units_gone_for_champion: int) -> float:
    """
    Calculates the probability of finding a specific unit in the next shop, using a simplified model.

    Args:
        level: The player's current level (1-10).
        tier: The cost of the unit (1-5).
        units_gone_for_champion: The number of copies of that specific unit already taken from the pool.

    Returns:
        The probability (0.0 to 1.0) of the unit appearing in the next 5 shop slots.
    """
    if not (1 <= level <= 10):
        raise ValueError("Level must be between 1 and 10.")
    if not (1 <= tier <= 5):
        raise ValueError("Tier must be between 1 and 5.")

    tier_odds = SHOP_ODDS.get(level, {}).get(tier, 0) / 100.0
    if tier_odds == 0:
        return 0.0
        
    total_copies = POOL_SIZES.get(tier, 0)
    remaining_copies = total_copies - units_gone_for_champion

    if remaining_copies <= 0:
        return 0.0

    # Simplified: Assume the pool of other champions in the tier is full.
    # This is not accurate, but it's a starting point for V1.
    # Total cards in tier = (Num champs in tier * copies per champ)
    # We are missing `total_units_gone_from_tier`
    # Let's just use a placeholder for now for the denominator.
    
    # Let's use the formula from the prompt, interpreting `Total_Tier_Pool` as the
    # total number of champion cards of that tier *remaining* in the pool.
    # This requires knowing all units gone from the tier.
    
    # For V1, let's make a simplifying assumption: the number of *other* champions
    # gone from the pool is negligible.
    
    # Note: This simple function doesn't know about the total pool size, so we approximate.
    approx_champs_in_tier = 13 
    total_cards_in_tier_at_start = approx_champs_in_tier * total_copies
    
    # Let's assume `units_gone_for_champion` is the only reduction in the pool size.
    current_total_cards_in_tier = total_cards_in_tier_at_start - units_gone_for_champion
    
    if current_total_cards_in_tier <= 0:
        return 0.0
        
    # Probability of a card from this tier being the one we want
    p_champ_if_tier_is_hit = remaining_copies / current_total_cards_in_tier
    
    # Probability of a single shop slot containing the desired champion
    p_slot = tier_odds * p_champ_if_tier_is_hit
    
    # Probability of hitting at least one in 5 slots
    p_hit_in_shop = 1 - (1 - p_slot) ** 5
    
    return p_hit_in_shop

# Example usage:
if __name__ == '__main__':
    # Ensure data is loaded
    load_configurations()

    # What are the odds of finding a 4-cost unit at level 7 if 3 copies are gone?
    level = 7
    tier = 4
    units_gone = 3
    
    odds = calculate_shop_odds_simple(level, tier, units_gone)
    # Test the new accurate function with dummy pool data
    # Assume 0 other units gone, and 12 champs in tier
    odds = calculate_shop_odds(level, tier, units_gone, units_gone, 12)
    print(f"Level: {level}, Tier: {tier}, Units Gone: {units_gone}")
    print(f"Probability of finding the unit in the next shop: {odds:.2%}")

    # Level 8, 5-cost, 0 gone
    level = 8
    tier = 5
    units_gone = 0
    odds = calculate_shop_odds_simple(level, tier, units_gone)
    print(f"\nLevel: {level}, Tier: {tier}, Units Gone: {units_gone}")
    print(f"Probability of finding the unit in the next shop: {odds:.2%}")