"""
Kolmogorov Engine — Trace & Compression
=========================================
This is the core contribution.

Records proof traces, finds recurring subtrace patterns,
and promotes them into reusable macros (learned lemmas)
when doing so reduces total description length.

The compression objective is MDL-based:
  total_cost = Σ proof_length + library_size

A macro is promoted when:
  compression_gain = old_total_cost - new_total_cost > 0
"""

from __future__ import annotations
from dataclasses import dataclass, field
from typing import Optional
from collections import Counter
import hashlib

from nexus_gate_pkg.engine.proof import ProofStep, ProofState, Goal


# ══════════════════════════════════════════════════════════════
# Proof Trace
# ══════════════════════════════════════════════════════════════

@dataclass
class ProofTrace:
    """
    A complete proof trace: the sequence of primitive applications
    that solved a theorem.
    """
    theorem_name: str
    steps: list[ProofStep] = field(default_factory=list)
    success: bool = False
    nodes_expanded: int = 0

    @property
    def length(self) -> int:
        return len(self.steps)

    @property
    def cost(self) -> int:
        """Total cost (each primitive = 1, each macro call = 1)."""
        return sum(s.cost for s in self.steps)

    def add_step(self, step: ProofStep):
        self.steps.append(step)

    def step_signature(self, i: int) -> str:
        """A hashable signature for a step — abstracts over specific variable names."""
        s = self.steps[i]
        # Abstract: keep primitive name and argument KEYS but normalize values
        # "rewrite_lhs with IH_xs→" and "rewrite_lhs with IH_a→" should match
        abstract_args = {}
        for k, v in sorted(s.args.items()):
            if isinstance(v, str):
                if v.startswith("IH_"):
                    abstract_args[k] = "IH_*"
                elif k == "var":
                    abstract_args[k] = "*"  # Abstract induction variable
                elif v.endswith("→") or v.endswith("←"):
                    abstract_args[k] = v
                else:
                    abstract_args[k] = v
            else:
                abstract_args[k] = str(v)
        return f"{s.primitive}({sorted(abstract_args.items())})"

    def subtrace_signature(self, start: int, end: int) -> str:
        """Signature for a contiguous subsequence of steps."""
        sigs = [self.step_signature(i) for i in range(start, end)]
        return "|".join(sigs)

    def __repr__(self):
        status = "✓" if self.success else "✗"
        lines = [f"Trace[{self.theorem_name}] {status} ({self.length} steps, cost={self.cost})"]
        for i, s in enumerate(self.steps):
            lines.append(f"  {i}: {s.primitive}({s.args})")
        return "\n".join(lines)


# ══════════════════════════════════════════════════════════════
# Macro (learned abstraction)
# ══════════════════════════════════════════════════════════════

@dataclass
class Macro:
    """
    A reusable proof fragment promoted from recurring subtraces.
    
    This is the analog of a lemma, but learned automatically.
    """
    name: str
    steps: list[ProofStep]           # The template steps
    signature: str                    # Pattern signature
    parameters: list[str] = field(default_factory=list)  # Variable parts
    use_count: int = 0               # How many times it's been applied
    definition_cost: int = 0         # Cost to store this macro

    @property
    def body_length(self) -> int:
        return len(self.steps)

    def __repr__(self):
        return f"Macro[{self.name}] ({self.body_length} steps, used {self.use_count}x)"


# ══════════════════════════════════════════════════════════════
# Trace Corpus
# ══════════════════════════════════════════════════════════════

