;;;; High-level interface for building quantum circuits
(in-package :cl-quantum/circuit)

(defparameter *circuit-measure-places* 5
  "The number of places to use with the :measure instruction in circuits.")

(defun circuit-arbitrary-gate (size operator gate target &optional control)
  (declare (ignorable operator))
  (if control
      (make-controlled-operator size gate target control)
      (make-operator size gate target)))

(defun circuit-matrix-operation (size operator target &optional control)
  (let ((matrix (case operator
                  (:* +identity-2x2+)
                  (:x +pauli-x-gate+)
                  (:y +pauli-y-gate+)
                  (:z +pauli-z-gate+)
                  (:h +hadamard-gate+)
                  (:p +phase-gate+)
                  (:t +pi/8-gate+)
                  (:cnot +pauli-x-gate+)
                  (:cz +pauli-z-gate+))))
    (if control
        (make-controlled-operator size matrix target control)
        (make-operator size matrix target))))

(defun circuit-n-controlled-gate (size operator target controls)
  (case operator
    (:ntoff (make-n-toffoli-operator size target controls))
    (:ncz (make-n-controlled-z-operator size target controls))))

(defparameter *circuit-operators*
  ;; Operator, # args, more args?, has output?, function
  ;; The output is always the last argument
  `((:gate 2 nil nil circuit-arbitrary-gate)
    (:cgate 3 nil nil circuit-arbitrary-gate)
    (:* 1 nil nil circuit-matrix-operation)
    (:x 1 nil nil circuit-matrix-operation)
    (:y 1 nil nil circuit-matrix-operation)
    (:z 1 nil nil circuit-matrix-operation)
    (:h 1 nil nil circuit-matrix-operation)
    (:p 1 nil nil circuit-matrix-operation)
    (:t 1 nil nil circuit-matrix-operation)
    (:cnot 2 nil nil circuit-matrix-operation)
    (:cz 2 nil nil circuit-matrix-operation)
    (:ntoff 2 nil nil circuit-n-controlled-gate)
    (:ncz 2 nil nil circuit-n-controlled-gate)
    (:measure 2 nil t ,(lambda (size operator bit target)
                         (declare (ignorable size operator target))
                         (lambda (state)
                           (nmeasure state bit
                                     :places *circuit-measure-places*))))))

(defun make-circuit (bits)
  "Create a new blank circuit of BITS bits."
  `(:circuit ,bits))

(defun circuitp (obj)
  "Return non-nil if OBJ is a quantum circuit."
  (and (eq (first obj) :circuit)
       (integerp (second obj))))

(defun add-to-circuit (circuit operator &rest args)
  "Add OPERATOR to CIRCUIT."
  (let ((bits (second circuit))
        (entry (assoc operator *circuit-operators*)))
    (unless entry
      (error "Unknown circuit operator: ~s" operator))
    (destructuring-bind (name arg-count has-rest has-output func) entry
      (declare (ignorable name))
      (unless (or (and has-rest (>= (length args) arg-count))
                  (= (length args) arg-count))
        (error "Operator ~s expects ~@[~*exactly ~]~s arg~:p, got ~s" operator
               (not has-rest) arg-count (length args)))
      (let ((action (apply func bits operator args))
            (prev-action (car (last circuit))))
        (cond
          ((and (matrixp action) (matrixp prev-action))
           (setf (car (last circuit)) (*mm action prev-action)))
          ((functionp action)
           (setf (cdr (last circuit)) (list (cons action
                                                  (and has-output
                                                       (car (last args)))))))
          (t (setf (cdr (last circuit)) (list action))))))
    circuit))

(defun build-circuit-add-action (bits actions args)
  "Add the action denoted by ARGS to the front of ACTIONS. The action is assumed
to operate on a state of BITS bits."
  (let ((entry (assoc (car args) *circuit-operators*)))
    (destructuring-bind (name arg-count has-rest has-output func) entry
      (declare (ignorable arg-count has-rest))
      (declare (ignorable name))
      (let ((action (apply func bits (car args) (cdr args)))
            (prev-action (car actions)))
        (cond
          ((and (matrixp action) (matrixp prev-action))
           (setf (car actions) (*mm action prev-action)))
          ((functionp action)
           (push (cons action (and has-output (car (last args))))
                 actions))
          (t (push action actions)))
        actions))))

(defmacro with-build-circuit (bits &body body)
  "Create a circuit using a simple DSL. BODY can be any valid Lisp forms, in
addition to function calls to functions named in `*circuit-operators*'."
  (let ((circuit-var (gensym))
        (bits-var (gensym)))
    `(let ((,circuit-var)
           (,bits-var ,bits))
       (macrolet
           (,@(mapcar (lambda (oper)
                        (let ((arg (gensym))
                              (arg-list (loop repeat (second oper)
                                              collect (gensym)))
                              (whole-arg (when (third oper)
                                           (gensym))))
                          `(,(car oper) (&whole ,arg
                                                ,@arg-list
                                                ,@(when (third oper)
                                                    (list '&rest whole-arg)))
                            (declare (ignorable ,@arg-list
                                                ,@(when (third oper)
                                                    (list whole-arg))))
                            `(setq ,',circuit-var
                                   (build-circuit-add-action
                                    ,',bits-var ,',circuit-var (list ,@,arg))))))
                *circuit-operators*))
         ,@body)
       (apply 'list :circuit ,bits-var (nreverse ,circuit-var)))))

(defun run-circuit (circuit &key (state :zero) (places 5))
  "Run the circuit CIRCUIT and return the final state. The initial STATE can be
specified in one of four ways:
 - :ZERO: a zero state is used -> #(1 0 0 0), etc.
 - :UNIFORM: a uniform state -> #(1/sqrt(2) 1/sqrt(2)), etc
 - a vector: the initial coefficients"
  (let ((real-state (cond
                 ((eq state :zero)
                  (make-zero-state (second circuit)))
                 ((eq state :uniform)
                  (make-uniform-normal-state (second circuit)))
                 (t (coerce state 'vector)))))
    (values
     real-state
     (loop with *circuit-measure-places* = places
           with outputs = ()
           for action in (cddr circuit) do
             (cond
               ((matrixp action)
                (replace-state real-state (*mv action real-state)))
               ((consp action)
                (destructuring-bind (func output) action
                  (let ((res (funcall func real-state)))
                    (when output
                      (push res outputs)
                      (push output outputs)))))
               (t (error "Unknown circuit element type: ~s" action)))))))