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

import os
import sys

# Set environment variables for testing
os.environ['SIMILARITY_TEMPERATURE'] = '0.7'  
os.environ['USE_SOFTMAX_SELECTION'] = 'true'

# Test the configuration loading
def test_config_loading():
    from thematic_word_generator import UnifiedThematicWordGenerator
    
    print("🧪 Testing configuration loading...")
    
    # Test default values
    generator = UnifiedThematicWordGenerator()
    print(f"   Similarity Temperature: {generator.similarity_temperature}")
    print(f"   Use Softmax Selection: {generator.use_softmax_selection}")
    
    # Test environment variable override
    os.environ['SIMILARITY_TEMPERATURE'] = '0.3'
    os.environ['USE_SOFTMAX_SELECTION'] = 'false'
    
    generator2 = UnifiedThematicWordGenerator()
    print(f"   After env change - Temperature: {generator2.similarity_temperature}")
    print(f"   After env change - Use Softmax: {generator2.use_softmax_selection}")
    
    print("✅ Configuration test passed!")

def test_softmax_logic():
    """Test just the softmax logic without full initialization"""
    import numpy as np
    
    print("\n🧪 Testing softmax selection logic...")
    
    # Mock data - candidates with (word, similarity, tier)
    candidates = [
        ("elephant", 0.85, "tier_5_common"),
        ("tiger", 0.75, "tier_6_moderately_common"),
        ("dog", 0.65, "tier_4_highly_common"),
        ("cat", 0.55, "tier_3_very_common"),
        ("fish", 0.45, "tier_5_common"),
        ("bird", 0.35, "tier_4_highly_common"),
        ("ant", 0.25, "tier_7_somewhat_uncommon"),
    ]
    
    # Test multiple runs to see randomness
    print("   Testing selection variability (temperature=0.7):")
    
    class MockGenerator:
        def __init__(self):
            self.similarity_temperature = 0.7
            
        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([score for _, score, _ 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]
    
    generator = MockGenerator()
    
    # Run selection multiple times to show variety
    for run in range(3):
        selected = generator._softmax_weighted_selection(candidates, 4)
        selected.sort(key=lambda x: x[1], reverse=True)  # Sort by similarity for display
        words = [f"{word}({sim:.2f})" for word, sim, _ in selected]
        print(f"   Run {run+1}: {', '.join(words)}")
    
    print("✅ Softmax selection logic test passed!")

if __name__ == "__main__":
    test_config_loading()
    test_softmax_logic()
    print("\n🎉 All tests completed successfully!")