"""
Text analysis example using the PPMI library.

This example demonstrates how to use PPMI for analyzing word co-occurrences
in text, which is useful for NLP tasks.
"""

from ppmi import PMI
from collections import defaultdict
import re


def extract_word_pairs(text, window_size=2):
    """
    Extract word pairs within a sliding window.
    
    Args:
        text: Input text
        window_size: Number of words to consider on each side
    
    Returns:
        List of (word, context_word) pairs
    """
    # Simple tokenization (lowercase and split)
    words = re.findall(r'\b[a-z]+\b', text.lower())
    
    pairs = []
    for i, word in enumerate(words):
        # Get context words within window
        start = max(0, i - window_size)
        end = min(len(words), i + window_size + 1)
        
        for j in range(start, end):
            if i != j:
                context_word = words[j]
                pairs.append((word, context_word))
    
    return pairs


def analyze_text_cooccurrence():
    """Analyze word co-occurrences in a sample text."""
    print("=" * 70)
    print("Text Co-occurrence Analysis with PPMI")
    print("=" * 70)
    print()
    
    # Sample text
    text = """
    Machine learning is a subset of artificial intelligence. 
    Deep learning is a type of machine learning that uses neural networks.
    Neural networks are inspired by biological neural networks.
    Artificial intelligence and machine learning are transforming technology.
    Deep learning models require large amounts of data for training.
    Machine learning algorithms can learn from data without explicit programming.
    """
    
    print("Sample Text:")
    print(text.strip())
    print()
    
    # Extract word pairs
    pairs = extract_word_pairs(text, window_size=2)
    
    # Calculate PPMI
    pmi_calc = PMI()
    pmi_calc.add_observations_batch(pairs)
    
    # Get top PPMI pairs
    all_ppmi = pmi_calc.get_all_ppmi()
    sorted_ppmi = sorted(all_ppmi.items(), key=lambda x: x[1], reverse=True)
    
    print("Top 20 Word Associations (by PPMI):")
    print(f"{'Rank':<6} {'Word Pair':<35} {'PPMI':<10} {'Count':<8}")
    print("-" * 70)
    
    for rank, ((word, context), ppmi_score) in enumerate(sorted_ppmi[:20], 1):
        count = pmi_calc.co_occurrence[word][context]
        print(f"{rank:<6} {f'({word}, {context})':<35} {ppmi_score:<10.4f} {count:<8}")
    
    print()
    
    # Analyze specific words
    target_words = ["learning", "neural", "intelligence"]
    
    print("Strongest associations for specific words:")
    print()
    
    for target_word in target_words:
        print(f"'{target_word}' associates most strongly with:")
        
        # Get all pairs involving the target word
        associations = {}
        for (word, context), score in all_ppmi.items():
            if word == target_word:
                associations[context] = score
        
        # Sort and display top 5
        top_assoc = sorted(associations.items(), key=lambda x: x[1], reverse=True)[:5]
        for i, (assoc_word, score) in enumerate(top_assoc, 1):
            count = pmi_calc.co_occurrence[target_word][assoc_word]
            print(f"  {i}. {assoc_word:<15} PPMI={score:.4f} (co-occurred {count} times)")
        print()


def compare_different_windows():
    """Compare PPMI scores with different window sizes."""
    print("=" * 70)
    print("Comparing Different Window Sizes")
    print("=" * 70)
    print()
    
    text = """
    The cat sat on the mat. The dog ran in the park.
    The cat and the dog are best friends.
    The cat likes to sleep. The dog likes to play.
    """
    
    window_sizes = [1, 2, 3]
    
    print("Target pair: ('cat', 'dog')")
    print()
    
    for window_size in window_sizes:
        pairs = extract_word_pairs(text, window_size=window_size)
        pmi_calc = PMI()
        pmi_calc.add_observations_batch(pairs)
        
        ppmi = pmi_calc.calculate_ppmi("cat", "dog")
        count = pmi_calc.co_occurrence["cat"]["dog"]
        
        print(f"Window size {window_size}:")
        print(f"  Co-occurrence count: {count}")
        print(f"  PPMI(cat, dog): {ppmi:.4f}")
        print()


def word_similarity_demo():
    """Demonstrate using PPMI for word similarity."""
    print("=" * 70)
    print("Word Similarity using PPMI Vectors")
    print("=" * 70)
    print()
    
    text = """
    The dog is a loyal animal. Dogs are friendly pets.
    The cat is an independent animal. Cats are curious pets.
    The bird can fly high. Birds have colorful feathers.
    The fish swims in water. Fish live in the ocean.
    """
    
    pairs = extract_word_pairs(text, window_size=2)
    pmi_calc = PMI()
    pmi_calc.add_observations_batch(pairs)
    
    # Get PPMI matrix
    matrix, words, contexts = pmi_calc.to_matrix(metric='ppmi', base=2)
    
    print(f"Vocabulary size: {len(words)} words")
    print(f"Context size: {len(contexts)} contexts")
    print(f"Matrix shape: {matrix.shape}")
    print()
    
    # Show PPMI vectors for specific words
    target_words = ["dog", "cat", "bird"]
    
    print("PPMI vectors (showing non-zero values):")
    for word in target_words:
        if word in words:
            idx = words.index(word)
            vector = matrix[idx]
            non_zero = [(contexts[i], vector[i]) for i in range(len(vector)) if vector[i] > 0]
            
            print(f"\n'{word}':")
            for context, value in sorted(non_zero, key=lambda x: x[1], reverse=True)[:5]:
                print(f"  {context:<15} {value:.4f}")


if __name__ == "__main__":
    analyze_text_cooccurrence()
    print("\n" + "=" * 70 + "\n")
    compare_different_windows()
    print("\n" + "=" * 70 + "\n")
    word_similarity_demo()
    
    print("\n" + "=" * 70)
    print("Text analysis examples completed!")
    print("=" * 70)
