(in-package :cl-quantum)

(defun pprint-complex (n &key parens (places 5))
  "Pretty-print the complex (or real, rational, etc.) number N. If PARENS is
non-nil, surround the output with parenthesis if it is multiple terms. PLACES is
the number of places to round each part before printing."
  (let* ((real (realpart n))
         (imag (imagpart n))
         ;; Also put parenthesis on fractions to make them easier to read
         (has-frac (and parens
                        (or (and (not (integerp real))
                                 (rationalp real))
                            (and (not (integerp imag))
                                 (rationalp imag))))))
    (when (and (not (rationalp real))
               (realp real))
      (setq real (round-to-place real places)))
    (when (and (not (rationalp imag))
               (realp imag))
      (setq imag (round-to-place imag places)))
    (cond
      ((and (zerop real)
            (zerop imag))
       "0")
      ((not (or (zerop real)
                (zerop imag)))
       (format nil "~@[~*(~]~a ~:[-~;+~] ~ai~@[~*)~]"
               parens real (>= imag 0) (abs imag) parens))
      (t
       (format nil "~@[~*(~]~a~@[~*i~]~@[~*)~]"
               has-frac (if (zerop real) imag real)
               (zerop real) has-frac)))))

(defun pprint-format-bits (index size)
  "A state formatter that converts the index to binary and pads it with zeros."
  (format nil "~v,,,'0<~b~>" (ceiling (log size 2)) index))

(defun pprint-format-linear (index size)
  "A state formatter that just returns the index +1 as a string."
  (declare (ignorable size))
  (format nil "~d" (1+ index)))

(defun pprint-state (state &key (formatter 'pprint-format-linear)
                             (places 5))
  "Pretty-print STATE, a quantum state represented as an array. FORMATTER is a
function which takes the index of the quantum state and the total size of the
state. It should convert these to a printable representation. This
representation will be put inside of a ket after each coefficient. PLACES is the
number of places to which to print the coefficients."
  (with-output-to-string (out)
    (loop with need-sign = nil
          for i below (length state)
          for coef = (aref state i)
          when (and need-sign (not (zerop coef)))
            if (>= (realpart coef) 0)
              do (format out " + ")
          else
            do (format out " - ")
            and do (setq coef (* -1 coef))
          end
          end
          unless (zerop coef)
            do (format out "~a|~a>" (pprint-complex coef :parens t
                                                         :places places)
                       (funcall formatter i (length state)))
            and do (setq need-sign t))))