"""
Core PMI calculation module.

This module provides functions and classes for calculating
Pointwise Mutual Information (PMI) and Positive PMI (PPMI)
between events or words.
"""

import math
import csv
from typing import Dict, Tuple, Union, Optional
from collections import defaultdict
import numpy as np


class PMI:
    """
    A class for calculating Pointwise Mutual Information between events.
    
    PMI measures the association between two events x and y:
    PMI(x, y) = log(P(x,y) / (P(x) * P(y)))
    
    Attributes:
        co_occurrence: Dictionary storing co-occurrence counts
        x_counts: Dictionary storing counts for x events
        y_counts: Dictionary storing counts for y events
        total: Total number of observations
    """
    
    def __init__(self):
        """Initialize PMI calculator with empty count dictionaries."""
        self.co_occurrence = defaultdict(lambda: defaultdict(int))
        self.x_counts = defaultdict(int)
        self.y_counts = defaultdict(int)
        self.total = 0
    
    def add_observation(self, x, y, count: int = 1):
        """
        Add an observation of x and y co-occurring.
        
        Args:
            x: First event/word
            y: Second event/word
            count: Number of times this co-occurrence was observed (default: 1)
        """
        self.co_occurrence[x][y] += count
        self.x_counts[x] += count
        self.y_counts[y] += count
        self.total += count
    
    def add_observations_batch(self, pairs: list):
        """
        Add multiple observations at once.
        
        Args:
            pairs: List of tuples (x, y) or (x, y, count)
        """
        for pair in pairs:
            if len(pair) == 2:
                self.add_observation(pair[0], pair[1])
            elif len(pair) == 3:
                self.add_observation(pair[0], pair[1], pair[2])
    
    def add_observations_from_text_csv(
            self,
            filepath: str,
            text_col: Union[int, str] = 0,
            window_size: int = 5,
            has_header: bool = True,
            delimiter: str = ',',
            min_word_length: int = 2):
        """
        Add observations from a CSV containing text documents using
        a sliding window. Tokenizes text into words and creates
        co-occurrence pairs based on proximity.
        
        This method is designed for continuous text (articles,
        documents, sentences) where word order and proximity are
        meaningful. Words appearing close together in the text are
        considered to co-occur.
        
        Args:
            filepath: Path to the CSV file
            text_col: Column index (int) or name (str) containing
                     text (default: 0)
            window_size: Number of words before/after target word to
                        consider as context (default: 5)
            has_header: Whether the CSV has a header row (default: True)
            delimiter: CSV delimiter (default: ',')
            min_word_length: Minimum word length to include (default: 2)
        
        Example CSV format:
            text
            "Brand awareness is the first stage of building a relationship."
            "Brand familiarity goes beyond mere awareness."
        
        With window_size=2, in the text "brand awareness is essential":
        - Creates pairs: (brand, awareness), (brand, is),
          (awareness, is), (awareness, essential), etc.
        
        Returns:
            Number of word pairs added
        """
        words = []
        
        # Read text from CSV and tokenize
        with open(filepath, 'r', encoding='utf-8') as f:
            reader = csv.reader(f, delimiter=delimiter)
            
            # Handle header
            if has_header:
                header = next(reader)
                # Convert column name to index if string provided
                if isinstance(text_col, str):
                    text_col = header.index(text_col)
            
            # Read and tokenize text from each row
            for row in reader:
                if not row:  # Skip empty rows
                    continue
                text = row[text_col].strip()
                if not text:
                    continue
                
                # Tokenize: split on whitespace and punctuation,
                # convert to lowercase
                import re
                tokens = re.findall(r'\b\w+\b', text.lower())
                
                # Filter by minimum word length
                tokens = [t for t in tokens if len(t) >= min_word_length]
                words.extend(tokens)
        
        # Create co-occurrence pairs using sliding window
        pairs_added = 0
        for i, word in enumerate(words):
            # Get context words within window
            # (only look forward to avoid duplicates)
            start = i + 1
            end = min(len(words), i + window_size + 1)
            
            for j in range(start, end):
                context_word = words[j]
                # Skip if same word text
                # (avoid self-pairs like "brand" with "brand")
                if word != context_word:
                    self.add_observation(word, context_word)
                    pairs_added += 1
        
        return pairs_added
    
    def calculate_pmi(
            self, x, y, base: Union[int, str] = 2
    ) -> Optional[float]:
        """
        Calculate PMI for a specific pair (x, y).
        
        Args:
            x: First event/word
            y: Second event/word
            base: Logarithm base (2, 10, 'e', or math.e). Default is 2.
        
        Returns:
            PMI value or None if the pair was never observed together
        """
        if self.co_occurrence[x][y] == 0:
            return None
        
        p_xy = self.co_occurrence[x][y] / self.total
        p_x = self.x_counts[x] / self.total
        p_y = self.y_counts[y] / self.total
        
        if p_x == 0 or p_y == 0:
            return None
        
        pmi_value = p_xy / (p_x * p_y)
        
        if base == 2:
            return math.log2(pmi_value)
        elif base == 10:
            return math.log10(pmi_value)
        elif base == 'e' or base == math.e:
            return math.log(pmi_value)
        else:
            return math.log(pmi_value) / math.log(base)
    
    def _calculate_smoothed_context_prob(self, y, alpha: float = 1.0) -> float:
        """
        Calculate smoothed context probability with
        Pα(c) = count(c)^α / Σ count(c)^α
        
        Args:
            y: Context word/event
            alpha: Smoothing parameter (default: 1.0 for no smoothing)
        
        Returns:
            Smoothed probability of context
        """
        if alpha == 1.0:
            return self.y_counts[y] / self.total
        
        # Calculate sum of all context counts raised to alpha
        total_alpha = sum(count ** alpha for count in self.y_counts.values())
        
        # Return smoothed probability
        return (self.y_counts[y] ** alpha) / total_alpha
    
    def calculate_ppmi(self, x, y, base: Union[int, str] = 2, alpha: float = 1.0) -> float:
        """
        Calculate Positive PMI (PPMI) for a specific pair (x, y).
        Supports context distribution smoothing with alpha parameter.
        
        PPMI_α(x, y) = max(0, log(P(x,y) / (P(x) * P_α(y))))
        where P_α(y) = count(y)^α / Σ count(c)^α
        
        Args:
            x: First event/word
            y: Second event/word (context)
            base: Logarithm base (2, 10, 'e', or math.e). Default is 2.
            alpha: Context smoothing parameter. Default is 1.0 (no smoothing).
                   Values < 1 (e.g., 0.75) increase probability of rare contexts.
        
        Returns:
            PPMI value (0 if PMI is negative or pair never observed)
        
        References:
            Levy et al. (2015): Improving Distributional Similarity with Lessons 
            Learned from Word Embeddings. Alpha=0.75 improves performance.
        """
        if self.co_occurrence[x][y] == 0:
            return 0
        
        p_xy = self.co_occurrence[x][y] / self.total
        p_x = self.x_counts[x] / self.total
        p_y_alpha = self._calculate_smoothed_context_prob(y, alpha)
        
        if p_x == 0 or p_y_alpha == 0:
            return 0
        
        pmi_value = p_xy / (p_x * p_y_alpha)
        
        if base == 2:
            pmi = math.log2(pmi_value)
        elif base == 10:
            pmi = math.log10(pmi_value)
        elif base == 'e' or base == math.e:
            pmi = math.log(pmi_value)
        else:
            pmi = math.log(pmi_value) / math.log(base)
        
        return max(0, pmi)
    
    def calculate_npmi(
            self, x, y, base: Union[int, str] = 2
    ) -> Optional[float]:
        """
        Calculate Normalized PMI (NPMI) for a specific pair (x, y).
        NPMI = PMI(x, y) / -log(P(x, y))
        NPMI ranges from -1 to 1.
        
        Args:
            x: First event/word
            y: Second event/word
            base: Logarithm base (2, 10, 'e', or math.e). Default is 2.
        
        Returns:
            NPMI value or None if the pair was never observed together
        """
        pmi = self.calculate_pmi(x, y, base)
        if pmi is None:
            return None
        
        p_xy = self.co_occurrence[x][y] / self.total
        
        if base == 2:
            log_p_xy = math.log2(p_xy)
        elif base == 10:
            log_p_xy = math.log10(p_xy)
        elif base == 'e' or base == math.e:
            log_p_xy = math.log(p_xy)
        else:
            log_p_xy = math.log(p_xy) / math.log(base)
        
        return pmi / (-log_p_xy)
    
    def get_all_pmi(self, base: Union[int, str] = 2) -> Dict[Tuple, float]:
        """
        Calculate PMI for all observed pairs.
        
        Args:
            base: Logarithm base (2, 10, 'e', or math.e). Default is 2.
        
        Returns:
            Dictionary mapping (x, y) tuples to their PMI values
        """
        pmi_dict = {}
        for x in self.co_occurrence:
            for y in self.co_occurrence[x]:
                pmi = self.calculate_pmi(x, y, base)
                if pmi is not None:
                    pmi_dict[(x, y)] = pmi
        return pmi_dict
    
    def get_all_ppmi(self, base: Union[int, str] = 2, alpha: float = 1.0) -> Dict[Tuple, float]:
        """
        Calculate PPMI for all observed pairs.
        
        Args:
            base: Logarithm base (2, 10, 'e', or math.e). Default is 2.
            alpha: Context smoothing parameter. Default is 1.0 (no smoothing).
                   Values < 1 (e.g., 0.75) increase probability of rare contexts.
        
        Returns:
            Dictionary mapping (x, y) tuples to their PPMI values
        """
        ppmi_dict = {}
        for x in self.co_occurrence:
            for y in self.co_occurrence[x]:
                ppmi = self.calculate_ppmi(x, y, base, alpha)
                ppmi_dict[(x, y)] = ppmi
        return ppmi_dict
    
    def to_matrix(
            self,
            metric: str = 'ppmi',
            base: Union[int, str] = 2,
            alpha: float = 1.0
    ) -> Tuple[np.ndarray, list, list]:
        """
        Convert PMI/PPMI scores to a matrix format.
        
        Args:
            metric: 'pmi', 'ppmi', or 'npmi'
            base: Logarithm base (2, 10, 'e', or math.e). Default is 2.
            alpha: Context smoothing parameter for PPMI (default: 1.0).
                   Only used when metric='ppmi'.
        
        Returns:
            Tuple of (matrix, x_labels, y_labels)
        """
        x_items = sorted(self.x_counts.keys())
        y_items = sorted(self.y_counts.keys())
        
        matrix = np.zeros((len(x_items), len(y_items)))
        
        for i, x in enumerate(x_items):
            for j, y in enumerate(y_items):
                if metric == 'pmi':
                    value = self.calculate_pmi(x, y, base)
                    matrix[i, j] = value if value is not None else 0
                elif metric == 'ppmi':
                    matrix[i, j] = self.calculate_ppmi(x, y, base, alpha)
                elif metric == 'npmi':
                    value = self.calculate_npmi(x, y, base)
                    matrix[i, j] = value if value is not None else 0
        
        return matrix, x_items, y_items


