"""
Kolmogorov Engine — Search
============================
Beam search over the primitive algebra to find proofs.

The searcher tries all applicable primitives at each step,
keeps the best candidates by a heuristic score, and
records complete traces for the compression engine.
"""

from __future__ import annotations
from dataclasses import dataclass
from typing import Optional
import time

from nexus_gate_pkg.engine.terms import RewriteRule, term_size, Term, TermKind
from nexus_gate_pkg.engine.proof import (
    ProofState, ProofStep, PrimitiveResult, Goal,
    prim_refl, prim_rewrite_lhs, prim_rewrite_rhs,
    prim_normalize_lhs, prim_normalize_rhs,
    prim_symmetry, prim_congruence, prim_induction,
    prim_substitute, replay_macro,
)
from nexus_gate_pkg.engine.compression import ProofTrace, Macro


@dataclass
class SearchNode:
    """A node in the search tree."""
    state: ProofState
    trace: list[ProofStep]
    score: float
    depth: int


def _shared_structure(t1: Term, t2: Term) -> int:
    """Count shared top-level structure between two terms (higher = more similar)."""
    if t1 == t2:
        return term_size(t1)
    if t1.kind == TermKind.APP and t2.kind == TermKind.APP and t1.name == t2.name and len(t1.args) == len(t2.args):
        return 1 + sum(_shared_structure(a, b) for a, b in zip(t1.args, t2.args))
    return 0


def heuristic_score(state: ProofState) -> float:
    """
    Score a proof state (lower is better).
    """
    unsolved = [g for g in state.goals if not g.is_trivial()]
    if not unsolved:
        return -1000.0

    goal_scores = []
    for goal in unsolved:
        size = term_size(goal.lhs) + term_size(goal.rhs)
        shared = _shared_structure(goal.lhs, goal.rhs)
        goal_scores.append(size - shared * 2.0)

    return max(goal_scores) + len(unsolved) * 2.0 + sum(goal_scores) * 0.1
    
    return score


