;; -*- Mode: Irken -*-

;; The '10958 Problem', from numberphile.
;; https://www.youtube.com/watch?v=-ruC5A9EzzE

(include "lib/basis.scm")
(include "lib/map.scm")

;; (product '(0 1) '(#\a #\b)) => ((0 #\a) (0 #\b) (1 #\a) (1 #\b))
(define (product al bl)
  (make-generator
   (lambda (consumer)
     (for-list a al
       (for-list b bl
         (consumer (maybe:yes (:tuple a b)))))
     (forever (consumer (maybe:no))))))

(define (listify xl)
  (map (lambda (x) (list:cons x (list:nil))) xl))

;; (productn '(0 1 2) 2) => ((0 0 0) (0 0 1) ... (2 2 2))
(define (productn xl n)
  (make-generator
   (lambda (consumer)
     (let recur ((acc (listify xl))
                 (n n))
       (if (= n 0)
           (for-list item acc
             (consumer (maybe:yes (reverse item))))
           (for-list b acc
             (recur (map (lambda (x) (list:cons x b)) xl) (- n 1)))
           ))
     (forever (consumer (maybe:no))))))

(define (generator->list gen)
  (let ((r '()))
    (for item gen
      (PUSH r item))
    (reverse r)))

;; https://en.wikipedia.org/wiki/Catalan_number
;; https://en.wikipedia.org/wiki/Dyck_language

(datatype group
  (:int int)
  (:pair group group)
  )

(define format-group
  (group:int n)    -> (format (int n))
  (group:pair l r) -> (format "(" (format-group l) " " (format-group r) ")")
  )

;; this computes all possible balanced n-parenthesis groupings.
;; http://stackoverflow.com/a/41310973
(define (dyck num)
  (let ((map (tree/make int-cmp (0 (LIST (group:int 0))))))
    (match (tree/member map int-cmp num) with
      (maybe:no)
      -> (let ((r '()))
           (for-range i num
             (let ((i1 (+ i 1)))
               (for (a b) (product (dyck (- i1 1)) (dyck (- num i1)))
                 (PUSH r (group:pair a b)))))
           (tree/insert! map int-cmp num r)
           r)
      (maybe:yes v)
      -> v
      )))

;; ((0)(0(0))) => ((1)(2(3)))
(define (renumber-dyck t)
  (let ((cell {val=1}))
    (let recur ((t t))
      (match t with
        (group:int n)
        -> (let ((val cell.val))
             (set! cell.val (+ val 1))
             (group:int val))
        (group:pair l r)
        -> (group:pair (recur l) (recur r))
        ))))

(datatype op
  (:add)
  (:sub)
  (:mul)
  (:div)
  (:cat)
  )

(define op-repr
  (op:add) -> "+"
  (op:sub) -> "-"
  (op:mul) -> "*"
  (op:div) -> "/"
  (op:cat) -> "."
  )

;; we restrict to integer divide.
(define (divop l r)
  (if (= r 0)
      (raise (:BadOp))
      (let ((quo (/ l r))
            (rem (remainder l r)))
        (if (= rem 0)
            quo
            (raise (:BadOp))))))

(define apply-op
  (op:add) l r -> (+ l r)
  (op:sub) l r -> (- l r)
  (op:mul) l r -> (* l r)
  (op:div) l r -> (divop l r)
  (op:cat) l r -> (raise (:BadOp)) ;; handled as a special case
  )

;; binary expressions
(datatype exp
  (:op op exp exp)
  (:int int)
  )

(define format-exp
  (exp:int n)           -> (format (int n))
  (exp:op (op:cat) l r) -> (format (format-exp l) (format-exp r))
  (exp:op op l r)       -> (format "(" (op-repr op) " " (format-exp l) " " (format-exp r) ")")
  )

(define infix-exp
  (exp:int n)           -> (format (int n))
  (exp:op (op:cat) l r) -> (format (infix-exp l) (infix-exp r))
  (exp:op op l r)       -> (format "(" (infix-exp l) " " (op-repr op) " " (infix-exp r) ")")
  )

;; only valid with a tree of cats.
;; Note: this restriction is the 'obvious' one, and is the restriction violated by
;;  Matt Parker's "solution".
(define (cat t)
  (match t with
    (exp:int n)           -> (LIST n)
    (exp:op (op:cat) l r) -> (append (cat l) (cat r))
    _                     -> (raise (:BadOp))
    ))

;; (1 2 3) => 123
(define render-cat
  acc ()               -> acc
  acc (digit . digits) -> (render-cat (+ (* 10 acc) digit) digits)
  )

(define (eval-exp exp)
  (match exp with
    (exp:int n)           -> n
    (exp:op (op:cat) l r) -> (render-cat 0 (cat exp))
    (exp:op op l r)       -> (apply-op op (eval-exp l) (eval-exp r))
    ))

(define ops-list (LIST (op:cat) (op:add) (op:sub) (op:mul) (op:div)))

;; given a tree and a list of ops,
;;  create an expression tree with the ops
;;  placed according to an inorder traversal.
(define (infill-ops tree ops)
  (let recur ((t tree))
    (match t with
      (group:int n)
      -> (exp:int n)
      (group:pair l r)
      -> (let ((op (pop ops))
               (lv (recur l))
               (rv (recur r)))
           (exp:op op lv rv)))))

;; we compute the set of all paren-groupings, then
;;  the set of all possible operators for each position.
;;   then we use a double for loop to iterate over the product of the two.
;; `infill-ops` combines a paren-grouping with a set of operators to
;;   create an expression tree.

(define (solve n)
  (let ((trees (map renumber-dyck (dyck 8)))
        (ops8 (generator->list (productn ops-list 7))))
    (make-generator
     (lambda (consumer)
       (for-list tree trees
         (for-list ops ops8
           (try
            (let ((exp (infill-ops tree ops))
                  (val (eval-exp exp)))
              (when (= val n)
                (consumer (maybe:yes (infix-exp exp)))
                ))
            except (:BadOp)
            -> #u
            )))
       (forever (consumer (maybe:no)))
       ))))

(define target
  (if (> sys.argc 1)
      (string->int sys.argv[1])
      814))

(set-verbose-gc #f)

;; next steps:
;; 1) make a pass over each solution to remove un-needed parens.
;;    this will require encoding knowledge of precedence rules.
;; 2) build a table of 11,111 entries and fill each result with
;;    its 'smallest' solution along with a count of how many solutions
;;    were found. Lower numbers have many more solutions, I bet it
;;    ends up looking like a zipf distribution.

(let ((last-solution ""))
  (for solution (solve target)
    (if (not (string=? last-solution solution))
        (printf solution "\n")
        #u)
    (set! last-solution solution)))