def calculate_pmi_from_probabilities(p_xy: float, p_x: float, p_y: float, base: Union[int, str] = 2) -> float:
    """
    Calculate PMI from probabilities.
    
    Args:
        p_xy: Joint probability P(x, y)
        p_x: Marginal probability P(x)
        p_y: Marginal probability P(y)
        base: Logarithm base (2, 10, 'e', or math.e). Default is 2.
    
    Returns:
        PMI value
    
    Raises:
        ValueError: If any probability is 0 or not in range [0, 1]
    """
    if not (0 < p_xy <= 1 and 0 < p_x <= 1 and 0 < p_y <= 1):
        raise ValueError("Probabilities must be in the range (0, 1]")
    
    pmi_value = p_xy / (p_x * p_y)
    
    if base == 2:
        return math.log2(pmi_value)
    elif base == 10:
        return math.log10(pmi_value)
    elif base == 'e' or base == math.e:
        return math.log(pmi_value)
    else:
        return math.log(pmi_value) / math.log(base)


def calculate_ppmi_from_probabilities(p_xy: float, p_x: float, p_y: float, base: Union[int, str] = 2) -> float:
    """
    Calculate Positive PMI from probabilities.
    
    Args:
        p_xy: Joint probability P(x, y)
        p_x: Marginal probability P(x)
        p_y: Marginal probability P(y)
        base: Logarithm base (2, 10, 'e', or math.e). Default is 2.
    
    Returns:
        PPMI value (max of 0 and PMI)
    """
    pmi = calculate_pmi_from_probabilities(p_xy, p_x, p_y, base)
    return max(0, pmi)


