"""
Kolmogorov Engine — Domain: Natural Numbers & Lists
=====================================================
Function definitions and theorems for Phase 1.

Natural numbers: Z, S(n)
Lists: Nil, Cons(h, t)

Functions:
  add(Z, y)      = y
  add(S(x), y)   = S(add(x, y))
  
  append(Nil, ys)      = ys
  append(Cons(x, xs), ys) = Cons(x, append(xs, ys))
  
  rev(Nil)             = Nil
  rev(Cons(x, xs))     = append(rev(xs), Cons(x, Nil))
  
  length(Nil)          = Z
  length(Cons(x, xs))  = S(length(xs))
  
  map(f, Nil)          = Nil
  map(f, Cons(x, xs))  = Cons(f(x), map(f, xs))
"""

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


# ══════════════════════════════════════════════════════════════
# Convenience constructors
# ══════════════════════════════════════════════════════════════

def Z():
    return Const("Z")

def S(n):
    return App("S", [n])

def Nil():
    return Const("Nil")

def Cons(h, t):
    return App("Cons", [h, t])

def add(x, y):
    return App("add", [x, y])

def append(xs, ys):
    return App("append", [xs, ys])

def rev(xs):
    return App("rev", [xs])

def length(xs):
    return App("length", [xs])

def map_f(f, xs):
    return App("map", [f, xs])

def snoc(xs, x):
    """append(xs, Cons(x, Nil)) — used in reverse proofs."""
    return append(xs, Cons(x, Nil()))


# ══════════════════════════════════════════════════════════════
# Function definitions (as rewrite rules)
# ══════════════════════════════════════════════════════════════

x, y, xs, ys, h, t, f, n, m = (
    Var("x"), Var("y"), Var("xs"), Var("ys"),
    Var("h"), Var("t"), Var("f"), Var("n"), Var("m"),
)

DEFINITIONS = [
    # add
    RewriteRule("add_Z",    add(Z(), y),      y),
    RewriteRule("add_S",    add(S(n), m),     S(add(n, m))),

    # append
    RewriteRule("append_Nil",  append(Nil(), ys),        ys),
    RewriteRule("append_Cons", append(Cons(h, t), ys),   Cons(h, append(t, ys))),

    # rev
    RewriteRule("rev_Nil",  rev(Nil()),        Nil()),
    RewriteRule("rev_Cons", rev(Cons(h, t)),   append(rev(t), Cons(h, Nil()))),

    # length
    RewriteRule("length_Nil",  length(Nil()),        Z()),
    RewriteRule("length_Cons", length(Cons(h, t)),   S(length(t))),

    # map
    RewriteRule("map_Nil",  map_f(f, Nil()),        Nil()),
    RewriteRule("map_Cons", map_f(f, Cons(h, t)),   Cons(App("apply", [f, h]), map_f(f, t))),
]


# ══════════════════════════════════════════════════════════════
# Theorems to prove
# ══════════════════════════════════════════════════════════════

def make_theorem(name: str, lhs: Term, rhs: Term, 
                 extra_context: list[Equation] = None) -> ProofState:
    """Create a proof state for a theorem."""
    state = ProofState(
        goals=[Goal(lhs, rhs, name)],
        context=list(extra_context or []),
        definitions=DEFINITIONS,
    )
    return state


# ─── Natural number theorems ─────────────────────────────────

def thm_add_zero_right():
    """∀n. add(n, Z) = n"""
    return make_theorem("add_zero_right",
        add(Var("n"), Z()), Var("n"))

def thm_add_succ_right():
    """∀n m. add(n, S(m)) = S(add(n, m))"""
    return make_theorem("add_succ_right",
        add(Var("n"), S(Var("m"))), S(add(Var("n"), Var("m"))))

def thm_add_comm():
    """∀n m. add(n, m) = add(m, n)"""
    return make_theorem("add_comm",
        add(Var("n"), Var("m")), add(Var("m"), Var("n")),
        # Need these as lemmas
        extra_context=[
            Equation("add_zero_right", add(Var("k"), Z()), Var("k")),
            Equation("add_succ_right", add(Var("k"), S(Var("j"))), S(add(Var("k"), Var("j")))),
        ])

def thm_add_assoc():
    """∀a b c. add(add(a, b), c) = add(a, add(b, c))"""
    return make_theorem("add_assoc",
        add(add(Var("a"), Var("b")), Var("c")),
        add(Var("a"), add(Var("b"), Var("c"))))


# ─── List theorems ───────────────────────────────────────────

def thm_append_nil_right():
    """∀xs. append(xs, Nil) = xs"""
    return make_theorem("append_nil_right",
        append(Var("xs"), Nil()), Var("xs"))

def thm_append_assoc():
    """∀xs ys zs. append(append(xs, ys), zs) = append(xs, append(ys, zs))"""
    return make_theorem("append_assoc",
        append(append(Var("xs"), Var("ys")), Var("zs")),
        append(Var("xs"), append(Var("ys"), Var("zs"))))

def thm_length_append():
    """∀xs ys. length(append(xs, ys)) = add(length(xs), length(ys))"""
    return make_theorem("length_append",
        length(append(Var("xs"), Var("ys"))),
        add(length(Var("xs")), length(Var("ys"))))

