'''
Teague Tomesh - 2/10/2020

Implementation of an n-bit ripple-carry adder.

Based on the specification given in Cuccaro, Draper, Kutin, Moulton.
(https://arxiv.org/abs/quant-ph/0410184v1)
'''

from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister


class RCAdder:
    """
    An n-bit ripple-carry adder can be generated using an instance of the
    RCAdder class by calling the gen_circuit() method.

    This adder circuit uses 1 ancilla qubit to add together two values
        a = a_(n-1)...a_0    and   b = b_(n-1)...a_0
    and store their sum
        s = s_n...s_0
    in the registers which initially held the b value.

    The adder circuit uses 2 + binary_len(a) + binary_len(b) qubits.
    The initial carry value is stored in the qubit at index = 0.
    The binary value of a_i is stored in the qubit at index = 2*i + 2
    The binary value of b_i is stored in the qubit at index = 2*i + 1
    The high bit, s_n, is stored in the last qubit at index = num_qubits - 1

    Attributes
    ----------
    nbits : int
        size, in bits, of the numbers the adder can handle
    nq : int
        number of qubits needed to construct the adder circuit
    a, b : int
        optional parameters to specify the numbers the adder should add.
        Will throw an exception if the length of the bitstring representations
        of a or b are greater than nbits.
    use_toffoli : bool
        Should the toffoli gate be used in the generated circuit or should it
        first be decomposed
    barriers : bool
        should barriers be included in the generated circuit
    regname : str
        optional string to name the quantum and classical registers. This
        allows for the easy concatenation of multiple QuantumCircuits.
    qr : QuantumRegister
        Qiskit QuantumRegister holding all of the quantum bits
    circ : QuantumCircuit
        Qiskit QuantumCircuit that represents the uccsd circuit
    """

    def __init__(self, nbits=None, a=0, b=0, use_toffoli=False, barriers=False,
                 measure=False, regname=None):

        # number of bits the adder can handle
        if nbits is None:
            raise Exception('Number of bits must be specified')
        else:
            self.nbits = nbits

        # given nbits, compute the number of qubits the adder will need
        # number of qubits = 1 ancilla for the initial carry value
        #                    + 2*nbits to hold both a and b
        #                    + 1 more qubit to hold the high bit, s_n
        self.nq = 1 + 2*nbits + 1

        # set flags for circuit generation
        if len('{0:b}'.format(a)) > nbits or len('{0:b}'.format(b)) > nbits:
            raise Exception('Binary representations of a and b must be less than or equal to nbits')

        self.a = a
        self.b = b
        self.use_toffoli = use_toffoli
        self.barriers = barriers
        self.measure = measure

        # create a QuantumCircuit object
        if regname is None:
            self.qr = QuantumRegister(self.nq)
        else:
            self.qr = QuantumRegister(self.nq, name=regname)
        self.circ = QuantumCircuit(self.qr)

        # add ClassicalRegister if measure is True
        if self.measure:
            self.cr = ClassicalRegister(self.nq)
            self.circ.add_register(self.cr)


    def _initialize_value(self, indices, value):
        """
        Initialize the qubits at indices to the given value

        Parameters
        ----------
        indices : List[int]
            List of qubit indices
        value : int
            The desired initial value
        """
        binstr = '{0:b}'.format(value)
        for index, val in enumerate(reversed(binstr)):
            if val is '1':
                self.circ.x(indices[index])


    def _toffoli(self, x, y, z):
        """
        Implement the toffoli gate using 1 and 2 qubit gates
        """
        self.circ.h(z)
        self.circ.cx(y,z)
        self.circ.tdg(z)
        self.circ.cx(x,z)
        self.circ.t(z)
        self.circ.cx(y,z)
        self.circ.t(y)
        self.circ.tdg(z)
        self.circ.cx(x,z)
        self.circ.cx(x,y)
        self.circ.t(z)
        self.circ.h(z)
        self.circ.t(x)
        self.circ.tdg(y)
        self.circ.cx(x,y)


    def _MAJ(self, x, y, z):
        """
        Implement the MAJ (Majority) gate described in Cuccaro, Draper, Kutin,
        Moulton.
        """
        self.circ.cx(z,y)
        self.circ.cx(z,x)

        if self.use_toffoli:
            self.circ.ccx(x,y,z)
        else:
            # use a decomposed version of toffoli
            self._toffoli(x,y,z)


    def _UMA(self, x, y, z):
        """
        Implement the UMA (UnMajority and Add) gate described in Cuccaro,
        Draper, Kutin, Moulton.
        """
        self.circ.x(y)
        self.circ.cx(x,y)
        if self.use_toffoli:
            self.circ.ccx(x,y,z)
        else:
            # use a decomposed version of toffoli
            self._toffoli(x,y,z)
        self.circ.x(y)
        self.circ.cx(z,x)
        self.circ.cx(z,y)


    def gen_circuit(self):
        """
        Create a circuit implementing the ripple-carry adder

        Returns
        -------
        QuantumCircuit
            QuantumCircuit object of size self.nq
        """
        high_bit_index = self.nq-1

        # initialize the a and b registers
        a_indices = [2*i+2 for i in range(self.nbits)]
        b_indices = [2*i+1 for i in range(self.nbits)]
        for index_list, value in zip([a_indices, b_indices], [self.a, self.b]):
            self._initialize_value(index_list, value)

        # compute the carry bits, c_i, in order using the MAJ ladder
        for a_i in a_indices:
            self._MAJ(a_i-2, a_i-1, a_i)

        # write the final carry bit value to the high bit register
        self.circ.cx(a_indices[-1], high_bit_index)

        # erase the carry bits in reverse order using the UMA ladder
        for a_i in reversed(a_indices):
            self._UMA(a_i-2, a_i-1, a_i)

        if self.measure:
            # measure the circuit
            self.circ.measure_all()

        return self.circ


if __name__ == '__main__':
    adder = RCAdder(nbits=6, a=0, b=0, use_toffoli=True, measure=True)
    circ = adder.gen_circuit()
    print(circ)