def generate_actions(state: ProofState, 
                     macros: list[Macro] = None,
                     lemma_rules: list[RewriteRule] = None) -> list[tuple[str, PrimitiveResult]]:
    """
    Generate all applicable primitive actions from a proof state.
    If macros are provided, also try replaying each macro.
    Returns list of (description, result) pairs.
    """
    actions = []
    goal = state.current_goal()
    if goal is None:
        return actions

    # 0. Try macros FIRST (they're cheaper if they work)
    if macros:
        for macro in macros:
            r = replay_macro(macro.name, macro.steps, state)
            if r.success:
                actions.append((f"MACRO:{macro.name}", r))

    # 1. Reflexivity (try to close trivial goal)
    r = prim_refl(state)
    if r.success:
        actions.append(("refl", r))

    # 2. Normalize LHS and RHS
    r = prim_normalize_lhs(state)
    if r.success:
        actions.append(("normalize_lhs", r))

    r = prim_normalize_rhs(state)
    if r.success:
        actions.append(("normalize_rhs", r))

    # 3. Congruence
    r = prim_congruence(state)
    if r.success:
        actions.append(("congruence", r))

    # 4. Symmetry
    r = prim_symmetry(state)
    if r.success:
        actions.append(("symmetry", r))

    # 5. Apply each rewrite rule from definitions (both sides)
    for rule in state.definitions:
        r = prim_rewrite_lhs(state, rule)
        if r.success:
            actions.append((f"rewrite_lhs({rule.name})", r))

        r = prim_rewrite_rhs(state, rule)
        if r.success:
            actions.append((f"rewrite_rhs({rule.name})", r))

    # 6. Apply context equations
    for eq in state.context:
        fwd = eq.forward()
        bwd = eq.backward()

        r = prim_rewrite_lhs(state, fwd)
        if r.success:
            actions.append((f"apply_ctx_lhs({eq.name}→)", r))

        r = prim_rewrite_rhs(state, fwd)
        if r.success:
            actions.append((f"apply_ctx_rhs({eq.name}→)", r))

        r = prim_rewrite_lhs(state, bwd)
        if r.success:
            actions.append((f"apply_ctx_lhs({eq.name}←)", r))

        r = prim_rewrite_rhs(state, bwd)
        if r.success:
            actions.append((f"apply_ctx_rhs({eq.name}←)", r))

    # 6b. Apply lemma rules (proven theorems from previous proofs)
    if lemma_rules:
        for rule in lemma_rules:
            r = prim_rewrite_lhs(state, rule)
            if r.success:
                actions.append((f"lemma_lhs({rule.name})", r))
            
            r = prim_rewrite_rhs(state, rule)
            if r.success:
                actions.append((f"lemma_rhs({rule.name})", r))

    # 7. Induction (only if goal has free variables)
    from nexus_gate_pkg.engine.terms import free_vars
    goal_vars = free_vars(goal.lhs) | free_vars(goal.rhs)

    # Infer likely type from variable name and what functions it appears in
    nat_hints = {"n", "m", "a", "b", "c", "k", "j", "i"}
    list_hints = {"xs", "ys", "zs", "l", "ls"}
    tree_hints = {"t", "t1", "t2", "tree", "lt", "rt"}
    # f is a function variable — not inductable
    non_induct = {"f", "v", "h", "x", "y"}

    for var_name in sorted(goal_vars):
        base = var_name.split("_")[0]  # handle xs_tail etc
        if base in non_induct:
            continue
        
        # Check custom type hints first
        from nexus_gate_pkg.engine.proof import CUSTOM_INDUCTION
        custom_match = None
        for ctype, scheme in CUSTOM_INDUCTION.items():
            if base in scheme.get("var_hints", set()):
                custom_match = ctype
                break
        
        if custom_match:
            r = prim_induction(state, var_name, custom_match)
            if r.success:
                actions.append((f"induction({var_name}, {custom_match})", r))
        elif base in list_hints:
            r = prim_induction(state, var_name, "List")
            if r.success:
                actions.append((f"induction({var_name}, List)", r))
        elif base in nat_hints:
            r = prim_induction(state, var_name, "Nat")
            if r.success:
                actions.append((f"induction({var_name}, Nat)", r))
        elif base in tree_hints:
            r = prim_induction(state, var_name, "Tree")
            if r.success:
                actions.append((f"induction({var_name}, Tree)", r))
        else:
            # Try all built-in types + custom registered types
            # Use definition analysis: after induction, does the base case normalize?
            # If it does, the type is likely correct.
            from nexus_gate_pkg.engine.proof import CUSTOM_INDUCTION
            all_types = ["Nat", "List", "Tree"] + list(CUSTOM_INDUCTION.keys())
            for dtype in all_types:
                r = prim_induction(state, var_name, dtype)
                if r.success:
                    # Check if base case can normalize — strong signal of correct type
                    base_goal = r.state.goals[0] if r.state.goals else None
                    if base_goal and not base_goal.is_trivial():
                        test_state = r.state.clone()
                        norm_r = prim_normalize_lhs(test_state)
                        if norm_r.success:
                            # Base case normalizes — this type is productive
                            actions.append((f"induction({var_name}, {dtype})", r))
                        # else: skip — base case can't normalize, wrong type
                    else:
                        actions.append((f"induction({var_name}, {dtype})", r))

    return actions