def thm_rev_append():
    """∀xs ys. rev(append(xs, ys)) = append(rev(ys), rev(xs))"""
    return make_theorem("rev_append",
        rev(append(Var("xs"), Var("ys"))),
        append(rev(Var("ys")), rev(Var("xs"))),
        extra_context=[
            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")),
        ])

def thm_rev_rev():
    """∀xs. rev(rev(xs)) = xs"""
    return make_theorem("rev_rev",
        rev(rev(Var("xs"))), Var("xs"),
        extra_context=[
            Equation("rev_append",
                rev(append(Var("a"), Var("b"))),
                append(rev(Var("b")), rev(Var("a")))),
            Equation("append_nil_right",
                append(Var("a"), Nil()), Var("a")),
            Equation("append_assoc",
                append(append(Var("a"), Var("b")), Var("c")),
                append(Var("a"), append(Var("b"), Var("c")))),
        ])

def thm_map_append():
    """∀f xs ys. map(f, append(xs, ys)) = append(map(f, xs), map(f, ys))"""
    return make_theorem("map_append",
        map_f(Var("f"), append(Var("xs"), Var("ys"))),
        append(map_f(Var("f"), Var("xs")), map_f(Var("f"), Var("ys"))))

def thm_length_rev():
    """∀xs. length(rev(xs)) = length(xs)"""
    return make_theorem("length_rev",
        length(rev(Var("xs"))), length(Var("xs")),
        extra_context=[
            Equation("length_append",
                length(append(Var("a"), Var("b"))),
                add(length(Var("a")), length(Var("b")))),
            Equation("add_comm",
                add(Var("a"), Var("b")),
                add(Var("b"), Var("a"))),
        ])


# ══════════════════════════════════════════════════════════════
# All theorems in order of difficulty
# ══════════════════════════════════════════════════════════════

# ══════════════════════════════════════════════════════════════
# All theorems in order of difficulty
# ══════════════════════════════════════════════════════════════

# Additional Nat theorems
def thm_add_zero_left():
    """∀n. add(Z, n) = n  (trivial by definition, but good baseline)"""
    return make_theorem("add_zero_left",
        add(Z(), Var("n")), Var("n"))

def thm_add_one_right():
    """∀n. add(n, S(Z)) = S(n)"""
    return make_theorem("add_one_right",
        add(Var("n"), S(Z())), S(Var("n")),
        extra_context=[
            Equation("add_zero_right", add(Var("k"), Z()), Var("k")),
        ])

def thm_add_one_left():
    """∀n. add(S(Z), n) = S(n)"""
    return make_theorem("add_one_left",
        add(S(Z()), Var("n")), S(Var("n")))

# Additional list theorems
def thm_append_nil_left():
    """∀xs. append(Nil, xs) = xs  (trivial by definition)"""
    return make_theorem("append_nil_left",
        append(Nil(), Var("xs")), Var("xs"))

def thm_length_nil():
    """length(Nil) = Z"""
    return make_theorem("length_nil",
        length(Nil()), Z())

def thm_length_cons():
    """∀x xs. length(Cons(x, xs)) = S(length(xs))"""
    return make_theorem("length_cons",
        length(Cons(Var("x"), Var("xs"))), S(length(Var("xs"))))

def thm_length_singleton():
    """∀x. length(Cons(x, Nil)) = S(Z)"""
    return make_theorem("length_singleton",
        length(Cons(Var("x"), Nil())), S(Z()))

def thm_map_nil():
    """∀f. map(f, Nil) = Nil"""
    return make_theorem("map_nil",
        map_f(Var("f"), Nil()), Nil())

def thm_rev_nil():
    """rev(Nil) = Nil"""
    return make_theorem("rev_nil",
        rev(Nil()), Nil())

def thm_rev_singleton():
    """∀x. rev(Cons(x, Nil)) = Cons(x, Nil)"""
    return make_theorem("rev_singleton",
        rev(Cons(Var("x"), Nil())),
        Cons(Var("x"), Nil()),
        extra_context=[
            Equation("append_nil_right", append(Var("a"), Nil()), Var("a")),
        ])

def thm_length_map():
    """∀f xs. length(map(f, xs)) = length(xs)"""
    return make_theorem("length_map",
        length(map_f(Var("f"), Var("xs"))), length(Var("xs")))


ALL_THEOREMS = [
    # Trivial (0-1 steps, baselines)
    ("add_zero_left",      thm_add_zero_left),
    ("append_nil_left",    thm_append_nil_left),
    ("length_nil",         thm_length_nil),
    ("length_cons",        thm_length_cons),
    ("map_nil",            thm_map_nil),
    ("rev_nil",            thm_rev_nil),
    # Easy induction (3-4 steps)
    ("add_zero_right",     thm_add_zero_right),
    ("add_succ_right",     thm_add_succ_right),
    ("add_one_left",       thm_add_one_left),
    ("add_one_right",      thm_add_one_right),
    ("append_nil_right",   thm_append_nil_right),
    ("length_singleton",   thm_length_singleton),
    ("rev_singleton",      thm_rev_singleton),
    # Medium induction (4-6 steps)
    ("add_assoc",          thm_add_assoc),
    ("add_comm",           thm_add_comm),
    ("append_assoc",       thm_append_assoc),
    ("length_append",      thm_length_append),
    ("map_append",         thm_map_append),
    ("length_map",         thm_length_map),
    # Hard (6+ steps, multiple lemmas)
    ("length_rev",         thm_length_rev),
    ("rev_rev",            thm_rev_rev),
]
