;;;; 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 (state operator gate target &optional control)
  (declare (ignorable operator))
  (if control
      (napply-controlled-gate state gate target control)
      (napply-gate state gate target)))

(defun circuit-matrix-operation (state 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
        (napply-controlled-gate state matrix target control)
        (napply-gate state matrix target))))

(defun circuit-n-controlled-gate (state operator target controls)
  (let ((func (case operator
                (:ntoff 'make-n-toffoli-operator)
                (:ncz 'make-n-controlled-z-operator))))
    (replace-state state (*mv (funcall func (state-bits state)
                                       target controls)
                              state))))

(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)
    (:measure 2 nil t ,(lambda (state operator &rest args)
                         (declare (ignorable operator))
                         (nmeasure state (car args)
                                   :places *circuit-measure-places*)))
    (:ntoff 2 nil nil circuit-n-controlled-gate)
    (:ncz 2 nil nil circuit-n-controlled-gate)))

(defun make-circuit ()
  "Create a new blank circuit."
  '(:circuit))

(defun add-to-circuit (circuit operator &rest args)
  "Add OPERATOR to CIRCUIT."
  (let ((entry (assoc operator *circuit-operators*)))
    (unless entry
      (error "Unknown circuit operator: ~s" operator))
    (destructuring-bind (name arg-count has-rest &rest r) entry
      (declare (ignorable name r))
      (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)))
      (nconc circuit (list (cons operator args))))
    circuit))

(defmacro with-build-circuit (&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)))
    `(let ((,circuit-var))
       (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))))
                            `(push (list ,@,arg) ,',circuit-var))))
                *circuit-operators*))
         ,@body)
       (cons :circuit (nreverse ,circuit-var)))))

(defun apply-circuit-operator-to-state (state operator args)
  "Apply the circuit operator OPERATOR to STATE by calling its function with
ARGS."
  (destructuring-bind (&optional name arg-count has-rest has-output function)
      (assoc operator *circuit-operators*)
    (declare (ignorable name arg-count has-rest))
    (assert function ()
            "Invalid circuit operator: ~s" operator)
    (let ((output (apply function state operator args)))
      (when has-output
        (cons (car (last args)) output)))))

(defun run-circuit (circuit &key bits uniform coefficients probabilities (places 5))
  "Run the circuit CIRCUIT and return the final state. The initial state can be
specified in one of three ways:
 - BITS: the number of qbits
 - COEFFICIENTS: the initial coefficients
 - PROBABILITES: the initial probabilities"
  (assert (= 1 (count-if 'identity (list bits coefficients probabilities uniform)))
          ()
          "Exactly one of BITS, UNIFORM, COEFFICIENTS, and PROBABILITIES can ~
be present")
  (let ((state (cond
                 (bits (make-zero-state bits))
                 (uniform (make-uniform-normal-state uniform))
                 (coefficients (coerce coefficients 'vector))
                 (probabilities (make-normal-state probabilities)))))
    (values
     state
     (loop with *circuit-measure-places* = places
           for element in (cdr circuit)
           for name = (car element)
           for args = (cdr element)
           for result = (apply-circuit-operator-to-state state name args)
           when result
             collect result))))