"""Decomposition forest for behavioral sequences.

Inspired by Crucio's decomposition forest (ICSE 2026).
Breaks down long sequences into shorter ones that still capture the pattern.

Why decomposition helps:
- Long sequences like ["if", "return", "if", "return", "if", "return"]
  → CRX sees 6 symbols, often produces flat bags like (if|return)*
- Decompose into shorter fragments:
  → ["if", "return"], ["if", "return"], ["if", "return"]
  → CRX sees clear pattern: if.return

Three decomposition strategies:
1. Prefix extraction: first N symbols
2. Suffix extraction: last N symbols  
3. Window extraction: sliding window of size N

All strategies preserve the original sequences (additive, not destructive).
"""

from collections import Counter


def decompose_sequence(seq, max_length=5):
    """Decompose a sequence into shorter fragments.
    
    Strategies:
    1. If seq <= max_length, return as-is (no decomposition needed)
    2. Extract prefixes of length 1..max_length
    3. Extract suffixes of length 1..max_length
    4. Extract windows of length max_length
    
    Args:
        seq: List of symbols
        max_length: Maximum fragment length
        
    Returns:
        List of fragments (shorter sequences)
    """
    if not seq:
        return []
    
    if len(seq) <= max_length:
        return [seq]
    
    fragments = []
    
    # Prefixes (1, 2, ..., max_length symbols from start)
    for i in range(1, min(max_length + 1, len(seq) + 1)):
        fragments.append(seq[:i])
    
    # Suffixes (1, 2, ..., max_length symbols from end)
    for i in range(1, min(max_length + 1, len(seq) + 1)):
        fragments.append(seq[-i:])
    
    # Windows (sliding window of max_length)
    for start in range(0, len(seq) - max_length + 1):
        fragments.append(seq[start:start + max_length])
    
    return fragments


def decompose_all(sequences, max_length=5):
    """Decompose all sequences in a list.
    
    Args:
        sequences: List of lists of symbols
        max_length: Maximum fragment length
        
    Returns:
        List of fragments (shorter sequences, may have duplicates)
    """
    all_fragments = []
    for seq in sequences:
        all_fragments.extend(decompose_sequence(seq, max_length))
    return all_fragments


def decompose_with_coverage(sequences, max_length=5, min_coverage=0.3):
    """Decompose sequences and filter by coverage.
    
    Keep only fragments that appear in at least min_coverage fraction
    of the original sequences. This ensures we keep common patterns,
    not rare edge cases.
    
    Args:
        sequences: List of lists of symbols
        max_length: Maximum fragment length
        min_coverage: Minimum fraction of sequences a fragment must appear in
        
    Returns:
        List of filtered fragments
    """
    if not sequences:
        return []
    
    # Decompose all sequences
    all_fragments = decompose_all(sequences, max_length)
    
    if not all_fragments:
        return []
    
    # Count how many original sequences each fragment appears in
    fragment_sources = Counter()
    for seq in sequences:
        # Get unique fragments from this sequence
        seq_fragments = set()
        for frag in decompose_sequence(seq, max_length):
            seq_fragments.add(tuple(frag))
        
        # Count each unique fragment once per source sequence
        for frag in seq_fragments:
            fragment_sources[frag] += 1
    
    # Keep fragments that appear in enough source sequences
    min_count = max(1, int(len(sequences) * min_coverage))
    filtered = [list(frag) for frag, count in fragment_sources.items() 
                if count >= min_count]
    
    return filtered


def get_decomposition_stats(sequences, max_length=5):
    """Get statistics about decomposition.
    
    Args:
        sequences: List of lists of symbols
        max_length: Maximum fragment length
        
    Returns:
        Dict with statistics
    """
    original_lengths = [len(s) for s in sequences]
    fragments = decompose_all(sequences, max_length)
    fragment_lengths = [len(f) for f in fragments]
    
    return {
        'n_original': len(sequences),
        'n_fragments': len(fragments),
        'expansion_ratio': len(fragments) / len(sequences) if sequences else 0,
        'avg_original_length': sum(original_lengths) / len(original_lengths) if original_lengths else 0,
        'avg_fragment_length': sum(fragment_lengths) / len(fragment_lengths) if fragment_lengths else 0,
        'max_original_length': max(original_lengths) if original_lengths else 0,
    }