burst_detector.py

"""
Comprehensive burst detection for topics using multiple algorithms.
"""
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from typing import List, Dict, Optional
from scipy import stats
from scipy.signal import find_peaks


class BurstDetector:
    """
    Comprehensive burst detection for topics using multiple algorithms:
    1. Kleinberg's burst detection
    2. Moving average burst detection
    3. Exponential weighted moving average (EWMA) detection
    4. Chi-square test for burst detection
    5. Poisson-based burst detection
    """

    def __init__(self, gamma: float = 1.0, sensitivity: float = 2.0, min_burst_length: int = 2,
                 ma_window: int = 5, ewma_alpha: float = 0.3):
        """
        Initialize burst detector

        Parameters:
        -----------
        gamma : float
            Cost parameter for Kleinberg's algorithm (higher = fewer bursts)
        sensitivity : float
            Sensitivity for threshold-based detection (in standard deviations)
        min_burst_length : int
            Minimum consecutive time points to consider as a burst
        ma_window : int
            Window size for moving average burst detection
        ewma_alpha : float
            Alpha parameter for EWMA burst detection
        """
        self.gamma = gamma
        self.sensitivity = sensitivity
        self.min_burst_length = min_burst_length
        self.ma_window = ma_window
        self.ewma_alpha = ewma_alpha

    def kleinberg_burst_detection(self, frequencies: np.ndarray, timestamps: List) -> List[Dict]:
        """
        Kleinberg's burst detection algorithm

        Returns list of burst periods with start/end indices and burst weight
        """
        n = len(frequencies)
        if n < 2:
            return []

        # Calculate rates
        total_freq = np.sum(frequencies)
        if total_freq == 0:
            return []

        avg_rate = total_freq / n

        # Two-state model: base state and burst state
        # Estimate burst rate as peaks above mean + std
        threshold = np.mean(frequencies) + np.std(frequencies)
        burst_indices = np.where(frequencies > threshold)[0]

        if len(burst_indices) == 0:
            return []

        burst_rate = np.mean(frequencies[burst_indices])

        # Cost function for state transitions
        transition_cost = -np.log(self.gamma)

        # Dynamic programming to find optimal state sequence
        states = self._viterbi_burst_detection(frequencies, avg_rate, burst_rate, transition_cost)

        # Extract burst periods
        bursts = []
        in_burst = False
        burst_start = 0

        for i, state in enumerate(states):
            if state == 1 and not in_burst:  # Burst starts
                burst_start = i
                in_burst = True
            elif state == 0 and in_burst:  # Burst ends
                if i - burst_start >= self.min_burst_length:
                    burst_weight = np.sum(frequencies[burst_start:i]) / (i - burst_start)
                    bursts.append({
                        'start_idx': burst_start,
                        'end_idx': i - 1,
                        'start_time': timestamps[burst_start] if timestamps else burst_start,
                        'end_time': timestamps[i-1] if timestamps else i-1,
                        'burst_weight': burst_weight,
                        'duration': i - burst_start
                    })
                in_burst = False

        # Handle burst that extends to the end
        if in_burst and n - burst_start >= self.min_burst_length:
            burst_weight = np.sum(frequencies[burst_start:]) / (n - burst_start)
            bursts.append({
                'start_idx': burst_start,
                'end_idx': n - 1,
                'start_time': timestamps[burst_start] if timestamps else burst_start,
                'end_time': timestamps[-1] if timestamps else n-1,
                'burst_weight': burst_weight,
                'duration': n - burst_start
            })

        return bursts

    def _viterbi_burst_detection(self, frequencies: np.ndarray, base_rate: float,
                                burst_rate: float, transition_cost: float) -> np.ndarray:
        """
        Viterbi algorithm for finding optimal state sequence
        """
        n = len(frequencies)
        states = np.zeros(n, dtype=int)

        # Initialize
        prob_base = np.zeros(n)
        prob_burst = np.zeros(n)

        # Emission probabilities (using Poisson approximation)
        for i in range(n):
            prob_base[i] = stats.poisson.pmf(round(frequencies[i]), base_rate)
            prob_burst[i] = stats.poisson.pmf(round(frequencies[i]), burst_rate)

        # Viterbi algorithm
        viterbi = np.zeros((2, n))
        path = np.zeros((2, n), dtype=int)

        # Initialize
        viterbi[0, 0] = prob_base[0]
        viterbi[1, 0] = prob_burst[0]

        for t in range(1, n):
            # Base state
            if viterbi[0, t-1] > viterbi[1, t-1] * np.exp(transition_cost):
                viterbi[0, t] = viterbi[0, t-1] * prob_base[t]
                path[0, t] = 0
            else:
                viterbi[0, t] = viterbi[1, t-1] * np.exp(transition_cost) * prob_base[t]
                path[0, t] = 1

            # Burst state
            if viterbi[1, t-1] > viterbi[0, t-1] * np.exp(transition_cost):
                viterbi[1, t] = viterbi[1, t-1] * prob_burst[t]
                path[1, t] = 1
            else:
                viterbi[1, t] = viterbi[0, t-1] * np.exp(transition_cost) * prob_burst[t]
                path[1, t] = 0

        # Backtrack
        states[-1] = 1 if viterbi[1, -1] > viterbi[0, -1] else 0
        for t in range(n-2, -1, -1):
            states[t] = path[states[t+1], t+1]

        return states

    def moving_average_burst_detection(self, frequencies: np.ndarray, window_size: int = 5,
                                      timestamps: Optional[List] = None) -> List[Dict]:
        """
        Detect bursts using moving average threshold
        """
        if len(frequencies) < window_size:
            return []

        # Calculate moving average
        ma = np.convolve(frequencies, np.ones(window_size)/window_size, mode='valid')

        # Calculate threshold
        threshold = np.mean(ma) + self.sensitivity * np.std(ma)

        # Detect burst points
        burst_mask = ma > threshold

        # Extract burst periods
        bursts = []
        in_burst = False
        burst_start = 0

        for i, is_burst in enumerate(burst_mask):
            if is_burst and not in_burst:
                burst_start = i
                in_burst = True
            elif not is_burst and in_burst:
                if i - burst_start >= self.min_burst_length:
                    bursts.append({
                        'start_idx': burst_start,
                        'end_idx': i - 1,
                        'start_time': timestamps[burst_start] if timestamps else burst_start,
                        'end_time': timestamps[i-1] if timestamps else i-1,
                        'burst_weight': np.mean(ma[burst_start:i]),
                        'duration': i - burst_start,
                        'method': 'moving_average'
                    })
                in_burst = False

        # Handle burst extending to end
        if in_burst and len(burst_mask) - burst_start >= self.min_burst_length:
            bursts.append({
                'start_idx': burst_start,
                'end_idx': len(burst_mask) - 1,
                'start_time': timestamps[burst_start] if timestamps else burst_start,
                'end_time': timestamps[len(burst_mask)-1] if timestamps else len(burst_mask)-1,
                'burst_weight': np.mean(ma[burst_start:]),
                'duration': len(burst_mask) - burst_start,
                'method': 'moving_average'
            })

        return bursts

    def ewma_burst_detection(self, frequencies: np.ndarray, alpha: float = 0.3,
                           timestamps: Optional[List] = None) -> List[Dict]:
        """
        Exponentially Weighted Moving Average burst detection
        """
        if len(frequencies) < 2:
            return []

        # Calculate EWMA
        ewma = np.zeros(len(frequencies))
        ewma[0] = frequencies[0]
        for i in range(1, len(frequencies)):
            ewma[i] = alpha * frequencies[i] + (1 - alpha) * ewma[i-1]

        # Calculate dynamic threshold
        ewma_std = pd.Series(frequencies).ewm(alpha=alpha).std().fillna(0).values
        threshold = ewma + self.sensitivity * ewma_std

        # Detect bursts
        burst_mask = frequencies > threshold

        # Extract burst periods
        bursts = []
        in_burst = False
        burst_start = 0

        for i, is_burst in enumerate(burst_mask):
            if is_burst and not in_burst:
                burst_start = i
                in_burst = True
            elif not is_burst and in_burst:
                if i - burst_start >= self.min_burst_length:
                    bursts.append({
                        'start_idx': burst_start,
                        'end_idx': i - 1,
                        'start_time': timestamps[burst_start] if timestamps else burst_start,
                        'end_time': timestamps[i-1] if timestamps else i-1,
                        'burst_weight': np.mean(frequencies[burst_start:i]),
                        'duration': i - burst_start,
                        'method': 'ewma'
                    })
                in_burst = False

        if in_burst and len(frequencies) - burst_start >= self.min_burst_length:
            bursts.append({
                'start_idx': burst_start,
                'end_idx': len(frequencies) - 1,
                'start_time': timestamps[burst_start] if timestamps else burst_start,
                'end_time': timestamps[-1] if timestamps else len(frequencies)-1,
                'burst_weight': np.mean(frequencies[burst_start:]),
                'duration': len(frequencies) - burst_start,
                'method': 'ewma'
            })

        return bursts

    def chi_square_burst_detection(self, frequencies: np.ndarray,
                                  timestamps: Optional[List] = None) -> List[Dict]:
        """
        Chi-square test based burst detection
        """
        if len(frequencies) < 3:
            return []

        bursts = []
        expected_freq = np.mean(frequencies)

        for i in range(len(frequencies)):
            # Calculate chi-square statistic for each point
            if expected_freq > 0:
                chi_sq = ((frequencies[i] - expected_freq) ** 2) / expected_freq

                # Critical value for significance level 0.05
                critical_value = stats.chi2.ppf(0.95, df=1)

                if chi_sq > critical_value * self.sensitivity:
                    # Check if this extends a previous burst
                    if bursts and bursts[-1]['end_idx'] == i - 1:
                        bursts[-1]['end_idx'] = i
                        bursts[-1]['end_time'] = timestamps[i] if timestamps else i
                        bursts[-1]['duration'] += 1
                        bursts[-1]['burst_weight'] = np.mean(
                            frequencies[bursts[-1]['start_idx']:i+1]
                        )
                    else:
                        # Start new burst
                        bursts.append({
                            'start_idx': i,
                            'end_idx': i,
                            'start_time': timestamps[i] if timestamps else i,
                            'end_time': timestamps[i] if timestamps else i,
                            'burst_weight': frequencies[i],
                            'duration': 1,
                            'method': 'chi_square',
                            'chi_square_stat': chi_sq
                        })

        # Filter by minimum burst length
        bursts = [b for b in bursts if b['duration'] >= self.min_burst_length]

        return bursts

    def poisson_burst_detection(self, frequencies: np.ndarray,
                              timestamps: Optional[List] = None) -> List[Dict]:
        """
        Poisson-based burst detection assuming count data
        """
        if len(frequencies) < 2:
            return []

        bursts = []

        # Estimate lambda (rate parameter)
        lambda_param = np.mean(frequencies)

        if lambda_param <= 0:
            return []

        # Calculate probability threshold
        threshold_prob = 1 - (0.05 / self.sensitivity)  # Adjusted significance level
        threshold_count = stats.poisson.ppf(threshold_prob, lambda_param)

        # Detect bursts
        burst_mask = frequencies > threshold_count

        # Extract burst periods
        in_burst = False
        burst_start = 0

        for i, is_burst in enumerate(burst_mask):
            if is_burst and not in_burst:
                burst_start = i
                in_burst = True
            elif not is_burst and in_burst:
                if i - burst_start >= self.min_burst_length:
                    # Calculate burst intensity
                    burst_frequencies = frequencies[burst_start:i]
                    burst_intensity = np.sum(burst_frequencies) / lambda_param

                    bursts.append({
                        'start_idx': burst_start,
                        'end_idx': i - 1,
                        'start_time': timestamps[burst_start] if timestamps else burst_start,
                        'end_time': timestamps[i-1] if timestamps else i-1,
                        'burst_weight': np.mean(burst_frequencies),
                        'burst_intensity': burst_intensity,
                        'duration': i - burst_start,
                        'method': 'poisson'
                    })
                in_burst = False

        if in_burst and len(frequencies) - burst_start >= self.min_burst_length:
            burst_frequencies = frequencies[burst_start:]
            burst_intensity = np.sum(burst_frequencies) / lambda_param

            bursts.append({
                'start_idx': burst_start,
                'end_idx': len(frequencies) - 1,
                'start_time': timestamps[burst_start] if timestamps else burst_start,
                'end_time': timestamps[-1] if timestamps else len(frequencies)-1,
                'burst_weight': np.mean(burst_frequencies),
                'burst_intensity': burst_intensity,
                'duration': len(frequencies) - burst_start,
                'method': 'poisson'
            })

        return bursts

    def detect_all_bursts(self, frequencies: np.ndarray,
                         timestamps: Optional[List] = None) -> Dict[str, List[Dict]]:
        """
        Apply all burst detection methods and return results
        """
        results = {
            'kleinberg': self.kleinberg_burst_detection(frequencies, timestamps),
            'moving_average': self.moving_average_burst_detection(frequencies, self.ma_window, timestamps),
            'ewma': self.ewma_burst_detection(frequencies, self.ewma_alpha, timestamps),
            'chi_square': self.chi_square_burst_detection(frequencies, timestamps),
            'poisson': self.poisson_burst_detection(frequencies, timestamps)
        }

        return results

    def consensus_bursts(self, frequencies: np.ndarray,
                        timestamps: Optional[List] = None,
                        min_methods: int = 3) -> List[Dict]:
        """
        Find burst periods detected by multiple methods (consensus approach)
        """
        all_results = self.detect_all_bursts(frequencies, timestamps)

        # Collect all burst periods from all methods
        all_bursts = []
        for method, bursts in all_results.items():
            for burst in bursts:
                burst['method'] = method
                all_bursts.append(burst)

        if not all_bursts:
            return []

        # Find overlapping bursts
        consensus_bursts = []
        processed = set()

        for i, burst1 in enumerate(all_bursts):
            if i in processed:
                continue

            overlapping_methods = [burst1['method']]
            overlapping_bursts = [burst1]

            for j, burst2 in enumerate(all_bursts[i+1:], i+1):
                if j in processed:
                    continue

                # Check for overlap
                if (burst1['start_idx'] <= burst2['end_idx'] and
                    burst2['start_idx'] <= burst1['end_idx']):
                    overlapping_methods.append(burst2['method'])
                    overlapping_bursts.append(burst2)
                    processed.add(j)

            # If enough methods agree, add to consensus
            if len(set(overlapping_methods)) >= min_methods:
                # Merge overlapping bursts
                start_idx = min(b['start_idx'] for b in overlapping_bursts)
                end_idx = max(b['end_idx'] for b in overlapping_bursts)

                consensus_bursts.append({
                    'start_idx': start_idx,
                    'end_idx': end_idx,
                    'start_time': timestamps[start_idx] if timestamps else start_idx,
                    'end_time': timestamps[end_idx] if timestamps else end_idx,
                    'burst_weight': np.mean(frequencies[start_idx:end_idx+1]),
                    'duration': end_idx - start_idx + 1,
                    'methods': list(set(overlapping_methods)),
                    'confidence': len(set(overlapping_methods)) / 5.0  # Normalized by total methods
                })

            processed.add(i)

        return consensus_bursts

    def plot_burst_detection_results(self, frequencies: np.ndarray,
                                    timestamps: Optional[List] = None,
                                    topic_id: Optional[int] = None,
                                    save_path: Optional[str] = None):
        """
        Visualize burst detection results from all methods
        """
        fig, axes = plt.subplots(3, 2, figsize=(15, 12))
        axes = axes.flatten()

        # Prepare x-axis (timestamps or indices)
        x_axis = timestamps if timestamps else range(len(frequencies))

        # Get all burst results
        all_results = self.detect_all_bursts(frequencies, timestamps)
        consensus = self.consensus_bursts(frequencies, timestamps)

        methods = ['kleinberg', 'moving_average', 'ewma', 'chi_square', 'poisson', 'consensus']
        titles = ['Kleinberg Burst Detection', 'Moving Average', 'EWMA',
                 'Chi-Square Test', 'Poisson-based', 'Consensus (≥3 methods)']

        for idx, (method, title) in enumerate(zip(methods, titles)):
            ax = axes[idx]

            # Plot frequency time series
            ax.plot(x_axis, frequencies, 'b-', alpha=0.6, label='Frequency')

            # Get bursts for this method
            if method == 'consensus':
                bursts = consensus
            else:
                bursts = all_results[method]

            # Highlight burst periods
            for burst in bursts:
                start = burst['start_idx']
                end = burst['end_idx']
                ax.axvspan(x_axis[start], x_axis[end], alpha=0.3, color='red',
                          label='Burst' if burst == bursts[0] else '')

                # Add burst weight annotation
                mid_point = (start + end) // 2
                ax.annotate(f'{burst["burst_weight"]:.1f}',
                          xy=(x_axis[mid_point], frequencies[mid_point]),
                          xytext=(0, 10), textcoords='offset points',
                          fontsize=8, alpha=0.7)

            ax.set_title(title, fontsize=10, fontweight='bold')
            ax.set_xlabel('Time', fontsize=9)
            ax.set_ylabel('Frequency', fontsize=9)
            ax.grid(True, alpha=0.3)

            # Add burst count to title
            burst_count = len(bursts)
            ax.text(0.02, 0.98, f'Bursts: {burst_count}',
                   transform=ax.transAxes, fontsize=9,
                   verticalalignment='top',
                   bbox=dict(boxstyle='round', facecolor='wheat', alpha=0.5))

            if idx == 0:
                ax.legend(loc='upper right', fontsize=8)

        # Main title
        topic_str = f' - Topic {topic_id}' if topic_id is not None else ''
        plt.suptitle(f'Burst Detection Results{topic_str}', fontsize=14, fontweight='bold')
        plt.tight_layout()

        if save_path:
            plt.savefig(save_path, dpi=300, bbox_inches='tight')
            print(f"Saved burst detection visualization to {save_path}")
        plt.show()

    def summarize_bursts(self, topics_over_time: pd.DataFrame,
                        top_n_topics: int = 10) -> pd.DataFrame:
        """
        Summarize burst detection results for top topics
        """
        summary_data = []

        # Get unique topics (excluding outliers)
        unique_topics = [t for t in topics_over_time['Topic'].unique() if t != -1][:top_n_topics]

        for topic in unique_topics:
            topic_data = topics_over_time[topics_over_time['Topic'] == topic].sort_values('Timestamp')
            frequencies = topic_data['Frequency'].values
            timestamps = topic_data['Timestamp'].tolist()

            if len(frequencies) < 3:
                continue

            # Get bursts from all methods
            all_results = self.detect_all_bursts(frequencies, timestamps)
            consensus = self.consensus_bursts(frequencies, timestamps)

            # Calculate statistics
            total_bursts = {method: len(bursts) for method, bursts in all_results.items()}

            # Calculate average burst duration and intensity
            avg_duration = {}
            avg_intensity = {}

            for method, bursts in all_results.items():
                if bursts:
                    avg_duration[method] = np.mean([b['duration'] for b in bursts])
                    avg_intensity[method] = np.mean([b['burst_weight'] for b in bursts])
                else:
                    avg_duration[method] = 0
                    avg_intensity[method] = 0

            summary_data.append({
                'Topic': topic,
                'Total_Bursts_Kleinberg': total_bursts.get('kleinberg', 0),
                'Total_Bursts_MA': total_bursts.get('moving_average', 0),
                'Total_Bursts_EWMA': total_bursts.get('ewma', 0),
                'Total_Bursts_ChiSquare': total_bursts.get('chi_square', 0),
                'Total_Bursts_Poisson': total_bursts.get('poisson', 0),
                'Consensus_Bursts': len(consensus),
                'Avg_Burst_Duration': np.mean(list(avg_duration.values())),
                'Avg_Burst_Intensity': np.mean(list(avg_intensity.values())),
                'Max_Burst_Intensity': max(list(avg_intensity.values())) if avg_intensity else 0,
                'Burst_Coverage': sum(b['duration'] for b in consensus) / len(frequencies) if consensus else 0
            })

        summary_df = pd.DataFrame(summary_data)
        summary_df = summary_df.sort_values('Consensus_Bursts', ascending=False)

        return summary_df