"""
Example demonstrating weighted PPMI with context distribution smoothing (alpha parameter).

The alpha parameter adjusts the context probability distribution:
- alpha = 1.0: Standard PPMI (no smoothing)
- alpha < 1.0 (e.g., 0.75): Increases probability of rare contexts, reducing their PPMI
- alpha > 1.0: Decreases probability of rare contexts, increasing their PPMI

Levy et al. (2015) found that alpha=0.75 improves performance on many tasks.
"""

from ppmi import PMI


def example_weighted_ppmi_comparison():
    """Compare standard PPMI with weighted PPMI (alpha=0.75)."""
    print("=" * 70)
    print("Weighted PPMI with Context Distribution Smoothing")
    print("=" * 70)
    print()
    
    pmi_calc = PMI()
    
    # Simulate word-context co-occurrences
    # "data" co-occurs frequently with common contexts and rarely with specific ones
    pmi_calc.add_observation("data", "the", count=100)      # very common context
    pmi_calc.add_observation("data", "analysis", count=20)   # specific context
    pmi_calc.add_observation("data", "science", count=15)    # specific context
    pmi_calc.add_observation("data", "mining", count=8)      # rare specific context
    
    # Add some other words to create a realistic distribution
    pmi_calc.add_observation("information", "the", count=80)
    pmi_calc.add_observation("research", "the", count=60)
    pmi_calc.add_observation("machine", "learning", count=25)
    pmi_calc.add_observation("statistics", "analysis", count=18)
    
    # Compare different alpha values
    contexts = ["the", "analysis", "science", "mining"]
    alphas = [1.0, 0.75, 0.5]
    
    print(f"Target word: 'data'")
    print()
    print(f"{'Context':<15} {'alpha=1.0':<12} {'alpha=0.75':<12} {'alpha=0.5':<12}")
    print("-" * 70)
    
    for context in contexts:
        ppmi_values = []
        for alpha in alphas:
            ppmi = pmi_calc.calculate_ppmi("data", context, alpha=alpha)
            ppmi_values.append(ppmi)
        
        count = pmi_calc.co_occurrence["data"][context]
        print(f"{context:<15} {ppmi_values[0]:<12.4f} {ppmi_values[1]:<12.4f} "
              f"{ppmi_values[2]:<12.4f} (count: {count})")
    
    print()
    print("Observations:")
    print("- Common contexts (e.g., 'the') have lower PPMI with smoothing")
    print("- Rare specific contexts (e.g., 'mining') show bigger reduction")
    print("- This helps reduce the bias toward rare events in word embeddings")
    print()


def example_effect_on_rare_contexts():
    """Demonstrate how alpha affects rare vs common contexts differently."""
    print("=" * 70)
    print("Effect of Alpha on Rare vs Common Contexts")
    print("=" * 70)
    print()
    
    pmi_calc = PMI()
    
    # Create a scenario with very common and very rare contexts
    pmi_calc.add_observation("word", "common", count=1000)
    pmi_calc.add_observation("word", "uncommon", count=50)
    pmi_calc.add_observation("word", "rare", count=5)
    
    # Add background for other words
    for i in range(10):
        pmi_calc.add_observation(f"other{i}", "common", count=100)
    
    contexts = ["common", "uncommon", "rare"]
    
    print(f"Target word: 'word'")
    print()
    print(f"{'Context':<12} {'Count':<8} {'Standard':<12} {'Smoothed':<12} {'% Change':<12}")
    print("-" * 70)
    
    for context in contexts:
        count = pmi_calc.co_occurrence["word"][context]
        ppmi_std = pmi_calc.calculate_ppmi("word", context, alpha=1.0)
        ppmi_smooth = pmi_calc.calculate_ppmi("word", context, alpha=0.75)
        
        if ppmi_std > 0:
            change = ((ppmi_smooth - ppmi_std) / ppmi_std) * 100
        else:
            change = 0
        
        print(f"{context:<12} {count:<8} {ppmi_std:<12.4f} {ppmi_smooth:<12.4f} "
              f"{change:<12.2f}%")
    
    print()
    print("Key insight: Rare contexts show larger percentage changes,")
    print("helping to reduce their inflated importance in the model.")
    print()


def example_recommended_alpha():
    """Show the recommended alpha=0.75 value in practice."""
    print("=" * 70)
    print("Using Recommended Alpha=0.75 (Levy et al., 2015)")
    print("=" * 70)
    print()
    
    pmi_calc = PMI()
    
    # Add realistic NLP data
    observations = [
        ("king", "queen", 15),
        ("king", "royal", 20),
        ("king", "the", 100),
        ("king", "and", 80),
        ("queen", "royal", 18),
        ("queen", "the", 90),
        ("royal", "palace", 12),
        ("palace", "the", 50),
    ]
    
    pmi_calc.add_observations_batch(observations)
    
    # Get all PPMI scores with different alpha values
    all_ppmi_std = pmi_calc.get_all_ppmi(alpha=1.0)
    all_ppmi_075 = pmi_calc.get_all_ppmi(alpha=0.75)
    
    print("Top associations by Standard PPMI:")
    sorted_std = sorted(all_ppmi_std.items(), key=lambda x: x[1], reverse=True)[:5]
    for (w, c), score in sorted_std:
        print(f"  {w:<10} + {c:<10} = {score:.4f}")
    
    print()
    print("Top associations by Weighted PPMI (alpha=0.75):")
    sorted_075 = sorted(all_ppmi_075.items(), key=lambda x: x[1], reverse=True)[:5]
    for (w, c), score in sorted_075:
        print(f"  {w:<10} + {c:<10} = {score:.4f}")
    
    print()
    print("Benefits of alpha=0.75:")
    print("- Reduces bias toward rare co-occurrences")
    print("- Improves quality of word embeddings")
    print("- Better performance on semantic similarity tasks")
    print()


if __name__ == "__main__":
    example_weighted_ppmi_comparison()
    example_effect_on_rare_contexts()
    example_recommended_alpha()
    
    print("=" * 70)
    print("All examples completed!")
    print("=" * 70)
