game.py

#!/usr/bin/env python3
"""
Tag Collector Game
A gamified version of the Image Tagger application where users collect tags
to earn currency that can be used to lower thresholds and discover rarer tags.
"""

import streamlit as st
import numpy as np
import os
import sys
import tempfile
import time
import math
import json
import random
from PIL import Image
from collections import Counter
import torch
import importlib
import traceback

import torch
import torch.nn as nn
from safetensors.torch import load_file
import timm

from huggingface_hub import hf_hub_download
from typing import Tuple, Dict, Any, Optional
from pathlib import Path
import numpy as np  # you use np in ONNXModelWrapper


# Import storage system
import tag_storage
import state_manager

# Import constants and utility modules
from game_constants import (
    TAG_CURRENCY_NAME, TAG_ANIMATIONS, TAG_POWER_BONUSES, ENKEPHALIN_CURRENCY_NAME, ENKEPHALIN_ICON, RARITY_LEVELS, ACHIEVEMENTS, 
    STARTING_THRESHOLD, MIN_THRESHOLD,
    THRESHOLD_UPGRADES
)
from tag_categories import (
    TAG_CATEGORIES, initialize_progression_system, get_unlocked_categories,
    get_max_detectable_tags, get_collection_power_level
)
from tag_mosaic import display_tag_mosaic, RevealMosaic
from series_mosaics import display_series_mosaics
from scan_handler import enhanced_scan_button_handler
from library_system import initialize_library_system, display_library_extraction
from dev_tools import display_dev_tools

# =============================================================================
# SET PAGE CONFIG AS THE FIRST STREAMLIT COMMAND
# =============================================================================
st.set_page_config(layout="wide", page_title="Camie Collector", page_icon="🎮")

# =============================================================================
# DATA LOADING FUNCTIONS
# =============================================================================

# ---- CONFIG ----
MODEL_REPO = "Camais03/camie-tagger-v2"
ONNX_MODEL_FILE = "camie-tagger-v2.onnx"
METADATA_FILE = "camie-tagger-v2-metadata.json"

@st.cache_data
def load_tag_rarity_metadata():
    """
    Load the tag rarity metadata from the JSON file.
    Returns the tag_rarity dictionary or None if the file doesn't exist.
    
    Returns:
        dict: Tag rarity metadata or None if not found
    """
    try:
        parent_dir = "data/tag_rarity_metadata.json"
        
        metadata_path = os.path.join("data/tag_rarity_metadata.json")
        if os.path.exists(metadata_path):
            with open(metadata_path, 'r') as f:
                metadata = json.load(f)
            print(f"Loaded tag rarity metadata for {len(metadata['tag_rarity'])} tags")
            return metadata["tag_rarity"]
        else:
            print(f"Warning: {metadata_path} not found.")
            print("No metadata file found.")
            return None
    except Exception as e:
        print(f"Error loading tag rarity metadata: {str(e)}")
        traceback.print_exc()
        return None

# =============================================================================
# MODEL LOADING UTILITIES
# =============================================================================

def is_windows():
    """Check if the system is Windows"""
    import platform
    return platform.system() == "Windows"

class TaggerTorch(nn.Module):
    def __init__(self, backbone_name="vit_base_patch16_384", img_size=512, num_tags=70527, normalize=True):
        super().__init__()
        # num_classes=0 -> return features; we add our own head
        self.backbone = timm.create_model(backbone_name, pretrained=False, num_classes=0, img_size=img_size)
        in_features = self.backbone.num_features  # 768 for vit_base_patch16_384
        self.head = nn.Linear(in_features, num_tags)

        # Most ViT taggers expect ImageNet normalization; keep it configurable
        self.normalize = normalize
        if self.normalize:
            self.register_buffer("mean", torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1))
            self.register_buffer("std",  torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1))

    def forward(self, x):
        if self.normalize:
            x = (x - self.mean) / self.std
        feats = self.backbone.forward_features(x)   # [B, C] or [B, tokens, C]
        if feats.ndim == 3:                          # if tokens, take CLS
            feats = feats[:, 0, :]
        return self.head(feats)

def build_torch_model_from_hub_safetensors(safetensors_path, metadata, backbone="vit_base_patch16_384"):
    """Build TaggerTorch model from safetensors file downloaded from Hub"""
    dataset_info = metadata.get("dataset_info", {})
    num_tags = int(dataset_info.get("total_tags", 70527))
    img_size = int(metadata.get("model_info", {}).get("img_size", 512))
    
    model = TaggerTorch(backbone_name=backbone, img_size=img_size, num_tags=num_tags, normalize=True)
    sd = load_file(safetensors_path)
    sd = _remap_backbone_keys(sd)

    # pull out tag embedding/bias if present and later copy into the linear head
    te_w = sd.pop("tag_embedding.weight", sd.pop("module.tag_embedding.weight", None))
    te_b = sd.pop("tag_bias", sd.pop("module.tag_bias", None))

    # load backbone etc.
    missing, unexpected = model.load_state_dict(sd, strict=False)
    print("[load] missing:", len(missing), "keys")
    print("[load] unexpected:", len(unexpected), "keys")

    # copy tag embedding → head
    with torch.no_grad():
        if te_w is not None and te_w.shape == model.head.weight.shape:
            model.head.weight.copy_(te_w)
            print("[load] copied tag_embedding.weight → head.weight")
        if te_b is not None and model.head.bias is not None and te_b.shape == model.head.bias.shape:
            model.head.bias.copy_(te_b)
            print("[load] copied tag_bias → head.bias")

    return model

def _remap_backbone_keys(sd):
    """Remap keys for safetensors loading"""
    out = {}
    for k, v in sd.items():
        if k.startswith("module."): k = k[7:]

        # collapse ImageTagger → TaggerTorch vit paths
        if k.startswith("backbone.vit."):
            k = "backbone." + k[len("backbone.vit."):]
        elif k.startswith("vit."):
            k = "backbone." + k[len("vit."):]
        elif k.startswith(("pos_embed","patch_embed.","blocks.","norm.","cls_token")):
            k = "backbone." + k

        out[k] = v
    return out

class ONNXModelWrapper:
    """
    - predict(image_path=..., threshold=...) -> {"refined_probabilities": np.ndarray[N, C]}
    - dataset.get_tag_info(idx) -> (tag, category)
    - Keeps your signature compatible, ignores 'threshold' in the wrapper.
    """
    def __init__(self, session, metadata: dict):
        self.session = session
        self.metadata = metadata or {}
        dataset_info = self.metadata.get("dataset_info", {})
        self.total_tags = dataset_info.get("total_tags", 0)

        # idx <-> tag mapping
        tag_mapping = dataset_info.get("tag_mapping", {})
        if "idx_to_tag" in tag_mapping:
            self.idx_to_tag = {int(k): v for k, v in tag_mapping["idx_to_tag"].items()}
        elif "tag_to_idx" in tag_mapping:
            t2i = tag_mapping["tag_to_idx"]
            self.idx_to_tag = {v: k for k, v in t2i.items()}
        else:
            self.idx_to_tag = {}

        self.tag_to_category = self.metadata.get("tag_to_category", {})
        if not self.tag_to_category:
            # Try to get from dataset_info.tag_mapping.tag_to_category (correct path)
            if "tag_mapping" in dataset_info and "tag_to_category" in dataset_info["tag_mapping"]:
                self.tag_to_category = dataset_info["tag_mapping"]["tag_to_category"]
            # Fallback to direct path in case structure varies
            elif "tag_to_category" in dataset_info:
                self.tag_to_category = dataset_info["tag_to_category"]

        # Compatibility shim for scan_handler: model.dataset.get_tag_info(...)
        self.dataset = self

        # Preprocess config (defaults match your prior function)
        pre = self.metadata.get("preprocess", {})
        self.input_h, self.input_w = _coerce_hw(pre.get("input_size", 512))  # default 512
        mean = pre.get("mean", [0.485, 0.456, 0.406])
        std  = pre.get("std",  [0.229, 0.224, 0.225])
        self.mean = np.array(mean, dtype="float32")
        self.std  = np.array(std,  dtype="float32")

        # Letterbox pad color from mean
        self.pad_rgb = tuple(int(round(c * 255)) for c in self.mean.tolist())

        # Cache input name
        self.input_name = self.session.get_inputs()[0].name

    # --- API expected by scan_handler ---
    def predict(self, image_tensor=None, image_path=None, pil_image=None, threshold=None):
        if image_tensor is not None:
            x = _ensure_numpy_float32(image_tensor)
        elif pil_image is not None:
            x = self._preprocess_pil(pil_image)
        elif image_path is not None:
            x = self._preprocess_path(image_path)
        else:
            raise ValueError("predict expects one of image_tensor, pil_image, or image_path")

        outputs = self.session.run(None, {self.input_name: x})
        logits = outputs[0].astype("float32")
        # Convert to probs (sigmoid) if not already [0,1]
        if np.nanmax(logits) > 1.0 or np.nanmin(logits) < 0.0:
            probs = 1.0 / (1.0 + np.exp(-logits))
        else:
            probs = logits
        return {"refined_probabilities": probs}

    def get_tag_info(self, idx: int):
        idx = int(idx)
        tag = self.idx_to_tag.get(idx, f"unknown_{idx}")
        category = self.tag_to_category.get(tag, "general")
        return tag, category

    # --- preprocessing (aspect ratio preserved via letterbox padding) ---
    def _preprocess_path(self, path: str):
        img = Image.open(path)
        return self._preprocess_pil(img)

    def _preprocess_pil(self, img: Image.Image):
        img = img.convert("RGB")
        orig_w, orig_h = img.size
        scale = min(self.input_w / max(1, orig_w), self.input_h / max(1, orig_h))
        new_w = max(1, int(round(orig_w * scale)))
        new_h = max(1, int(round(orig_h * scale)))
        img = img.resize((new_w, new_h), Image.Resampling.LANCZOS)

        canvas = Image.new("RGB", (self.input_w, self.input_h), self.pad_rgb)
        paste_x = (self.input_w - new_w) // 2
        paste_y = (self.input_h - new_h) // 2
        canvas.paste(img, (paste_x, paste_y))

        arr = np.array(canvas).astype("float32") / 255.0
        arr = (arr - self.mean) / self.std
        arr = arr.transpose(2, 0, 1)[None, ...]  # NCHW
        return arr

def _ensure_numpy_float32(x):
    if hasattr(x, "detach"):
        x = x.detach().cpu().numpy()
    elif hasattr(x, "numpy"):
        x = x.numpy()
    return x.astype("float32")

def _coerce_hw(size):
    # accepts int, [H,W], (H,W)
    if isinstance(size, int):
        return size, size
    if isinstance(size, (list, tuple)) and len(size) == 2:
        h, w = int(size[0]), int(size[1])
        return h, w
    return 512, 512

def _pick_cache_root() -> str:
    """
    Pick a writable cache directory for Hugging Face files in Spaces.
    """
    # Prefer an explicit env var if set; otherwise use /tmp (always writable in Spaces)
    root = os.environ.get("HF_HOME") or "/tmp/huggingface"
    Path(root).mkdir(parents=True, exist_ok=True)
    return root

def _load_json(path: str) -> dict:
    with open(path, "r", encoding="utf-8") as f:
        return json.load(f)


@st.cache_resource(show_spinner="Loading ONNX model from Hugging Face…")
def load_models_from_hub() -> Tuple[ONNXModelWrapper, Optional[TaggerTorch], Dict[str, Any]]:
    """
    Load both ONNX (for inference) and PyTorch safetensors (for essence generation) models
    Returns: (onnx_model_wrapper, torch_model_or_none, metadata)
    """
    cache_dir = _pick_cache_root()
    os.environ.setdefault("HF_HOME", cache_dir)
    os.environ.setdefault("TRANSFORMERS_CACHE", cache_dir)

    # Download files
    onnx_path = hf_hub_download(
        repo_id=MODEL_REPO, filename=ONNX_MODEL_FILE, cache_dir=cache_dir, resume_download=True
    )
    metadata_path = hf_hub_download(
        repo_id=MODEL_REPO, filename=METADATA_FILE, cache_dir=cache_dir, resume_download=True
    )
    
    # Try to download safetensors for essence generation
    torch_model = None
    safetensors_path = None
    try:
        safetensors_path = hf_hub_download(
            repo_id=MODEL_REPO, filename=SAFETENSORS_MODEL_FILE, cache_dir=cache_dir, resume_download=True
        )
        print(f"Downloaded safetensors model: {safetensors_path}")
    except Exception as e:
        print(f"Could not download safetensors model: {e}")
        print("Essence generation will not be available")
    
    metadata = _load_json(metadata_path)

    # Create ONNX model for inference
    import onnxruntime as ort
    providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
    try:
        session = ort.InferenceSession(onnx_path, providers=providers)
        print(f"ONNX model loaded with providers: {session.get_providers()}")
    except Exception:
        session = ort.InferenceSession(onnx_path, providers=["CPUExecutionProvider"])
        print("ONNX model loaded with CPU provider only")

    onnx_model = ONNXModelWrapper(session, metadata)

    # Create PyTorch model for essence generation if safetensors available
    if safetensors_path:
        try:
            torch_model = build_torch_model_from_hub_safetensors(safetensors_path, metadata)
            device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
            torch_model = torch_model.to(device).eval()
            print(f"PyTorch model loaded on {device} for essence generation")
        except Exception as e:
            print(f"Failed to load PyTorch model: {e}")
            torch_model = None

    return onnx_model, torch_model, metadata

