"""CRX Refined — Cluster-then-infer CRX with over-approximation detection.

Standard CRX produces over-approximated grammars when the Hasse diagram is
non-linear (branching). This module fixes that by:
1. Clustering similar sequences before inference
2. Inferring per-cluster (tighter grammars)
3. Reporting confidence metrics

Key insight: code call sequences have branching patterns that CHAREs
(linear chain expressions) can't represent. Clustering reduces branching
within each group, making CRX's linear assumption more valid.

Tradeoff: cluster granularity.
- Too coarse → over-approximation (standard CRX)
- Too fine → no generalization (one grammar per sequence)
- Sweet spot → cluster by structural features
"""

from collections import defaultdict
from .crx import CRX


def _cluster_by_structure(sequences):
    """Cluster sequences by structural features.

    Uses (first_symbol, last_symbol, length_bucket) as the cluster key.
    This groups sequences that start and end the same way, which
    typically means they follow the same calling pattern.

    Args:
        sequences: list of token lists

    Returns:
        dict mapping cluster_key → list of sequences
    """
    groups = defaultdict(list)
    for seq in sequences:
        if not seq:
            continue
        # Length bucket: short (1-3), medium (4-8), long (9+)
        length = len(seq)
        if length <= 3:
            length_bucket = 'short'
        elif length <= 8:
            length_bucket = 'med'
        else:
            length_bucket = 'long'
        key = (seq[0], seq[-1], length_bucket)
        groups[key].append(seq)
    return dict(groups)


def crx_refined(sequences, min_cluster=2):
    """Cluster-then-infer CRX.

    Clusters sequences by (first, last, length), infers CRX per cluster,
    and returns the grammar from the largest cluster. Falls back to standard
    CRX if no cluster has enough sequences.

    Args:
        sequences: list of token lists
        min_cluster: minimum cluster size to infer from (default: 2)

    Returns:
        CHARE expression string
    """
    S = [list(s) for s in sequences if s]
    if not S:
        return 'ε'

    clusters = _cluster_by_structure(S)

    # Find clusters large enough to infer from
    valid_clusters = {k: v for k, v in clusters.items() if len(v) >= min_cluster}

    if not valid_clusters:
        # No cluster large enough — fall back to standard CRX
        return CRX().infer(S)

    # Pick the largest cluster
    best_key = max(valid_clusters, key=lambda k: len(valid_clusters[k]))
    best_cluster = valid_clusters[best_key]

    return CRX().infer(best_cluster)


def crx_with_confidence(sequences, min_cluster=2):
    """Cluster-then-infer CRX with confidence metrics.

    Returns:
        dict with keys:
            grammar: str — the refined grammar
            confidence: float — fraction of training pairs captured by grammar structure
            n_clusters: int — number of clusters
            largest_cluster: int — size of largest cluster
            n_sequences: int — total sequences
    """
    S = [list(s) for s in sequences if s]
    if not S:
        return {
            'grammar': 'ε',
            'confidence': 1.0,
            'n_clusters': 0,
            'largest_cluster': 0,
            'n_sequences': 0,
        }

    clusters = _cluster_by_structure(S)
    valid_clusters = {k: v for k, v in clusters.items() if len(v) >= min_cluster}

    if valid_clusters:
        best_key = max(valid_clusters, key=lambda k: len(valid_clusters[k]))
        best_cluster = valid_clusters[best_key]
        grammar = CRX().infer(best_cluster)
    else:
        best_cluster = S
        grammar = CRX().infer(S)

    # Compute confidence: what fraction of consecutive pairs in the training
    # data are "captured" by the grammar's structure
    all_pairs = set()
    for w in S:
        for i in range(len(w) - 1):
            all_pairs.add((w[i], w[i + 1]))

    # Pairs in the best cluster that appear in grammar structure
    cluster_pairs = set()
    for w in best_cluster:
        for i in range(len(w) - 1):
            cluster_pairs.add((w[i], w[i + 1]))

    # How many training pairs are cluster pairs? (coverage)
    captured = len(all_pairs & cluster_pairs)
    confidence = captured / len(all_pairs) if all_pairs else 1.0

    return {
        'grammar': grammar,
        'confidence': confidence,
        'n_clusters': len(clusters),
        'largest_cluster': len(best_cluster),
        'n_sequences': len(S),
    }