(in-package :cl-quantum)

(defun pprint-complex (n &key parens)
  "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."
  (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))))))
    (cond
      ((zerop n) "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))
  "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."
  (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)
                       (funcall formatter i (length state)))
            and do (setq need-sign t))))

(defconstant +parse-state-regexp+
  (ppcre:create-scanner
   "^\\s*([-+])?\\s*(\\()?\\s*([-+0-9ei./]*)\\s*(\\))?\\s*\\|([^>]*)>"
   :extended-mode t)
  "The regexp scanner used in `parse-state'.")

(defun parse-bits-state (state)
  "A `parse-state' parser that parses its state as a binary string."
  (parse-integer state :radix 2))

(defun parse-state (str &key (parser 'parse-integer))
  "Try to parse STR into a quantum state. PARSER should be a function of one
argument that will take the string inside each ket and return the index of the
state."
  (loop for start = 0 then (+ start (length whole))
        for (whole matches) = (multiple-value-list
                               (ppcre:scan-to-strings +parse-state-regexp+
                                                      str
                                                      :sharedp t
                                                      :start start))
        while whole
        for coef = (if (zerop (length (aref matches 2)))
                       1
                       (parse-complex (aref matches 2)))
        for index = (funcall parser (aref matches 4))
        unless (eq (not (aref matches 1))
                   (not (aref matches 3)))
          do (error "Mismatches parenthesis: ~s" whole)
        when (and (complexp coef)
                  (not (aref matches 1)))
          do (error "Coefficient without matching state: ~s" whole)
        collect (if (equal (aref matches 0) "-")
                    (* -1 coef)
                    coef)
          into coefs
        collect index into indecies
        maximizing (1+ index) into state-size
        finally
           (return
             (let ((state (make-array state-size)))
               (loop for index in indecies
                     for coef in coefs
                     do (incf (aref state index) coef))
               state))))

(defun normal-state-p (state &key (places 5))
  "Return non-nil if state is normalized. PLACES is the number of places to
round the norm of STATE before checking."
  (= (round-to-place (dot state state) places) 1))

(defun normalize-state (state)
  "Return a copy of STATE that is normalized."
  (/vs state (norm state)))

(defun make-uniform-normal-state (bits)
  "Make a uniform normalized quantum state of BITS qbits."
  (let ((size (ash 1 bits)))
    (make-array (ash 1 bits) :initial-element (/ (sqrt size)))))

(defun bit-unset-index (bit n &key (period (ash 1 bit)))
  "Return the Nth index in a state in which BIT is 0."
  (multiple-value-bind (quo rem)
      (floor n period)
    (+ (* 2 period quo) rem)))

(defun bit-set-index (bit n &key (period (ash 1 bit)))
  "Return the Nth index in a state in which BIT is 1."
  (+ (bit-unset-index bit n :period period)))

(defun bit-probability (state bit)
  "Return the probability that BIT is set in STATE."
  (setq state (normalize-state state))
  (loop with period = (ash 1 bit)
        for i below (/ (length state) 2)
        for index = (bit-set-index bit i :period period)
        for coef = (aref state index)
        summing (* coef coef)))

(defun nmeasure (state bit)
  "Collapse BIT in STATE by measuring it. This will return t or nil depending
on the state the bit collapsed to. Note that this will also modify STATE."
  (loop with prob = (round-to-place (bit-probability state bit) 5)
        with limit = (* most-positive-fixnum prob)
        with rnum = (random most-positive-fixnum)
        with result = (>= rnum limit)
        with period = (ash 1 bit)
        for i below (/ (length state) 2)
        for unset-index = (bit-unset-index bit i :period period)
        for set-index = (+ period unset-index)
        for unset-coef = (aref state unset-index)
        for set-coef = (aref state set-index)
        for new-coef = (sqrt (+ (* set-coef set-coef)
                                (* unset-coef unset-coef)))
        if result
          do (setf (aref state unset-index) 0
                   (aref state set-index) new-coef)
        else
          do (setf (aref state unset-index) new-coef
                   (aref state set-index) 0)
        finally (return (values result state))))

(defun make-operator (bits operator target)
  "Create an operator matrix that can act on a state with BITS bits and will
apply OPERATOR to TARGET."
  (loop with out = (if (= (1- bits) target)
                       operator
                       identity-2x2)
        for i from (- bits 2) downto 0
        do (setq out (tensor-mm out (if (= i target)
                                        operator
                                        identity-2x2)))
        finally (return out)))

(defun make-controlled-operator (bits operator target controls)
  "Create an operator matrix that can act on a state with BITS bits and will
apply OPERATOR to TARGET if CONTROLS are all set."
  (labels ((matrix-for (bit target-operator control-operator)
             (cond
               ((= bit target) target-operator)
               ((member bit controls :test '=) control-operator)
               (t identity-2x2)))
           (tensor-chain (target-operator control-operator)
             (loop with out = (matrix-for (1- bits) target-operator
                                          control-operator)
                   for i from (- bits 2) downto 0
                   do (setq out (tensor-mm out (matrix-for i target-operator
                                                           control-operator)))
                   finally (return out))))
    (+mm (tensor-chain identity-2x2 unset-projector)
         (tensor-chain operator set-projector))))

;;; Gates and Operators:
(defconstant unset-projector
  #2A((1 0)
      (0 0)))

(defconstant set-projector
  #2A((0 0)
      (0 1)))

(defconstant identity-2x2
  (make-identity-matrix 2))

(defconstant pauli-x-gate
  #2A((0 1)
      (1 0)))

(defconstant pauli-y-gate
  #2A((0 #C(0 -1))
      (#C(0 -1) 0)))

(defconstant pauli-z-gate
  #2A((1 0)
      (0 -1)))

(defconstant hadamard-gate
  (let ((oort (/ (sqrt 2))))
    (make-array '(2 2) :initial-contents
                `((,oort ,oort)
                  (,oort ,(- oort))))))

(defconstant phase-gate
  #2A((1 0)
      (0 #C(0 1))))

(defconstant cnot-gate
  #2A((1 0 0 0)
      (0 1 0 0)
      (0 0 0 1)
      (0 0 1 0)))

(defconstant cz-gate
  #2A((1 0 0 0)
      (0 1 0 0)
      (0 0 1 0)
      (0 0 0 -1)))

(defconstant swap-gate
  #2A((1 0 0 0)
      (0 0 1 0)
      (0 1 0 0)
      (0 0 0 1)))

(defconstant ccnot-gate
  (nswap-rows (make-identity-matrix 9) 7 8))