#!/usr/bin/env python3
"""
Test script for softmax-based word selection in ThematicWordService.
"""

import os
import sys

# Add src directory to path
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

def test_config_loading():
    """Test configuration loading from environment variables"""
    print("🧪 Testing ThematicWordService configuration loading...")
    
    # Set test environment variables
    os.environ['SIMILARITY_TEMPERATURE'] = '0.5' 
    os.environ['USE_SOFTMAX_SELECTION'] = 'true'
    
    from services.thematic_word_service import ThematicWordService
    
    service = ThematicWordService()
    
    print(f"   Similarity Temperature: {service.similarity_temperature}")
    print(f"   Use Softmax Selection: {service.use_softmax_selection}")
    
    # Test environment variable changes
    os.environ['SIMILARITY_TEMPERATURE'] = '1.2'
    os.environ['USE_SOFTMAX_SELECTION'] = 'false'
    
    service2 = ThematicWordService()
    print(f"   After env change - Temperature: {service2.similarity_temperature}")
    print(f"   After env change - Use Softmax: {service2.use_softmax_selection}")
    
    print("✅ Configuration test passed!")

def test_softmax_logic():
    """Test just the softmax logic without full service initialization"""
    print("\n🧪 Testing softmax selection logic...")
    
    import numpy as np
    
    # Mock word data similar to what ThematicWordService uses
    mock_words = [
        {"word": "ELEPHANT", "similarity": 0.85, "clue": "Large African mammal", "tier": "tier_5_common"},
        {"word": "TIGER", "similarity": 0.75, "clue": "Striped big cat", "tier": "tier_6_moderately_common"}, 
        {"word": "DOG", "similarity": 0.65, "clue": "Domestic pet", "tier": "tier_4_highly_common"},
        {"word": "CAT", "similarity": 0.55, "clue": "Feline pet", "tier": "tier_3_very_common"},
        {"word": "FISH", "similarity": 0.45, "clue": "Aquatic animal", "tier": "tier_5_common"},
        {"word": "BIRD", "similarity": 0.35, "clue": "Flying animal", "tier": "tier_4_highly_common"},
        {"word": "ANT", "similarity": 0.25, "clue": "Small insect", "tier": "tier_7_somewhat_uncommon"},
    ]
    
    # Test the actual ThematicWordService softmax logic
    class MockService:
        def __init__(self, temperature=0.7):
            self.similarity_temperature = temperature
        
        def _softmax_with_temperature(self, scores, temperature=1.0):
            if temperature <= 0:
                temperature = 0.01
            scaled_scores = scores / temperature
            max_score = np.max(scaled_scores)
            exp_scores = np.exp(scaled_scores - max_score)
            probabilities = exp_scores / np.sum(exp_scores)
            return probabilities
            
        def _softmax_weighted_selection(self, candidates, num_words, temperature=None):
            if len(candidates) <= num_words:
                return candidates
                
            if temperature is None:
                temperature = self.similarity_temperature
            
            similarities = np.array([word_data['similarity'] for word_data in candidates])
            probabilities = self._softmax_with_temperature(similarities, temperature)
            
            selected_indices = np.random.choice(
                len(candidates), 
                size=min(num_words, len(candidates)),
                replace=False,
                p=probabilities
            )
            
            return [candidates[i] for i in selected_indices]
    
    service = MockService(temperature=0.7)
    
    print("   Testing selection variability (temperature=0.7):")
    for run in range(3):
        selected = service._softmax_weighted_selection(mock_words, 4)
        # Sort by similarity for consistent display
        selected.sort(key=lambda x: x['similarity'], reverse=True)
        words = [f"{word['word']}({word['similarity']:.2f})" for word in selected]
        print(f"   Run {run+1}: {', '.join(words)}")
    
    print("✅ Softmax selection logic test passed!")

def test_environment_integration():
    """Test environment variable integration"""
    print("\n🧪 Testing backend environment integration...")
    
    # Test configuration scenarios
    scenarios = [
        {"SIMILARITY_TEMPERATURE": "0.3", "USE_SOFTMAX_SELECTION": "true", "desc": "Deterministic"},
        {"SIMILARITY_TEMPERATURE": "0.7", "USE_SOFTMAX_SELECTION": "true", "desc": "Balanced"},
        {"SIMILARITY_TEMPERATURE": "1.5", "USE_SOFTMAX_SELECTION": "true", "desc": "Random"},
        {"SIMILARITY_TEMPERATURE": "0.7", "USE_SOFTMAX_SELECTION": "false", "desc": "Disabled"},
    ]
    
    for scenario in scenarios:
        # Set environment variables
        for key, value in scenario.items():
            if key != "desc":
                os.environ[key] = value
        
        # Import fresh service (without initialization to avoid long loading times)
        if 'services.thematic_word_service' in sys.modules:
            del sys.modules['services.thematic_word_service']
        
        from services.thematic_word_service import ThematicWordService
        service = ThematicWordService()
        
        print(f"   {scenario['desc']}: T={service.similarity_temperature}, Enabled={service.use_softmax_selection}")
    
    print("✅ Environment integration test passed!")

if __name__ == "__main__":
    test_config_loading()
    test_softmax_logic()
    test_environment_integration()
    print("\n🎉 All ThematicWordService tests completed successfully!")
    print("\n📝 Usage in production:")
    print("   export SIMILARITY_TEMPERATURE=0.7")
    print("   export USE_SOFTMAX_SELECTION=true")
    print("   # Backend will automatically use these settings")