(in-package :cl-quantum/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 nnormalize-state (state)
  "Normalize STATE by mutating it."
  (let ((norm (norm state)))
    (dotimes (i (length state) state)
      (setf (aref state i) (/ (aref state i) norm)))))

(defun state-bits (state)
  "Return the number of bits in STATE."
  (values (ceiling (log (length state) 2))))

(defun make-normal-state (probabilities)
  "Create a new normalized state with the probability of each state
corresponding to an element in PROBABILITIES."
  (map 'vector 'sqrt probabilities))

(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."
  (+ period (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 &key (places 5))
  "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. The
probability will be rounded to PLACES before calculations are carried out."
  (loop with prob = (round-to-place (bit-probability state bit) places)
        ;; Theocratically there are multiple numbers that could be equivalent
        ;; to 100% here, but the chance of those failing is so low anyway that
        ;; it doesn't matter.
        with limit = (* most-positive-fixnum prob)
        with rnum = (random most-positive-fixnum)
        with result = (< rnum limit)
        with period = (ash 1 bit)
        with 1-prob = (- 1 prob)
        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)
        if result
          do (setf (aref state unset-index) 0
                   (aref state set-index)
                   (sqrt (/ (* set-coef set-coef) prob)))
        else
          do (setf (aref state unset-index)
                   (sqrt (/ (* unset-coef unset-coef) 1-prob))
                   (aref state set-index) 0)
        finally (return (values result state))))

(defun measure (state bit &key (places 5))
  "Like `nmeasure', but don't modify STATE. Note that the new state is returned
as the second value, with the first value being the result of the
measurement. You will probably want to save the second value as the first value
too as the second value alone means pretty much nothing."
  (nmeasure (copy-seq state) bit :places places))

(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 control)
  "Create an operator matrix that can act on a state with BITS bits and will
apply OPERATOR to TARGET if CONTROL is set."
  (labels ((matrix-for (bit target-operator control-operator)
             (cond
               ((= bit target) target-operator)
               ((= bit control) 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))))

(defun replace-state (target template)
  "Replace each element of TARGET with the corresponding element in TEMPLATE."
  (assert (= (length target) (length template))
          (target template)
          "Vectors must be of the same length: ~s and ~s" target template)
  (dotimes (i (length target) target)
    (setf (aref target i) (aref template i))))

(defun apply-operator (state operator target)
  "Apply OPERATOR to the bit numbered TARGET in STATE."
  (*mv (make-operator (state-bits state) operator target) state))

(defun napply-operator (state operator target)
  "Apply OPERATOR to the bit numbered TARGET in STATE. This modifies state."
  (replace-state state (apply-operator state operator target)))

(defun apply-controlled-operator (state operator target control)
  "Apply OPERATOR to the bit numbered TARGET in STATE if CONTROL is set."
  (*mv (make-controlled-operator (state-bits state) operator target control)
       state))

(defun napply-controlled-operator (state operator target control)
  "Apply OPERATOR to the bit numbered TARGET in STATE if CONTROL is set. This
modified STATE."
  (replace-state state (apply-controlled-operator state operator target control)))

;;; 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 pi/8-gate
  (make-array '(2 2) :initial-contents
              `((1 0)
                (0 ,(exp (complex 0 (/ pi 4)))))))

(defconstant cnot-gate
  (make-controlled-operator 2 pauli-x-gate 0 1))

(defconstant cz-gate
  (make-controlled-operator 2 pauli-z-gate 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 8) 6 7))