def load_model():
    try:
        onnx_model, torch_model, metadata = load_models_from_hub()
        return onnx_model, torch_model, metadata
    except Exception as e:
        st.error(f"Error loading models from HF Hub: {e}")
        st.code(traceback.format_exc())
        st.stop()

# =============================================================================
# GAME STATE MANAGEMENT
# =============================================================================

def initialize_game_state():
    """Initialize the game state in session state if not already present."""
    if 'game_initialized' not in st.session_state:
        st.session_state.game_initialized = True
        
        # Initialize state version counter
        st.session_state.state_version = 0
        
        # Try to load from storage first
        loaded_tags, tag_history, _ = tag_storage.load_tag_collection()
        
        if loaded_tags:
            # If we have saved data, use it
            st.session_state.collected_tags = loaded_tags
            st.session_state.tag_history = tag_history
        else:
            # Otherwise initialize with defaults
            st.session_state.collected_tags = {}  # {tag_name: {"count": int, "rarity": str}}
            st.session_state.tag_history = []  # List of recent tag acquisitions
            
        # Initialize other state variables with defaults
        initialize_default_state()

    if 'tag_rarity_metadata' not in st.session_state:
        st.session_state.tag_rarity_metadata = load_tag_rarity_metadata()

def initialize_default_state():
    """Initialize default state variables"""
    # Core game mechanics
    st.session_state.threshold = STARTING_THRESHOLD
    st.session_state.tag_currency = 0
    st.session_state.enkephalin = 0
    st.session_state.purchased_upgrades = []
    st.session_state.achievements = set()
    st.session_state.current_scan = None
    
    # Game statistics
    st.session_state.game_stats = {
        "images_processed": 0,
        "total_tags_found": 0,
        "total_currency_earned": 0,
        "currency_spent": 0,
        "enkephalin_generated": 0,
        "enkephalin_spent": 0,
        "tags_sacrificed": 0,
        "essences_generated": 0
    }
    
    # Tag power features
    st.session_state.tag_power_bonus = 0
    st.session_state.coin_multiplier = 1.0
    st.session_state.unlocked_combinations = set()
    st.session_state.combination_bonuses = {"threshold_reduction": 0, "coin_bonus": 0}
    
    # Sacrificed tags tracking
    st.session_state.sacrificed_tags = {}

def save_game():
    """Save the game state to files using the enhanced tag storage system."""
    try:
        # Use the enhanced tag storage system to save all game data
        import tag_storage
        success = tag_storage.save_game(st.session_state)
        return success
    except Exception as e:
        st.error(f"Failed to save game: {str(e)}")
        return False

def load_game():
    """
    Load the game state from files using the enhanced tag storage system.
    
    Returns:
        bool: True if load was successful, False otherwise
    """
    try:
        # Use the enhanced tag storage system to load all game data
        import tag_storage
        success = tag_storage.load_game(st.session_state)
        return success
    except Exception as e:
        st.error(f"Failed to load game: {str(e)}")
        return False

# =============================================================================
# USER INTERFACE FUNCTIONS
# =============================================================================

def get_discovered_rarities():
    """
    Helper function to get a set of all rarities the player has discovered
    Returns:
        set: Set of discovered rarity names
    """
    # Start with Canard as always discovered
    discovered_rarities = set(["Canard"])
    
    # Add rarities from collected tags
    if hasattr(st.session_state, 'collected_tags') and st.session_state.collected_tags:
        for tag_info in st.session_state.collected_tags.values():
            if "rarity" in tag_info and tag_info["rarity"]:
                discovered_rarities.add(tag_info["rarity"])
    
    return discovered_rarities

def apply_tag_animations():
    """Apply CSS animations for special tag rarities from the game constants"""
    # Create CSS styles for animations
    css = ""
    
    # Add animations from TAG_ANIMATIONS
    for rarity, animation_info in TAG_ANIMATIONS.items():
        css += animation_info["animation"]
    
    # Apply the CSS
    st.markdown(f"<style>{css}</style>", unsafe_allow_html=True)

def create_sidebar():
    """Create a simplified sidebar without game stats"""
    with st.sidebar:
        # Save/Load controls
        display_save_load_controls()
        
        # Rarity legend
        display_rarity_legend()
        
        # # Developer mode toggle
        # display_developer_mode_toggle()
        
        # Support info
        display_support_info()

def display_save_load_controls():
    """Display save/load game controls"""
    st.subheader("Save/Load Game")
    save_col, load_col = st.columns(2)
    with save_col:
        if st.button("Save Game"):
            if save_game():
                st.success("Game saved!")
                # Force refresh
                st.rerun()
    with load_col:
        if st.button("Load Game"):
            if load_game():
                st.success("Game loaded!")
                # Force refresh
                st.rerun()
            else:
                st.info("No saved game found.")

def display_rarity_legend():
    """Display a legend explaining the rarity levels with TagCoins and Enkephalin rewards, with undiscovered rarities blacked out and special effects for rare rarities"""
    st.sidebar.subheader("Tag Rarity Legend")
    
    # Add animations for the sidebar
    st.sidebar.markdown("""
    <style>
    /* Star of the City animation */
    @keyframes sidebar-glow {
        0% { text-shadow: 0 0 2px gold; }
        50% { text-shadow: 0 0 6px gold; }
        100% { text-shadow: 0 0 2px gold; }
    }
    
    .sidebar-star {
        animation: sidebar-glow 2s infinite;
        font-weight: bold;
    }
    
    /* Impuritas Civitas animation */
    @keyframes sidebar-rainbow {
        0% { color: red; }
        14% { color: orange; }
        28% { color: yellow; }
        42% { color: green; }
        57% { color: blue; }
        71% { color: indigo; }
        85% { color: violet; }
        100% { color: red; }
    }
    
    .sidebar-impuritas {
        animation: sidebar-rainbow 4s linear infinite;
        font-weight: bold;
    }
    
    /* Urban Nightmare animation */
    @keyframes sidebar-pulse {
        0% { opacity: 0.8; }
        50% { opacity: 1; }
        100% { opacity: 0.8; }
    }
    
    .sidebar-nightmare {
        animation: sidebar-pulse 3s infinite;
        font-weight: bold;
    }
    
    /* Urban Plague subtle effect */
    .sidebar-plague {
        text-shadow: 0 0 2px #9C27B0;
        font-weight: bold;
    }
    
    /* Special Effect for Enkephalin */
    @keyframes enkephalin-glow {
        0% { text-shadow: 0 0 2px cyan; }
        50% { text-shadow: 0 0 5px cyan; }
        100% { text-shadow: 0 0 2px cyan; }
    }
    
    .enkephalin-value {
        color: #00BCD4;
        animation: enkephalin-glow 2s infinite;
        font-weight: bold;
    }
    
    /* Reward row styling */
    .rewards-row {
        display: flex;
        align-items: center;
        margin-bottom: 8px;
        padding: 5px;
        border-radius: 4px;
    }
    
    .rewards-row:hover {
        background-color: rgba(0,0,0,0.05);
    }
    
    .rarity-name {
        flex: 0 0 35%;
        font-weight: bold;
    }
    
    .coin-value {
        flex: 0 0 30%;
        color: #FFD700;
    }
    
    .enkephalin-value {
        flex: 0 0 35%;
        color: #00BCD4;
    }
    </style>
    """, unsafe_allow_html=True)
    
    # Get the rarities the player has discovered
    discovered_rarities = get_discovered_rarities()
    
    # Create a header explaining the values
    st.sidebar.markdown("""
    <div style="display: flex; margin-bottom: 10px; font-weight: bold;">
        <div style="flex: 0 0 35%;">Rarity</div>
        <div style="flex: 0 0 30%;">TagCoins</div>
        <div style="flex: 0 0 35%;">Enkephalin</div>
    </div>
    """, unsafe_allow_html=True)
    
    # Display in order from most common to most rare
    for rarity, info in RARITY_LEVELS.items():
        # Get the enkephalin reward from TAG_POWER_BONUSES
        enkephalin_reward = TAG_POWER_BONUSES.get(rarity, {}).get("enkephalin_reward", 0)
        
        # Check if this rarity has been discovered
        if rarity in discovered_rarities:
            # Apply special styling based on rarity
            if rarity == "Impuritas Civitas":
                # Rainbow animation for Impuritas Civitas
                st.sidebar.markdown(
                    f"""
                    <div class="rewards-row">
                        <div class="rarity-name">
                            <span class='sidebar-impuritas'>{rarity}</span>
                        </div>
                        <div class="coin-value">
                            {info['value']} {TAG_CURRENCY_NAME}
                        </div>
                        <div class="enkephalin-value">
                            +{enkephalin_reward} {ENKEPHALIN_ICON}
                        </div>
                    </div>
                    """,
                    unsafe_allow_html=True
                )
            elif rarity == "Star of the City":
                # Glowing effect for Star of the City
                st.sidebar.markdown(
                    f"""
                    <div class="rewards-row">
                        <div class="rarity-name">
                            <span class='sidebar-star' style='color:{info['color']}'>{rarity}</span>
                        </div>
                        <div class="coin-value">
                            {info['value']} {TAG_CURRENCY_NAME}
                        </div>
                        <div class="enkephalin-value">
                            +{enkephalin_reward} {ENKEPHALIN_ICON}
                        </div>
                    </div>
                    """,
                    unsafe_allow_html=True
                )
            elif rarity == "Urban Nightmare":
                # Pulsing effect for Urban Nightmare
                st.sidebar.markdown(
                    f"""
                    <div class="rewards-row">
                        <div class="rarity-name">
                            <span class='sidebar-nightmare' style='color:{info['color']}'>{rarity}</span>
                        </div>
                        <div class="coin-value">
                            {info['value']} {TAG_CURRENCY_NAME}
                        </div>
                        <div class="enkephalin-value">
                            +{enkephalin_reward} {ENKEPHALIN_ICON}
                        </div>
                    </div>
                    """,
                    unsafe_allow_html=True
                )
            elif rarity == "Urban Plague":
                # Subtle glow for Urban Plague
                st.sidebar.markdown(
                    f"""
                    <div class="rewards-row">
                        <div class="rarity-name">
                            <span class='sidebar-plague' style='color:{info['color']}'>{rarity}</span>
                        </div>
                        <div class="coin-value">
                            {info['value']} {TAG_CURRENCY_NAME}
                        </div>
                        <div class="enkephalin-value">
                            +{enkephalin_reward} {ENKEPHALIN_ICON}
                        </div>
                    </div>
                    """,
                    unsafe_allow_html=True
                )
            else:
                # Standard display for common rarities
                st.sidebar.markdown(
                    f"""
                    <div class="rewards-row">
                        <div class="rarity-name">
                            <span style='color:{info['color']}'>{rarity}</span>
                        </div>
                        <div class="coin-value">
                            {info['value']} {TAG_CURRENCY_NAME}
                        </div>
                        <div class="enkephalin-value">
                            {'+' + str(enkephalin_reward) + ' ' + ENKEPHALIN_ICON if enkephalin_reward > 0 else '-'}
                        </div>
                    </div>
                    """,
                    unsafe_allow_html=True
                )
        else:
            # Show blacked out version for undiscovered rarities
            st.sidebar.markdown(
                f"""
                <div class="rewards-row">
                    <div class="rarity-name">
                        <span style='color:#333333;'>?????</span>
                    </div>
                    <div class="coin-value">
                        ??? {TAG_CURRENCY_NAME}
                    </div>
                    <div class="enkephalin-value">
                        ??? {ENKEPHALIN_ICON}
                    </div>
                </div>
                """,
                unsafe_allow_html=True
            )

# def display_developer_mode_toggle():
#     """Display developer mode toggle"""
#     if 'show_dev_tools' not in st.session_state:
#         st.session_state.show_dev_tools = False

#     if st.sidebar.checkbox("Developer Mode", value=st.session_state.show_dev_tools, key="dev_mode_toggle"):
#         st.session_state.show_dev_tools = True
#     else:
#         st.session_state.show_dev_tools = False

def display_support_info():
    # Add separator for visual clarity
    st.markdown("---")
    
    # Support information
    st.subheader("💡 Notes")
    
    st.markdown("""
    This tagger was trained on a subset of the available data due to hardware limitations.
    
    A more comprehensive model trained on the full 3+ million image dataset would provide:
    - More recent characters and tags.
    - Improved accuracy.
    
    If you find this tool useful and would like to support future development:
    """)
    
    # Add Buy Me a Coffee button with Star of the City-like glow effect
    st.markdown("""
    <style>
    @keyframes coffee-button-glow {
        0% { box-shadow: 0 0 5px #FFD700; }
        50% { box-shadow: 0 0 15px #FFD700; }
        100% { box-shadow: 0 0 5px #FFD700; }
    }
    
    .coffee-button {
        display: inline-block;
        animation: coffee-button-glow 2s infinite;
        border-radius: 5px;
        transition: transform 0.3s ease;
    }
    
    .coffee-button:hover {
        transform: scale(1.05);
    }
    </style>
    
    <a href="https://ko-fi.com/camais" target="_blank" class="coffee-button">
        <img src="https://cdn.buymeacoffee.com/buttons/v2/default-yellow.png" 
            alt="Buy Me A Coffee" 
            style="height: 45px; width: 162px; border-radius: 5px;" />
    </a>
    """, unsafe_allow_html=True)
    
    st.markdown("""
    Your support helps with:
    - GPU costs for training
    - Storage for larger datasets
    - Development of new features
    - Future projects
    
    Thank you! 🙏
                
    Full Details: https://huggingface.co/Camais03/camie-tagger-v2
    """)

