(defpackage :fourfours
  (:use :common-lisp
        :iterate))
(in-package :fourfours)

(defparameter *base-operators* (list #'+ #'- #'* #'/)
  "The basic mathematical operators.")

(defun ^ (base exponent)
  "Hack of an exponent function to try to avoid the slowness of large numbers"
  (if (and (integerp exponent)
           (< exponent (log 512 base)))
      (expt base exponent)
      nil))

;;; This is for reference, mostly.  I was playing with the functions here to see
;;; if I could get the results I wanted with fewer than four operators.
(defparameter *all-operators*
  (list #'ceiling
        ;#'expt  ; sloooow
        #'floor
        #'gcd
        #'lcm
        #'ldb  ; esoteric
        #'logand
        #'logandc1
        #'logandc2
        #'logeqv
        #'logior
        #'lognand
        #'lognor
        #'logorc1
        #'logorc2
        #'logxor
        #'mask-field
        #'max
        #'min
        #'mod
        #'rem
        #'round
        #'truncate
        #'^  ; custom func
        #'*
        #'+
        #'-
        #'/)
  "All CL functions which make sense in the context of this problem.")

(defun list-trees (item-count)
  "Returns a list of all possible binary trees with ITEM-COUNT nodes."
  (if (= item-count 1)
      (list t)
      (iter outer
            (for first-count from 1 to (1- item-count))
            (iter (for first-elts in (list-trees first-count))
                  (iter (for second-elts in (list-trees (- item-count first-count)))
                        (in outer (collecting (list first-elts second-elts))))))))

(defun function-name (function)
  "Convenience function for pulling a function's name out of its function object."
  (multiple-value-bind (lambda-expression closure-p name)
      (function-lambda-expression function)
    (declare (ignore lambda-expression closure-p))
    name))

(defun walk-tree (action tree params)
  "Walks a binary tree, taking the specified action:
   :replace-vals - replace the nodes in the tree with the nodes in PARAMS.  Nodes
                   are processed inorder.
   :add-funcalls - PARAMS is a list of functions.  'funcall <function>' is added at
                   each inner node.  The result is not a binary tree.
   :format-infix - PARAMS is a list of functions.  The return value is a string
                   representing the binary tree as the infix expansion equivalent
                   to the result of :add-funcalls.

   The second value returned is a list of the elements of PARAMS that were not
   used in whatever expansion was done."
  (if (atom tree)
      (if (eql action :replace-vals)
          (values (car params) (cdr params))
          (values tree params))
      (let ((remaining-params params))
        (labels ((handle-node (node)
                   (multiple-value-bind (new-tree new-params)
                       (walk-tree action node remaining-params)
                     (setf remaining-params new-params)
                     new-tree)))
          (values (destructuring-bind (left right) tree
                    (ecase action
                      (:add-funcalls
                       (list 'funcall (pop remaining-params)
                             (handle-node left) (handle-node right)))
                      (:add-ops
                       (list (function-name (pop remaining-params))
                             (handle-node left) (handle-node right)))
                      (:format-infix
                       (let ((op (function-name (pop remaining-params))))
                         (format nil "(~A ~A ~A)"
                                 (handle-node left) op (handle-node right))))
                      (:replace-vals
                       (list (handle-node left) (handle-node right)))))
                  remaining-params)))))

(defun count-elements (tr)
  "Returns the number of leaf nodes in the binary tree TR."
  (iter (for e in tr)
        (if (atom e)
            (summing 1 into val-count)
            (multiple-value-bind (o-c v-c)
                (count-elements e)
              (summing o-c into op-count)
              (summing v-c into val-count)))
        (finally (return (values (1+ op-count) val-count)))))

(defun expand-tree (tree)
  "Turns the binary tree TREE into a function that takes a list of functions
   and a list of parameters, applies the functions to the parameters in the
   order indicated by the structure of TREE, and returns the result."
  (declare (values (function (list list))))
  (multiple-value-bind (op-count val-count)
      (count-elements tree)
    (let ((op-vars (iter (repeat op-count)
                         (collecting (gensym))))
          (val-vars (iter (repeat val-count)
                          (collecting (gensym)))))
      (compile
       nil
       `(lambda (ops vals)
          (declare (type list ops vals))
          (destructuring-bind ,op-vars ops
            (declare (type (function (rational rational) rational) ,@op-vars))
            (destructuring-bind ,val-vars vals
              (declare (type rational ,@val-vars))
              ,(walk-tree :add-funcalls (walk-tree :replace-vals tree val-vars)
                          op-vars))))))))

(defun visit-combinations (source-set set-size visit-func)
  "Calls VISIT-FUNC for each combination of SET-SIZE elements from SOURCE-SET,
   including conbinations with duplicate elements."
  (labels ((visit-combinations-r (cur-set elts-remaining)
             (if (zerop elts-remaining)
                 (funcall visit-func cur-set)
                 (iter (for e in source-set)
                       (visit-combinations-r (cons e cur-set)
                                             (1- elts-remaining))))))
    (visit-combinations-r nil set-size))
  (values))

(defun val-weight (val-source vals)
  "Returns an integer giving the relative weight of the elements of VALS, where
   elements that appear earlier in VAL-SOURCE are given a lower weight."
  (iter (for v in vals)
        (summing (expt 2 (position v val-source)))))

(defun calculate (ops vals val-count)
  "The meat of the algorithm.  Tries every combination of VAL-COUNT VALS (and
   (VAL-COUNT - 1) OPS) and stores the lowest-weighted valid results in a
   hashtable.  That hashtable is the return value.
   The hashtable's keys are the calculated integers.  Its values are lists of
   the operators, parameters, and binary-tree that resulted in the key value."
  (let* ((trees (list-trees val-count))
         (funcs (iter (for tree in trees)
                      (collecting (expand-tree tree))))
         (results (make-hash-table)))
    (labels ((update-result (r op-list val-list tree)
               (when (and (integerp r)
                          (or (null (gethash r results nil))
                              (< (val-weight vals val-list)
                                 (val-weight vals (second
                                                   (gethash r results))))))
                 (setf (gethash r results) (list op-list val-list tree))))
             (do-calculation (op-list val-list)
               (iter (for tree in trees)
                     (for func in funcs)
                     (for r = (handler-case (funcall func op-list val-list)
                                (division-by-zero () nil)
                                (type-error () nil)))
                     (update-result r op-list val-list tree))))
      (visit-combinations
       ops (1- val-count)
       (lambda (op-list)
         (visit-combinations vals val-count
                             (lambda (val-list)
                               (do-calculation op-list val-list))))))
    results))

(defun list-consecutive (results &key (start 1) end)
  "Takes a hashtable with numberic keys and returns the values of sequential
   keys from START to END (or until the first key that doesn't exist)."
  (iter (for i from start)
        (for val = (gethash i results nil))
        (while (if end
                   (<= i end)
                   val))
        (collecting val)))

(defun create-display-filter (num-alist)
  (lambda (num) (cdr (assoc num num-alist))))

(defun format-result (ops vals tree display-filter)
  (let ((display-str (walk-tree :format-infix
                                (walk-tree :replace-vals tree
                                           (mapcar display-filter vals))
                                ops)))
    (subseq display-str 1 (1- (length display-str)))))

(defun print-results (results display-filter &key (start 1))
  (iter (for (ops vals tree) in results)
        (for i from start)
        (format t "~A: ~A~%" i
                (format-result ops vals tree display-filter)))
  (values))

(defun generate-val-alist (n)
  "Generate an alist of the possible values and their representations based off
   the digit N."
  (mapcar (lambda (e) (cons (second e) (format nil (first e) n)))
          (remove nil (list (list "~A" n)
                            (list "~A!" (iter (for i from n downto 1)
                                              (multiplying i)))
                            (list ".~A" (/ n 10))
                            (when (= n (expt (isqrt n) 2))
                              (list "sqrt(~A)" (isqrt n)))
                            (list ".~Abar" (/ n 9))
                            (when (= n (expt (isqrt n) 2))
                              (list "sqrt(.~Abar)" (/ (isqrt n) 3)))))))

(defun n-ns (n &key vals (ops *base-operators*))
  (let* ((val-alist (or vals (generate-val-alist n)))
         (results (list-consecutive (calculate ops (mapcar #'car val-alist) n))))
    (print-results results (create-display-filter val-alist))))