scan_handler.py

#!/usr/bin/env python3
"""
Modified scan handler for Tag Collector Game with Enkephalin rewards for rare tags
"""

import streamlit as st
import torch
import time
import math
from tag_categories import (
    get_unlocked_categories,
    get_collection_power_level
)
from game_constants import TAG_CURRENCY_NAME, ENKEPHALIN_CURRENCY_NAME, ENKEPHALIN_ICON, RARITY_LEVELS, TAG_POWER_BONUSES
from state_manager import update_game_state 

def enhanced_scan_button_handler(image_path):
    """
    Enhanced scan button handler with Enkephalin rewards for rare tag discoveries
    """
    try:
        # Run inference
        results = st.session_state.model.predict(
            image_path=image_path,
            threshold=st.session_state.threshold
        )
        
        # Get probabilities
        probs = results['refined_probabilities'][0]  # Remove batch dimension
        
        # All categories are now unlocked
        unlocked_categories = get_unlocked_categories(st.session_state)
        
        # No limit on tags
        max_tags = float('inf')  # Unlimited
        
        # Get collection power bonus
        power_info = get_collection_power_level(st.session_state)
        coin_multiplier = power_info['current_level']['coin_bonus']
        
        # Store all detected tags for display
        all_tags = {}  # Format: {category: [(tag, probability, rarity), ...]}
        found_tags = []  # Tags that pass the threshold
        total_currency_earned = 0
        total_enkephalin_earned = 0
        new_tag_count = 0
        
        # Minimum probability to display (can be lower than threshold)
        min_display_prob = 0.1
        
        # Track any newly completed combinations
        previously_unlocked = set(st.session_state.unlocked_combinations) if hasattr(st.session_state, 'unlocked_combinations') else set()
        
        # Collect all tags that are above threshold
        candidates = []
        for idx in range(len(probs)):
            prob_value = probs[idx].item()
            
            # Only process tags above threshold
            if prob_value >= st.session_state.threshold:
                tag, category = st.session_state.model.dataset.get_tag_info(idx)
                candidates.append({
                    "idx": idx,
                    "tag": tag,
                    "probability": prob_value,
                    "category": category
                })
        
        # Sort by probability (highest first)
        candidates.sort(key=lambda x: x["probability"], reverse=True)
        
        # Add all tags to appropriate categories (for display purposes)
        for idx in range(len(probs)):
            prob_value = probs[idx].item()
            
            # Skip if below display threshold
            if prob_value < min_display_prob:
                continue
            
            # Get tag information
            tag, category = st.session_state.model.dataset.get_tag_info(idx)
            
            # Determine rarity
            rarity = determine_tag_rarity(tag, prob_value, category)
            
            # Add to category in all_tags
            if category not in all_tags:
                all_tags[category] = []
            
            # Check if this tag is above threshold (collected)
            is_collected = prob_value >= st.session_state.threshold
            
            # Add to all_tags with appropriate status
            if is_collected:
                all_tags[category].append((tag, prob_value, rarity, "collected"))
            else:
                all_tags[category].append((tag, prob_value, rarity, "displayed"))
        
        # Process all candidates for collection and currency
        for candidate in candidates:
            tag = candidate["tag"]
            prob_value = candidate["probability"]
            category = candidate["category"]
            
            # Determine rarity
            rarity = determine_tag_rarity(tag, prob_value, category)
            
            # Add tag to collection and get currency
            currency_earned, enkephalin_earned, is_new_tag = add_tag_to_collection(tag, prob_value, category, coin_multiplier)
            
            # Count new tags
            if is_new_tag:
                new_tag_count += 1
            
            # Add to total currency earned
            if currency_earned > 0:
                total_currency_earned += currency_earned
                
            # Add to total enkephalin earned
            if enkephalin_earned > 0:
                total_enkephalin_earned += enkephalin_earned
            
            # Add to found tags list
            found_tags.append({
                "tag": tag,
                "probability": prob_value,
                "category": category,
                "currency": currency_earned,
                "enkephalin": enkephalin_earned,
                "rarity": rarity,
                "is_new": is_new_tag
            })
        
        # Sort each category by probability
        for category in all_tags:
            all_tags[category].sort(key=lambda x: x[1], reverse=True)
        
        # Update game stats - IMPORTANT: Direct modification to ensure updates 
        # are properly tracked by Streamlit's session state mechanism
        st.session_state.game_stats["images_processed"] += 1
        st.session_state.game_stats["total_tags_found"] += len(found_tags)
        
        # Update enkephalin stats if we earned any
        if total_enkephalin_earned > 0:
            if "enkephalin_generated" not in st.session_state.game_stats:
                st.session_state.game_stats["enkephalin_generated"] = 0
            st.session_state.game_stats["enkephalin_generated"] += total_enkephalin_earned
            
        # Update tag power ONCE
        threshold_bonus, coin_bonus = calculate_tag_power()
        st.session_state.tag_power_bonus = threshold_bonus
        st.session_state.coin_multiplier = coin_bonus

        # Force update of state_version to trigger UI refresh
        if 'state_version' not in st.session_state:
            st.session_state.state_version = 0
        st.session_state.state_version += 1
        
        # Store results in session state
        st.session_state.current_scan = {
            "all_tags": all_tags,
            "found_tags": found_tags,
            "threshold": st.session_state.threshold,
            "total_currency_earned": total_currency_earned,
            "total_enkephalin_earned": total_enkephalin_earned,
            "new_tag_count": new_tag_count,
        }
        
        # # Save game state explicitly
        # save_game_state()
        
        return True
            
    except Exception as e:
        st.error(f"Error scanning image: {str(e)}")
        import traceback
        st.code(traceback.format_exc())
        return False