def display_scan_interface():
    """Display the scan interface tab"""
    st.subheader("Scan Images for Tags")
    
    # Create two columns
    col1, col2 = st.columns([1, 1])
    
    with col1:
        # Image uploader
        uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])
        
        image_path = None
        if uploaded_file:
            # Create a temporary file
            with tempfile.NamedTemporaryFile(delete=False, suffix='.jpg') as tmp_file:
                tmp_file.write(uploaded_file.getvalue())
                image_path = tmp_file.name
            
            # Display the image
            image = Image.open(uploaded_file)
            st.image(image, use_container_width=True)
            
            # Scan button with special key to avoid rerun loop
            button_key = f"scan_btn_{st.session_state.get('state_version', 0)}"
            if st.button("Scan for Tags", key=button_key):
                # Run scan
                with st.spinner("Scanning image..."):
                    success = enhanced_scan_button_handler(image_path)
                    if success:
                        # Save state
                        import tag_storage
                        tag_storage.update_tag_storage_from_session(st.session_state)
                        
                        # Force rerun to update UI everywhere
                        st.rerun()
    
    with col2:
        display_scanner_settings()
    
    # Display scan results if available
    if hasattr(st.session_state, 'current_scan') and st.session_state.current_scan:
        display_scan_results(st.session_state.current_scan)

def display_scan_results(scan_data):
    """Display scan results from session state with improved tag categorization and special effects for rare tags"""
    if not scan_data:
        return
        
    found_tags = scan_data.get("found_tags", [])
    all_tags = scan_data.get("all_tags", {})
    total_currency_earned = scan_data.get("total_currency_earned", 0)
    total_enkephalin_earned = scan_data.get("total_enkephalin_earned", 0)
    new_tag_count = scan_data.get("new_tag_count", 0)
    
    # Check for extremely rare new discoveries and celebrate them first
    new_rare_tags = [t for t in found_tags if t.get("is_new", False) and 
                     t.get("rarity") in ["Star of the City", "Impuritas Civitas"]]
    
    if new_rare_tags:
        # Sort by rarity (most rare first)
        rarity_order = ["Impuritas Civitas", "Star of the City"]
        new_rare_tags.sort(key=lambda x: rarity_order.index(x["rarity"]) if x["rarity"] in rarity_order else 999)
        
        # Display a celebration for the rarest tag found
        rarest_tag = new_rare_tags[0]
        if rarest_tag["rarity"] == "Impuritas Civitas":
            st.markdown(f"""
            <style>
            @keyframes rainbow-bg {{
                0% {{ background-color: rgba(255,0,0,0.1); }}
                14% {{ background-color: rgba(255,165,0,0.1); }}
                28% {{ background-color: rgba(255,255,0,0.1); }}
                42% {{ background-color: rgba(0,128,0,0.1); }}
                57% {{ background-color: rgba(0,0,255,0.1); }}
                71% {{ background-color: rgba(75,0,130,0.1); }}
                85% {{ background-color: rgba(238,130,238,0.1); }}
                100% {{ background-color: rgba(255,0,0,0.1); }}
            }}
            
            .rare-discovery-banner {{
                background-color: rgba(0,0,0,0.1);
                border: 3px solid red;
                border-radius: 10px;
                padding: 20px;
                margin: 20px 0;
                text-align: center;
                animation: rainbow-bg 4s linear infinite;
            }}
            </style>
            
            <div class="rare-discovery-banner">
                <h2>✨ EXTRAORDINARY DISCOVERY! ✨</h2>
                <h3>You found the ultra-rare Impuritas Civitas tag</h3>
                <p>This is one of the rarest tags in the game!</p>
                <p>Rewards: +{RARITY_LEVELS[rarest_tag["rarity"]]["value"]} {TAG_CURRENCY_NAME} and +25 {ENKEPHALIN_ICON}</p>
            </div>
            """, unsafe_allow_html=True)
            # # Show balloons for the extraordinary discovery
            # st.balloons()
        elif rarest_tag["rarity"] == "Star of the City":
            st.markdown(f"""
            <style>
            @keyframes glow-bg {{
                0% {{ box-shadow: 0 0 10px #FFD700; }}
                50% {{ box-shadow: 0 0 30px #FFD700; }}
                100% {{ box-shadow: 0 0 10px #FFD700; }}
            }}
            
            .star-discovery-banner {{
                background-color: rgba(255,215,0,0.1);
                border: 2px solid gold;
                border-radius: 10px;
                padding: 20px;
                margin: 20px 0;
                text-align: center;
                box-shadow: 0 0 10px #FFD700;
                animation: glow-bg 2s infinite;
            }}
            </style>
            
            <div class="star-discovery-banner">
                <h2>🌟 EXCEPTIONAL DISCOVERY! 🌟</h2>
                <h3>You found a Star of the City tag</h3>
                <p>This is a very rare and valuable tag!</p>
                <p>Rewards: +{RARITY_LEVELS[rarest_tag["rarity"]]["value"]} {TAG_CURRENCY_NAME} and +10 {ENKEPHALIN_ICON}</p>
            </div>
            """, unsafe_allow_html=True)
    
    if found_tags:
        st.success(f"Found {len(found_tags)} tags! {f'({new_tag_count} new discoveries)' if new_tag_count > 0 else ''}")
        
        # Count tags by rarity
        rarity_counts = {}
        for tag_info in found_tags:
            rarity = tag_info.get("rarity", "Unknown")
            is_new = tag_info.get("is_new", False)
            
            if rarity not in rarity_counts:
                rarity_counts[rarity] = {"total": 0, "new": 0}
            
            rarity_counts[rarity]["total"] += 1
            if is_new:
                rarity_counts[rarity]["new"] += 1
        
        # Get valid rarity order
        rarity_order = list(RARITY_LEVELS.keys())
        
        # Display rarity distribution
        for rarity, counts in sorted(rarity_counts.items(), 
                             key=lambda x: rarity_order.index(x[0]) if x[0] in rarity_order else 999):
            color = RARITY_LEVELS.get(rarity, {}).get("color", "#AAAAAA")
            total = counts["total"]
            new_count = counts["new"]
            
            # Create two columns for showing the rarity count and new count
            col1, col2 = st.columns([3, 1])
            with col1:
                st.markdown(f"<span style='color:{color};font-weight:bold;'>{rarity}:</span> {total} tags", unsafe_allow_html=True)
            with col2:
                if new_count > 0:
                    st.write(f"({new_count} new)")
        
        st.divider()
        
        # Show earned currency and enkephalin
        col1, col2, col3 = st.columns([1, 1, 1])
        with col1:
            st.write("**Earned:**")
        with col2:
            st.write(f"**{total_currency_earned} {TAG_CURRENCY_NAME}**")
        with col3:
            if total_enkephalin_earned > 0:
                st.write(f"**{total_enkephalin_earned} {ENKEPHALIN_ICON}**")
        
        # Show current stats after this scan
        st.divider()
        col1, col2 = st.columns(2)
        with col1:
            st.write("**Images Processed:**")
            st.write("**Total Tags Found:**")
            st.write("**Current TagCoins:**")
            st.write(f"**Current {ENKEPHALIN_CURRENCY_NAME}:**")
        with col2:
            st.write(f"**{st.session_state.game_stats['images_processed']}**")
            st.write(f"**{st.session_state.game_stats['total_tags_found']}**")
            st.write(f"**{st.session_state.tag_currency}**")
            st.write(f"**{st.session_state.enkephalin} {ENKEPHALIN_ICON}**")
        
        # Display found tags by category
        st.subheader("Found Tags by Category")
        
        # Check if we have any rare tags (rarer than Urban Myth)
        rare_rarities = ["Urban Legend", "Urban Plague", "Urban Nightmare", "Star of the City", "Impuritas Civitas"]
        has_rare_tags = any(tag_info.get("rarity") in rare_rarities for tag_info in found_tags)
        
        # Display category expanders
        for category, tags in all_tags.items():
            # Only display categories with collected tags
            collected_tags = [(tag, prob, rarity, any(t["tag"] == tag and t["is_new"] for t in found_tags)) 
                             for tag, prob, rarity, status in tags if status == "collected"]
            
            if not collected_tags:
                continue
                
            # Split into new and existing tags
            new_tags = [(tag, prob, rarity, True) for tag, prob, rarity, is_new in collected_tags if is_new]
            existing_tags = [(tag, prob, rarity, False) for tag, prob, rarity, is_new in collected_tags if not is_new]
            
            # Sort each group by rarity (rarest first)
            new_tags.sort(key=lambda x: rarity_order.index(x[2]) if x[2] in rarity_order else 999, reverse=True)
            existing_tags.sort(key=lambda x: rarity_order.index(x[2]) if x[2] in rarity_order else 999, reverse=True)
            
            # Combine with new tags first
            ordered_tags = new_tags + existing_tags
            
            # Check if this category has any rare tags
            category_has_rare = any(rarity in rare_rarities for _, _, rarity, _ in collected_tags)
            
            # Determine if this category should be expanded by default
            default_expanded = category == None or (category_has_rare and has_rare_tags) or len(new_tags) > 0
            
            # Create category title with badge for new tags
            category_title = f"{category.capitalize()} ({len(collected_tags)} tags)"
            if new_tags:
                category_title += f" 🆕 {len(new_tags)} new!"
            
            # Add special indicator for categories with rare tags
            if any(tag[2] in ["Star of the City", "Impuritas Civitas"] for tag in ordered_tags):
                category_title += " ✨ RARE!"
            
            with st.expander(category_title, expanded=default_expanded):
                # Display new tags section if any
                if new_tags:
                    st.markdown("### ✨ New Discoveries")
                    for tag, prob, rarity, _ in new_tags:
                        # Find the corresponding tag in found_tags to get enkephalin reward
                        tag_info = next((t for t in found_tags if t["tag"] == tag), None)
                        enkephalin_reward = tag_info.get("enkephalin", 0) if tag_info else 0
                        
                        # Get enkephalin reward if not available in tag_info
                        if enkephalin_reward == 0 and rarity in TAG_POWER_BONUSES:
                            enkephalin_reward = TAG_POWER_BONUSES[rarity]["enkephalin_reward"]
                            
                        display_tag_with_effects(tag, prob, rarity, is_new=True, enkephalin=enkephalin_reward)
                    
                    # Add separator if we have both new and existing tags
                    if existing_tags:
                        st.markdown("---")
                
                # Display existing tags
                if existing_tags:
                    if new_tags:  # Only show this header if we have new tags above
                        st.markdown("### Previously Discovered")
                    for tag, prob, rarity, _ in existing_tags:
                        display_tag_with_effects(tag, prob, rarity, is_new=False)
    else:
        st.warning("No tags found above the current threshold.")

