import numpy as np
from scipy.optimize import minimize
import qiskit
from qiskit import Aer, QuantumCircuit
from qiskit.utils import QuantumInstance
from matplotlib import pyplot as plt

from qiskit_machine_learning.neural_networks import CircuitQNN

# linear entangler (as in scales linearly with trash qubits)
def get_entangler_map(L, num_trash, i_permut=1):
    result = []
    nums = list(range(L)) # here was the problem, it doesnt like when list elements are taken from numpy
    nums_compressed = nums.copy()[:L-num_trash]
    nums_trash = nums.copy()[-num_trash:]
    #print(nums, nums_compressed, nums_trash)
    # combine all trash qubits with themselves
    for trash_q in nums_trash[:-1]:
        result.append((trash_q+1,trash_q))
    # combine each of the trash qubits with every n-th
    repeated = list(nums_trash) * (L-num_trash) # repeat the list of trash indices cyclicly
    for i in range(L-num_trash):
        result.append((repeated[i_permut + i], nums_compressed[i]))
    return result

def QAE_Ansatz(L,num_trash, insert_barriers=False, parametrized_gate = "ry", entangling_gate = "cz"):
    entanglement = [get_entangler_map(L,num_trash,i_permut) for i_permut in range(num_trash)]
    return qiskit.circuit.library.TwoLocal(L, 
                                           parametrized_gate,
                                           entangling_gate,
                                           entanglement,
                                           reps=num_trash,
                                           insert_barriers=insert_barriers,
                                          )

QAE_Ansatz(6,2,insert_barriers=True).draw("mpl")

QAE_Ansatz(5,2,insert_barriers=True).draw("mpl")

QAE_Ansatz(7,2,insert_barriers=True).draw("mpl")

QAE_Ansatz(6,3,insert_barriers=True).draw("mpl")

# more complicated and as far as i can tell exponential
def get_entangler_map_exp(L, num_trash, i_permut=0):
    result = []
    nums = list(range(L))
    nums_compressed = nums.copy()[:L-num_trash]
    nums_trash = nums.copy()[-num_trash:]
    #print(nums, nums_compressed, nums_trash)
    # combine all trash qubits with themselves
    for trash_q in nums_trash[:-1]:
        result.append((trash_q+1,trash_q))
    # combine each of the trash qubits with every n-th
    permutations = list(itertools.permutations(nums_trash))
    for _,trash_q in enumerate(permutations[i_permut]):
        for comp_q in nums_compressed[_%num_trash::num_trash]: # combine each trash_q with even or odd site
            result.append((trash_q, comp_q))
    return result

# Use Qiskit layer generator thingy because now it has everything it needs to get a gradient!
# Note that it has redundant Ry rotations at the end on the compression qubits, didnt find a clever way how to get rid of them, but this shouldnt change too much

# variational ansatz sub-circuit
ansatz = QAE_Ansatz(6,2)

# measurement circuit
qreg = qiskit.QuantumRegister(L, 'q')
creg = qiskit.ClassicalRegister(2, 'c')
measurement_circ = qiskit.QuantumCircuit(qreg, creg)
measurement_circ.measure(qreg[4], creg[0])
measurement_circ.measure(qreg[5], creg[1])
circ2 = ansatz + measurement_circ
circ2.decompose()
#print(circ2)
#measurement_circ.draw("mpl")
circ2.draw("mpl")

qi_qasm = QuantumInstance(Aer.get_backend('qasm_simulator'), shots=10)
qi_sv = QuantumInstance(Aer.get_backend('statevector_simulator'))

qnn4 = CircuitQNN(circ2, [], circ2.parameters, sparse=False, quantum_instance=qi_qasm)

input4 = np.random.rand(qnn4.num_inputs)
weights4 = np.random.rand(qnn4.num_weights)

qnn4.forward(input4, weights4)

qnn4.backward(input4, weights4)

from qiskit.algorithms.variational_algorithm import VariationalAlgorithm

VariationalAlgorithm(ansatz = circ2, optimizer = None)

qiskit.__version__

import qiskit
import numpy as np

L = 6

def create_map(L):
    return [(i,i+1) for i in np.arange(L-1)]

map1 = create_map(L)
print(map1)

map_by_hand1 = [(0, 1), (1, 2), (2, 3), (3, 4), (4, 5)]
print(map_by_hand1)

map1 == map_by_hand1

ansatz = qiskit.circuit.library.TwoLocal(L,"ry","cz",[map_by_hand1],reps=2, insert_barriers=True)
ansatz.draw("mpl")

ansatz = qiskit.circuit.library.TwoLocal(L,"ry","cz",[map1],reps=2, insert_barriers=True)
ansatz.draw("mpl")

map1 = get_entangler_map(6,2)
map2 = get_entangler_map(6,2,1)

print(map1)
print(map2)

ansatz = qiskit.circuit.library.TwoLocal(L,"ry","cz",[map1,map2],reps=2, insert_barriers=True)
ansatz.draw("mpl")



ansatz = qiskit.circuit.library.TwoLocal(L,"ry","cz",[map_by_hand],reps=2, insert_barriers=True)
ansatz.draw("mpl")

