"""Ensemble grammar inference — run multiple algorithms, pick best by scoring."""

from .crx import CRX
from .idregex import idregex
from .grammar import alphabet, match as grammar_match, Empty, Epsilon
from .mdl import score_grammar


def _matches(grammar, sequence):
    """Check if a sequence matches the grammar."""
    if grammar is None or isinstance(grammar, (Empty, Epsilon)):
        return not sequence if isinstance(grammar, Epsilon) else False
    try:
        return grammar_match(grammar, sequence)
    except Exception:
        return False


def _fit_score(grammar, seq):
    """Score how tightly a sequence fits the grammar core."""
    if not seq:
        return 0.0
    if grammar is None or isinstance(grammar, Empty):
        return 0.0
    try:
        if not grammar_match(grammar, seq):
            return 0.0
        alpha = alphabet(grammar)
        if not alpha:
            return 0.5
        unique_syms = len(set(seq))
        total_syms = len(seq)
        return max(0.0, 1.0 - (total_syms - unique_syms) / max(total_syms, 1))
    except Exception:
        return 0.0


def _symbol_rarity_score(seq, all_sequences):
    """Score a sequence by how rare its symbols are across the dataset.
    1.0 = all symbols are common, 0.0 = mostly rare symbols.
    """
    from collections import Counter
    all_syms = Counter()
    for s in all_sequences:
        all_syms.update(s)
    n = len(all_sequences)
    scores = []
    for sym in seq:
        freq = all_syms.get(sym, 0) / n
        scores.append(min(freq, 1.0))
    return sum(scores) / len(scores) if scores else 0.0


def _find_core(sequences, min_coverage=0.8):
    """Find the core subset of sequences by iterative CRX + outlier removal."""
    if not sequences or min_coverage >= 1.0:
        crx_g = CRX().infer(sequences)
        return crx_g, sequences, [], []

    from collections import Counter
    all_syms = Counter()
    for s in sequences:
        all_syms.update(s)
    n = len(sequences)

    def _rarity(seq):
        rare_count = sum(1 for sym in seq if all_syms.get(sym, 0) / n < 0.3)
        return rare_count / max(len(seq), 1)

    working = list(sequences)
    removed_indices = []
    crx = CRX()

    for _ in range(50):
        if len(working) < 3:
            break
        target = max(int(len(sequences) * min_coverage), 1)
        if len(working) <= target:
            break
        scores = [(i, _rarity(seq)) for i, seq in enumerate(working)]
        scores.sort(key=lambda x: -x[1])
        if len(scores) < 2 or scores[0][1] == scores[-1][1]:
            break
        worst_idx = scores[0][0]
        removed_indices.append(working[worst_idx])
        working = [s for i, s in enumerate(working) if i != worst_idx]

    core_g = crx.infer(working) if working else None
    return core_g, working, removed_indices, []


def mdl_score_simple(grammar, sequences, method='langsize'):
    """Score a grammar. Default: Language Size (Bex et al., arXiv:1004.2372)."""
    return score_grammar(grammar, sequences, method=method)


def _run_idregex(sequences, kmax, N, method='langsize'):
    """Run standalone iDRegEx, return (grammar, score) or (None, inf)."""
    g = idregex(sequences, kmax=kmax, N=N)
    if g and not isinstance(g, Empty):
        return g, mdl_score_simple(g, sequences, method=method)
    return None, float('inf')


_ALGO_NAMES = {
    'crx': 'CRX',
    'idregex': 'iDRegEx',
}


_ALGORITHMS = {
    'crx': lambda s, k, n, m='langsize': (CRX().infer(s), mdl_score_simple(CRX().infer(s), s, method=m)),
    'idregex': _run_idregex,
}


def _run_kore(sequences, kmax, N, method='langsize'):
    """Run kOREInference, return (grammar, score) or (None, inf)."""
    from .kore import kOREInference
    kore = kOREInference(k_max=kmax, N=N)
    result = kore.infer(sequences)
    if result:
        _, expr, _ = result
        if not isinstance(expr, Empty):
            return expr, mdl_score_simple(expr, sequences, method=method)
    return None, float('inf')


