# %matplotlib inline
# Importing standard Qiskit libraries and configuring account
from qiskit import QuantumCircuit, QuantumRegister, ClassicalRegister, execute, Aer, IBMQ
from qiskit.compiler import transpile, assemble
from qiskit.tools.jupyter import *
from qiskit.visualization import *
# import math
# Loading your IBM Q account(s)
# IBMQ.load_account()

backend = Aer.get_backend("qasm_simulator")

from ArithmaticFunctions import *

def convolute_shift_func(circuit, control_reg, acc_reg, power):
    """Implement bit shift.
    @param circuit: QuantumCircuit containing operands
    @param control_reg: QuantumRegister containing qubit that defines power (all operations conditional on this)
    @param acc_reg: QuantumRegister containing integer to shift
    @param power: int number of places to shift (if control_reg[power])"""
    ## NOTE: Power is the number of places to shift. Power 1 is controlled by
    ## control register's 2nd bit, which is control_reg[1]
    ## and refers to shift of 2^1, or shift one place
    
    # Flip all qubits in accumulator to allow ripple carry
#     for qubit in range(len(acc_reg)):
#         circuit.x(acc_reg[qubit])
        
    for target_bit in range(len(acc_reg) - 1):
        # shift bit by correct number of places
#         circuit.ccx(control_reg[power], acc_reg[target_bit], acc_reg[target_bit + power])

        # Circuit will flip target bit before evaluating. Preempt that first flip.
        circuit.cx(control_reg[power], acc_reg[target_bit])
        # Flip next bit if carry occurs
        for flipped_bit in range(target_bit, len(acc_reg) - 1):
            ## Carry occurs if target now 0 (if flip makes it 1).
            circuit.cx(control_reg[power], acc_reg[flipped_bit])
            circuit.ccx(control_reg[power], acc_reg[flipped_bit], acc_reg[flipped_bit + 1])
            circuit.cx(control_reg[power], acc_reg[flipped_bit])
            # Flip bit that's being shifted
            # circuit.cx(control_reg[power], acc_reg[target_bit])

def ripple_carry_bits(circuit, control_bit, acc_reg, scratch_reg, start_shifting_at):
    """Continue a carry operation
    This assumes that whatever simple op triggered the ripple has already occurred.
    For example, if the ripple started with addition of 1 to the least significant bit,
    manually flip and then call function.
    
    @param circuit: QuantumCircuit containing operands
    @param control_bit: single bit controling whether op begins.
    @param acc_reg: QuantumRegister containing integer to shift
    @param scratch_reg: Quantum registered in initial |0..0> state used to store carry result of each simple addition.
    @param start_shifting_at: index in acc_reg of bit to shift
    
    TESTING - Passed minimal test cases July 18 at 8:24am Pacific Time
    """
    
    """
    NOTE: Power is the number of places to shift. Power 1 is controlled by
    control register's 2nd bit, which is control_reg[1]
    and refers to shift of 2^1, or shift one place
    """
    
    circuit.ccx(control_bit, acc_reg[start_shifting_at], scratch_reg[start_shifting_at])
    circuit.cx(control_bit, acc_reg[start_shifting_at])
    
    # After flipping target, state |0> means we should have carried... =>
    # before flipping target, state |1> means we should carry.
    
    # Flip each bit, if it's 1, carry to the next bit, and flip it back.
    for bit in range(start_shifting_at + 1, len(acc_reg) - 1):
        circuit.ccx(acc_reg[bit], scratch_reg[bit - 1], scratch_reg[bit])
        circuit.cx(scratch_reg[bit - 1], acc_reg[bit])
        
    

def c_ripple_subtract(circuit, control_bit, min_reg, scratch_reg, start_shifting_at):
    """Continue a subtract operation
    
    @param circuit: QuantumCircuit containing operands
    @param control_bit: single bit controling whether op begins.
    @param min_reg: QuantumRegister containing integer to transform
    @param scratch_reg: Quantum register in initial |0..0> state used to store carry result of each simple addition.
    @param start_shifting_at: index in acc_reg of bit to shift
    """
    
    circuit.x(min_reg[start_shifting_at])
    circuit.ccx(min_reg[start_shifting_at], control_bit, scratch_reg[start_shifting_at])
    circuit.x(min_reg[start_shifting_at])
    
    # Check how deep the borrow goes. Flip a scratch bit for every borrow op.
    for bit in range(start_shifting_at + 1, len(min_reg)):
        circuit.x(min_reg[bit])
        circuit.ccx(min_reg[bit], scratch_reg[bit - 1], scratch_reg[bit])
        circuit.x(min_reg[bit])
        
    # Perform every borrow indicated in scratch, and then flip targeted bit.
    for bit in range(len(scratch_reg) - 2, -1, -1):
        circuit.cx(scratch_reg[bit], min_reg[bit+1])
        
    # Finally flip original target bit
    circuit.x(min_reg[start_shifting_at])