def beam_search(initial_state: ProofState,
                beam_width: int = 20,
                max_depth: int = 30,
                max_nodes: int = 5000,
                timeout: float = 30.0,
                verbose: bool = False,
                macros: list[Macro] = None,
                lemma_rules: list[RewriteRule] = None) -> Optional[ProofTrace]:
    """
    Beam search for a proof.
    
    Returns a ProofTrace if successful, None if not.
    """
    start_time = time.time()

    initial_node = SearchNode(
        state=initial_state,
        trace=[],
        score=heuristic_score(initial_state),
        depth=0,
    )

    beam = [initial_node]
    nodes_expanded = 0
    best_score = float('inf')
    visited: set[int] = set()  # Hashes of goal states we've already explored

    def _state_hash(state: ProofState) -> int:
        """Hash the goal list for cycle detection."""
        parts = []
        for g in state.goals:
            if not g.is_trivial():
                parts.append((str(g.lhs), str(g.rhs)))
        return hash(tuple(sorted(parts)))

    theorem_name = initial_state.goals[0].name if initial_state.goals else "unknown"

    for depth in range(max_depth):
        if time.time() - start_time > timeout:
            if verbose:
                print(f"  Timeout at depth {depth}")
            break

        candidates = []

        for node in beam:
            if node.state.is_complete():
                # Found a proof!
                trace = ProofTrace(
                    theorem_name=theorem_name,
                    steps=node.trace,
                    success=True,
                    nodes_expanded=nodes_expanded,
                )
                elapsed = time.time() - start_time
                if verbose:
                    print(f"  ✓ Proof found! depth={depth}, "
                          f"steps={len(node.trace)}, "
                          f"nodes={nodes_expanded}, "
                          f"time={elapsed:.2f}s")
                return trace

            actions = generate_actions(node.state, macros=macros, lemma_rules=lemma_rules)
            nodes_expanded += len(actions)

            if nodes_expanded > max_nodes:
                if verbose:
                    print(f"  Node limit at depth {depth}")
                return None

            for desc, result in actions:
                sh = _state_hash(result.state)
                if sh in visited:
                    continue
                visited.add(sh)
                new_score = heuristic_score(result.state)
                new_node = SearchNode(
                    state=result.state,
                    trace=node.trace + result.steps,
                    score=new_score + len(node.trace) * 0.1,  # Slight preference for shorter proofs
                    depth=depth + 1,
                )
                candidates.append(new_node)

        if not candidates:
            if verbose:
                print(f"  No candidates at depth {depth}")
            break

        # Stratified beam: reserve slots for post-induction paths
        # so they don't get drowned out by cheaper non-induction alternatives
        has_induction = []
        no_induction = []
        for c in candidates:
            if any(s.primitive == "induction" for s in c.trace):
                has_induction.append(c)
            else:
                no_induction.append(c)

        has_induction.sort(key=lambda n: n.score)
        no_induction.sort(key=lambda n: n.score)

        half = beam_width // 2
        if has_induction and no_induction:
            beam = no_induction[:half] + has_induction[:half]
            # Fill remaining slots from whichever pool has leftovers
            remaining = beam_width - len(beam)
            if remaining > 0:
                leftover = sorted(
                    no_induction[half:] + has_induction[half:],
                    key=lambda n: n.score
                )
                beam.extend(leftover[:remaining])
        else:
            all_sorted = sorted(candidates, key=lambda n: n.score)
            beam = all_sorted[:beam_width]

        if verbose and depth % 5 == 0:
            print(f"  depth={depth}, candidates={len(candidates)}, "
                  f"best_score={beam[0].score:.1f}, nodes={nodes_expanded}")

    if verbose:
        print(f"  ✗ No proof found (nodes={nodes_expanded})")
    return None


def prove_theorem(state: ProofState, verbose: bool = True, **kwargs) -> Optional[ProofTrace]:
    """
    Attempt to prove a theorem. Convenience wrapper around beam_search.
    """
    name = state.goals[0].name if state.goals else "?"
    if verbose:
        print(f"\nProving: {state.goals[0] if state.goals else '?'}")

    trace = beam_search(state, verbose=verbose, **kwargs)

    if trace and verbose:
        print(f"  Proof trace:")
        for i, step in enumerate(trace.steps):
            print(f"    {i}: {step.primitive}({step.args})")

    return trace