def infer_ensemble(sequences, kmax=2, N=3, prefer=None, min_coverage=1.0, include_kore=False, include_idregex=False, method='langsize'):
    """Run all applicable algorithms and return the best by scoring.

    Args:
        sequences: List of sequences, each a list of strings.
        kmax: Maximum k for k-ORE inference (iDRegEx, kOREInference).
        N: Number of random trials for k-ORE inference.
        prefer: Optional — 'crx', 'idregex', or 'koreinference' to skip
                ensemble and return only that algorithm's result.
        min_coverage: When < 1.0, also runs CRX on the tightest core subset.
        include_idregex: Run iDRegEx (slow, opt-in).
        method: Scoring method — 'langsize' (default) or 'mdl' (fallback).

    Returns:
        dict with keys: best, all, why, core (optional)
    """
    if prefer and prefer.lower() in _ALGORITHMS:
        key = prefer.lower()
        fn = _ALGORITHMS[key]
        algo_name = _ALGO_NAMES.get(key, key)
        g, score = fn(sequences, kmax, N, method)
        if g and not isinstance(g, Empty):
            return {
                'best': {'algorithm': algo_name, 'grammar': g, 'mdl_score': round(score, 2)},
                'all': [{'algorithm': algo_name, 'grammar': g, 'mdl_score': round(score, 2)}],
                'why': f"Requested {algo_name} only.",
            }
        return {
            'best': None,
            'all': [],
            'why': f"{algo_name} returned empty (no grammar found).",
        }

    results = []

    crx_g = CRX().infer(sequences)
    crx_score = mdl_score_simple(crx_g, sequences, method=method) if crx_g and not isinstance(crx_g, Empty) else float('inf')
    results.append(('CRX', crx_g if crx_g and not isinstance(crx_g, Empty) else Empty(), crx_score))

    if include_idregex:
        idr_g, idr_score = _run_idregex(sequences, kmax, N, method=method)
        if idr_g:
            results.append(('iDRegEx', idr_g, idr_score))

    if include_kore:
        kore_g, kore_score = _run_kore(sequences, kmax, N, method=method)
        if kore_g:
            results.append(('kOREInference', kore_g, kore_score))

    results = [r for r in results if r[1] and not isinstance(r[1], Empty)]
    if not results:
        base = {
            'best': None,
            'all': [],
            'why': "No algorithm produced a non-empty grammar.",
        }
        if min_coverage < 1.0:
            core_g, core_seqs, outliers, _ = _find_core(sequences, min_coverage)
            base['core'] = {
                'grammar': core_g,
                'coverage': round(len(core_seqs) / max(len(sequences), 1), 2) if sequences else 0,
                'outliers': outliers,
            }
        return base

    results.sort(key=lambda x: x[2])
    best = results[0]
    all_results = [
        {'algorithm': a, 'grammar': g, 'mdl_score': round(s, 2)}
        for a, g, s in results
    ]

    why_parts = []
    if len(results) == 1:
        why_parts.append(f"Only {results[0][0]} produced a result.")
    else:
        scores_str = ', '.join(f"{r[0]}={r[2]:.1f}" for r in results)
        why_parts.append(f"Scores: {scores_str}.")

    match_strs = []
    for r_algo, r_grammar, _ in results:
        if r_grammar and not isinstance(r_grammar, Empty):
            m = sum(1 for s in sequences if _matches(r_grammar, s))
            match_strs.append(f"{r_algo}={m}/{len(sequences)}")
    if match_strs:
        why_parts.append(f"Match rates: {', '.join(match_strs)}.")

    why_parts.append(f"{best[0]} selected (MDL score {best[2]}).")

    result = {
        'best': {
            'algorithm': best[0],
            'grammar': best[1],
            'mdl_score': round(best[2], 2),
        },
        'all': all_results,
        'why': ' '.join(why_parts),
    }

    if min_coverage < 1.0:
        core_g, core_seqs, outliers, _ = _find_core(sequences, min_coverage)
        result['core'] = {
            'grammar': core_g,
            'coverage': round(len(core_seqs) / max(len(sequences), 1), 2) if sequences else 0,
            'outlier_count': len(outliers),
            'outliers': outliers,
        }
        result['why'] += f' Core CRX ({min_coverage:.0%} coverage, {len(outliers)} outliers): {core_g}'

    return result
