"""Reduce — Algorithm 4 (TODS 2010).

Paper: Bex et al. "Inference of Concise Regular Expressions and DTDs"
ACM Transactions on Database Systems, 2010. Section 5.

When iLocal runs on an incomplete corpus, it may infer separate types for what is
actually the same type in the target schema. Reduce fixes this by measuring
similarity between inferred types and merging those that are close enough.

Key definitions:

  dist(A, B) — SOA edit distance (Definition 14):
    dist(A, B) = (∑_{(a,b)∈E-F} suppA(a,b) / ∑_{(a,b)∈E} suppA(a,b))
               + (∑_{(a,b)∈F-E} suppB(a,b) / ∑_{(a,b)∈F} suppB(a,b))

  distD(s, t) — type distance in inferred XSD D (Definition 15):
    distD(s, t) = max_{(s',t')∈reachD(s,t)} dist(soa(s'), soa(t'))

  Adjunction A ∪ B — merge SOAs by unioning edges, summing support:
    supp_{A∪B}(a,b) = supp_A(a,b) + supp_B(a,b)

  ToSore(SOA) — convert SOA to SORE (Algorithm 6: RWR₀)
  Minimize(D, r) — unify equivalent types in XSD D

For our use case (flat sequences, no XSD type hierarchy):
  - "types" are context groupings (context_key -> sequences)
  - reachD(s, t) = {(s, t)} (no hierarchy, so distD = dist)
  - elemsD(t) = ∅ (no child elements, so lines 8-11 are no-ops)
  - ToSore = rwr0 (our REWRITE implementation)
  - Minimize = merge contexts with identical SOREs

Important: SOA state IDs are not comparable across different SOAs.
We compare labeled transitions (symbol pairs) instead of state ID pairs.
"""

from collections import defaultdict
from itertools import combinations

from .soa import SOA
from .twotinf import build_soa
from .rwr0 import rwr0


def _labeled_transitions(soa):
    """Extract labeled transitions from SOA.

    Returns a dict mapping (from_label, to_label) -> set of state pairs.
    Labels are: None for src/sink, symbol string for other states.
    We use ('SRC',) for src and ('SINK',) for sink to make them distinguishable.
    """
    transitions = defaultdict(set)
    for f, targets in soa._succ.items():
        for t in targets:
            # Get labels: src -> ('SRC',), sink -> ('SINK',), others -> symbol
            if f == soa.src:
                f_label = ('SRC',)
            elif f == soa.sink:
                f_label = ('SINK',)
            else:
                f_label = (soa.label(f),)

            if t == soa.src:
                t_label = ('SRC',)
            elif t == soa.sink:
                t_label = ('SINK',)
            else:
                t_label = (soa.label(t),)

            transitions[(f_label, t_label)].add((f, t))
    return dict(transitions)


def build_soa_with_support(sequences):
    """Build a SOA with support annotations.

    For each edge (a, b), track suppA(a, b) = number of strings in S
    for which (a, b) needed to be added to the edges of A.

    Args:
        sequences: list of lists of symbols

    Returns:
        (SOA, support_dict) where support_dict maps (from_label, to_label) -> count
    """
    G = SOA()
    # Support keyed by labeled transitions, not state IDs
    support = {}
    symbol_states = {}

    for seq in sequences:
        if not seq:
            if not G.has_edge(G.src, G.sink):
                G.add_edge(G.src, G.sink)
            key = (('SRC',), ('SINK',))
            support[key] = support.get(key, 0) + 1
            continue
        for i, token in enumerate(seq):
            if token not in symbol_states:
                symbol_states[token] = G.add_state(token)
            if i == 0:
                if not G.has_edge(G.src, symbol_states[token]):
                    G.add_edge(G.src, symbol_states[token])
                key = (('SRC',), (token,))
                support[key] = support.get(key, 0) + 1
            if i == len(seq) - 1:
                if not G.has_edge(symbol_states[token], G.sink):
                    G.add_edge(symbol_states[token], G.sink)
                key = ((token,), ('SINK',))
                support[key] = support.get(key, 0) + 1
            if i + 1 < len(seq):
                nxt = seq[i + 1]
                if nxt not in symbol_states:
                    symbol_states[nxt] = G.add_state(nxt)
                if not G.has_edge(symbol_states[token], symbol_states[nxt]):
                    G.add_edge(symbol_states[token], symbol_states[nxt])
                key = ((token,), (nxt,))
                support[key] = support.get(key, 0) + 1

    return G, support


