"""
Basic usage examples for the PPMI library.
"""

from ppmi import PMI, calculate_pmi_from_probabilities, calculate_ppmi_from_probabilities
from ppmi.pmi import calculate_pmi_from_counts


def example1_basic_pmi_calculation():
    """Example 1: Basic PMI calculation from probabilities."""
    print("=" * 60)
    print("Example 1: Basic PMI calculation from probabilities")
    print("=" * 60)
    
    # Calculate PMI from known probabilities
    p_xy = 0.1   # P(x and y occur together)
    p_x = 0.3    # P(x occurs)
    p_y = 0.4    # P(y occurs)
    
    pmi = calculate_pmi_from_probabilities(p_xy, p_x, p_y, base=2)
    ppmi = calculate_ppmi_from_probabilities(p_xy, p_x, p_y, base=2)
    
    print(f"P(x,y) = {p_xy}")
    print(f"P(x) = {p_x}")
    print(f"P(y) = {p_y}")
    print(f"PMI(x, y) = {pmi:.4f}")
    print(f"PPMI(x, y) = {ppmi:.4f}")
    print()


def example2_using_pmi_class():
    """Example 2: Using PMI class to track observations."""
    print("=" * 60)
    print("Example 2: Using PMI class to track observations")
    print("=" * 60)
    
    # Create PMI calculator
    pmi_calc = PMI()
    
    # Add observations
    pmi_calc.add_observation("cat", "animal")
    pmi_calc.add_observation("dog", "animal")
    pmi_calc.add_observation("cat", "pet")
    pmi_calc.add_observation("dog", "pet")
    pmi_calc.add_observation("cat", "meow")
    pmi_calc.add_observation("dog", "bark")
    pmi_calc.add_observation("cat", "animal")
    pmi_calc.add_observation("dog", "animal")
    
    # Calculate PMI for specific pairs
    pairs = [
        ("cat", "animal"),
        ("dog", "animal"),
        ("cat", "meow"),
        ("dog", "bark"),
    ]
    
    print("PMI Scores:")
    for x, y in pairs:
        pmi = pmi_calc.calculate_pmi(x, y)
        ppmi = pmi_calc.calculate_ppmi(x, y)
        print(f"  PMI({x}, {y}) = {pmi:.4f}, PPMI = {ppmi:.4f}")
    print()


def example3_batch_observations():
    """Example 3: Adding observations in batch."""
    print("=" * 60)
    print("Example 3: Adding observations in batch")
    print("=" * 60)
    
    pmi_calc = PMI()
    
    # Add multiple observations at once
    observations = [
        ("apple", "fruit"),
        ("banana", "fruit"),
        ("carrot", "vegetable"),
        ("apple", "red"),
        ("banana", "yellow"),
        ("carrot", "orange"),
        ("apple", "fruit"),  # duplicate
        ("fruit", "healthy", 5),  # with count
    ]
    
    pmi_calc.add_observations_batch(observations)
    
    # Get all PPMI scores
    all_ppmi = pmi_calc.get_all_ppmi()
    
    print("All PPMI Scores:")
    for (x, y), score in sorted(all_ppmi.items(), key=lambda item: item[1], reverse=True):
        print(f"  PPMI({x}, {y}) = {score:.4f}")
    print()


def example4_from_counts():
    """Example 4: Calculate PMI from raw counts."""
    print("=" * 60)
    print("Example 4: Calculate PMI from raw counts")
    print("=" * 60)
    
    # Scenario: In a corpus of 1000 word pairs
    # "big" appears 100 times
    # "data" appears 80 times
    # "big data" appears together 30 times
    
    co_count = 30
    x_count = 100
    y_count = 80
    total = 1000
    
    pmi = calculate_pmi_from_counts(co_count, x_count, y_count, total, base=2)
    
    print(f"Co-occurrence count: {co_count}")
    print(f"'big' count: {x_count}")
    print(f"'data' count: {y_count}")
    print(f"Total observations: {total}")
    print(f"PMI('big', 'data') = {pmi:.4f}")
    print()


def example5_normalized_pmi():
    """Example 5: Normalized PMI (NPMI)."""
    print("=" * 60)
    print("Example 5: Normalized PMI (NPMI)")
    print("=" * 60)
    
    pmi_calc = PMI()
    
    # Add observations
    observations = [
        ("coffee", "hot", 10),
        ("coffee", "drink", 15),
        ("tea", "hot", 8),
        ("tea", "drink", 12),
        ("water", "drink", 20),
        ("ice", "cold", 10),
    ]
    
    pmi_calc.add_observations_batch(observations)
    
    # Calculate different metrics for the same pairs
    pairs = [("coffee", "hot"), ("tea", "hot"), ("water", "drink")]
    
    print(f"{'Pair':<20} {'PMI':<10} {'PPMI':<10} {'NPMI':<10}")
    print("-" * 50)
    for x, y in pairs:
        pmi = pmi_calc.calculate_pmi(x, y)
        ppmi = pmi_calc.calculate_ppmi(x, y)
        npmi = pmi_calc.calculate_npmi(x, y)
        print(f"{f'({x}, {y})':<20} {pmi:.4f}    {ppmi:.4f}    {npmi:.4f}")
    print()


def example6_different_log_bases():
    """Example 6: Using different logarithm bases."""
    print("=" * 60)
    print("Example 6: Using different logarithm bases")
    print("=" * 60)
    
    p_xy = 0.1
    p_x = 0.3
    p_y = 0.4
    
    bases = [2, 10, 'e']
    
    print(f"Probabilities: P(x,y)={p_xy}, P(x)={p_x}, P(y)={p_y}")
    print()
    for base in bases:
        pmi = calculate_pmi_from_probabilities(p_xy, p_x, p_y, base=base)
        print(f"PMI with base {base}: {pmi:.4f}")
    print()


def example7_matrix_export():
    """Example 7: Export to matrix format."""
    print("=" * 60)
    print("Example 7: Export to matrix format")
    print("=" * 60)
    
    pmi_calc = PMI()
    
    # Simulate word-context matrix
    observations = [
        ("the", "article"),
        ("a", "article"),
        ("run", "verb"),
        ("jump", "verb"),
        ("the", "common"),
        ("a", "common"),
        ("run", "action"),
        ("jump", "action"),
    ]
    
    pmi_calc.add_observations_batch(observations)
    
    # Get PPMI matrix
    matrix, words, contexts = pmi_calc.to_matrix(metric='ppmi', base=2)
    
    print("PPMI Matrix:")
    print(f"Words: {words}")
    print(f"Contexts: {contexts}")
    print(matrix)
    print()


if __name__ == "__main__":
    example1_basic_pmi_calculation()
    example2_using_pmi_class()
    example3_batch_observations()
    example4_from_counts()
    example5_normalized_pmi()
    example6_different_log_bases()
    example7_matrix_export()
    
    print("=" * 60)
    print("All examples completed!")
    print("=" * 60)
