(define-module (d-)
  #:use-module (srfi srfi-1)
  #:use-module (language tree-il)
  #:export
  (~>
   ~>>
   as~>
   if-not
   if-let
   for
   partial
   argmin
   iterate
   upply
   conjoin
   distinct?
   generator
   macro-expand
   amb
   amb-reset
   amb-require))

(define-syntax ~>
  (syntax-rules ()
    [(~> v) v]
    [(~> v (fn args ...) more ...) (~> (fn v args ...) more ...)]
    [(~> v fn more ...) (~> (fn v) more ...)]))

(define-syntax ~>>
  (syntax-rules ()
    [(~>> v) v]
    [(~>> v (fn args ...) more ...) (~>> (fn args ... v) more ...)]
    [(~>> v fn more ...) (~>> (fn v) more ...)]))

(define-syntax as~>
  (syntax-rules ()
    [(_ as v) v]
    [(_ as v (fn args ...) more ...)
     (let [(as v)]
       (as~> as (fn args ...) more ...))]))

(define-syntax-rule (if-not pred body ...)
  (if (not pred) body ...))

(define-syntax-rule (if-let ([ident test]) expr ...)
  (let ([ident test])
    (if ident
	(begin expr ...)
	#f)))

(define-syntax for
  (syntax-rules ()
    [(for () expr ...) (list (begin expr ...))]
    [(for ([ident lst] bindings ...) expr ...)
     (let iter ([rest lst])
       (if (pair? rest)
	   (let ([ident (car rest)])
	     (append (for (bindings ...) expr ...) (iter (cdr rest))))
	   '()))]))

(define (partial fn . args)
  (lambda x (apply fn (append args x))))

(define (argmin arg lt? . vals)
  (reduce (lambda (val min)
	    (if (lt? (arg val) (arg min))
		val
		min))
	  #f
	  vals))

(define (iterate n f v)
  "Repeatedly call f on values returned from (f v)"
  (if (zero? n)
      v
      (iterate (1- n) f (f v))))

(define (upply a b cmp . fs)
  "U-shapped apply, apply fs to a and b as in compose, then apply cmp to both results"
  (let ([arm-f (apply compose fs)])
    (cmp (arm-f a) (arm-f b))))

(define (conjoin . preds)
  "Returns a procedure that applies each pred to a single value, and returns #t if all return a truthy value. With no preds returns a one arg function that always returns true"
  (if (null? preds)
      (lambda (x) #t)
      (lambda (x)
	(if ((car preds) x)
	    ((apply conjoin (cdr preds)) x)
	    #f))))

(define (distinct? . a)
  "Are all values distinct, as in equal?"
  (if (or (null? a) (null? (cdr a)))
      #t
      (and (not (any (lambda (x) (equal? (car a) x)) (cdr a)))
	   (apply distinct? (cdr a)))))

;; Shamelessly ripped from https://wingolog.org/archives/2013/02/25/on-generators
(define (make-generator f)
  (define tag (make-prompt-tag))
  (define (thunk)
    (f (lambda (val) (abort-to-prompt tag val))))
  (lambda ()
    (call-with-prompt tag
      thunk
      (lambda (k value)
	(set! thunk k)
	value))))

(define-syntax generator
  ;; generator with an anaphoric yield
  (lambda (x)
    (syntax-case x ()
      [(generator expr ...)
       (with-syntax ([yield (datum->syntax x 'yield)])
	 #'(make-generator
	    (lambda (yield)
	      expr ...)))])))

;; Why wasn't this included?
(define macro-expand (compose tree-il->scheme macroexpand))

;; This is the "classic" implementation of McCarthy's `amb'
;; Would rather it use delimited continuations, but I'm too dumb
(define (*fail*) (error "Could not satisfy amb"))
(define-syntax amb
  (syntax-rules ()
    [(amb) (*fail*)]
    [(amb a) a]
    [(amb a b ...)
     (let ([fail0 *fail*])
       (call/cc
	(lambda (cc)
	  (set! *fail*
		(lambda ()
		  (set! *fail* fail0)
		  (cc (amb b ...))))
	  (cc a))))]))
(define (amb-reset)
  ;; Shouldn't be necessary, but prompts are hard and I'm dumb dumb dummy
  (set! *fail* (lambda () (error "Could not satisfy amb"))))
(define (amb-require pred)
  (or pred (amb)))