def calculate_pmi_from_counts(
    co_count: int,
    x_count: int,
    y_count: int,
    total: int,
    base: Union[int, str] = 2
) -> float:
    """
    Calculate PMI from raw counts.
    
    Args:
        co_count: Co-occurrence count of x and y
        x_count: Total count of x
        y_count: Total count of y
        total: Total number of observations
        base: Logarithm base (2, 10, 'e', or math.e). Default is 2.
    
    Returns:
        PMI value
    
    Raises:
        ValueError: If any count is 0 or invalid
    """
    if co_count <= 0 or x_count <= 0 or y_count <= 0 or total <= 0:
        raise ValueError("All counts must be positive")
    
    p_xy = co_count / total
    p_x = x_count / total
    p_y = y_count / total
    
    return calculate_pmi_from_probabilities(p_xy, p_x, p_y, base)


def calculate_ppmi_from_counts(
    co_count: int,
    x_count: int,
    y_count: int,
    total: int,
    base: Union[int, str] = 2
) -> float:
    """
    Calculate PPMI from raw counts.
    
    Args:
        co_count: Co-occurrence count of x and y
        x_count: Total count of x
        y_count: Total count of y
        total: Total number of observations
        base: Logarithm base (2, 10, 'e', or math.e). Default is 2.
    
    Returns:
        PPMI value
    """
    pmi = calculate_pmi_from_counts(co_count, x_count, y_count, total, base)
    return max(0, pmi)