def display_tag_with_effects(tag, prob, rarity, is_new=False, enkephalin=0):
    """
    Display a tag with special effects based on rarity, including animations for the rarest tags,
    sample count, and both TagCoins and Enkephalin rewards for new discoveries.
    """
    if rarity is None:
        rarity = "Unknown"
        
    # Get color and determine if this is a rare tag
    color = RARITY_LEVELS.get(rarity, {}).get("color", "#AAAAAA")
    
    # Get rewards for this tag's rarity
    coin_value = RARITY_LEVELS.get(rarity, {}).get("value", 0)
    
    # Get sample count if available
    sample_count = None
    if hasattr(st.session_state, 'tag_rarity_metadata') and st.session_state.tag_rarity_metadata:
        if tag in st.session_state.tag_rarity_metadata:
            tag_info = st.session_state.tag_rarity_metadata[tag]
            if isinstance(tag_info, dict) and "sample_count" in tag_info:
                sample_count = tag_info["sample_count"]
    
    # Format sample count for display
    sample_display = ""
    if sample_count is not None:
        if sample_count >= 1000000:
            sample_display = f"({sample_count/1000000:.1f}M samples)"
        elif sample_count >= 1000:
            sample_display = f"({sample_count/1000:.1f}K samples)"
        else:
            sample_display = f"({sample_count} samples)"
    
    # Add custom CSS for animations
    st.markdown("""
    <style>
    @keyframes glowing {
        0% { box-shadow: 0 0 5px #FFD700; }
        50% { box-shadow: 0 0 20px #FFD700; }
        100% { box-shadow: 0 0 5px #FFD700; }
    }
    
    @keyframes rainbow-border {
        0% { border-color: red; }
        14% { border-color: orange; }
        28% { border-color: yellow; }
        42% { border-color: green; }
        57% { border-color: blue; }
        71% { border-color: indigo; }
        85% { border-color: violet; }
        100% { border-color: red; }
    }
    
    @keyframes rainbow-text {
        0% { color: red; }
        14% { color: orange; }
        28% { color: yellow; }
        42% { color: green; }
        57% { color: blue; }
        71% { color: indigo; }
        85% { color: violet; }
        100% { color: red; }
    }
    
    @keyframes pulse-border {
        0% { border-color: #FF9800; }
        50% { border-color: #FF5722; }
        100% { border-color: #FF9800; }
    }
    
    @keyframes enkephalin-glow {
        0% { text-shadow: 0 0 2px cyan; }
        50% { text-shadow: 0 0 5px cyan; }
        100% { text-shadow: 0 0 2px cyan; }
    }
    
    @keyframes coin-glow {
        0% { text-shadow: 0 0 2px gold; }
        50% { text-shadow: 0 0 5px gold; }
        100% { text-shadow: 0 0 2px gold; }
    }
    
    .star-of-city {
        background-color: rgba(255, 215, 0, 0.2);
        padding: 8px;
        border-radius: 5px;
        border: 2px solid gold;
        animation: glowing 2s infinite;
    }
    
    .impuritas-civitas {
        background-color: rgba(0, 0, 0, 0.1);
        padding: 10px;
        border-radius: 5px;
        border: 3px solid red;
        animation: rainbow-border 4s linear infinite;
    }
    
    .impuritas-text {
        font-weight: bold;
        animation: rainbow-text 4s linear infinite;
    }
    
    .urban-nightmare {
        background-color: rgba(255, 152, 0, 0.1);
        padding: 6px;
        border-radius: 5px;
        border: 2px solid #FF9800;
        animation: pulse-border 3s infinite;
    }
    
    .urban-plague {
        background-color: rgba(156, 39, 176, 0.08);
        padding: 5px;
        border-radius: 5px;
        border: 1px solid #9C27B0;
        box-shadow: 0 0 3px #9C27B0;
    }
    
    .sample-count {
        font-size: 0.8em;
        color: #666;
        margin-left: 5px;
    }
    
    .enkephalin-reward {
        font-size: 0.9em;
        color: #00BCD4;
        margin-left: 5px;
        animation: enkephalin-glow 2s infinite;
        font-weight: bold;
    }
    
    .coin-reward {
        font-size: 0.9em;
        color: #FFD700;
        margin-left: 5px;
        animation: coin-glow 2s infinite;
        font-weight: bold;
    }
    
    .rewards-container {
        display: inline-block;
        background-color: rgba(0, 0, 0, 0.05);
        border-radius: 4px;
        padding: 2px 6px;
        margin-left: 5px;
    }
    </style>
    """, unsafe_allow_html=True)
    
    # Style based on rarity
    rare_rarities = {
        "Urban Legend": {"style": "background-color:rgba(33, 150, 243, 0.1); padding:5px; border-radius:5px;", "emoji": "🔮"},
        "Urban Plague": {"style": "class='urban-plague'", "emoji": "⚔️"},
        "Urban Nightmare": {"style": "class='urban-nightmare'", "emoji": "👑"},
        "Star of the City": {"style": "class='star-of-city'", "emoji": "🌟"},
        "Impuritas Civitas": {"style": "class='impuritas-civitas'", "emoji": "✨"}
    }
    
    # Determine if this is a rare tag
    is_rare = rarity in rare_rarities
    style = ""
    emoji_prefix = ""
    
    if is_rare:
        style = rare_rarities[rarity]["style"]
        emoji_prefix = rare_rarities[rarity]["emoji"] + " "
    
    # Create rewards display if new discovery
    rewards_display = ""
    if is_new:
        coin_display = f"<span class='coin-reward'>+{coin_value} {TAG_CURRENCY_NAME}</span>"
        enkephalin_display = ""
        if enkephalin > 0:
            enkephalin_display = f"<span class='enkephalin-reward'>+{enkephalin} {ENKEPHALIN_ICON}</span>"
            
        if enkephalin > 0:
            rewards_display = f"<span class='rewards-container'>{coin_display} {enkephalin_display}</span>"
        else:
            rewards_display = f"<span class='rewards-container'>{coin_display}</span>"
    
    # Create the tag display
    if is_new:
        if rarity == "Impuritas Civitas":
            tag_html = f"""
            <div {style}>
                {emoji_prefix}<span class="impuritas-text">{tag}</span> - 
                <span style="color:{color};font-weight:bold;">{rarity}</span> ({prob:.2f}) 
                <span style="background-color:#4CAF50;color:white;padding:2px 6px;border-radius:10px;font-size:0.8em;">NEW</span>
                {rewards_display}
                <span class="sample-count">{sample_display}</span>
            </div>
            """
        else:
            tag_html = f"""
            <div {style}>
                {emoji_prefix}<strong>{tag}</strong> - 
                <span style="color:{color};font-weight:bold;">{rarity}</span> ({prob:.2f}) 
                <span style="background-color:#4CAF50;color:white;padding:2px 6px;border-radius:10px;font-size:0.8em;">NEW</span>
                {rewards_display}
                <span class="sample-count">{sample_display}</span>
            </div>
            """
    else:
        if rarity == "Impuritas Civitas":
            tag_html = f"""
            <div {style}>
                {emoji_prefix}<span class="impuritas-text">{tag}</span> - 
                <span style="color:{color};font-weight:bold;">{rarity}</span> ({prob:.2f})
                <span class="sample-count">{sample_display}</span>
            </div>
            """
        else:
            tag_html = f"""
            <div {style}>
                {emoji_prefix}{tag} - 
                <span style="color:{color};font-weight:bold;">{rarity}</span> ({prob:.2f})
                <span class="sample-count">{sample_display}</span>
            </div>
            """
    
    # Add notification sound/confetti for very rare tags
    if is_new and rarity in ["Star of the City", "Impuritas Civitas"]:
        # Add confetti effect for extremely rare tags
        # st.balloons()
        
        # Add special notification
        if rarity == "Impuritas Civitas":
            st.markdown(f"""
            <div style="border:3px solid red; padding:10px; margin:10px 0; text-align:center; background-color:rgba(0,0,0,0.1); animation:rainbow-border 4s linear infinite;">
                <h3>🎉 EXTRAORDINARY DISCOVERY! 🎉</h3>
                <p>You found an ultra-rare Impuritas Civitas tag!</p>
                <p>Rewards: +{RARITY_LEVELS[rarity]["value"]} {TAG_CURRENCY_NAME} and +25 {ENKEPHALIN_ICON}</p>
            </div>
            """, unsafe_allow_html=True)
        elif rarity == "Star of the City":
            st.markdown(f"""
            <div style="border:2px solid gold; padding:10px; margin:10px 0; text-align:center; background-color:rgba(255,215,0,0.2); animation:glowing 2s infinite;">
                <h3>🌟 EXCEPTIONAL DISCOVERY! 🌟</h3>
                <p>You found a very rare Star of the City tag!</p>
                <p>Rewards: +{RARITY_LEVELS[rarity]["value"]} {TAG_CURRENCY_NAME} and +10 {ENKEPHALIN_ICON}</p>
            </div>
            """, unsafe_allow_html=True)
    
    st.markdown(tag_html, unsafe_allow_html=True)

def display_scanner_settings():
    """Display scanner settings and help"""
    st.write("### Scanner Settings")
    
    # Get threshold values
    base_threshold = st.session_state.threshold
    tag_power = st.session_state.tag_power_bonus if hasattr(st.session_state, 'tag_power_bonus') else 0
    effective_threshold = max(MIN_THRESHOLD, base_threshold - tag_power)
    
    # Display threshold info
    if tag_power > 0:
        st.write(f"Base Threshold: **{base_threshold:.3f}**")
        st.write(f"Tag Power Bonus: **-{tag_power:.4f}**")
        st.write(f"Effective Threshold: **{effective_threshold:.3f}**")
    else:
        st.write(f"Current Threshold: **{base_threshold:.3f}**")
                  
    if hasattr(st.session_state, 'coin_multiplier') and st.session_state.coin_multiplier > 1.0:
        st.info(f"Your Tag Power gives a coin multiplier of {st.session_state.coin_multiplier:.2f}x for new discoveries!")
    
    # Instruction box
    st.markdown(
        """
        ### How to Play
        1. Upload an image
        2. Scan for tags to discover them
        3. Earn TagCoins for new discoveries
        4. Spend TagCoins on upgrades to lower the threshold
        5. Lower thresholds reveal rarer tags!
        6. Collect sets of related tags for bonuses and reveal unique mosaics!
        7. Visit the Library System to discover unique tags (not collect)
        8. Use collected tags to either inspire new searches or generate essence
        9. Use Enkephalin to generate Tag Essences
        10. Use the Tag Essence Generator to collect the tag and related tags to it.
        """
    )

def display_game_stats_panel():
    """Display a prominent game stats panel at the top of the main content area"""
    # Force read latest values from session state
    tag_currency = st.session_state.tag_currency
    enkephalin = st.session_state.enkephalin
    images_processed = st.session_state.game_stats.get('images_processed', 0)
    total_tags_found = st.session_state.game_stats.get('total_tags_found', 0)
    total_currency_earned = st.session_state.game_stats.get('total_currency_earned', 0)
    unique_tags = len(st.session_state.collected_tags) if hasattr(st.session_state, 'collected_tags') else 0
    
    # Get threshold values
    base_threshold = st.session_state.threshold
    tag_power = st.session_state.tag_power_bonus if hasattr(st.session_state, 'tag_power_bonus') else 0
    effective_threshold = max(MIN_THRESHOLD, base_threshold - tag_power)
    
    # Prepare the container
    st.markdown('<div class="stats-panel">', unsafe_allow_html=True)
    
    # Primary stats row
    cols = st.columns([1, 1, 1, 1])
    
    with cols[0]:
        st.markdown(f'<div class="stats-value">{tag_currency}</div>', unsafe_allow_html=True)
        st.markdown(f'<div class="stats-label">{TAG_CURRENCY_NAME}</div>', unsafe_allow_html=True)
    
    with cols[1]:
        st.markdown(f'<div class="stats-value">{enkephalin} {ENKEPHALIN_ICON}</div>', unsafe_allow_html=True)
        st.markdown(f'<div class="stats-label">{ENKEPHALIN_CURRENCY_NAME}</div>', unsafe_allow_html=True)
    
    with cols[2]:
        if tag_power > 0:
            st.markdown(f'<div class="stats-value">{effective_threshold:.3f}</div>', unsafe_allow_html=True)
            st.markdown(f'<div class="stats-label">Effective Threshold (Base: {base_threshold:.3f})</div>', unsafe_allow_html=True)
        else:
            st.markdown(f'<div class="stats-value">{base_threshold:.3f}</div>', unsafe_allow_html=True)
            st.markdown(f'<div class="stats-label">Threshold</div>', unsafe_allow_html=True)
    
    with cols[3]:
        st.markdown(f'<div class="stats-value">{unique_tags}</div>', unsafe_allow_html=True)
        st.markdown(f'<div class="stats-label">Unique Tags</div>', unsafe_allow_html=True)
    
    # Add a secondary row with collapsible additional stats
    with st.expander("More Game Stats"):
        cols2 = st.columns(3)
        
        with cols2[0]:
            st.markdown("### Collection Stats")
            st.write(f"Images Processed: **{images_processed}**")
            st.write(f"Total Tags Found: **{total_tags_found}**")
            st.write(f"Unique Tags: **{unique_tags}**")
        
        with cols2[1]:
            st.markdown("### Economy Stats")
            st.write(f"Total Earned: **{total_currency_earned}** {TAG_CURRENCY_NAME}")
            if hasattr(st.session_state.game_stats, 'currency_spent'):
                st.write(f"Currency Spent: **{st.session_state.game_stats.get('currency_spent', 0)}** {TAG_CURRENCY_NAME}")
            if hasattr(st.session_state.game_stats, 'enkephalin_generated'):
                st.write(f"Enkephalin Generated: **{st.session_state.game_stats.get('enkephalin_generated', 0)}**")
        
        with cols2[2]:
            st.markdown("### Bonuses")
            if hasattr(st.session_state, 'tag_power_bonus') and st.session_state.tag_power_bonus > 0:
                st.write(f"Tag Power Bonus: **-{st.session_state.tag_power_bonus:.4f}**")
            if hasattr(st.session_state, 'coin_multiplier') and st.session_state.coin_multiplier > 1.0:
                st.write(f"Coin Multiplier: **{st.session_state.coin_multiplier:.2f}x**")
            if hasattr(st.session_state, 'unlocked_combinations'):
                st.write(f"Unlocked Combinations: **{len(st.session_state.unlocked_combinations)}**")
    
    # Close the container
    st.markdown('</div>', unsafe_allow_html=True)

def display_tag_collection():
    """Display the user's tag collection with stats and analysis."""
    st.subheader("Your Tag Collection")
    
    # Create tabs for different views
    list_tab, stats_tab, analysis_tab, mosaic_tab = st.tabs([
        "Tag List", 
        "Collection Stats",
        "Category Analysis",
        "Collection Mosaic"
    ])
     
    with list_tab:
        display_tag_list()
    
    with stats_tab:
        display_collection_stats()
    
    with analysis_tab:
        display_category_stats()

    with mosaic_tab:
        # Display the tag mosaic
        display_tag_mosaic()

