(require '[ansatz.core :as a])
(a/init! "/var/tmp/ansatz-cslib" "cslib")

;; Verified merge sort — kernel-checked, compiles to Clojure
(a/defn merge [xs :- (List Nat), ys :- (List Nat)] (List Nat)
  :termination-by (+ (List.length xs) (List.length ys))
  (match xs (List Nat) (List Nat)
    (nil ys)
    (cons [x xs']
      (match ys (List Nat) (List Nat)
        (nil (cons x xs'))
        (cons [y ys']
          (if (<= x y)
            (cons x (merge xs' (cons y ys')))
            (cons y (merge (cons x xs') ys'))))))))

;; Run it — it's a normal Clojure function
(merge '(1 3 5) '(2 4 6)) ;; => (1 2 3 4 5 6)

;; Structures compile to Clojure defrecord
(a/structure Point [] (x Nat) (y Nat))
(->Point 3 4)  ;; => #user.Point{:x 3, :y 4}
(:x (->Point 3 4))  ;; => 3

;; Define your own types (like Clojure records, but type-checked)
(a/inductive RBColor [] (red) (black))
(a/inductive RBTree [α Type]
  (leaf)
  (node [color RBColor] [left (RBTree α)] [key α] [right (RBTree α)]))

;; Pattern matching with recursion — the kernel checks termination
;; ih_left and ih_right are induction hypotheses (recursive results)
(a/defn rb-size [t (RBTree Nat)] Nat
  (match t (RBTree Nat) Nat
    (leaf 0)
    (node [color left key right] (+ 1 (+ ih_left ih_right)))))

;; Nested match for BST lookup — references outer param k
(a/defn rb-member [t (RBTree Nat) k Nat] Bool
  (match t (RBTree Nat) Bool
    (leaf false)
    (node [color left key right]
      (match (< k key) Bool Bool
        (true ih_left)                    ;; recurse into left subtree
        (false (match (== k key) Bool Bool
                 (true true)              ;; found it
                 (false ih_right)))))))   ;; recurse into right subtree

;; All verified functions compile to ordinary Clojure and run on the JVM
(rb-size tree)           ;; => 10
(rb-member tree 4)       ;; => true
(rb-member tree 42)      ;; => false

;; Simplest proof: rb-size(leaf) computes to 0 by definition.
;; (rfl) = "reflexivity" — both sides reduce to the same thing.
(a/theorem leaf-size-zero []
  (= Nat (rb-size (RBTree.leaf Nat)) 0)
  (rfl))

;; (omega) solves linear arithmetic over natural numbers automatically.
;; It sees: goal is (rb-size l) ≤ 1 + (rb-size l) + (rb-size r)
;; and closes it because x ≤ 1 + x + y for all naturals.
(a/theorem left-le-size [c :- RBColor, l :- (RBTree Nat), k :- Nat, r :- (RBTree Nat)]
  (<= Nat (rb-size l) (+ 1 (+ (rb-size l) (rb-size r))))
  (omega))

;; Induction + grind: the standard pattern for list properties.
;; (induction l) splits into two goals:
;;   1. Base case:  llen (lmap f []) = llen []        — both sides are 0
;;   2. Inductive:  llen (lmap f (x::xs)) = llen (x::xs)
;;      with IH:    llen (lmap f xs) = llen xs
;; (grind "lmap" "llen") tells grind to use the equation theorems
;; for lmap and llen (e.g. llen.eq_2: llen(x::xs) = 1 + llen xs).
;; It unfolds one step, sees 1 + llen(lmap f xs) = 1 + llen xs,
;; and closes via congruence from the IH.
(a/theorem map-preserves-len [f :- (arrow Nat Nat), l :- (List Nat)]
  (= Nat (llen (lmap f l)) (llen l))
  (induction l) (all_goals (grind "lmap" "llen")))

;; Sorted is an indexed inductive: Sorted [] | Sorted [a] | Sorted (a::b::tl) when a≤b.
;; (induction h) on h : Sorted l gives 3 goals (one per constructor).
;; (grind "insertSorted") uses insertSorted's equation theorems to unfold it,
;; case-splits on ≤ comparisons, applies constructors, uses omega
;; for arithmetic, and the IH for the recursive case.
(a/theorem insert-preserves
  [x :- Nat, l :- (List Nat), h :- (Sorted l)]
  (Sorted (insertSorted x l))
  (induction h) (grind "insertSorted"))

;; balance1 has 7 branches (from nested pattern matching on color + subtree shape).
;; (cases hl) splits ValidRB(l) into leaf/node.
;; (simp "balance1") uses balance1's equation theorems to unfold it for each branch.
;; (grind) then applies ValidRB constructors and matches sub-proofs from context.
;; The middle lines split further on color, subtree shape, and inner color
;; to expose the left-left rotation case (the only one that restructures the tree).
(a/theorem balance1-preserves-valid
  [l :- (RBTree Nat), v :- Nat, r :- (RBTree Nat),
   hl :- (ValidRB l), hr :- (ValidRB r)]
  (ValidRB (balance1 l v r))
  (cases hl)                           ;; leaf or node?
  (all_goals (try (simp "balance1")))  ;; unfold via equation theorems
  (all_goals (try (grind)))            ;; close the easy cases
  (all_goals (try (cases c)))          ;; red or black?
  (all_goals (try (cases l)))          ;; left subtree shape
  (all_goals (try (cases color)))      ;; inner node color (detects LL rotation)
  (all_goals (try (cases hl)))         ;; decompose inner ValidRB proof
  (all_goals (try (simp "balance1")))  ;; unfold the rotation case
  (all_goals (try (grind))))           ;; close all remaining goals

{:deps {org.replikativ/ansatz {:mvn/version "0.1.0-SNAPSHOT"}}}

;; Works out of the box for Nat proofs (no Real/analysis)
(require '[ansatz.export.storage :as s] '[ansatz.export.replay :as r]
         '[ansatz.export.parser :as p])
(def env (:env (r/replay (:decls (p/parse-ndjson-file "test-data/init-medium.ndjson")))))

(require '[ansatz.core :as a])

;; Load Mathlib environment
(a/init! "/var/tmp/ansatz-mathlib" "mathlib")

;; Define and prove
(a/defn double [n :- Nat] Nat (+ n n))

(a/theorem add-zero [n :- Nat]
  (= Nat (+ n 0) n)
  (simp "Nat.add_zero"))

(double 21) ;; => 42

Nat                          ;; natural numbers
Int                          ;; integers
Real                         ;; real numbers
Bool                         ;; booleans
Prop                         ;; propositions
Type                         ;; types
(List Nat)                   ;; lists

(+ a b)                      ;; Nat addition (default)
(add Real a b)               ;; typed addition
(mul Real a b)               ;; typed multiplication
(sub Real a b)               ;; subtraction
(div Real a b)               ;; division
(pow Real k n)               ;; power (base^Nat)
(= Nat a b)                  ;; equality (Prop)
(le Real a b)                ;; ≤ (Prop)
(<= Real a b)                ;; ≤ sugar (Prop, 3-arg form)
(<= a b)                     ;; ≤ (Bool, Nat default, 2-arg form)

;; Verified function (structural recursion)
(a/defn name [param :- Type, ...] ReturnType body)

;; Well-founded recursion (non-structural, e.g. merge sort)
(a/defn merge [xs :- (List Nat), ys :- (List Nat)] (List Nat)
  :termination-by (+ (List.length xs) (List.length ys))
  body)

;; Typeclass params
(a/defn sort [α :- Type, inst :- (Ord α) :inst, xs :- (List α)] (List α) ...)

;; Theorem
(a/theorem name [param :- Type, hyp :- (le Real 0 x), ...]
  proposition
  (tactic1 arg1 arg2)
  (tactic2)
  ...)

;; Structure (compiles to defrecord)
(a/structure Point [] (x Nat) (y Nat))
;; Creates: (->Point 3 4) => #user.Point{:x 3, :y 4}
;; Access: (:x p), (:y p)

;; Inductive type
(a/inductive Color [] (red) (green) (blue))
(a/inductive MyList [α Type] (nil) (cons [head α] [tail (MyList α)]))

;; Indexed inductive family (length-indexed vector)
(a/inductive Vec [α Type] :indices [n Nat]
  (nil :where [0])
  (cons [a α] [m Nat] [tail (Vec α m)] :where [(+ m 1)]))

;; Match on inductive types — uses CIC recursor under the hood
(match expr InductiveType ReturnType
  (ctor1 body1)                          ;; leaf case
  (ctor2 [field1 field2 ...] body2))     ;; node case with field bindings

;; Induction hypotheses are auto-generated for recursive fields:
;;   ih_left, ih_right (named after fields), or ih0, ih1
(a/defn size [t (MyList Nat)] Nat
  (match t (MyList Nat) Nat
    (nil 0)
    (cons [head tail] (+ 1 ih_tail))))

(require '[ansatz.config :as config])

;; Increase fuel for complex proofs
(binding [config/*high-fuel* 100000000]
  (a/theorem ...))

;; Enable grind debugging
(binding [config/*grind-debug* true]
  (a/theorem ... (grind)))

;; Register custom simplification procedure
(binding [config/*simprocs* [(fn [st expr] ...)]]
  (a/theorem ... (simp)))

;; Register a tactic that tries rfl, then assumption, then omega
(a/register-tactic! 'auto
  (fn [ps args]
    (or (try (basic/rfl ps) (catch Exception _ nil))
        (try (basic/assumption ps) (catch Exception _ nil))
        (try (omega/omega ps) (catch Exception _ nil)))))

;; Use it in proofs:
(a/theorem foo [n :- Nat] (<= Nat 0 n) (auto))

;; Register a custom syntax form that elaborates to kernel Expr
(a/register-elaborator! 'double
  (fn [env scope depth args lctx]
    (let [x (ansatz.core/sexp->ansatz env scope depth (first args) lctx)]
      ;; Build: x + x
      ...)))

;; Use in definitions:
(a/defn f [n :- Nat] Nat (double n))

;; Register a simplification procedure (always active during simp)
(a/register-simproc! 'my-domain/simplify
  (fn [st expr]
    ;; Return {:expr simplified :proof eq-proof} or nil
    nil))