def soa_distance(supp_a, supp_b):
    """SOA edit distance (Definition 14, TODS 2010).

    dist(A, B) = (∑_{(a,b)∈E-F} suppA(a,b) / ∑_{(a,b)∈E} suppA(a,b))
               + (∑_{(a,b)∈F-E} suppB(a,b) / ∑_{(a,b)∈F} suppB(a,b))

    Support-weighted: edges that appear in many strings contribute more.
    dist = 0 means the SOAs accept the same language.

    Args:
        supp_a: Support dict for SOA A, maps (from_label, to_label) -> count
        supp_b: Support dict for SOA B, maps (from_label, to_label) -> count

    Returns:
        Float in [0, 1]. 0 = identical, 1 = completely disjoint.
    """
    edges_a = frozenset(supp_a.keys())
    edges_b = frozenset(supp_b.keys())

    if not edges_a and not edges_b:
        return 0.0

    total_a = sum(supp_a.values())
    total_b = sum(supp_b.values())

    if total_a == 0 and total_b == 0:
        return 0.0

    # ∑_{(a,b)∈E-F} suppA(a,b) / ∑_{(a,b)∈E} suppA(a,b)
    only_a = sum(supp_a[e] for e in (edges_a - edges_b))
    dist_a = only_a / total_a if total_a > 0 else 0

    # ∑_{(a,b)∈F-E} suppB(a,b) / ∑_{(a,b)∈F} suppB(a,b)
    only_b = sum(supp_b[e] for e in (edges_b - edges_a))
    dist_b = only_b / total_b if total_b > 0 else 0

    return dist_a + dist_b


def adjunct_support(supp_a, supp_b):
    """Adjunction of support-annotated SOAs (line 6, Algorithm 4).

    supp_{A∪B}(a, b) = supp_A(a, b) + supp_B(a, b)
    (assuming supp_A(a, b) = 0 if (a,b) ∉ E_A, and similarly for B)
    """
    combined = dict(supp_a)
    for edge, count in supp_b.items():
        combined[edge] = combined.get(edge, 0) + count
    return combined


def reduce_contexts(contexts, threshold):
    """Reduce algorithm for context groupings (Algorithm 4, TODS 2010).

    Faithful implementation of Algorithm 4, adapted for flat sequences
    (no XSD type hierarchy):

      Line 1: let (T, ρ, τ) = D
        T = set of context keys
        ρ(context) = SORE (built from sequences)
        τ = ∅ (no type hierarchy for flat contexts)

      Line 2: M := {(s, t) ∈ T² | 0 < distD(s, t) < ε}
        For flat contexts: distD(s, t) = dist(soa(s), soa(t))

      Lines 3-12: while M is non-empty
        For each (s, t) ∈ M:
          Line 6: soa(s) := soa(s) ∪ soa(t)  [adjunction]
          Line 7: soa(t) := soa(s)
          Lines 8-11: copy elems (no-op: elemsD = ∅)
        Line 12: recompute M

      Lines 13-14: for each t, replace ρ(t) by ToSore(soa(t))
        ToSore = rwr0 (our REWRITE implementation)

      Line 15: Minimize(D, r)
        Merge contexts with identical SOREs

    Args:
        contexts: dict mapping context_key -> list of sequences
        threshold: similarity threshold ε in [0, 1]
            - 0.05: very conservative
            - 0.15: moderate (Bex's recommended starting point)
            - 0.30: aggressive

    Returns:
        (merged_contexts, merge_info)
    """
    # Line 1: Build support-annotated SOAs for contexts with >= 2 sequences
    ctx_data = {}  # ctx -> (soa, support_dict, sequences)
    for ctx, seqs in contexts.items():
        if len(seqs) >= 2:
            soa, supp = build_soa_with_support(seqs)
            ctx_data[ctx] = (soa, supp, seqs)

    ctx_list = list(ctx_data.keys())
    merge_log = []
    total_comparisons = 0
    iteration = 0

    # Lines 3-12: while M is non-empty
    while True:
        iteration += 1

        # Line 2: M := {(s, t) ∈ T² | 0 < distD(s, t) < ε}
        M = []
        for i in range(len(ctx_list)):
            for j in range(i + 1, len(ctx_list)):
                ctx_a, ctx_b = ctx_list[i], ctx_list[j]
                if ctx_a not in ctx_data or ctx_b not in ctx_data:
                    continue

                _, supp_a, _ = ctx_data[ctx_a]
                _, supp_b, _ = ctx_data[ctx_b]

                total_comparisons += 1
                dist = soa_distance(supp_a, supp_b)

                if 0 < dist < threshold:
                    M.append((ctx_a, ctx_b, dist))

        # Line 3: while M is non-empty
        if not M:
            break

        # Line 4: for each (s, t) ∈ M
        for ctx_a, ctx_b, dist in M:
            if ctx_a not in ctx_data or ctx_b not in ctx_data:
                continue  # Already merged in this iteration

            _, supp_a, seqs_a = ctx_data[ctx_a]
            _, supp_b, seqs_b = ctx_data[ctx_b]

            # Line 6: soa(s) := soa(s) ∪ soa(t) [adjunction]
            combined_supp = adjunct_support(supp_a, supp_b)

            # Merge sequences (equivalent to adjunction for SOA building)
            merged_seqs = seqs_a + seqs_b

            # Rebuild SOA from merged sequences
            new_soa, new_supp = build_soa_with_support(merged_seqs)

            # Line 7: soa(t) := soa(s)
            ctx_data[ctx_a] = (new_soa, new_supp, merged_seqs)
            del ctx_data[ctx_b]

            merge_log.append({
                "iteration": iteration,
                "merged_into": str(ctx_a),
                "removed": str(ctx_b),
                "distance": round(dist, 4),
                "new_size": len(merged_seqs),
            })

        # Line 12: recompute M (loop continues)

    # Lines 13-14: for each type t, replace ρ(t) by ToSore(soa(t))
    # This happens implicitly when we infer SOREs later

    # Build final merged contexts dict
    merged = {}
    for ctx, (soa, supp, seqs) in ctx_data.items():
        merged[ctx] = seqs
    # Add contexts with < 2 sequences (no SOA built)
    for ctx, seqs in contexts.items():
        if ctx not in merged:
            merged[ctx] = seqs

    merge_info = {
        "iterations": iteration,
        "merges": len(merge_log),
        "comparisons": total_comparisons,
        "contexts_before": len(contexts),
        "contexts_after": len(merged),
        "threshold": threshold,
        "merge_log": merge_log,
    }

    return merged, merge_info