def display_tag_list():
    """Display a simple list view of tags with improved sorting options"""
    # Show total unique tags
    unique_tags = len(st.session_state.collected_tags)
    st.write(f"You have collected {unique_tags} unique tags.")
    
    # Count tags by rarity
    rarity_counts = get_rarity_counts()
    
    # Only display rarity categories that have tags
    active_rarities = {r: c for r, c in rarity_counts.items() if c > 0}
    
    # If there are active rarities to display
    if active_rarities:
        display_rarity_distribution(active_rarities)
    
    # Add a sorting option
    sort_options = ["Category (rarest first)", "Rarity"]
    selected_sort = st.selectbox("Sort tags by:", sort_options)
    
    # Group tags by the selected method
    if selected_sort == "Category (rarest first)":
        categories = group_tags_by_category()
        
        # Display tags by category in expanders
        for category, tags in sorted(categories.items()):
            # Get rarity order for sorting
            rarity_order = list(RARITY_LEVELS.keys())
            
            # Sort tags by rarity (rarest first)
            def get_rarity_index(tag_tuple):
                tag, info = tag_tuple
                rarity = info.get("rarity", "Unknown")
                if rarity in rarity_order:
                    return len(rarity_order) - rarity_order.index(rarity)
                return 0
            
            sorted_tags = sorted(tags, key=get_rarity_index, reverse=True)
            
            # Check if category has any rare tags
            has_rare_tags = any(info.get("rarity") in ["Impuritas Civitas", "Star of the City"] 
                               for _, info in sorted_tags)
            
            # Get category info if available
            category_display = category.capitalize()
            if category in TAG_CATEGORIES:
                category_info = TAG_CATEGORIES[category]
                category_icon = category_info.get("icon", "")
                category_color = category_info.get("color", "#888888")
                category_display = f"<span style='color:{category_color};'>{category_icon} {category.capitalize()}</span>"
            
            # Create header with information about rare tags if present
            header = f"{category_display} ({len(tags)} tags)"
            if has_rare_tags:
                header += " ✨ Contains rare tags!"
                
            # Display the header and expander
            st.markdown(header, unsafe_allow_html=True)
            with st.expander("Show/Hide"):
                # Group by rarity within category
                rarity_groups = {}
                for tag, info in sorted_tags:
                    rarity = info.get("rarity", "Unknown")
                    if rarity not in rarity_groups:
                        rarity_groups[rarity] = []
                    rarity_groups[rarity].append((tag, info))
                
                # Display each rarity group in order (rarest first)
                for rarity in reversed(rarity_order):
                    if rarity in rarity_groups:
                        tags_in_rarity = rarity_groups[rarity]
                        if tags_in_rarity:
                            color = RARITY_LEVELS[rarity]["color"]
                            
                            # Special styling for rare rarities
                            if rarity == "Impuritas Civitas":
                                rarity_style = f"animation:rainbow-text 4s linear infinite;font-weight:bold;"
                            elif rarity == "Star of the City":
                                rarity_style = f"color:{color};text-shadow:0 0 3px gold;font-weight:bold;"
                            elif rarity == "Urban Nightmare":
                                rarity_style = f"color:{color};text-shadow:0 0 1px #FF5722;font-weight:bold;"
                            else:
                                rarity_style = f"color:{color};font-weight:bold;"
                            
                            st.markdown(f"<span style='{rarity_style}'>{rarity.capitalize()}</span> ({len(tags_in_rarity)} tags)", unsafe_allow_html=True)
                            display_tag_grid(tags_in_rarity)
                            st.markdown("---")
                
    elif selected_sort == "Rarity":
        # Group tags by rarity level
        rarity_groups = {}
        for tag, info in st.session_state.collected_tags.items():
            rarity = info.get("rarity", "Unknown")
            if rarity not in rarity_groups:
                rarity_groups[rarity] = []
            rarity_groups[rarity].append((tag, info))
        
        # Get ordered rarities (rarest first)
        ordered_rarities = list(RARITY_LEVELS.keys())
        ordered_rarities.reverse()  # Reverse to show rarest first
        
        # Display tags by rarity
        for rarity in ordered_rarities:
            if rarity in rarity_groups:
                tags = rarity_groups[rarity]
                color = RARITY_LEVELS[rarity]["color"]
                
                # Add special styling for rare rarities
                rarity_html = f"<span style='color:{color};font-weight:bold;'>{rarity.capitalize()}</span>"
                if rarity == "Impuritas Civitas":
                    rarity_html = f"<span style='animation:rainbow-text 4s linear infinite;font-weight:bold;'>{rarity.capitalize()}</span>"
                elif rarity == "Star of the City":
                    rarity_html = f"<span style='color:{color};text-shadow:0 0 3px gold;font-weight:bold;'>{rarity.capitalize()}</span>"
                elif rarity == "Urban Nightmare":
                    rarity_html = f"<span style='color:{color};text-shadow:0 0 1px #FF5722;font-weight:bold;'>{rarity.capitalize()}</span>"
                
                # First create the title with HTML, then use it in the expander
                st.markdown(f"### {rarity_html} ({len(tags)} tags)", unsafe_allow_html=True)
                with st.expander("Show/Hide"):
                    # Group by category within rarity
                    category_groups = {}
                    for tag, info in tags:
                        category = info.get("category", "unknown")
                        if category not in category_groups:
                            category_groups[category] = []
                        category_groups[category].append((tag, info))
                    
                    # Display each category within this rarity level
                    for category, category_tags in sorted(category_groups.items()):
                        # Get category info if available
                        category_display = category.capitalize()
                        if category in TAG_CATEGORIES:
                            category_info = TAG_CATEGORIES[category]
                            category_icon = category_info.get("icon", "")
                            category_color = category_info.get("color", "#888888")
                            category_display = f"<span style='color:{category_color};'>{category_icon} {category.capitalize()}</span>"
                        
                        st.markdown(f"#### {category_display} ({len(category_tags)} tags)", unsafe_allow_html=True)
                        display_tag_grid(category_tags)
                        st.markdown("---")

def get_rarity_counts():
    """Count tags by rarity levels"""
    rarity_counts = {
        "Canard": 0, 
        "Urban Myth": 0, 
        "Urban Legend": 0, 
        "Urban Plague": 0, 
        "Urban Nightmare": 0, 
        "Star of the City": 0, 
        "Impuritas Civitas": 0,
        "Unknown": 0  # Add an "Unknown" category for handling None or invalid rarities
    }
    
    for tag_info in st.session_state.collected_tags.values():
        rarity = tag_info.get("rarity")
        if rarity is None:
            rarity = "Unknown"
        
        if rarity in rarity_counts:
            rarity_counts[rarity] += 1
        else:
            # Handle any unexpected rarity values
            rarity_counts["Unknown"] += 1
            
    return rarity_counts

def display_rarity_distribution(active_rarities):
    """Display distribution of tags by rarity with themed animations for rare tags"""
    # Add the necessary CSS for animations
    st.markdown("""
    <style>
    @keyframes grid-glow {
        0% { text-shadow: 0 0 2px gold; }
        50% { text-shadow: 0 0 6px gold; }
        100% { text-shadow: 0 0 2px gold; }
    }
    
    @keyframes grid-rainbow {
        0% { color: red; }
        14% { color: orange; }
        28% { color: yellow; }
        42% { color: green; }
        57% { color: blue; }
        71% { color: indigo; }
        85% { color: violet; }
        100% { color: red; }
    }
    
    @keyframes grid-pulse {
        0% { opacity: 0.8; }
        50% { opacity: 1; }
        100% { opacity: 0.8; }
    }
    
    .grid-star {
        text-shadow: 0 0 3px gold;
        animation: grid-glow 2s infinite;
    }
    
    .grid-impuritas {
        animation: grid-rainbow 4s linear infinite;
    }
    
    .grid-nightmare {
        text-shadow: 0 0 1px #FF5722;
        animation: grid-pulse 3s infinite;
    }
    
    .grid-plague {
        text-shadow: 0 0 1px #9C27B0;
    }
    </style>
    """, unsafe_allow_html=True)
    
    rarity_cols = st.columns(len(active_rarities))
    for i, (rarity, count) in enumerate(active_rarities.items()):
        with rarity_cols[i]:
            # Get color with fallback
            color = RARITY_LEVELS.get(rarity, {}).get("color", "#888888")
            
            # Apply special styling based on rarity
            style = f"color:{color};font-weight:bold;"
            class_name = ""
            
            if rarity == "Impuritas Civitas":
                class_name = "grid-impuritas"
            elif rarity == "Star of the City":
                class_name = "grid-star"
            elif rarity == "Urban Nightmare":
                class_name = "grid-nightmare"
            elif rarity == "Urban Plague":
                class_name = "grid-plague"
            
            if class_name:
                st.markdown(
                    f"<div style='text-align:center;'><span class='{class_name}' style='font-weight:bold;'>{rarity.capitalize()}</span><br>{count}</div>",
                    unsafe_allow_html=True
                )
            else:
                st.markdown(
                    f"<div style='text-align:center;'><span style='{style}'>{rarity.capitalize()}</span><br>{count}</div>",
                    unsafe_allow_html=True
                )

def group_tags_by_category():
    """Group tags by their categories"""
    categories = {}
    for tag, info in st.session_state.collected_tags.items():
        category = info.get("category", "unknown")
        if category not in categories:
            categories[category] = []
        categories[category].append((tag, info))
    return categories

def display_tag_grid(tags):
    """Display tags in a grid layout with sample count information"""
    # Create a grid layout for tags
    cols = st.columns(3)
    for i, (tag, info) in enumerate(sorted(tags)):
        col_idx = i % 3
        with cols[col_idx]:
            rarity = info.get("rarity")
            if rarity is None:
                rarity = "Unknown"
            
            count = info.get("count", 1)
            color = RARITY_LEVELS.get(rarity, {}).get("color", "#888888")
            
            # Get sample count if available
            sample_count = None
            if hasattr(st.session_state, 'tag_rarity_metadata') and st.session_state.tag_rarity_metadata:
                if tag in st.session_state.tag_rarity_metadata:
                    tag_info = st.session_state.tag_rarity_metadata[tag]
                    if isinstance(tag_info, dict) and "sample_count" in tag_info:
                        sample_count = tag_info["sample_count"]
            
            # Format sample count
            sample_display = ""
            if sample_count is not None:
                if sample_count >= 1000000:
                    sample_display = f"<span style='font-size:0.8em;color:#666;'>({sample_count/1000000:.1f}M)</span>"
                elif sample_count >= 1000:
                    sample_display = f"<span style='font-size:0.8em;color:#666;'>({sample_count/1000:.1f}K)</span>"
                else:
                    sample_display = f"<span style='font-size:0.8em;color:#666;'>({sample_count})</span>"
            
            # Apply special styling for rare tags
            tag_html = tag
            if rarity == "Impuritas Civitas":
                tag_html = f"<span style='animation: rainbow-text 4s linear infinite;'>{tag}</span>"
            elif rarity == "Star of the City":
                tag_html = f"<span style='text-shadow: 0 0 3px gold;'>{tag}</span>"
            elif rarity == "Urban Nightmare":
                tag_html = f"<span style='text-shadow: 0 0 1px #FF9800;'>{tag}</span>"
            
            # Display tag with rarity badge, count, and sample count
            st.markdown(
                f"{tag_html} <span style='background-color:{color};color:white;padding:2px 6px;border-radius:10px;font-size:0.8em;'>{rarity.capitalize()}</span> (×{count}) {sample_display}",
                unsafe_allow_html=True
            )

def display_collection_stats():
    """Display overall collection statistics with themed animations"""
    st.subheader("Collection Overview")
    
    # Check if we have collected tags
    if not hasattr(st.session_state, 'collected_tags') or not st.session_state.collected_tags:
        st.info("Start scanning images to collect tags!")
        return
    
    # Add the necessary CSS for animations
    st.markdown("""
    <style>
    @keyframes grid-glow {
        0% { text-shadow: 0 0 2px gold; }
        50% { text-shadow: 0 0 6px gold; }
        100% { text-shadow: 0 0 2px gold; }
    }
    
    @keyframes grid-rainbow {
        0% { color: red; }
        14% { color: orange; }
        28% { color: yellow; }
        42% { color: green; }
        57% { color: blue; }
        71% { color: indigo; }
        85% { color: violet; }
        100% { color: red; }
    }
    
    @keyframes grid-pulse {
        0% { opacity: 0.8; }
        50% { opacity: 1; }
        100% { opacity: 0.8; }
    }
    
    @keyframes gradient-shift {
        0% { background-position: 0% 50%; }
        50% { background-position: 100% 50%; }
        100% { background-position: 0% 50%; }
    }
    
    .grid-star {
        text-shadow: 0 0 3px gold;
        animation: grid-glow 2s infinite;
    }
    
    .grid-impuritas {
        animation: grid-rainbow 4s linear infinite;
    }
    
    .grid-nightmare {
        text-shadow: 0 0 1px #FF5722;
        animation: grid-pulse 3s infinite;
    }
    
    .grid-plague {
        text-shadow: 0 0 1px #9C27B0;
    }
    
    /* Custom styles for collection metrics */
    .metric-container {
        background-color: #f8f9fa;
        border-radius: 5px;
        padding: 10px;
        text-align: center;
        border: 1px solid #dee2e6;
    }
    
    .metric-value {
        font-size: 24px;
        font-weight: bold;
    }
    
    .metric-label {
        font-size: 14px;
        color: #6c757d;
    }
    </style>
    """, unsafe_allow_html=True)
    
    # Rarity breakdown
    display_rarity_breakdown()
    
    # Most common tags
    display_common_tags()
    
    # Recently discovered tags
    display_recent_discoveries()

