(use fugue)

(defproto Stack () data {:default @[]})

(defproto Element ()
  doc {:default ""})

(defproto Number Element
  out-arity {:allocate-value 1})

(defproto Int Number
  type {:allocate-value :integer}
  value {:init? true})

(defproto Float Number
  type {:allocate-value :float}
  value {:init? true})

(defproto Operation Element
  type {}
  arity {}
  out-arity {}
  fun-ref {})

(defproto Adverb Element
  type {}
  arity {}
  out-arity {})

(defproto Quotation Element
  out-arity {:allocate-value 1}
  data {:init? true}
  type {:allocate-value :quotation})

(defproto Vector Quotation
  shape {:default [] :init? true}
  data {:init? true}
  type {:allocate-value :vector})

(defmethod get-shape Number [self] [])
(defmethod get-shape Vector [self] (self :shape))
(defmethod get-shape Quotation [_]
  (error "Shape error: attempted vector operation on a quotation"))

(defgeneric get-arity [_] 0)
(defmethod get-arity Adverb [a] (a :arity))
(defmethod get-arity Operation [a] (a :arity))
(defmethod get-arity Quotation [q]
  (def arities @[])
  (loop [i :range [0 (length (q :data))]]
    (if (zero? i)
      (array/push arities (get-arity ((q :data) i)))
      (array/push arities
                  (+ (last arities)
                     (- (get-arity ((q :data) i))
                        (((q :data) (dec i)) :out-arity))))))
  (max ;arities))

(defmethod fill Number
  [element shape-to-fill]
  (if (empty? shape-to-fill)
    element
    (let [y (array ;shape-to-fill)
          yi (last y)
          v (:new Vector [yi] (array/new-filled yi element))]
      (fill v y))))

(defmethod fill Vector
  [self shape-to-fill]
  (let [x (array ;(shape self))
        y (array ;shape-to-fill)]

    (while (not (empty? x))
      (let [xi (array/pop x)
            yi (array/pop y)]
        (unless (= xi yi)
          (errorf "Shape error: can't fill vector with shape %q to %q"
                  (shape self) shape-to-fill))))

    (reduce (fn [acc length]
              (let [new-shape (tuple length ;(shape acc))
                    new-data (array/new-filled length acc)]
                (:new Vector new-shape new-data)))
            self
            (reverse y))))

# TODO: Do we want to do something more sophisticated here?
(defmulti make-element [:number] [n] (:new Float n
                                           :doc (string n)))
(defmulti make-element [_] [v] v)

(defn- inner-apply
  [op & args]
  (let [x (first args)]
    (if (Vector? x)
      (:new Vector
            (x :shape)
            (map (partial inner-apply op) ;(map |($ :data) args)))
      (let [f (op :fun-ref)
            unwrapped (map |($ :value) args)]
        (->> (f ;unwrapped)
             (make-element))))))

(defmethod apply-operation Operation
  [op & args]
  (let [cmpr (comp length get-shape)
        largest (get-shape (extreme |(> (cmpr $0) (cmpr $1)) args))
        filled (map |(fill $ largest) args)]
    (inner-apply op ;filled)))

(defn quote-wrap
  [& data]
  (:new Quotation data))

(defn wrap
  [& data]
  (let [hd (first data)
        inner-shape (if hd (get-shape hd) ())]
    (each datum data
      (unless (deep= (get-shape datum) inner-shape)
        (errorf "Vector error: can't wrap heterogeneous data. Expected %j, got %j."
                inner-shape
                (get-shape datum))))
    (:new Vector (tuple (length data) ;inner-shape) data)))

(defmethod pop Stack
  [self]
  (-> (self :data) (array/pop)))

(defmethod peek Stack
  [self]
  (-> (self :data) (last)))

(defmethod size Stack
  [self]
  (-> (self :data) (length)))

(defn check-arity
  [stack elem]
  (let [arity (get-arity elem)]
    (when (> arity (size stack))
      (errorf "Not enough stack: %s has arity %i; stack has size %i"
              (string (elem :type))
              arity
              (size stack)))))

(defmulti push [Stack Operation]
  [self op]
  (check-arity self op)

  (def buf @[])
  (loop [_ :range [0 (get-arity op)]]
    (array/insert buf 0 (pop self)))

  (->> (apply-operation op ;buf)
       (make-element)
       (push self)))

(defmulti push [Stack Adverb]
  [self adv]
  (check-arity self adv)
  ((adv :fun-ref) self))

(defmulti push [Stack _]
  [self item]
  (array/push (self :data) item))

(defn shape?
  [obj]
  (and (Vector? obj) (all Number*? (obj :data))))