def determine_tag_rarity(tag, probability, category):
    """
    Determine the rarity of a tag based on the metadata if available,
    otherwise use a simplified probability-based approach.
    Always returns a valid rarity level.
    """
    # Set a default rarity (safeguard against returning None)
    default_rarity = "Canard"
    
    try:
        # If we have the rarity metadata loaded, use it
        if hasattr(st.session_state, 'tag_rarity_metadata') and st.session_state.tag_rarity_metadata:
            # Check if this tag exists in our metadata
            if tag in st.session_state.tag_rarity_metadata:
                tag_info = st.session_state.tag_rarity_metadata[tag]
                
                # Handle both new and old format
                if isinstance(tag_info, dict) and "rarity" in tag_info:
                    # New format with rarity and sample_count
                    rarity = tag_info["rarity"]
                else:
                    # Old format where tag_info is the rarity string directly
                    rarity = tag_info
                    
                # Verify the rarity is valid
                if rarity in RARITY_LEVELS:
                    return rarity
            
            # Special handling for rating_* and year_* tags
            if tag.startswith("rating_") or tag.startswith("year_"):
                for metadata_tag in st.session_state.tag_rarity_metadata:
                    if metadata_tag == tag:
                        tag_info = st.session_state.tag_rarity_metadata[metadata_tag]
                        if isinstance(tag_info, dict) and "rarity" in tag_info:
                            rarity = tag_info["rarity"]
                        else:
                            rarity = tag_info
                            
                        if rarity in RARITY_LEVELS:
                            return rarity
    
    except Exception as e:
        # Log the error but don't crash
        print(f"Error determining rarity for tag {tag}: {str(e)}")
    
    # Return the default rarity as a fallback
    return default_rarity

def get_enkephalin_reward(rarity):
    """
    Determine the Enkephalin reward based on tag rarity
    
    Args:
        rarity (str): The tag rarity level
        
    Returns:
        int: Amount of Enkephalin to award
    """
    # Get enkephalin reward from TAG_POWER_BONUSES
    if rarity in TAG_POWER_BONUSES:
        return TAG_POWER_BONUSES[rarity]["enkephalin_reward"]
    return 0