def display_rarity_breakdown():
    """Display rarity breakdown with progress bar visualization showing collected vs total tags"""
    st.subheader("Rarity Collection Progress")
    
    # Get counts of collected tags by rarity
    collected_counts = {}
    for tag_info in st.session_state.collected_tags.values():
        rarity = tag_info["rarity"]
        if rarity not in collected_counts:
            collected_counts[rarity] = 0
        collected_counts[rarity] += 1
    
    # Get total possible tags by rarity from metadata
    total_counts = {}
    if hasattr(st.session_state, 'tag_rarity_metadata') and st.session_state.tag_rarity_metadata:
        for tag, info in st.session_state.tag_rarity_metadata.items():
            if isinstance(info, dict) and "rarity" in info:
                rarity = info["rarity"]
                if rarity not in total_counts:
                    total_counts[rarity] = 0
                total_counts[rarity] += 1
    
    # Get ordered rarities
    ordered_rarities = list(RARITY_LEVELS.keys())
    
    # Create a progress bar for each rarity
    for rarity in ordered_rarities:
        if rarity in collected_counts:
            collected = collected_counts[rarity]
            total = total_counts.get(rarity, 0)
            
            # If total is not available or 0, use collected as total
            if total <= 0:
                total = collected
            
            # Calculate percentage
            percentage = (collected / total) * 100 if total > 0 else 0
            
            # Get color for this rarity
            color = RARITY_LEVELS[rarity]["color"]
            
            # Special styling based on rarity
            rarity_span = f"<span style='color:{color};font-weight:bold;'>{rarity}</span>"
            bar_style = f"background-color: {color};"
            
            if rarity == "Impuritas Civitas":
                rarity_span = f"<span class='grid-impuritas' style='font-weight:bold;'>{rarity}</span>"
                bar_style = "background: linear-gradient(to right, red, orange, yellow, green, blue, indigo, violet); background-size: 400% 400%; animation: gradient-shift 4s linear infinite;"
            elif rarity == "Star of the City":
                rarity_span = f"<span class='grid-star' style='color:{color};font-weight:bold;'>{rarity}</span>"
                bar_style = f"background-color: {color}; box-shadow: 0 0 5px #FFD700;"
            elif rarity == "Urban Nightmare":
                rarity_span = f"<span class='grid-nightmare' style='color:{color};font-weight:bold;'>{rarity}</span>"
                bar_style = f"background-color: {color}; animation: grid-pulse 3s infinite;"
            elif rarity == "Urban Plague":
                rarity_span = f"<span class='grid-plague' style='color:{color};font-weight:bold;'>{rarity}</span>"
            
            # Create progress bar with appropriate styling
            st.markdown(
                f"""
                <div style="display: flex; align-items: center; margin-bottom: 10px;">
                    <div style="width: 150px;">{rarity_span}</div>
                    <div style="flex-grow: 1; background-color: #f0f0f0; border-radius: 5px; height: 20px;">
                        <div style="width: {percentage}%; {bar_style} height: 20px; border-radius: 5px;"></div>
                    </div>
                    <div style="width: 120px; text-align: right;">{collected}/{total} ({percentage:.1f}%)</div>
                </div>
                """,
                unsafe_allow_html=True
            )
            
    # Add the necessary CSS for animations if not already added
    st.markdown("""
    <style>
    @keyframes gradient-shift {
        0% { background-position: 0% 50%; }
        50% { background-position: 100% 50%; }
        100% { background-position: 0% 50%; }
    }
    </style>
    """, unsafe_allow_html=True)

def display_common_tags():
    """Display the most common tags with themed animations"""
    st.subheader("Most Common Tags")
    
    # Sort tags by count
    sorted_tags = sorted(
        st.session_state.collected_tags.items(),
        key=lambda x: x[1]["count"],
        reverse=True
    )
    
    # Display top 10
    for i, (tag, info) in enumerate(sorted_tags[:10]):
        rarity = info["rarity"]
        count = info["count"]
        color = RARITY_LEVELS[rarity]["color"]
        
        # Special styling based on rarity
        rarity_span = f"<span style='color:{color};'>{rarity}</span>"
        
        if rarity == "Impuritas Civitas":
            rarity_span = f"<span class='grid-impuritas'>{rarity}</span>"
        elif rarity == "Star of the City":
            rarity_span = f"<span class='grid-star' style='color:{color};'>{rarity}</span>"
        elif rarity == "Urban Nightmare":
            rarity_span = f"<span class='grid-nightmare' style='color:{color};'>{rarity}</span>"
        elif rarity == "Urban Plague":
            rarity_span = f"<span class='grid-plague' style='color:{color};'>{rarity}</span>"
        
        # Get sample count if available
        sample_display = ""
        if hasattr(st.session_state, 'tag_rarity_metadata') and st.session_state.tag_rarity_metadata:
            if tag in st.session_state.tag_rarity_metadata:
                tag_info = st.session_state.tag_rarity_metadata[tag]
                if isinstance(tag_info, dict) and "sample_count" in tag_info:
                    sample_count = tag_info["sample_count"]
                    if sample_count >= 1000000:
                        sample_display = f"<span style='font-size:0.8em;color:#666;'>({sample_count/1000000:.1f}M samples)</span>"
                    elif sample_count >= 1000:
                        sample_display = f"<span style='font-size:0.8em;color:#666;'>({sample_count/1000:.1f}K samples)</span>"
                    else:
                        sample_display = f"<span style='font-size:0.8em;color:#666;'>({sample_count} samples)</span>"
        
        st.markdown(
            f"**{i+1}.** {tag} - {rarity_span} ({count} occurrences) {sample_display}",
            unsafe_allow_html=True
        )

def display_recent_discoveries():
    """Display recently discovered tags with themed animations"""
    st.subheader("Recent Discoveries")
    
    # Get recently discovered tags from tag history
    if hasattr(st.session_state, 'tag_history') and st.session_state.tag_history:
        # Get last 10 new discoveries
        new_discoveries = [entry for entry in st.session_state.tag_history if entry.get("is_new", False)]
        for entry in new_discoveries[-10:]:
            tag = entry["tag"]
            rarity = entry["rarity"]
            time = entry["time"]
            color = RARITY_LEVELS[rarity]["color"]
            
            # Special styling based on rarity
            rarity_span = f"<span style='color:{color};'>{rarity}</span>"
            
            if rarity == "Impuritas Civitas":
                rarity_span = f"<span class='grid-impuritas'>{rarity}</span>"
            elif rarity == "Star of the City":
                rarity_span = f"<span class='grid-star' style='color:{color};'>{rarity}</span>"
            elif rarity == "Urban Nightmare":
                rarity_span = f"<span class='grid-nightmare' style='color:{color};'>{rarity}</span>"
            elif rarity == "Urban Plague":
                rarity_span = f"<span class='grid-plague' style='color:{color};'>{rarity}</span>"
            
            st.markdown(
                f"{tag} - {rarity_span} (at {time})",
                unsafe_allow_html=True
            )

def display_category_stats():
    """Display stats and details about tag categories"""
    st.subheader("Category Analysis")
    
    # Check if we have collected tags
    if not hasattr(st.session_state, 'collected_tags') or not st.session_state.collected_tags:
        st.info("Start scanning images to collect tags!")
        return
    
    # Count tags by category
    category_counts = {}
    for tag, info in st.session_state.collected_tags.items():
        category = info.get("category", "general")
        if category not in category_counts:
            category_counts[category] = {"count": 0, "tags": []}
        category_counts[category]["count"] += 1
        category_counts[category]["tags"].append(tag)
    
    # Create category bar visualization
    display_category_bars(category_counts)
    
    st.divider()

def display_category_bars(category_counts):
    """Display category distribution with bar visualization"""
    total_tags = len(st.session_state.collected_tags)
    st.subheader("Category Distribution")
    
    # Sort categories by count
    sorted_categories = sorted(
        category_counts.items(),
        key=lambda x: x[1]["count"],
        reverse=True
    )
    
    max_count = max(cat_info["count"] for _, cat_info in sorted_categories) if sorted_categories else 1
    
    for category, cat_info in sorted_categories:
        count = cat_info["count"]
        percentage = (count / total_tags) * 100
        
        if category in TAG_CATEGORIES:
            color = TAG_CATEGORIES[category]["color"]
            icon = TAG_CATEGORIES[category]["icon"]
            name = TAG_CATEGORIES[category]["name"]
        else:
            color = "#888888"
            icon = "❓"
            name = category.capitalize()
        
        # Create bar
        st.markdown(
            f"""
            <div style="display: flex; align-items: center; margin-bottom: 10px;">
                <div style="width: 150px;">{icon} <span style="color:{color};font-weight:bold;">{name}</span></div>
                <div style="flex-grow: 1; background-color: #f0f0f0; border-radius: 5px; height: 20px;">
                    <div style="width: {(count/max_count)*100}%; background-color: {color}; height: 20px; border-radius: 5px;"></div>
                </div>
                <div style="width: 100px; text-align: right;">{count} ({percentage:.1f}%)</div>
            </div>
            """,
            unsafe_allow_html=True
        )

def display_achievements_tab():
    """Display the user's achievements in a dedicated tab"""
    st.subheader("🏆 Achievements")
    
    # Count unlocked achievements
    unlocked_count = len(st.session_state.achievements)
    total_count = len(ACHIEVEMENTS)
    
    # Create a progress bar for overall achievement completion
    progress_percentage = unlocked_count / total_count
    st.progress(progress_percentage, text=f"Unlocked {unlocked_count} of {total_count} achievements ({int(progress_percentage * 100)}%)")
    
    # Group achievements by category
    achievement_categories = group_achievements_by_category()
    
    # Create tabs for different achievement categories
    category_tabs = st.tabs(list(achievement_categories.keys()))
    
    # Display achievements in each category tab
    for i, (category, tab) in enumerate(zip(achievement_categories.keys(), category_tabs)):
        with tab:
            display_category_achievements(category, achievement_categories[category])

def group_achievements_by_category():
    """Group achievements into categories"""
    achievement_categories = {
        "Collection": [],
        "Rarity": [],
        "Progression": [],
        "Special": []
    }
    
    # Categorize achievements
    for achievement_id, achievement in ACHIEVEMENTS.items():
        if "collector" in achievement_id or "Collector" in achievement.get("name", ""):
            category = "Collection"
        elif any(rarity in achievement_id for rarity in ["canard", "myth", "legend", "plague", "nightmare", "star", "impuritas"]):
            category = "Rarity"
        elif any(term in achievement_id for term in ["milestone", "master", "threshold"]):
            category = "Progression"
        else:
            category = "Special"
        
        achievement_categories[category].append((achievement_id, achievement))
        
    return achievement_categories

def display_category_achievements(category, achievements_in_category):
    """Display achievements for a specific category"""
    # If no achievements in this category, show a message
    if not achievements_in_category:
        st.info(f"No {category} achievements available yet.")
        return
    
    # Otherwise, display achievements in a grid
    cols_per_row = 3
    for j in range(0, len(achievements_in_category), cols_per_row):
        cols = st.columns(cols_per_row)
        
        for k in range(cols_per_row):
            idx = j + k
            if idx < len(achievements_in_category):
                achievement_id, achievement = achievements_in_category[idx]
                with cols[k]:
                    display_achievement_card(achievement_id, achievement)

def display_achievement_card(achievement_id, achievement):
    """Display a single achievement card"""
    # Create a card-like container for each achievement
    is_unlocked = achievement_id in st.session_state.achievements
    
    # Set the card style based on unlock status
    if is_unlocked:
        card_style = "border:2px solid #4CAF50; border-radius:10px; padding:10px; margin:5px; background-color:rgba(76, 175, 80, 0.1);"
        icon = "✅"
    else:
        card_style = "border:2px solid #cccccc; border-radius:10px; padding:10px; margin:5px; background-color:rgba(0, 0, 0, 0.05);"
        icon = "🔒"
    
    # Display achievement card
    st.markdown(f"""
    <div style="{card_style}">
        <p style="font-size:18px; margin:0;">{icon} <b>{achievement.get('name', 'Unknown')}</b></p>
        <p style="font-size:14px; margin:5px 0;">{achievement.get('description', '')}</p>
        {f"<p style='font-size:12px; color:#666;'>Requires: {achievement.get('requirement', '')}</p>" if 'requirement' in achievement else ""}
        {f"<p style='font-size:12px; color:#4CAF50;'>Reward: {', '.join([f'{k}: {v}' for k, v in achievement.get('reward', {}).items()])}</p>" if 'reward' in achievement and is_unlocked else ""}
    </div>
    """, unsafe_allow_html=True)