def carry_op(circuit, a_reg, acc_reg, scratch_reg):
    """Used as 'CARRY' operator for circuit illustated in arxiv.org/quant-ph/9511018
    
    @param circuit: QuantumCircuit
    @param a_reg: QuantumRegister containing addend
    @param acc_reg: QuantumRegister 'b' starts with addend and accumulates sum
    @param scratch_reg: QuantumRegister holds carry bits at each step"""
    for bit in range(len(a_reg)):
        circuit.cx(a_reg[bit], acc_reg[bit])
        circuit.x(acc_reg[bit])
        circuit.cx(acc_reg[bit], scratch_reg[bit])
        circuit.x(acc_reg[bit])

def add_to_b_in_place(circuit, a_reg, b_reg, scratch_reg):
    """|a > | b > => |a > |a+b >
    
    TESTING - Passed minimal test cases July 18 at 8:24am Pacific Time
    """
        
    for bit in range(len(a_reg)):
        circuit.reset(scratch_reg)
        ripple_carry_bits(
            circuit=circuit,
            control_bit=a_reg[bit],
            acc_reg=b_reg,
            scratch_reg=scratch_reg,
            start_shifting_at=bit)

def sub_b_from_a_in_place(circuit, minnd_reg, subtrhnd_reg, scratch_reg):
    """Subtract subtrahend integer b from minuend integer a in register a
    
    @param circuit: QuantumCircuit containing other parameters
    @param minnd_reg: QuantumRegister transformed from minuend to difference
    @param subtrhnd_reg: QuantumRegister containing subtrahend
    @param scratch_reg: QuantumRegister in initial |0...0> state used as borrow flags
    for each primitive op (final bit indicates negative difference)
    """
    
    

def bit_shift_left(circuit, register, places=1):
    """
    TESTING - Passed minimal test cases July 21 at 2:30 Pacific Time
    """
    # zero out all trailing bits
    num_bits = len(register)
    for rollover in range(num_bits - places, num_bits):
        circuit.reset(register[rollover])
        
    # swap every bit 'places' forward, with last bits wrapping around to beginning
    for bit in range(num_bits - 1,places-1,-1):
        circuit.swap(register[bit], register[bit - places])

def c_copy_register(circuit, control_bit, origin, dest):
    
    """sets contents of dest with contents of origin if control_bit
    WARNING - perform a reset before use
    """
    circuit.reset(dest)
    for bit in range(len(dest)):
        circuit.ccx(control_bit, origin[bit], dest[bit])

def multiply(circuit, multiplicand, multiplier, scratch_zero_check, scratch_carrier, prod_accumulator):
    
    c_copy_register(circuit=circuit, control_bit=multiplier[0], origin=multiplicand, dest=prod_accumulator)
    
    for bit in range(1, len(multiplier)):
        # free up scratch space
        circuit.reset(scratch_carrier)
        circuit.reset(scratch_zero_check)
        
        # shift multiplicand one space left, to match magnitude of current multiplier bit
        bit_shift_left(circuit=circuit, register=multiplicand, places=1)
        
        # copy multiplicand into scratch register only if multiplicand bit, else keep scratch register |0>
        c_copy_register(circuit=circuit, control_bit=multiplier[bit], origin=multiplicand, dest=scratch_zero_check)
        
        # add that scratch term (shifted multiplicand or zero) to accumulated product
        add_to_b_in_place(circuit=circuit, a_reg=scratch_zero_check, b_reg=prod_accumulator, scratch_reg=scratch_carrier)

n_reg = QuantumRegister(3, 'n')
m_reg = QuantumRegister(3, 'm')
scratch_a = QuantumRegister(3, 'sca')
scratch_b = QuantumRegister(3, 'scb')
scratch_y = QuantumRegister(3, 'accm')
class_reg = ClassicalRegister(3, 'y')
qc = QuantumCircuit(n_reg, m_reg, scratch_a, scratch_b, scratch_y, class_reg)

# Register n set to 5
qc.x(n_reg[0])
qc.x(n_reg[1])
# qc.x(n_reg[2])
# Register m set to
# qc.x(m_reg[0])
qc.x(m_reg[1])