def add_tag_to_collection(tag, probability, category, coin_multiplier=1.0):
    """
    Add a tag to the user's collection and award currency and enkephalin for new discoveries.
    """
    try:
        # Get the tag's rarity
        rarity = determine_tag_rarity(tag, probability, category)
        
        # Check if this is a new tag discovery
        is_new_tag = tag not in st.session_state.collected_tags
        
        # Check if tag was previously sacrificed
        was_sacrificed = hasattr(st.session_state, 'sacrificed_tags') and tag in st.session_state.sacrificed_tags
        
        # Get base currency value for this rarity
        base_currency_value = RARITY_LEVELS[rarity]["value"]
        
        # Basic debug info - always print this
        print(f"Processing tag: {tag}, rarity: {rarity}, is_new: {is_new_tag}, base_value: {base_currency_value}")
        
        # Initialize currency earned
        currency_earned = 0
        enkephalin_earned = 0
        
        # Record the tag acquisition
        timestamp = time.strftime("%H:%M:%S")
        
        # Update tag collection
        if tag in st.session_state.collected_tags:
            # If already collected, just increment the count, no new currency
            st.session_state.collected_tags[tag]["count"] += 1
        else:
            # New tag discovery - create entry in collection
            st.session_state.collected_tags[tag] = {
                "count": 1,
                "rarity": rarity,
                "category": category,
                "discovery_time": timestamp
            }
            
            # Award currency and enkephalin only if not previously sacrificed and is a new tag
            if not was_sacrificed and is_new_tag:
                # Calculate tag power bonuses
                tag_power_bonus, base_coin_multiplier = calculate_tag_power()
                
                # Now print the debug info AFTER calculating base_coin_multiplier
                print(f"  Base multiplier: {base_coin_multiplier}, Collection multiplier: {coin_multiplier}")
                
                # Apply collection power multiplier
                total_multiplier = base_coin_multiplier * coin_multiplier
                
                # Check if category is mastered for additional bonus
                if hasattr(st.session_state, 'mastered_categories') and category in st.session_state.mastered_categories:
                    from tag_categories import CATEGORY_MASTERY_BONUS
                    total_multiplier += CATEGORY_MASTERY_BONUS.get('coin_multiplier', 0)
                
                # Award currency with multiplier for new discoveries
                currency_earned = int(base_currency_value * total_multiplier)
                
                # Award enkephalin for new tag discoveries based on rarity
                enkephalin_earned = get_enkephalin_reward(rarity)
                
                # Apply any enkephalin bonus from achievements if it exists
                if hasattr(st.session_state, 'enkephalin_bonus') and st.session_state.enkephalin_bonus > 0:
                    enkephalin_earned = int(enkephalin_earned * (1 + st.session_state.enkephalin_bonus))
                
                print(f"  Final multiplier: {total_multiplier}, Final award: {currency_earned} {TAG_CURRENCY_NAME}, {enkephalin_earned} {ENKEPHALIN_CURRENCY_NAME}")
                
                # Add to player's currency
                st.session_state.tag_currency += currency_earned
                
                # Add to player's enkephalin if earned
                if enkephalin_earned > 0:
                    st.session_state.enkephalin += enkephalin_earned
                
                # Track in stats
                st.session_state.game_stats["total_currency_earned"] += currency_earned
                
                # Add to history for new discoveries
                if not hasattr(st.session_state, 'tag_history'):
                    st.session_state.tag_history = []
                    
                st.session_state.tag_history.append({
                    "tag": tag,
                    "rarity": rarity,
                    "value": currency_earned,
                    "enkephalin": enkephalin_earned,
                    "time": timestamp,
                    "is_new": True
                })
        
        # Always increment total tags found count
        st.session_state.game_stats["total_tags_found"] += 1
        
        # Update tag power calculations with the new tag
        threshold_bonus, coin_bonus = calculate_tag_power()
        st.session_state.tag_power_bonus = threshold_bonus
        st.session_state.coin_multiplier = coin_bonus

        update_game_state(tag_currency=st.session_state.tag_currency, enkephalin=st.session_state.enkephalin)

        # Return the earned currency, enkephalin, and new tag status
        return currency_earned, enkephalin_earned, is_new_tag
    
    except Exception as e:
        # Log the error
        print(f"Error adding tag '{tag}' to collection: {str(e)}")
        import traceback
        traceback.print_exc()
        
        # Return safe defaults to prevent crashing
        return 0, 0, False

def calculate_tag_power():
    """
    Calculate tag power bonuses based on collected tags (not sacrificed ones).
    Tags that are sacrificed no longer contribute to tag power until recollected.
    
    Returns:
        (threshold_reduction, coin_multiplier)
    """
    if not hasattr(st.session_state, 'collected_tags'):
        return 0, 1.0
    
    total_threshold_reduction = 0
    total_coin_multiplier = 1.0
    
    # Calculate bonuses from individual tags
    for tag, info in st.session_state.collected_tags.items():
        rarity = info["rarity"]
        if rarity in TAG_POWER_BONUSES:
            # Add threshold reduction (scales with tag count, but with diminishing returns)
            count = info["count"]
            # Use logarithmic scaling to prevent excessive bonuses from farming
            scaling_factor = 1 + (math.log(count) / 2) if count > 1 else 1
            
            # Apply the bonus
            total_coin_multiplier += TAG_POWER_BONUSES[rarity]["coin_multiplier"] * scaling_factor
    
    # Apply collection power coin bonus
    if hasattr(st.session_state, 'collection_power_coin_bonus'):
        total_coin_multiplier *= st.session_state.collection_power_coin_bonus
    
    return total_threshold_reduction, total_coin_multiplier