"""
Kolmogorov Engine — Domain: Binary Trees
==========================================
Extends the engine to a new datatype for research benchmarks.

Binary trees: Leaf, Node(l, v, r)

Functions:
  mirror(Leaf)          = Leaf
  mirror(Node(l, v, r)) = Node(mirror(r), v, mirror(l))

  size(Leaf)            = Z
  size(Node(l, v, r))   = S(add(size(l), size(r)))

  flatten(Leaf)         = Nil
  flatten(Node(l, v, r)) = append(flatten(l), Cons(v, flatten(r)))

  height(Leaf)          = Z
  height(Node(l, v, r)) = S(max(height(l), height(r)))
"""

from nexus_gate_pkg.engine.terms import (
    Term, Var, Const, App, RewriteRule, Equation,
)
from nexus_gate_pkg.engine.proof import ProofState, Goal

# Import Nat/List constructors from the base domain
from nexus_gate_pkg.corpus.domain import (
    Z, S, Nil, Cons, add, append, length, rev,
    DEFINITIONS as BASE_DEFINITIONS,
)


# ══════════════════════════════════════════════════════════════
# Tree constructors
# ══════════════════════════════════════════════════════════════

def Leaf():
    return Const("Leaf")

def Node(l, v, r):
    return App("Node", [l, v, r])

def mirror(t):
    return App("mirror", [t])

def tree_size(t):
    return App("tsize", [t])

def flatten(t):
    return App("flatten", [t])


# ══════════════════════════════════════════════════════════════
# Variables
# ══════════════════════════════════════════════════════════════

l, r, v, t = Var("l"), Var("r"), Var("v"), Var("t")


# ══════════════════════════════════════════════════════════════
# Function definitions
# ══════════════════════════════════════════════════════════════

TREE_DEFINITIONS = [
    # mirror
    RewriteRule("mirror_Leaf", mirror(Leaf()), Leaf()),
    RewriteRule("mirror_Node", mirror(Node(l, v, r)),
                Node(mirror(r), v, mirror(l))),

    # tsize (tree size)
    RewriteRule("tsize_Leaf", tree_size(Leaf()), Z()),
    RewriteRule("tsize_Node", tree_size(Node(l, v, r)),
                S(add(tree_size(l), tree_size(r)))),

    # flatten
    RewriteRule("flatten_Leaf", flatten(Leaf()), Nil()),
    RewriteRule("flatten_Node", flatten(Node(l, v, r)),
                append(flatten(l), Cons(v, flatten(r)))),
]

ALL_DEFINITIONS = BASE_DEFINITIONS + TREE_DEFINITIONS


# ══════════════════════════════════════════════════════════════
# Induction scheme for Tree
# ══════════════════════════════════════════════════════════════
# We need to register Tree induction in the proof engine.
# For now, we add Tree theorems that use manual proof states
# with the right induction structure.
# The search engine's induction primitive needs extending — see below.


def make_tree_theorem(name, lhs, rhs, extra_context=None):
    """Create a proof state for a tree theorem."""
    return ProofState(
        goals=[Goal(lhs, rhs, name)],
        context=list(extra_context or []),
        definitions=ALL_DEFINITIONS,
    )


# ══════════════════════════════════════════════════════════════
# Tree theorems
# ══════════════════════════════════════════════════════════════

def thm_mirror_leaf():
    """mirror(Leaf) = Leaf"""
    return make_tree_theorem("mirror_leaf",
        mirror(Leaf()), Leaf())

def thm_mirror_mirror():
    """∀t. mirror(mirror(t)) = t"""
    return make_tree_theorem("mirror_mirror",
        mirror(mirror(Var("t"))), Var("t"))

def thm_size_mirror():
    """∀t. tsize(mirror(t)) = tsize(t)"""
    return make_tree_theorem("size_mirror",
        tree_size(mirror(Var("t"))), tree_size(Var("t")),
        extra_context=[
            Equation("add_comm",
                add(Var("a"), Var("b")),
                add(Var("b"), Var("a"))),
        ])

def thm_flatten_leaf():
    """flatten(Leaf) = Nil"""
    return make_tree_theorem("flatten_leaf",
        flatten(Leaf()), Nil())

def thm_size_leaf():
    """tsize(Leaf) = Z"""
    return make_tree_theorem("size_leaf",
        tree_size(Leaf()), Z())

def thm_mirror_node():
    """mirror(Node(l, v, r)) = Node(mirror(r), v, mirror(l))"""
    return make_tree_theorem("mirror_node",
        mirror(Node(Var("l"), Var("v"), Var("r"))),
        Node(mirror(Var("r")), Var("v"), mirror(Var("l"))))


def thm_length_flatten():
    """∀t. length(flatten(t)) = tsize(t)"""
    return make_tree_theorem("length_flatten",
        length(flatten(Var("t"))), tree_size(Var("t")),
        extra_context=[
            Equation("length_append",
                length(append(Var("a"), Var("b"))),
                add(length(Var("a")), length(Var("b")))),
            Equation("add_succ_right",
                add(Var("k"), S(Var("j"))),
                S(add(Var("k"), Var("j")))),
        ])

def thm_flatten_mirror():
    """∀t. flatten(mirror(t)) = rev(flatten(t))
    This is the hardest theorem — requires rev_append, append_assoc on RHS.
    Currently beyond search budget; serves as a research target.
    """
    return make_tree_theorem("flatten_mirror",
        flatten(mirror(Var("t"))), rev(flatten(Var("t"))),
        extra_context=[
            Equation("rev_append",
                rev(append(Var("a"), Var("b"))),
                append(rev(Var("b")), rev(Var("a")))),
            Equation("append_assoc",
                append(append(Var("a"), Var("b")), Var("c")),
                append(Var("a"), append(Var("b"), Var("c")))),
            Equation("append_nil_right",
                append(Var("a"), Nil()), Var("a")),
        ])


TREE_THEOREMS = [
    # Trivial
    ("mirror_leaf",       thm_mirror_leaf),
    ("flatten_leaf",      thm_flatten_leaf),
    ("size_leaf",         thm_size_leaf),
    ("mirror_node",       thm_mirror_node),
    # Induction required
    ("mirror_mirror",     thm_mirror_mirror),
    ("size_mirror",       thm_size_mirror),
    ("length_flatten",    thm_length_flatten),
    # Research frontier (beyond current search)
    ("flatten_mirror",    thm_flatten_mirror),
]
