"""
Example of loading observations from a single-column CSV (text as word list).
"""

from ppmi import PMI
import os

# Get the directory where this script is located
script_dir = os.path.dirname(os.path.abspath(__file__))
text_csv = os.path.join(script_dir, 'text_words.csv')


def example_text_csv_default_window():
    """Load words from CSV and create co-occurrences with default window."""
    print("=" * 70)
    print("Single Column CSV - Default Window Size (2)")
    print("=" * 70)
    print()
    
    pmi_calc = PMI()
    
    # Load from single column CSV
    pairs_added = pmi_calc.add_observations_from_text_csv(
        text_csv,
        text_col='word',
        window_size=2,
        has_header=True
    )
    
    print(f"Added {pairs_added} word pairs from text")
    print(f"Total co-occurrence count: {pmi_calc.total}")
    print()
    
    # Show top PPMI associations
    print("Top word associations (PPMI):")
    all_ppmi = pmi_calc.get_all_ppmi()
    sorted_ppmi = sorted(all_ppmi.items(), key=lambda x: x[1], reverse=True)
    
    for (word, context), score in sorted_ppmi[:15]:
        count = pmi_calc.co_occurrence[word][context]
        print(f"  {word:<10} + {context:<10} = {score:.4f} "
              f"(co-occurred {count} times)")


def example_compare_window_sizes():
    """Compare results with different window sizes."""
    print()
    print("=" * 70)
    print("Comparing Different Window Sizes")
    print("=" * 70)
    print()
    
    window_sizes = [1, 2, 3]
    
    for window_size in window_sizes:
        pmi_calc = PMI()
        pairs_added = pmi_calc.add_observations_from_text_csv(
            text_csv,
            text_col='word',
            window_size=window_size,
            has_header=True
        )
        
        print(f"Window size {window_size}:")
        print(f"  Total pairs added: {pairs_added}")
        
        # Check association between 'cat' and 'dog'
        if pmi_calc.co_occurrence['cat']['dog'] > 0:
            ppmi = pmi_calc.calculate_ppmi('cat', 'dog')
            count = pmi_calc.co_occurrence['cat']['dog']
            print(f"  PPMI(cat, dog): {ppmi:.4f} "
                  f"(co-occurred {count} times)")
        else:
            print("  'cat' and 'dog' did not co-occur within this window")
        print()


def example_specific_word_context():
    """Show context words for specific target words."""
    print("=" * 70)
    print("Context Analysis for Specific Words")
    print("=" * 70)
    print()
    
    pmi_calc = PMI()
    pmi_calc.add_observations_from_text_csv(
        text_csv,
        text_col='word',
        window_size=2,
        has_header=True
    )
    
    target_words = ['cat', 'dog', 'the']
    
    for target in target_words:
        print(f"\nContexts for '{target}':")
        
        # Get all contexts for this word
        contexts = {}
        all_ppmi = pmi_calc.get_all_ppmi()
        for (word, context), score in all_ppmi.items():
            if word == target:
                contexts[context] = score
        
        # Show top contexts
        top_contexts = sorted(
            contexts.items(), key=lambda x: x[1], reverse=True
        )[:5]
        for context, score in top_contexts:
            count = pmi_calc.co_occurrence[target][context]
            print(f"  → {context:<10} PPMI={score:.4f} (count={count})")


def example_create_custom_text():
    """Create and analyze a custom text."""
    print()
    print("=" * 70)
    print("Custom Text Example")
    print("=" * 70)
    print()
    
    # Create a custom CSV with a sentence
    custom_csv = os.path.join(script_dir, 'custom_text.csv')
    with open(custom_csv, 'w', encoding='utf-8') as f:
        f.write("word\n")
        # "machine learning is amazing machine learning models"
        words = ['machine', 'learning', 'is', 'amazing', 'machine',
                 'learning', 'models', 'are', 'powerful']
        for word in words:
            f.write(f"{word}\n")
    
    pmi_calc = PMI()
    pairs_added = pmi_calc.add_observations_from_text_csv(
        custom_csv,
        text_col='word',
        window_size=2
    )
    
    text_desc = (
        "Text: \"machine learning is amazing machine learning "
        "models are powerful\""
    )
    print(text_desc)
    print(f"Pairs added: {pairs_added}")
    print()
    
    # Show all PPMI scores
    print("Word associations:")
    all_ppmi = pmi_calc.get_all_ppmi()
    for (word, context), score in sorted(
        all_ppmi.items(), key=lambda x: x[1], reverse=True
    ):
        count = pmi_calc.co_occurrence[word][context]
        if score > 0:  # Only show positive associations
            print(f"  ({word}, {context}): PPMI={score:.4f}, count={count}")
    
    # Clean up
    os.remove(custom_csv)


if __name__ == "__main__":
    example_text_csv_default_window()
    example_compare_window_sizes()
    example_specific_word_context()
    example_create_custom_text()
    
    print("\n" + "=" * 70)
    print("Single-column CSV examples completed!")
    print("=" * 70)