class TraceCorpus:
    """
    Collection of all proof traces.
    Used for subtrace mining and compression analysis.
    """

    def __init__(self):
        self.traces: list[ProofTrace] = []
        self.macros: list[Macro] = []
        self._macro_counter = 0

    def add_trace(self, trace: ProofTrace):
        self.traces.append(trace)

    @property
    def total_cost(self) -> int:
        """
        Total description length = 
            sum of all proof costs + sum of all macro definition costs
        """
        proof_cost = sum(t.cost for t in self.traces)
        library_cost = sum(m.definition_cost for m in self.macros)
        return proof_cost + library_cost

    @property
    def total_steps(self) -> int:
        return sum(t.length for t in self.traces)

    # ─── Subtrace Mining ──────────────────────────────────────

    def find_recurring_subtraces(self, min_length: int = 2, 
                                  max_length: int = 8,
                                  min_count: int = 2) -> list[dict]:
        """
        Find contiguous step subsequences that appear in multiple traces.
        Returns candidates sorted by compression gain.
        """
        # Count all subtrace signatures
        sig_counts: Counter = Counter()
        sig_examples: dict[str, list[tuple[int, int, int]]] = {}  # sig -> [(trace_idx, start, end)]

        for t_idx, trace in enumerate(self.traces):
            for length in range(min_length, min(max_length + 1, trace.length + 1)):
                for start in range(trace.length - length + 1):
                    end = start + length
                    sig = trace.subtrace_signature(start, end)
                    sig_counts[sig] += 1
                    if sig not in sig_examples:
                        sig_examples[sig] = []
                    sig_examples[sig].append((t_idx, start, end))

        # Filter by minimum count
        candidates = []
        for sig, count in sig_counts.items():
            if count < min_count:
                continue

            examples = sig_examples[sig]
            # Get the actual steps from first example
            t_idx, start, end = examples[0]
            steps = self.traces[t_idx].steps[start:end]
            length = end - start

            # Calculate compression gain
            gain = self._compression_gain(length, count)

            if gain > 0:
                candidates.append({
                    "signature": sig,
                    "length": length,
                    "count": count,
                    "gain": gain,
                    "examples": examples,
                    "steps": steps,
                })

        # Sort by gain (highest first)
        candidates.sort(key=lambda c: c["gain"], reverse=True)
        return candidates

    def _compression_gain(self, subtrace_length: int, count: int) -> int:
        """
        MDL-based compression gain from promoting a subtrace to a macro.
        
        Before: each occurrence costs `subtrace_length` steps
        After:  each occurrence costs 1 (macro call), plus `subtrace_length` for definition
        
        Gain = count * subtrace_length - (count * 1 + subtrace_length)
             = (count - 1) * subtrace_length - count
             = count * (subtrace_length - 1) - subtrace_length
        """
        # Cost before: count * length
        cost_before = count * subtrace_length

        # Cost after: count * 1 (call) + length (definition)
        cost_after = count * 1 + subtrace_length

        return cost_before - cost_after

    # ─── Macro Promotion ──────────────────────────────────────

    def promote_best_macro(self) -> Optional[Macro]:
        """
        Find the best subtrace candidate and promote it to a macro.
        Returns the new macro, or None if no beneficial promotion exists.
        """
        candidates = self.find_recurring_subtraces()
        if not candidates:
            return None

        best = candidates[0]
        if best["gain"] <= 0:
            return None

        # Create macro
        self._macro_counter += 1
        macro = Macro(
            name=f"macro_{self._macro_counter}",
            steps=best["steps"],
            signature=best["signature"],
            definition_cost=best["length"],
            use_count=best["count"],
        )

        # Rewrite traces to use the macro
        self._apply_macro_to_corpus(macro, best["examples"])

        self.macros.append(macro)
        return macro

    def promote_all_beneficial(self) -> list[Macro]:
        """Keep promoting macros until no more beneficial ones exist."""
        promoted = []
        while True:
            macro = self.promote_best_macro()
            if macro is None:
                break
            promoted.append(macro)
        return promoted

    def _apply_macro_to_corpus(self, macro: Macro, 
                                examples: list[tuple[int, int, int]]):
        """
        Replace subtrace occurrences with a single macro-call step.
        
        Must process in reverse order within each trace to preserve indices.
        """
        # Group by trace
        by_trace: dict[int, list[tuple[int, int]]] = {}
        for t_idx, start, end in examples:
            if t_idx not in by_trace:
                by_trace[t_idx] = []
            by_trace[t_idx].append((start, end))

        for t_idx, ranges in by_trace.items():
            trace = self.traces[t_idx]
            # Sort by start position, reverse (so we replace from back to front)
            ranges.sort(key=lambda r: r[0], reverse=True)

            for start, end in ranges:
                # Check the range is still valid (might overlap with prior replacement)
                if end > len(trace.steps):
                    continue

                # Verify signature still matches
                current_sig = trace.subtrace_signature(start, end)
                if current_sig != macro.signature:
                    continue

                # Replace with single macro-call step
                macro_step = ProofStep(
                    primitive=f"MACRO:{macro.name}",
                    args={"macro": macro.name},
                    goal_before=trace.steps[start].goal_before,
                    goal_after=trace.steps[end-1].goal_after,
                    cost=1,  # Macro call costs 1
                )
                trace.steps = trace.steps[:start] + [macro_step] + trace.steps[end:]

    # ─── Analysis ─────────────────────────────────────────────

    def compression_report(self) -> dict:
        """Generate a report on the state of compression."""
        total_raw = sum(t.length for t in self.traces)
        total_cost = self.total_cost
        library_cost = sum(m.definition_cost for m in self.macros)
        macro_uses = sum(
            1 for t in self.traces for s in t.steps 
            if s.primitive.startswith("MACRO:")
        )

        return {
            "num_traces": len(self.traces),
            "total_raw_steps": total_raw,
            "total_cost": total_cost,
            "library_cost": library_cost,
            "num_macros": len(self.macros),
            "macro_uses_in_proofs": macro_uses,
            "macros": [
                {
                    "name": m.name,
                    "body_length": m.body_length,
                    "use_count": m.use_count,
                    "definition_cost": m.definition_cost,
                }
                for m in self.macros
            ],
        }

    def print_report(self):
        r = self.compression_report()
        print(f"\n{'='*50}")
        print(f"  COMPRESSION REPORT")
        print(f"{'='*50}")
        print(f"  Traces:        {r['num_traces']}")
        print(f"  Total steps:   {r['total_raw_steps']}")
        print(f"  Total cost:    {r['total_cost']}")
        print(f"  Library cost:  {r['library_cost']}")
        print(f"  Macros:        {r['num_macros']}")
        print(f"  Macro uses:    {r['macro_uses_in_proofs']}")
        if r['macros']:
            print(f"\n  Learned macros:")
            for m in r['macros']:
                print(f"    {m['name']}: {m['body_length']} steps, used {m['use_count']}x")
        print(f"{'='*50}\n")