def _extract_text_from_seqs(seqs):
    """Extract text from sequences that may be tuples or plain strings.

    Sequences from preprocessing are [(capture_name, text, line_number), ...].
    Sequences from reduce may be plain strings.
    """
    result = []
    for seq in seqs:
        if seq and isinstance(seq[0], tuple):
            # Sequence of tuples: extract text
            result.append([text for _, text, _ in seq])
        else:
            # Sequence of strings: use as-is
            result.append(list(seq))
    return result


def minimize_contexts(merged_contexts):
    """Minimize: merge contexts with identical SOREs (Line 15, Algorithm 4).

    After Reduce, some contexts may have identical SOREs. Minimize unifies them.

    For our use case: group contexts by their inferred SORE, merge those with
    the same SORE into a single context.
    """
    # Build SORE for each context
    ctx_sores = {}
    for ctx, seqs in merged_contexts.items():
        if len(seqs) < 2:
            ctx_sores[ctx] = ("∅", seqs)
            continue

        clean = _extract_text_from_seqs(seqs)
        clean = [s for s in clean if s]

        if len(clean) < 2:
            ctx_sores[ctx] = ("∅", seqs)
            continue

        soa = build_soa(clean)
        sore = rwr0(soa)
        ctx_sores[ctx] = (sore, seqs)

    # Group by SORE
    sore_groups = defaultdict(list)
    for ctx, (sore, seqs) in ctx_sores.items():
        sore_groups[sore].append((ctx, seqs))

    # Merge contexts with same SORE
    minimized = {}
    for sore, items in sore_groups.items():
        if len(items) == 1:
            ctx, seqs = items[0]
            minimized[ctx] = seqs
        else:
            # Multiple contexts with same SORE -> merge into one
            merged_seqs = []
            for ctx, seqs in items:
                merged_seqs.extend(seqs)
            # Use the shortest context key as the representative
            rep = min(items, key=lambda x: len(x[0]))[0]
            minimized[rep] = merged_seqs

    return minimized


def reduce_and_infer(contexts, threshold, min_methods=3):
    """Full Reduce + Minimize + Infer pipeline.

    Args:
        contexts: dict mapping context_key -> list of sequences
        threshold: similarity threshold for Reduce
        min_methods: minimum methods per context to attempt inference

    Returns:
        dict with keys:
            merged: final context->sequences mapping (after minimize)
            infer_results: list of (context, sore, methods_count) tuples
            merge_info: stats from reduce step
            minimize_info: stats from minimize step
    """
    merged, merge_info = reduce_contexts(contexts, threshold)

    minimize_info = {
        "contexts_before": merge_info["contexts_after"],
        "contexts_after": len(merged),
        "merged_by_sore": 0,
    }

    # Line 14: ToSore(soa(t)) — infer SOREs for each context
    infer_results = []
    for ctx, seqs in sorted(merged.items(), key=lambda x: -len(x[1])):
        n = len(seqs)
        if n < min_methods:
            continue
        clean = _extract_text_from_seqs(seqs)
        clean = [s for s in clean if s]
        if len(clean) < 2:
            continue

        soa = build_soa(clean)
        sore = rwr0(soa)
        if sore != "∅":
            infer_results.append((ctx, sore, n))

    return {
        "merged": merged,
        "infer_results": infer_results,
        "merge_info": merge_info,
        "minimize_info": minimize_info,
        "coverage_count": sum(n for _, _, n in infer_results),
    }
