(in-package :cl-quantum/circuit)

(defun circuit-matrix-operation (state operator &rest args)
  (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))))
    (destructuring-bind (target &optional control) args
      (if control
          (napply-controlled-operator state matrix target control)
          (napply-operator state matrix target)))))

(defparameter *circuit-operators*
  ;; Operator, # args, has output?, function
  ;; The output is always the last argument
  `((:* 1 nil circuit-matrix-operation)
    (:x 1 nil circuit-matrix-operation)
    (:y 1 nil circuit-matrix-operation)
    (:z 1 nil circuit-matrix-operation)
    (:h 1 nil circuit-matrix-operation)
    (:p 1 nil circuit-matrix-operation)
    (:t 1 nil circuit-matrix-operation)
    (:cnot 2 nil circuit-matrix-operation)
    (:cz 2 nil circuit-matrix-operation)
    (:measure 2 t ,(lambda (state operator &rest args)
                     (declare (ignorable operator))
                     (nmeasure state (car args))))))

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

(defun add-to-circuit (circuit operator &rest args)
  "Add OPERATOR to CIRCUIT."
  (let ((entry (assoc operator *circuit-operators*))
        (largest-arg (apply 'max (remove-if-not 'integerp args))))
    (unless entry
      (error "Unknown circuit operator: ~s" operator))
    (destructuring-bind (name arg-count &rest r) entry
      (declare (ignorable name r))
      (unless (= arg-count (length args))
        (error "Operator ~s expects ~s args, got ~s" operator
               arg-count (length args)))
      (when (> (1+ largest-arg) (car circuit))
        (setf (car circuit) (1+ largest-arg)))
      (nconc circuit (list (cons operator args))))
    circuit))

(defmacro with-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))
        (size-var (gensym))
        (ao-func (gensym))
        (ao-entry-var (gensym))
        (ao-i-var (gensym)))
    `(let ((,circuit-var)
           (,size-var 0))
       (flet ((,ao-func (,ao-entry-var)
                (dolist (,ao-i-var ,ao-entry-var)
                  (when (and (integerp ,ao-i-var)
                             (< ,size-var ,ao-i-var))
                    (setq ,size-var ,ao-i-var)))
                (push ,ao-entry-var ,circuit-var)))
        (macrolet
            (,@(mapcar (lambda (oper)
                         (let ((arg (gensym)))
                           `(,(car oper) (&rest ,arg)
                             (assert (= (length ,arg) ,(second oper))
                                     ()
                                     "~s expects ~s arguments, got ~s"
                                     ,(car oper) ,(second oper) (length ,arg))
                             `(,',ao-func (list ,',(car oper) ,@,arg)))))
                 *circuit-operators*))
          ,@body))
       (cons (1+ ,size-var) (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-output function)
      (assoc operator *circuit-operators*)
    (declare (ignorable name arg-count))
    (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 coefficients probabilities)
  "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)))
          ()
          "Exactly one of BITS, COEFFICIENTS, and PROBABILITIES can be present")
  (let ((state (cond
                 (bits (make-uniform-normal-state bits))
                 (coefficients (coerce coefficients 'vector))
                 (probabilities (make-normal-state probabilities)))))
    (destructuring-bind (circuit-size &rest elements) circuit
      (when (> circuit-size (state-bits state))
        (error "Circuit needs at least ~s bits, got ~s" circuit-size
               (state-bits state)))
      (values
       state
       (loop for element in elements
             for name = (car element)
             for args = (cdr element)
             for result = (apply-circuit-operator-to-state state name args)
             when result
               collect result)))))

(let ((circuit
        (with-circuit
          (:h 0)
          (:cnot 0 1)
          (:measure 0 :v1))))
  (run-circuit circuit :bits 2))