n_reg = QuantumRegister(3, 'n')
m_reg = QuantumRegister(2, 'm')
scratch_a = QuantumRegister(3, 'sca')
class_reg = ClassicalRegister(3, 'y')

backend = Aer.get_backend('qasm_simulator')

n_reg = QuantumRegister(3, 'n')
scratch_a = QuantumRegister(3, 'sca')
qc = QuantumCircuit(n_reg, m_reg, scratch_a, class_reg)

# Set control bits
qc.x(m_reg[1])

# 1 - 1 = 0

"""
PASSED July 29 10:58am Pacific Time
"""

qc.x(n_reg[0])
qc.barrier()
c_ripple_subtract(circuit=qc, min_reg=n_reg, control_bit=m_reg[1], scratch_reg=scratch_a, start_shifting_at=0)
qc.draw()



# 2 - 1 = 1
"""
PASSED July 29 10:58am Pacific Time
"""

qc.x(n_reg[1])
qc.barrier()
c_ripple_subtract(circuit=qc, min_reg=n_reg, control_bit=m_reg[True], scratch_reg=scratch_a, start_shifting_at=0)
# qc.draw()

# 4 - 1 = 3
"""
PASSED July 29 10:58am Pacific Time
"""

qc.x(n_reg[2])

c_ripple_subtract(circuit=qc, min_reg=n_reg, control_bit=m_reg[True], scratch_reg=scratch_a, start_shifting_at=0)
qc.draw()

# 6 - 1 = 5
"""
PASSED July 29 10:58am Pacific Time
"""

qc.x(n_reg[1])
qc.x(n_reg[2])

c_ripple_subtract(circuit=qc, min_reg=n_reg, control_bit=n_reg[True], scratch_reg=scratch_a, start_shifting_at=0)
qc.draw()

# 5 - 1 = 4
"""
PASSED July 29 10:58am Pacific Time
"""

qc.x(n_reg[0])
qc.x(n_reg[2])

c_ripple_subtract(circuit=qc, min_reg=n_reg, control_bit=m_reg[True], scratch_reg=scratch_a, start_shifting_at=0)
# c_draw()

# 5 - 2 = 3
"""
PASSED July 29 10:58am Pacific Time
"""

qc.x(n_reg[0])
qc.x(n_reg[2])

c_ripple_subtract(circuit=qc, min_reg=n_reg, control_bit=m_reg[True], scratch_reg=scratch_a, start_shifting_at=1)
# qc.draw()

# 7 - 2 = 5
qc.x(n_reg[0])
qc.x(n_reg[1])
qc.x(n_reg[2])

c_ripple_subtract(circuit=qc, min_reg=n_reg, control_bit=m_reg[True], scratch_reg=scratch_a, start_shifting_at=1)
qc.draw()

# 7 - 0 = 7
qc.x(n_reg[0])
qc.x(n_reg[1])
qc.x(n_reg[2])
qc.barrier()
c_ripple_subtract(circuit=qc, min_reg=n_reg, control_bit=m_reg[False], scratch_reg=scratch_a, start_shifting_at=1)
qc.draw()

# 1 - 2 = ?
qc.x(n_reg[0])

c_ripple_subtract(circuit=qc, min_reg=n_reg, control_bit=m_reg[True], scratch_reg=scratch_a, start_shifting_at=1)
# c_draw()

for bit in range(len(n_reg)):
    qc.measure(n_reg[bit], class_reg[bit])
    
execute(qc, backend=backend, shots=1024).result().get_counts()

multiply(
    circuit=qc,
    multiplier=n_reg,
    multiplicand=m_reg,
    scratch_zero_check=scratch_b,
    scratch_carrier=scratch_a,
    prod_accumulator=scratch_y)

# add_to_b_in_place(circuit=qc, a_reg=n_reg, b_reg=m_reg, scratch_reg=scratch_a)
# bit_shift_left(circuit=qc, register=m_reg, places=1)

# ripple_carry_bits(circuit=qc, control_bit=n_reg[0],acc_reg=m_reg,scratch_reg=scratch_reg, start_shifting_at=0)
# bit_shift_left(circuit=qc, register=m_reg, places=1)
for index in range(len(m_reg)):
    qc.measure(scratch_y[index], class_reg[index])

qc.draw()

qc_opt = transpile(qc, optimization_level=1, backend=backend)

qc_opt.draw()



simulate = execute(qc_opt, backend=backend, shots=1024).result()

simulate.get_counts()

qc.reset(n_reg[0])

print(max_memory_mb)