def check_achievements(session_state):
    """
    Check for achievement completion based on the current game state.
    Adds newly completed achievements to the session_state.achievements set
    and applies their rewards.
    
    Args:
        session_state: Streamlit session state containing game data
    
    Returns:
        list: List of newly unlocked achievement names (empty if none)
    """
    from game_constants import ACHIEVEMENTS, RARITY_LEVELS
    
    # Ensure achievements set exists
    if not hasattr(session_state, 'achievements'):
        session_state.achievements = set()
    
    newly_unlocked = []
    
    # Helper function to apply achievement rewards
    def apply_reward(achievement_id):
        reward = ACHIEVEMENTS[achievement_id].get("reward", {})
        
        # Apply tagcoin reward
        if "tagcoins" in reward:
            session_state.tag_currency += reward["tagcoins"]
            
        # Apply coin bonus
        if "coin_bonus" in reward:
            if not hasattr(session_state, 'achievement_coin_bonus'):
                session_state.achievement_coin_bonus = 0
            session_state.achievement_coin_bonus += reward["coin_bonus"]
            
        # Apply enkephalin reward
        if "enkephalin" in reward:
            session_state.enkephalin += reward["enkephalin"]
            
        # Apply essence cost reduction
        if "essence_cost_reduction" in reward:
            if not hasattr(session_state, 'essence_cost_reduction'):
                session_state.essence_cost_reduction = 0
            session_state.essence_cost_reduction += reward["essence_cost_reduction"]
            
        # Apply library cost reduction
        if "library_cost_reduction" in reward:
            if not hasattr(session_state, 'library_cost_reduction'):
                session_state.library_cost_reduction = 0
            session_state.library_cost_reduction += reward["library_cost_reduction"]
            
        # Apply enkephalin bonus
        if "enkephalin_bonus" in reward:
            if not hasattr(session_state, 'enkephalin_bonus'):
                session_state.enkephalin_bonus = 0
            session_state.enkephalin_bonus += reward["enkephalin_bonus"]
    
    # Check collection achievements
    tag_count = len(session_state.collected_tags) if hasattr(session_state, 'collected_tags') else 0
    
    for achievement_id, achievement in ACHIEVEMENTS.items():
        # Skip already unlocked achievements
        if achievement_id in session_state.achievements:
            continue
            
        # Check if this is a collection size achievement
        if achievement_id.startswith("tag_collector_") or achievement_id.startswith("collection_milestone_") or achievement_id == "tag_master":
            requirement = achievement.get("requirement", 0)
            if tag_count >= requirement:
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
        
        # Check rarity-based achievements
        elif achievement_id.endswith("_collector") or achievement_id == "legendary_hunter" or achievement_id == "multi_legendary":
            # Count tags by rarity
            rarity_counts = {}
            for tag_info in session_state.collected_tags.values():
                rarity = tag_info.get("rarity", "Unknown")
                if rarity not in rarity_counts:
                    rarity_counts[rarity] = 0
                rarity_counts[rarity] += 1
            
            # Check for specific rarity achievements
            if achievement_id == "canard_collector" and rarity_counts.get("Canard", 0) >= achievement.get("requirement", 0):
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
                
            elif achievement_id == "urban_myth_collector" and rarity_counts.get("Urban Myth", 0) >= achievement.get("requirement", 0):
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
                
            elif achievement_id == "urban_legend_collector" and rarity_counts.get("Urban Legend", 0) >= achievement.get("requirement", 0):
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
                
            elif achievement_id == "urban_plague_collector" and rarity_counts.get("Urban Plague", 0) >= achievement.get("requirement", 0):
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
                
            elif achievement_id == "urban_nightmare_collector" and rarity_counts.get("Urban Nightmare", 0) >= achievement.get("requirement", 0):
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
                
            elif achievement_id == "star_collector" and rarity_counts.get("Star of the City", 0) >= achievement.get("requirement", 0):
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
                
            elif achievement_id == "impuritas_collector" and rarity_counts.get("Impuritas Civitas", 0) >= achievement.get("requirement", 0):
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
                
            elif achievement_id == "legendary_hunter" and rarity_counts.get("Impuritas Civitas", 0) >= achievement.get("requirement", 0):
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
                
            elif achievement_id == "multi_legendary" and rarity_counts.get("Impuritas Civitas", 0) >= achievement.get("requirement", 0):
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
                
        # Check threshold achievements
        elif achievement_id == "perfect_scanner":
            effective_threshold = session_state.threshold
            if hasattr(session_state, 'tag_power_bonus'):
                effective_threshold -= session_state.tag_power_bonus
            if effective_threshold <= 0.1:  # MIN_THRESHOLD
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
                
        elif achievement_id == "optimal_threshold":
            if abs(session_state.threshold - 0.32857141) < 0.001:  # Within 0.001 of optimal F1 score
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
        
        # Check essence achievements
        elif achievement_id == "essence_creator" or achievement_id == "essence_master":
            essence_count = session_state.game_stats.get("essences_generated", 0)
            if essence_count >= achievement.get("requirement", 0):
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
        
        # Check tag explorer achievement
        elif achievement_id == "tag_explorer":
            explored_tiers = session_state.explored_library_tiers if hasattr(session_state, 'explored_library_tiers') else set()
            if len(explored_tiers) >= achievement.get("requirement", 0):
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
        
        # Check enkephalin master achievement
        elif achievement_id == "enkephalin_master" or achievement_id == "enkephalin_harvester":
            enkephalin_generated = session_state.game_stats.get("enkephalin_generated", 0)
            if enkephalin_generated >= achievement.get("requirement", 0):
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
        
        # Check sacrifice devotee achievement
        elif achievement_id == "sacrifice_devotee":
            tags_sacrificed = session_state.game_stats.get("tags_sacrificed", 0)
            if tags_sacrificed >= achievement.get("requirement", 0):
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
        
        # Check category explorer achievement
        elif achievement_id == "category_explorer":
            # Count unique categories
            categories = set()
            for tag_info in session_state.collected_tags.values():
                category = tag_info.get("category", "unknown")
                categories.add(category)
            
            if len(categories) >= achievement.get("requirement", 0):
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
        
        # Check series collector achievement
        elif achievement_id == "series_collector":
            completed_series = session_state.completed_series if hasattr(session_state, 'completed_series') else set()
            if len(completed_series) >= achievement.get("requirement", 0):
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
        
        # Check rapid tagger achievement
        elif achievement_id == "rapid_tagger":
            images_processed = session_state.game_stats.get("images_processed", 0)
            if images_processed >= achievement.get("requirement", 0):
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
        
        # Check library scholar achievement
        elif achievement_id == "library_scholar":
            extracted_tags = session_state.tags_extracted if hasattr(session_state, 'tags_extracted') else 0
            if extracted_tags >= achievement.get("requirement", 0):
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
        
        # Check rarity hunter achievement
        elif achievement_id == "rarity_hunter":
            # Count tags by rarity
            rarity_counts = {}
            for tag_info in session_state.collected_tags.values():
                rarity = tag_info.get("rarity", "Unknown")
                if rarity not in rarity_counts:
                    rarity_counts[rarity] = 0
                rarity_counts[rarity] += 1
            
            # Check if player has at least one tag of each rarity
            has_all_rarities = all(rarity in rarity_counts and rarity_counts[rarity] > 0 
                                  for rarity in RARITY_LEVELS.keys())
            
            if has_all_rarities:
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
        
        # Check legendary librarian achievement
        elif achievement_id == "legendary_librarian":
            extracted_legendary = session_state.extracted_impuritas if hasattr(session_state, 'extracted_impuritas') else False
            if extracted_legendary:
                session_state.achievements.add(achievement_id)
                apply_reward(achievement_id)
                newly_unlocked.append(achievement["name"])
    
    return newly_unlocked


def display_achievement_notifications(newly_unlocked):
    """
    Display notifications for newly unlocked achievements.
    
    Args:
        newly_unlocked: List of newly unlocked achievement names
    """
    import streamlit as st
    
    if not newly_unlocked:
        return
    
    # Add CSS for achievement notifications
    st.markdown("""
    <style>
    @keyframes slide-in {
        0% { transform: translateX(100%); opacity: 0; }
        10% { transform: translateX(0); opacity: 1; }
        90% { transform: translateX(0); opacity: 1; }
        100% { transform: translateX(100%); opacity: 0; }
    }
    
    .achievement-notification {
        position: fixed;
        top: 70px;
        right: 20px;
        background-color: #4CAF50;
        color: white;
        padding: 15px;
        border-radius: 5px;
        box-shadow: 0 4px 8px rgba(0,0,0,0.2);
        z-index: 9999;
        animation: slide-in 5s ease-in-out;
        animation-fill-mode: forwards;
    }
    </style>
    """, unsafe_allow_html=True)
    
    # Generate notification for each achievement
    for i, achievement_name in enumerate(newly_unlocked):
        delay = i * 0.5  # Stagger notifications
        notification_html = f"""
        <div class="achievement-notification" style="animation-delay: {delay}s; top: {70 + i*70}px;">
            <div>🏆 Achievement Unlocked!</div>
            <div><strong>{achievement_name}</strong></div>
        </div>
        """
        st.markdown(notification_html, unsafe_allow_html=True)


def update_tag_power_bonuses():
    """
    Update tag power bonuses based on the player's collection.
    This affects threshold reduction and coin multiplier.
    """
    import streamlit as st
    from game_constants import TAG_POWER_BONUSES, RARITY_LEVELS
    
    # Initialize tag power variables if they don't exist
    if not hasattr(st.session_state, 'tag_power_bonus'):
        st.session_state.tag_power_bonus = 0
    
    if not hasattr(st.session_state, 'coin_multiplier'):
        st.session_state.coin_multiplier = 1.0
    
    # Reset to defaults
    st.session_state.tag_power_bonus = 0
    coin_multiplier_bonus = 0
    
    # Calculate bonuses from rare tags
    if hasattr(st.session_state, 'collected_tags'):
        for tag, info in st.session_state.collected_tags.items():
            rarity = info.get("rarity")
            if rarity in TAG_POWER_BONUSES:
                bonus = TAG_POWER_BONUSES[rarity]
                
                # Accumulate coin multiplier bonus
                coin_multiplier_bonus += bonus["coin_multiplier"]
    
    # Apply coin multiplier (base 1.0 + bonuses)
    st.session_state.coin_multiplier = 1.0 + coin_multiplier_bonus
    
    # Apply achievement bonuses if they exist
    if hasattr(st.session_state, 'achievement_coin_bonus'):
        st.session_state.coin_multiplier += st.session_state.achievement_coin_bonus
    
    return st.session_state.tag_power_bonus, st.session_state.coin_multiplier

def display_upgrade_shop():
    """Display the upgrade shop with preset threshold levels and visual enhancements"""
    # Display header with visually appealing styling
    st.markdown("""
    <style>
    .upgrade-shop-header {
        background: linear-gradient(90deg, #0d6efd, #6610f2);
        color: white;
        padding: 15px;
        border-radius: 8px;
        margin-bottom: 20px;
        display: flex;
        flex-direction: column;
        align-items: center;
    }
    
    .currency-display {
        background-color: rgba(255, 255, 255, 0.1);
        border-radius: 8px;
        padding: 8px 15px;
        margin-top: 10px;
        font-size: 1.2em;
    }
    </style>
    
    <div class="upgrade-shop-header">
        <h2>🧪 Neural Scanner Upgrade Lab 🧪</h2>
        <div>Enhance your scanner's capabilities to discover rarer tags!</div>
    </div>
    """, unsafe_allow_html=True)
    
    # Add introduction if player has no purchased upgrades
    purchased = set(st.session_state.purchased_upgrades) if hasattr(st.session_state, 'purchased_upgrades') else set()

    # Get base threshold and tag power bonus
    base_threshold = st.session_state.threshold
    tag_power_bonus = st.session_state.tag_power_bonus if hasattr(st.session_state, 'tag_power_bonus') else 0
    effective_threshold = max(MIN_THRESHOLD, base_threshold - tag_power_bonus)
    
    st.divider()
    
    # Create upgrade cards
    display_available_upgrades(base_threshold, tag_power_bonus, purchased)

def display_available_upgrades(base_threshold, tag_power_bonus, purchased):
    """Display available upgrade options with progressive unlocking and visual effects"""
    # Add CSS for upgrade card effects
    st.markdown("""
    <style>
    @keyframes glow-effect {
        0% { box-shadow: 0 0 5px rgba(0,123,255,0.5); }
        50% { box-shadow: 0 0 15px rgba(0,123,255,0.8); }
        100% { box-shadow: 0 0 5px rgba(0,123,255,0.5); }
    }
    
    .upgrade-card {
        border: 1px solid #dee2e6;
        border-radius: 8px;
        padding: 15px;
        margin-bottom: 15px;
        background-color: #f8f9fa;
        transition: transform 0.2s, box-shadow 0.2s;
    }
    
    .upgrade-card:hover {
        transform: translateY(-2px);
        box-shadow: 0 4px 8px rgba(0,0,0,0.1);
    }
    
    /* Rarity-based animations - matching the tag rarities */
    @keyframes rainbow-border {
        0% { border-color: red; }
        14% { border-color: orange; }
        28% { border-color: yellow; }
        42% { border-color: green; }
        57% { border-color: blue; }
        71% { border-color: indigo; }
        85% { border-color: violet; }
        100% { border-color: red; }
    }
    
    @keyframes star-glow {
        0% { box-shadow: 0 0 5px #FFD700; }
        50% { box-shadow: 0 0 20px #FFD700; }
        100% { box-shadow: 0 0 5px #FFD700; }
    }
    
    @keyframes nightmare-pulse {
        0% { border-color: #FF9800; }
        50% { border-color: #FF5722; }
        100% { border-color: #FF9800; }
    }
    
    /* Upgrade card variants */
    .upgrade-card-available {
        border: 1px solid #2196F3;  /* Urban Legend blue */
        background-color: rgba(33, 150, 243, 0.05);
    }
    
    .upgrade-card-affordable {
        border: 2px solid #9C27B0;  /* Urban Plague purple */
        background-color: rgba(156, 39, 176, 0.08);
        box-shadow: 0 0 3px #9C27B0;
    }
    
    .upgrade-card-purchased {
        border: 2px solid #5D9C59;  /* Urban Myth green */
        background-color: rgba(93, 156, 89, 0.08);
    }
    
    .upgrade-card-locked {
        border: 1px solid #AAAAAA;  /* Canard gray */
        background-color: rgba(170, 170, 170, 0.05);
        filter: grayscale(100%);
        opacity: 0.7;
    }
    
    /* Special cards for high-tier upgrades */
    .upgrade-card-nightmare {
        border: 2px solid #FF9800;  /* Urban Nightmare orange */
        background-color: rgba(255, 152, 0, 0.08);
        animation: nightmare-pulse 3s infinite;
    }
    
    .upgrade-card-star {
        border: 2px solid #FFEB3B;  /* Star of the City yellow/gold */
        background-color: rgba(255, 235, 59, 0.08);
        animation: star-glow 2s infinite;
    }
    
    .upgrade-card-impuritas {
        border: 3px solid red;  /* Impuritas Civitas red/rainbow */
        background-color: rgba(0, 0, 0, 0.1);
        animation: rainbow-border 4s linear infinite;
    }
    
    .upgrade-name {
        font-weight: bold;
        font-size: 1.1em;
        margin-bottom: 5px;
    }
    
    .level-indicator {
        display: flex;
        margin-bottom: 10px;
    }
    
    .level-dot {
        width: 12px;
        height: 12px;
        border-radius: 50%;
        margin-right: 5px;
        background-color: #dee2e6;
    }
    
    .level-dot-filled {
        background-color: #0d6efd;
    }
    
    .level-dot-current {
        background-color: #28a745;
        box-shadow: 0 0 5px #28a745;
    }
    
    .locked-overlay {
        font-size: 1.2em;
        color: #6c757d;
    }
    </style>
    """, unsafe_allow_html=True)
    
    # Display purchased upgrades first
    if purchased:
        with st.expander("Your Purchased Upgrades", expanded=False):
            for i, upgrade in enumerate(THRESHOLD_UPGRADES):
                if upgrade["name"] in purchased:
                    # Create a completed upgrade card with appropriate rarity styling
                    card_class = "upgrade-card-purchased"
                    
                    # Apply special styling based on upgrade level
                    if i >= 6:  # Technological Singularity
                        card_class = "upgrade-card-impuritas"
                    elif i >= 5:  # Consciousness Emulation
                        card_class = "upgrade-card-star"
                    elif i >= 4:  # Recursive Self-Improvement
                        card_class = "upgrade-card-nightmare"
                    
                    upgrade_html = f"""
                    <div class="upgrade-card {card_class}">
                        <div class="level-indicator">
                            {create_level_dots(i, len(THRESHOLD_UPGRADES), purchased)}
                        </div>
                        <div class="upgrade-name">{upgrade["name"]} ✓</div>
                        <p>{upgrade["description"]}</p>
                        <p><strong>Threshold setting:</strong> {upgrade["threshold_setting"]:.4f}</p>
                    </div>
                    """
                    st.markdown(upgrade_html, unsafe_allow_html=True)
    
    # Find the next available upgrade (first unpurchased)
    next_available_index = None
    for i, upgrade in enumerate(THRESHOLD_UPGRADES):
        if upgrade["name"] not in purchased:
            next_available_index = i
            break
    
    # Display if all upgrades are purchased
    if next_available_index is None:
        st.success("🎉 Congratulations! You've unlocked all available scanner upgrades!")
        return
    
    # Display available and locked upgrades
    st.subheader("Available Upgrades")
    
    # Show only next available upgrade and a preview of what's next
    for i, upgrade in enumerate(THRESHOLD_UPGRADES):
        # Skip if already purchased
        if upgrade["name"] in purchased:
            continue
        
        # Only show the current upgrade and the next one
        if i > next_available_index + 1:
            continue
        
        # Determine if this is the next available upgrade or a future one
        is_next_upgrade = (i == next_available_index)
        can_afford = st.session_state.tag_currency >= upgrade["cost"]
        
        if is_next_upgrade:
            # Create a card-like container for the next available upgrade
            with st.container():
                # Use the direct threshold setting
                new_threshold = upgrade["threshold_setting"]
                
                # Determine threshold change direction and create text
                threshold_change = abs(new_threshold - base_threshold)
                if new_threshold < base_threshold:
                    change_text = f"Lowers threshold to {new_threshold:.4f} (↓ {threshold_change:.4f})"
                    change_emoji = "⬇️"
                else:
                    change_text = f"Raises threshold to {new_threshold:.4f} (↑ {threshold_change:.4f})"
                    change_emoji = "⬆️"
                
                # Determine card class based on upgrade tier and affordability
                card_class = ""
                
                # First set base affordability class
                base_class = "upgrade-card-affordable" if can_afford else "upgrade-card-available"
                
                # Then add special rarity class based on upgrade level
                if i >= 6:  # Technological Singularity (highest level)
                    card_class = "upgrade-card-impuritas"  # Impuritas Civitas style
                elif i >= 5:  # Consciousness Emulation
                    card_class = "upgrade-card-star"  # Star of the City style
                elif i >= 4:  # Recursive Self-Improvement
                    card_class = "upgrade-card-nightmare"  # Urban Nightmare style
                else:
                    card_class = base_class
                
                # Create the upgrade card HTML
                upgrade_html = f"""
                <div class="upgrade-card {card_class}">
                    <div class="level-indicator">
                        {create_level_dots(i, len(THRESHOLD_UPGRADES), purchased)}
                    </div>
                    <div class="upgrade-name">{upgrade["name"]}</div>
                    <p>{upgrade["description"]}</p>
                    <p><strong>{change_emoji} {change_text}</strong></p>
                </div>
                """
                st.markdown(upgrade_html, unsafe_allow_html=True)
                
                # Show effective threshold with tag power bonus
                if tag_power_bonus > 0:
                    effective_upgrade_threshold = max(MIN_THRESHOLD, new_threshold - tag_power_bonus)
                    st.info(f"Effective threshold: {effective_upgrade_threshold:.4f} with your Tag Power bonus of {tag_power_bonus:.4f}")
                
                # Purchase button
                col1, col2 = st.columns([3, 1])
                with col1:
                    st.write(f"**Cost:** {upgrade['cost']} {TAG_CURRENCY_NAME}")
                with col2:
                    if st.button("Purchase", key=f"upgrade_{i}", disabled=not can_afford, use_container_width=True):
                        purchase_upgrade(i)
        else:
            # Create a preview of the next locked upgrade
            # Add a hint of the upgrade's rarity styling for locked cards
            locked_class = "upgrade-card-locked"
            rarity_hint = ""
            
            if i >= 6:  # Technological Singularity
                rarity_hint = f'<div style="height: 3px; width: 50px; background: linear-gradient(to right, red, orange, yellow, green, blue, indigo, violet); margin: 5px 0;"></div>'
            elif i >= 5:  # Consciousness Emulation
                rarity_hint = f'<div style="height: 3px; width: 50px; background: #FFEB3B; box-shadow: 0 0 3px #FFD700; margin: 5px 0;"></div>'
            elif i >= 4:  # Recursive Self-Improvement
                rarity_hint = f'<div style="height: 3px; width: 50px; background: #FF9800; margin: 5px 0;"></div>'
            
            locked_html = f"""
            <div class="upgrade-card {locked_class}">
                <div class="level-indicator">
                    {create_level_dots(i, len(THRESHOLD_UPGRADES), purchased)}
                </div>
                <div class="upgrade-name">{upgrade["name"]}</div>
                {rarity_hint}
                <p style="color: #6c757d;">Complete previous upgrade to unlock</p>
                <p><strong>Cost:</strong> {upgrade['cost']} {TAG_CURRENCY_NAME}</p>
            </div>
            """
            st.markdown(locked_html, unsafe_allow_html=True)
    
    # Add a hint about next upgrades
    if next_available_index + 2 < len(THRESHOLD_UPGRADES):
        remaining = len(THRESHOLD_UPGRADES) - (next_available_index + 2)
        st.markdown(f"<div style='text-align:center;color:#6c757d;'>{remaining} more upgrades will be revealed as you progress</div>", unsafe_allow_html=True)

def create_level_dots(current_level, total_levels, purchased):
    """Create HTML for level indicator dots"""
    dots_html = ""
    for i in range(total_levels):
        if f"Pattern Recognition Module" in purchased and i == 0:
            # First upgrade is purchased
            dot_class = "level-dot level-dot-filled"
        elif i < current_level:
            # Previous levels are filled
            dot_class = "level-dot level-dot-filled"
        elif i == current_level:
            # Current level is highlighted
            dot_class = "level-dot level-dot-current"
        else:
            # Future levels are empty
            dot_class = "level-dot"
        
        dots_html += f'<div class="{dot_class}"></div>'
    
    return dots_html

def purchase_upgrade(upgrade_index):
    """Purchase a threshold upgrade"""
    upgrade = THRESHOLD_UPGRADES[upgrade_index]
    
    # Check if we have enough currency
    if st.session_state.tag_currency >= upgrade["cost"]:
        # Purchase the upgrade
        st.session_state.tag_currency -= upgrade["cost"]
        st.session_state.game_stats["currency_spent"] += upgrade["cost"]
        
        # Apply new threshold setting directly instead of reduction
        st.session_state.threshold = upgrade["threshold_setting"]
        
        # Add to purchased upgrades
        st.session_state.purchased_upgrades.append(upgrade["name"])
        
        # Save immediately
        save_game()
        
        # Update state version to force a refresh
        if 'state_version' in st.session_state:
            st.session_state.state_version += 1
        
        # Force refresh the UI to show updated state
        st.rerun()
    else:
        st.error("Not enough currency!")

def display_neural_specialization():
    """Display neural specialization grid"""
    st.subheader("Neural Specialization")
    
    # Count tags by category
    categories_count = {}
    if hasattr(st.session_state, 'collected_tags'):
        for tag, info in st.session_state.collected_tags.items():
            category = info.get("category", "general")
            if category not in categories_count:
                categories_count[category] = 0
            categories_count[category] += 1
    
    # Check for mastered categories
    mastered = st.session_state.mastered_categories if hasattr(st.session_state, 'mastered_categories') else {}
    
    # Make a grid of category stats
    cols = st.columns(3)
    col_idx = 0
    
    for category, info in TAG_CATEGORIES.items():
        with cols[col_idx]:
            count = categories_count.get(category, 0)
            is_mastered = category in mastered
            
            # Get color and status
            if is_mastered:
                color = info['color']
                status = "MASTERED"
            else:
                color = "#AAAAAA"  # Gray
                status = "Active"
            
            # Display category stats
            st.markdown(f"""
            <div style="border: 1px solid {color}; border-radius: 5px; padding: 10px; margin-bottom: 10px;">
                <p style="margin: 0; font-weight: bold;">{info['icon']} {info['name']}</p>
                <p style="margin: 0; font-size: 0.9em;">Tags: {count}</p>
                <p style="margin: 0; font-size: 0.9em; color: {color};">{status}</p>
            </div>
            """, unsafe_allow_html=True)
            
            # Move to next column
            col_idx = (col_idx + 1) % 3

def display_discovery_tips():
    """Display tag discovery tips"""
    with st.expander("💡 Tag Discovery Tips"):
        st.write("""
        **How to improve your model's cognitive abilities:**
        
        1. **Collect diverse tags** across many categories to increase your collection level
        2. **Master categories** by collecting many tags within each category
        3. **Find tag combinations** by collecting specific sets of related tags
        4. **Keep rare tags** in your collection to maintain tag power
        5. **Extract tags** from the Library System to discover new tags
        6. **Generate tag essences** to gain insight into what the AI recognizes
        """)

# =============================================================================
# MAIN APPLICATION FUNCTION
# =============================================================================

def is_essence_generation_available():
    """Check if PyTorch model is available for essence generation"""
    return (hasattr(st.session_state, 'model_torch') and 
            st.session_state.model_torch is not None and
            isinstance(st.session_state.model_torch, nn.Module))

def create_app_tabs():
    # Determine if essence generation is available
    essence_available = is_essence_generation_available()
    
    if essence_available:
        tab_names = [
            "Scan Images", 
            "Tag Collection",
            "Series Collections", 
            "Upgrade Shop", 
            "Library", 
            "Essence Generator"
        ]
    else:
        tab_names = [
            "Scan Images", 
            "Tag Collection",
            "Series Collections", 
            "Upgrade Shop", 
            "Library"
        ]
    
    tabs = st.tabs(tab_names)
    
    if essence_available:
        return tabs
    else:
        # Return tabs + None for the essence tab
        return tabs + (None,)

def main():
    """Main application function."""
    # Set up title with emoji
    st.title("🎮 Tag Collector Game")
    
    # Apply tag animations and themes
    apply_tag_animations()

    load_game()
    
    # Initialize game state
    initialize_game_state()
    initialize_library_system()  # Initialize the library system
    
    # Try to load the models
    if 'model' not in st.session_state:
        try:
            onnx_model, torch_model, metadata = load_model()
            st.session_state.model = onnx_model  # ONNX model for inference
            st.session_state.model_torch = torch_model  # PyTorch model for essence generation
            st.session_state.metadata = metadata
            
            # Show model availability status
            if torch_model is not None:
                st.success("✅ Both inference and essence generation models loaded successfully!")
            else:
                st.warning("⚠️ Only inference model loaded. Essence generation unavailable.")
                
        except Exception as e:
            st.error(f"Error loading model: {str(e)}")
            st.code(traceback.format_exc())
            st.stop()
    
    # Create game stats panel at the top
    display_game_stats_panel()
    
    # Create a simplified sidebar without game stats
    create_sidebar()
    
    # Create app tabs
    if is_essence_generation_available():
        tab1, tab2, tab3, tab4, tab5, tab6 = create_app_tabs()
    else:
        tab1, tab2, tab3, tab4, tab5, tab6 = create_app_tabs()
        tab6 = None
    
    # Fill tab content
    with tab1:
        display_scan_interface()
    
    with tab2:
        display_tag_collection()
    
    with tab3:
        display_series_mosaics()
        
    with tab4:
        display_upgrade_shop()
    
    with tab5:
        display_library_extraction()
            
    if tab6 is not None:
        with tab6:
            # Display the essence generator in its own tab
            from essence_generator import display_essence_generator
            display_essence_generator()

if __name__ == "__main__":
    main()